fixed checkpointing for tensors

argonne-lcf · Nov 29, 2023 · 8a1fb5a · 8a1fb5a
1 parent 6221b33
commit 8a1fb5a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/dlio_benchmark/framework/tf_framework.py b/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()

diff --git a/dlio_benchmark/framework/torch_framework.py b/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()