Adding support for saving and loading train step info properties in the state dict and checkpoint file. (#10569)

* Adding optimization step and step parameter to the ORTTrainer constructor

* Added ORTTrainerOptions for optimization step

* Adding Train Step Info Settings to State Dictionary

* Adding train step info key

* Updating comments

* Reverting changes

* Updating test case for new state dict entry train_step_info

orttraining/orttraining/python/training/_utils.py

@@ -299,3 +299,18 @@ def state_dict_trainer_options_optimizer_name_key():
     """Returns the trainer options optimizer_name key name in the state dictionary"""
 
     return 'optimizer_name'
+
+def state_dict_train_step_info_key():
+    """Returns the train step info key name in the state dictionary"""
+
+    return 'train_step_info'
+
+def state_dict_train_step_info_optimization_step_key():
+    """Returns the train step info optimization step key name in the state dictionary"""
+
+    return 'optimization_step'
+
+def state_dict_train_step_info_step_key():
+    """Returns the train step info step key name in the state dictionary"""
+
+    return 'step'

orttraining/orttraining/python/training/orttrainer.py

@@ -968,13 +968,24 @@ class ORTTrainer(object):
         state_dict[_utils.state_dict_trainer_options_key()][D_size] = self.options.distributed.data_parallel_size
         state_dict[_utils.state_dict_trainer_options_key()][H_size] = self.options.distributed.horizontal_parallel_size
 
+    def _extract_train_step_info(self, state_dict):
+        """Extract train step info settings and save it into the state_dict"""
+
+        optimization_step = _utils.state_dict_train_step_info_optimization_step_key()
+        step = _utils.state_dict_train_step_info_step_key()
+
+        state_dict[_utils.state_dict_train_step_info_key()] = {}
+        state_dict[_utils.state_dict_train_step_info_key()][optimization_step] = self._train_step_info.optimization_step
+        state_dict[_utils.state_dict_train_step_info_key()][step] = self._train_step_info.step
+
     def state_dict(self, pytorch_format=False):
-        """Returns a dictionary with model, and optionally, optimizer states
+        """Returns a dictionary with model, train step info and optionally, optimizer states
 
         The returned dictionary contains the following information:
         - Model and optimizer states
         - Required ORTTrainerOptions settings
         - Distributed training information, such as but not limited to ZeRO
+        - Train step info settings
 
         Structure of the returned dictionary:
         - When `pytorch_format = False`
@@ -1094,6 +1105,21 @@ class ORTTrainer(object):
                         }
                     }
                 }
+            },
+            "train_step_info":
+            {
+                type: dict,
+                schema:
+                {
+                    "optimization_step":
+                    {
+                        type: int
+                    },
+                    "step":
+                    {
+                        type: int
+                    }
+                }
             }
         }
         - When `pytorch_format = True`
@@ -1136,6 +1162,9 @@ class ORTTrainer(object):
         # extract the relevant training configuration from the trainer and load them into the state_dict
         self._extract_trainer_options(state_dict)
 
+        # Extract train step info settings and load it into the state_dict
+        self._extract_train_step_info(state_dict)
+
         # add partition information in case of a distributed run
         if self.options.distributed.deepspeed_zero_optimization.stage > 0 or self.options.distributed.horizontal_parallel_size > 1:
             state_dict[_utils.state_dict_partition_info_key()] = self._training_session.get_partition_info_map()
@@ -1294,6 +1323,19 @@ class ORTTrainer(object):
         return current_state_dict[_utils.state_dict_optimizer_key()] if \
             _utils.state_dict_optimizer_key() in current_state_dict else {}
 
+    def _load_train_step_info(self, state_dict):
+        """Load the train step info settings from state dict"""
+
+        if _utils.state_dict_train_step_info_key() not in state_dict:
+            warnings.warn("Missing key: train_step_info in state_dict", UserWarning)
+            return
+
+        optimization_step = _utils.state_dict_train_step_info_optimization_step_key()
+        step = _utils.state_dict_train_step_info_step_key()
+
+        self._train_step_info.optimization_step = state_dict[_utils.state_dict_train_step_info_key()][optimization_step]
+        self._train_step_info.step = state_dict[_utils.state_dict_train_step_info_key()][step]
+
     def load_state_dict(self, state_dict, strict=True):
         """Loads state_dict containing model/optimizer states into ORTTrainer
 
@@ -1315,6 +1357,9 @@ class ORTTrainer(object):
             self._load_state_dict = partial(self._load_state_dict_impl, state_dict, strict=strict)
             return
 
+        # load the train step info settings
+        self._load_train_step_info(state_dict)
+
         # load states onto the frontend onnx graph
         optimizer_state_dict = self._load_state_dict_impl(state_dict, strict=strict)
 

orttraining/orttraining/test/python/orttraining_test_orttrainer_checkpoint_functions.py

@@ -325,22 +325,50 @@ def test_load_state_dict_holds_when_training_session_not_initialized():
 
 @pytest.mark.parametrize("state_dict, input_state_dict, error_key", [
     ({
-        'optimizer':{},
+        'model':{},
+        'optimizer':{}
     },
     {
+        'model':{},
         'optimizer':{},
         'trainer_options': {
             'optimizer_name': 'LambOptimizer'
         }
     },
+    'train_step_info'),
+    ({
+        'optimizer':{},
+        'train_step_info': {
+            'optimization_step': 0,
+            'step': 0
+        }
+    },
+    {
+        'optimizer':{},
+        'trainer_options': {
+            'optimizer_name': 'LambOptimizer'
+        },
+        'train_step_info': {
+            'optimization_step': 0,
+            'step': 0
+        }
+    },
     'model'),
     ({
-        'model':{}
+        'model':{},
+        'train_step_info': {
+            'optimization_step': 0,
+            'step': 0
+        }
     },
     {
         'model':{},
         'trainer_options': {
             'optimizer_name': 'LambOptimizer'
+        },
+        'train_step_info': {
+            'optimization_step': 0,
+            'step': 0
         }
     },
     'optimizer')])