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

Author: hariharan-devarajan

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
 

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()