Merge remote-tracking branch 'origin/feature/checkpoint' into feature…

…/checkpoint
argonne-lcf · Jan 9, 2024 · 3727e5a · 3727e5a
2 parents a397d11 + 63ff8c4
commit 3727e5a
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 7 deletions.
diff --git a/dlio_benchmark/configs/workload/megatron_deepspeed.yaml b/dlio_benchmark/configs/workload/megatron_deepspeed.yaml
@@ -1,4 +1,5 @@
-model: unet3d
+# 8 node run with 4 GPUs per node and TPSIZE=4 and PPSIZE=8
+model: megatron_deepspeed
 
 framework: pytorch
 

diff --git a/dlio_benchmark/data_generator/hdf5_generator.py b/dlio_benchmark/data_generator/hdf5_generator.py
@@ -44,7 +44,7 @@ def generate(self):
         """
         super().generate()
         np.random.seed(10)
-        samples_per_iter=max(1, int(32*1024*1024/self._args.record_length))
+        samples_per_iter=max(1, int(self._args.generation_buffer_size/self._args.record_length))
         record_labels = [0] * self.num_samples
         for i in dlp.iter(range(self.my_rank, int(self.total_files_to_generate), self.comm_size)):
             progress(i, self.total_files_to_generate, "Generating HDF5 Data")

diff --git a/dlio_benchmark/framework/tf_framework.py b/dlio_benchmark/framework/tf_framework.py
@@ -58,11 +58,8 @@ def __init__(self, profiling):
         if self.args.checkpoint_type == CheckpointLocationType.RANK_ZERO:
             rank_to_checkpoint = 0
         if rank_to_checkpoint == self.args.my_rank:
-            num_ranks = 1
-            if self.args.checkpoint_type == CheckpointLocationType.RANK_ZERO:
-                num_ranks = self.args.comm_size
             if self.args.model_size > 0:
-                self.model_state = {"a": self._get_tensor(self.args.model_size*num_ranks)}
+                self.model_state = {"a": self._get_tensor(self.args.model_size)}
             self.optimization_state = None
             if len(self.args.optimization_groups) > 0:
                 self.optimization_state = dict()
@@ -78,7 +75,7 @@ def __init__(self, profiling):
                 self.layer_state = dict()
                 for index, state in enumerate(self.args.layer_parameters):
                     if state > 0:
-                        self.layer_state[str(index)] = self._get_tensor(state*num_ranks)
+                        self.layer_state[str(index)] = self._get_tensor(state)
 
     def _get_tensor(self, size):
         return tf.random.uniform((int(size / 4),), maxval=100, dtype=tf.dtypes.int32)