diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index e726a62105..8c54e93ddb 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -158,11 +158,6 @@ class GraphExecutionManager(ABC): self._optimized_onnx_model = onnx.load_model_from_string(self._graph_builder.get_model()) self._graph_info = self._graph_builder.get_graph_info() - # TODO: Explore ways to make self._graph_info.initializer_names and self._graph_info.initializer_names_to_train - # a set (unordered_set in the backend) that does not require a copy on each reference. - self._graph_initializer_names = set(self._graph_info.initializer_names) - self._graph_initializer_names_to_train = set(self._graph_info.initializer_names_to_train) - def _get_session_config(self): """Creates and returns the session configuration to be used for the ExecutionAgent""" providers = None @@ -202,6 +197,12 @@ class GraphExecutionManager(ABC): # 3. Export the user model under self._export_training_flag mode # Return True if the model needed to be exported, False if no export was required. + # Note: Model is only exported when: + # 1. Model has never been exported before. + # 2. Model input schema has changed (changes in inputs requiring gradient, shape, boolean inputs values change, etc) + # Model is not re-exported when the model parameters change. This can happen when the model is a stateful model, + # or the user explicitly changed model parameters after the onnx export. + schema = _io._extract_schema({'args': copy.copy(inputs), 'kwargs': copy.copy(kwargs)}) if self._onnx_model and schema == self._input_info.schema: # All required models have already been exported previously @@ -314,3 +315,8 @@ class GraphExecutionManager(ABC): # It is assumed here that the order and names of the inputs and outputs are not modified by the backend in any way # and are kept as they appear in the exported onnx model. self._graph_builder.initialize(self._onnx_model.SerializeToString(), grad_builder_config) + + # TODO: Explore ways to make self._graph_info.initializer_names and self._graph_info.initializer_names_to_train + # a set (unordered_set in the backend) that does not require a copy on each reference. + self._graph_initializer_names = set(initializer_names) + self._graph_initializer_names_to_train = set(initializer_names_to_train) diff --git a/orttraining/orttraining/python/training/ortmodule/_training_manager.py b/orttraining/orttraining/python/training/ortmodule/_training_manager.py index 97d9802467..96c634105a 100644 --- a/orttraining/orttraining/python/training/ortmodule/_training_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_training_manager.py @@ -241,8 +241,12 @@ class TrainingManager(GraphExecutionManager): def _reinitialize_graph_builder(self, input_info): """Return true if the module graph builder was reinitialized""" + # Model could have unused parameters which are dropped after export and so not a part of self._graph_initializer_names_to_train. + # To see if any trainable initializers changed, compare self._graph_initializer_names_to_train + # with initializers in module named_parameters that are known to the onnx graph. initializer_names_to_train_set_user_model = {name for name, param in - self._flattened_module.named_parameters() if param.requires_grad} + self._flattened_module.named_parameters() + if param.requires_grad and name in self._graph_initializer_names} # If inputs requiring gradient change from forward to the next, the module_gradient_graph_builder # needs to be reinitialized so it can compute the backward output for the new inputs that require_grad diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index 8ebb9c1c61..3a6607895d 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -16,7 +16,7 @@ from collections import OrderedDict from collections import namedtuple from inspect import signature -from onnxruntime.training.ortmodule import ORTModule, _utils +from onnxruntime.training.ortmodule import ORTModule, _utils, _io import _test_helpers # Import autocasting libs @@ -2597,3 +2597,25 @@ def test_stateless_model_unspecified_device(): ort_y = ort_model(ort_x) _test_helpers.assert_values_are_close(pt_y, ort_y) + +@pytest.mark.parametrize("model", + [(UnusedBeginParameterNet(784, 500, 400, 10)), + (UnusedMiddleParameterNet(784, 500, 400, 10)), + (UnusedEndParameterNet(784, 500, 400, 10))]) +def test_unused_parameters_does_not_unnecssarily_reinitilize(model): + device = 'cuda' + + N, D_in, H1, H2, D_out = 64, 784, 500, 400, 10 + model = model.to(device) + ort_model = ORTModule(copy.deepcopy(model)) + training_manager = ort_model._execution_manager(ort_model._is_training()) + + x = torch.randn(N, D_in, device=device) + _ = ort_model(x) + + input_info = _io.parse_inputs_for_onnx_export(training_manager._module_parameters, + training_manager._onnx_model, + x, + {}) + + assert not training_manager._reinitialize_graph_builder(input_info)