Support for models having partially non trainable parameters (#7058)

* Support for models having partially non trainable parameters

orttraining/orttraining/core/framework/module_gradient_graph_builder.cc

@@ -22,15 +22,15 @@ Status ModuleGradientGraphBuilder::Initialize(std::istream& model_istream,
   config_ = config;
 
   // Handle original model inputs, outputs and trainable initializers.
-  // We need to move the trainable initializers to graph inputs and keep the order in config,
-  // it's possible that the graph already has some trainable initializers in graph inputs,
-  // so we need to NOT assign these trainable initializers to the user inputs list.
+  // We need to move all the initializers to graph inputs and keep the order in config,
+  // it's possible that the graph already has some initializers in graph inputs,
+  // so we need to NOT assign these initializers to the user inputs list.
   Graph& graph = model_->MainGraph();
-  std::unordered_set<std::string> initializer_names_to_train_set(config.initializer_names_to_train.begin(),
-                                                                 config.initializer_names_to_train.end());
+  std::unordered_set<std::string> initializer_names(config.initializer_names.begin(),
+                                                    config.initializer_names.end());
   const std::vector<const NodeArg*>& graph_inputs = graph.GetInputsIncludingInitializers();
   for (auto& node_arg : graph_inputs) {
-    if (initializer_names_to_train_set.find(node_arg->Name()) == initializer_names_to_train_set.end()) {
+    if (initializer_names.find(node_arg->Name()) == initializer_names.end()) {
       training_graph_info_.user_input_names.emplace_back(node_arg->Name());
     }
   }
@@ -42,14 +42,16 @@ Status ModuleGradientGraphBuilder::Initialize(std::istream& model_istream,
 
   training_graph_info_.initializer_names_to_train.assign(config.initializer_names_to_train.begin(),
                                                          config.initializer_names_to_train.end());
+  training_graph_info_.initializer_names.assign(config.initializer_names.begin(),
+                                                config.initializer_names.end());
 
   std::vector<const NodeArg*> input_args;
   for (const auto& input_name : training_graph_info_.user_input_names) {
     input_args.emplace_back(graph.GetNodeArg(input_name));
   }
 
-  // Remove the training initializers from the graph and move them to graph inputs.
-  for (const auto& initializer_name : training_graph_info_.initializer_names_to_train) {
+  // Remove all the initializers from the graph and move them to graph inputs.
+  for (const auto& initializer_name : training_graph_info_.initializer_names) {
     const NodeArg* node_arg = graph.GetNodeArg(initializer_name);
     ORT_ENFORCE(node_arg != nullptr);
     input_args.emplace_back(node_arg);

orttraining/orttraining/core/framework/module_gradient_graph_builder.h

@@ -16,6 +16,8 @@ namespace training {
  * The training configuration options.
  */
 struct ModuleGradientGraphBuilderConfiguration {
+  // The names of the weights.
+  std::vector<std::string> initializer_names{};
   // The names of the weights to train.
   std::vector<std::string> initializer_names_to_train{};
   // The names of inputs that require gradient.
@@ -35,6 +37,8 @@ struct TrainingGraphInfo {
   std::vector<std::string> user_input_names{};
   // Map from user input names to corresponding user input grad names for those user inputs that require grad.
   std::unordered_map<std::string, std::string> user_input_grad_names{};
+  // All initializers (trainable as well as non trainable).
+  std::vector<std::string> initializer_names{};
   // Trainable initializers.
   std::vector<std::string> initializer_names_to_train{};
   // Trainable initializer grad names, ordered according to initializer_names_to_train.
@@ -60,9 +64,9 @@ class ModuleGradientGraphBuilder {
   Status Initialize(std::istream& model_istream, const ModuleGradientGraphBuilderConfiguration& config);
 
   /**
-   * Build the gradient graph and split it to forward and backward graphs.
+   * Build the gradient graph.
    * @param input_shapes_ptr The pointer to vector of concrete shapes of the user inputs.
-   * @return The status of the gradient graph building and forward/backward graphs splitting.
+   * @return The status of the gradient graph building.
    */
   Status Build(const std::vector<std::vector<int64_t>>* input_shapes_ptr = nullptr);
 
@@ -73,8 +77,8 @@ class ModuleGradientGraphBuilder {
   std::string GetGradientModel() const;
 
   /**
-   * Get the split graphs information.
-   * @return The split graphs information.
+   * Get the training graphs information.
+   * @return The training graphs information.
    */
   TrainingGraphInfo GetTrainingGraphInfo() const { return training_graph_info_; }
 

orttraining/orttraining/python/orttraining_pybind_state.cc

@@ -501,6 +501,7 @@ py::class_<TrainingAgent>(m, "TrainingAgent", R"pbdoc(This is the main class use
       m, "ModuleGradientGraphBuilderConfiguration",
       R"pbdoc(Configuration information for module gradient graph builder.)pbdoc");
   module_gradient_graph_builder_config.def(py::init())
+      .def_readwrite("initializer_names", &ModuleGradientGraphBuilderConfiguration::initializer_names)
       .def_readwrite("initializer_names_to_train", &ModuleGradientGraphBuilderConfiguration::initializer_names_to_train)
       .def_readwrite("input_names_require_grad", &ModuleGradientGraphBuilderConfiguration::input_names_require_grad)
       .def_readwrite("use_invertible_layernorm_grad",
@@ -511,6 +512,7 @@ py::class_<TrainingAgent>(m, "TrainingAgent", R"pbdoc(This is the main class use
   training_graph_info.def(py::init())
       .def_readwrite("user_input_names", &TrainingGraphInfo::user_input_names)
       .def_readwrite("user_input_grad_names", &TrainingGraphInfo::user_input_grad_names)
+      .def_readwrite("initializer_names", &TrainingGraphInfo::initializer_names)
       .def_readwrite("initializer_names_to_train", &TrainingGraphInfo::initializer_names_to_train)
       .def_readwrite("initializer_grad_names_to_train", &TrainingGraphInfo::initializer_grad_names_to_train)
       .def_readwrite("user_output_names", &TrainingGraphInfo::user_output_names)

orttraining/orttraining/python/training/ortmodule.py

@@ -107,7 +107,9 @@ class ORTModule(torch.nn.Module):
             Finally, we instantiate the ONNX Runtime InferenceSession.
             '''
             # TODO: using pytorch for evaluation for now. We will use ORT for evaluation later.
-            if not self._is_training:
+            # TODO: If the model is being executed with the gradient disabled (inside torch.no_grad() context for example),
+            # leverage pytorch model for now.
+            if not self._is_training():
                 return self._original_module(*inputs, **kwargs)
 
             # Exporting module to ONNX for the first time
@@ -122,16 +124,25 @@ class ORTModule(torch.nn.Module):
                 self._get_inference_graph_and_init_gradient_graph_builder(
                     *inputs, **kwargs)
 
+            # Flag to indicate whether the gradient_graph needs to be built
+            build_gradient_graph = self._current_input_shape is None
             _, _, input_names_require_grad, new_input_shape = \
                 _ortmodule_io.parse_inputs_for_onnx_export(
                     self._original_module_parameters, self._onnx_inference, *inputs, **kwargs)
+            initializer_names_to_train_set_user_model = {name for name, param in
+                self._flattened_output_module.named_parameters() if param.requires_grad}
+            initializer_names_to_train_set_onnx_graph = set(self._onnx_graphs_info.initializer_names_to_train) \
+                if self._onnx_graphs_info else None
             # 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
-            if input_names_require_grad != self._input_names_require_grad:
+            if input_names_require_grad != self._input_names_require_grad or \
+                initializer_names_to_train_set_user_model != initializer_names_to_train_set_onnx_graph:
                 self._input_names_require_grad = input_names_require_grad
                 self._initialize_module_gradient_graph_builder()
+                # Trigger the rebuilding of the gradient graph
+                build_gradient_graph = True
 
-            if self._current_input_shape is None:
+            if build_gradient_graph:
                 self._current_input_shape = new_input_shape
                 self._build_training_graph()
                 self._create_training_session()
@@ -238,9 +249,19 @@ class ORTModule(torch.nn.Module):
                             # input_name is not found in the self._input_names_require_grad list
                             # Append None to results for each input that did not require grad
                             results.append(None)
+
                     # Append gradients of initializer to results
-                    results += [_ortvalue_to_torch_tensor(backward_output)
-                                for backward_output in backward_outputs[num_user_input_grads:]]
+                    # Go over each initializer, check if it required grad and append to results accordingly
+                    initializer_names_to_train_set = set(self._onnx_graphs_info.initializer_names_to_train) \
+                        if self._onnx_graphs_info else None
+                    initializer_index = num_user_input_grads
+                    for initializer_name in self._onnx_graphs_info.initializer_names:
+                        if initializer_name in initializer_names_to_train_set:
+                            results.append(_ortvalue_to_torch_tensor(backward_outputs[initializer_index]))
+                            initializer_index += 1
+                        else:
+                            results.append(None)
+
                     # The OrtValue has a shared_ptr to the data.
                     # At this point there are two shared_ptrs to the data, one through the
                     # OrtValue in the output iobinding, and the other through the copy in OrtDLManagedTensor.
@@ -287,7 +308,6 @@ class ORTModule(torch.nn.Module):
                     "The model's forward method has **kwargs parameter which is currently not supported.")
 
         self._onnx_inference = None
-        self._is_training = True
 
         # Related to training graph shape inference
         self._current_input_shape = None
@@ -297,6 +317,7 @@ class ORTModule(torch.nn.Module):
         self._module_gradient_graph_builder = None
         self._input_names_require_grad = None
         self._original_module_output_schema = None
+        self._onnx_graphs_info = None
 
         # Training model
         self._onnx_training = None
@@ -325,18 +346,26 @@ class ORTModule(torch.nn.Module):
             self._torch_alloc = self._torch_cuda_allocator.cuda_caching_allocator_raw_alloc_address()
             self._torch_free = self._torch_cuda_allocator.cuda_caching_allocator_raw_delete_address()
 
+    def _is_training(self):
+        return self._flattened_output_module.training and torch.is_grad_enabled()
+
     def _initialize_module_gradient_graph_builder(self):
         # TODO: PyTorch exporter bug: changes the initializer order in ONNX model
-        initializer_names = [p[0]
-                             for p in self._flattened_output_module.named_parameters()]
+        initializer_names = [name
+                             for name, _ in self._flattened_output_module.named_parameters()]
+        initializer_names_to_train = []
+        if self._is_training():
+            initializer_names_to_train = [name
+                for name, param in self._flattened_output_module.named_parameters() if param.requires_grad]
         onnx_initializer_names = {
             p.name for p in self._onnx_inference.graph.initializer}
-        initializer_names = [
-            p for p in initializer_names if p in onnx_initializer_names]
+        initializer_names_to_train = [
+            p for p in initializer_names_to_train if p in onnx_initializer_names]
 
         # Build full training graph
         grad_builder_config = C.ModuleGradientGraphBuilderConfiguration()
-        grad_builder_config.initializer_names_to_train = initializer_names
+        grad_builder_config.initializer_names = initializer_names
+        grad_builder_config.initializer_names_to_train = initializer_names_to_train
         grad_builder_config.input_names_require_grad = self._input_names_require_grad
         self._module_gradient_graph_builder = C.ModuleGradientGraphBuilder()
         self._module_gradient_graph_builder.initialize(
@@ -402,11 +431,9 @@ class ORTModule(torch.nn.Module):
                       self._save_onnx_prefix + '_training.onnx')
 
     def eval(self: T) -> T:
-        self._is_training = False
         self._flattened_output_module.eval()
 
     def train(self: T, mode: bool = True) -> T:
-        self._is_training = mode
         self._flattened_output_module.train(mode)
 
     def _convert_training_graph_input_to_list(self, *inputs, **kwargs):

orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py

@@ -174,6 +174,19 @@ class NeuralNetNonDifferentiableOutput(torch.nn.Module):
         
         return out1, mask1, out2, mask2     # intentionally place the non-differentiable output in the middle
 
+class NeuralNetPartialNoGradModel(torch.nn.Module):
+    def __init__(self, input_size, hidden_size, num_classes):
+        super(NeuralNetPartialNoGradModel, self).__init__()
+
+        self.fc1 = torch.nn.Linear(input_size, hidden_size).requires_grad_(False)
+        self.relu = torch.nn.ReLU()
+        self.fc2 = torch.nn.Linear(hidden_size, num_classes)
+
+    def forward(self, model_input):
+        out = self.relu(self.fc1(model_input))
+        out = self.fc2(out)
+        return out
+
 # TODO: This is a workaround for the problem that pytest is still cleaning up the previous test
 # while the next task already start. 
 @pytest.fixture(autouse=True)
@@ -1404,3 +1417,59 @@ def test_uint8_input_and_output():
 
     assert y1 is not None
     assert y2 is not None and y2.dtype == torch.uint8
+
+def test_model_partially_requires_grad():
+    device = 'cuda'
+    N, D_in, H, D_out = 64, 784, 500, 10
+    model = NeuralNetPartialNoGradModel(D_in, H, D_out).to(device)
+    model = ORTModule(model)
+    x = torch.randn(N, D_in, device=device)
+
+    # Make sure no exception is raised
+    output = model(x)
+
+    loss = torch.sum(output)
+    loss.backward()
+
+def test_model_wrapped_inside_torch_no_grad():
+    device = 'cuda'
+    N, D_in, H, D_out = 64, 784, 500, 10
+    model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device)
+    model = ORTModule(model)
+    x = torch.randn(N, D_in, device=device)
+
+    # Make sure no exception is raised
+    with torch.no_grad():
+        output = model(x)
+
+def test_model_initializer_requires_grad_changes_from_one_forward_to_next():
+    device = 'cuda'
+    N, D_in, H, D_out = 64, 784, 500, 10
+    model = NeuralNetPartialNoGradModel(D_in, H, D_out).to(device)
+    model.fc1.requires_grad_(True)
+    model = ORTModule(model)
+    x = torch.randn(N, D_in, device=device)
+    assert model._original_module.fc1.weight.grad is None
+    assert model._original_module.fc1.bias.grad is None
+
+    # Make sure no exception is raised
+    output = model(x)
+    loss = torch.sum(output)
+    loss.backward()
+    training_session1 = model._training_session
+    weight_grad_2 = model._original_module.fc1.weight.grad
+    bias_grad_2 = model._original_module.fc1.bias.grad
+    assert weight_grad_2 is not None
+    assert bias_grad_2 is not None
+
+    model._original_module.fc1.requires_grad_(False)
+    output = model(x)
+    loss = torch.sum(output)
+    loss.backward()
+    training_session2 = model._training_session
+    weight_grad_3 = model._original_module.fc1.weight.grad
+    bias_grad_3 = model._original_module.fc1.bias.grad
+
+    assert training_session1 != training_session2
+    assert torch.equal(weight_grad_2, weight_grad_3)
+    assert torch.equal(bias_grad_2, bias_grad_3)