Prevent unnecessary re-initialization of the graph when model has unused parameters (#7799)

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)

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

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)