fixed checkpointing for tensors

Author: hariharan-devarajan

dlio_benchmark/configs/workload/megatron_deepspeed.yaml

@@ -17,7 +17,7 @@ dataset:
 reader: 
   data_loader: pytorch
   batch_size: 1024
-  read_threads: 8
+  read_threads: 1
   file_shuffle: seed
   sample_shuffle: seed
 
@@ -27,7 +27,7 @@ train:
 
 checkpoint:
   checkpoint_folder: checkpoints/megatron-deepspeed
-  checkpoint_after_epoch: 1000
+  steps_between_checkpoints: 1000
   model_size: 30102
   type: independent
   optimization_groups: [1009254400, 865075200, 793600]

dlio_benchmark/framework/tf_framework.py

@@ -66,13 +66,13 @@ def __init__(self, profiling):
             self.optimization_state = None
             if len(self.args.optimization_groups) > 0:
                 self.optimization_state = dict()
-                tensor_array = []
+                tensor_array_size = 0
                 for index, state in enumerate(self.args.optimization_groups):
                     if state > 0:
-                        self.optimization_state[str(index)] = {'a': self._get_tensor(state*num_ranks),
-                                                               'b': self._get_tensor(state*num_ranks)}
-                        tensor_array.append(self._get_tensor(state*num_ranks))
-                self.optimization_state["combined"] = tensor_array
+                        self.optimization_state[str(index)] = {'a': self._get_tensor(state),
+                                                               'b': self._get_tensor(state)}
+                        tensor_array_size += state
+                self.optimization_state["combined"] = self._get_tensor(tensor_array_size)
             self.layer_state = None
             if len(self.args.layer_parameters) > 0:
                 self.layer_state = dict()

dlio_benchmark/framework/torch_framework.py

@@ -71,12 +71,12 @@ def __init__(self, profiling):
             self.optimization_state = None
             if len(self.args.optimization_groups) > 0:
                 self.optimization_state = dict()
-                tensor_array = []
+                tensor_array_size = 0
                 for index, state in enumerate(self.args.optimization_groups):
                     if state > 0:
                         self.optimization_state[str(index)] = {'a': self._get_tensor(state), 'b': self._get_tensor(state)}
-                        tensor_array.append(self._get_tensor(state))
-                self.optimization_state["combined"] = tensor_array
+                        tensor_array_size += state
+                self.optimization_state["combined"] = self._get_tensor(tensor_array_size)
             self.layer_state = None
             if len(self.args.layer_parameters) > 0:
                 self.layer_state = dict()