From 1e917c879e3875d25413ac6586fa07f12262b709 Mon Sep 17 00:00:00 2001 From: mpapdiwala <99221944+mpapdiwala@users.noreply.github.com> Date: Thu, 24 Mar 2022 11:50:45 -0700 Subject: [PATCH] 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/python/training/_utils.py | 15 ++++++ .../orttraining/python/training/orttrainer.py | 47 ++++++++++++++++++- ...ng_test_orttrainer_checkpoint_functions.py | 32 ++++++++++++- 3 files changed, 91 insertions(+), 3 deletions(-) diff --git a/orttraining/orttraining/python/training/_utils.py b/orttraining/orttraining/python/training/_utils.py index 91ee336b10..0f882279a0 100644 --- a/orttraining/orttraining/python/training/_utils.py +++ b/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' diff --git a/orttraining/orttraining/python/training/orttrainer.py b/orttraining/orttraining/python/training/orttrainer.py index 93d2981432..af95713b47 100644 --- a/orttraining/orttraining/python/training/orttrainer.py +++ b/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) diff --git a/orttraining/orttraining/test/python/orttraining_test_orttrainer_checkpoint_functions.py b/orttraining/orttraining/test/python/orttraining_test_orttrainer_checkpoint_functions.py index ae67b72500..35e3d18c04 100644 --- a/orttraining/orttraining/test/python/orttraining_test_orttrainer_checkpoint_functions.py +++ b/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')])