Added RequiredGrad attribute to YieldOp (#6657)

* Added required_grad attribute to YieldOp

* Chagened YieldOp attribute to hold the indices of the required gradient outputs from the count, and removed the code reordering the outputs.

* Changed backward_output_grad_names to a map from backward output gradient name to the corresponding output index.

orttraining/orttraining/core/framework/module_gradient_graph_builder.cc

@@ -50,7 +50,9 @@ Status ModuleGradientGraphBuilder::Initialize(std::istream& model_istream,
 
   // 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) {
-    input_args.emplace_back(graph.GetNodeArg(initializer_name));
+    const NodeArg* node_arg = graph.GetNodeArg(initializer_name);
+    ORT_ENFORCE(node_arg != nullptr);
+    input_args.emplace_back(node_arg);
     graph.RemoveInitializedTensor(initializer_name);
   }
 
@@ -188,42 +190,35 @@ void ModuleGradientGraphBuilder::AddYieldOp() {
     }
   }
 
-  // Yield inputs include all user outputs, those require output gradients come first, so Yield Op can use their shapes
-  // to infer Op output shapes.
-  std::vector<std::string> user_output_names_require_grad;
-  std::vector<std::string> user_output_names_no_grad;
-  training_graph_info_.backward_output_grad_names.clear();
-  for (const auto& name : training_graph_info_.user_output_names) {
+  // YieldOps required_grad attribute specifies the indices of the required gradients.
+  ONNX_NAMESPACE::AttributeProto required_grad;
+  const std::string attribute_name = "required_grad";
+  required_grad.set_name(attribute_name);
+  required_grad.set_type(ONNX_NAMESPACE::AttributeProto::INTS);
+
+  training_graph_info_.backward_output_grad_names_map.clear();
+  for (std::size_t i = 0; i < training_graph_info_.user_output_names.size(); ++i) {
+    const auto& name = training_graph_info_.user_output_names[i];
     std::string grad_name = name + "_grad";
     if (non_backward_user_output_grad_names.find(grad_name) == non_backward_user_output_grad_names.end()) {
-      user_output_names_require_grad.emplace_back(name);
-      training_graph_info_.backward_output_grad_names.emplace_back(grad_name);
-    } else {
-      user_output_names_no_grad.emplace_back(name);
+      training_graph_info_.backward_output_grad_names_map.insert(std::make_pair(grad_name, i));
+      required_grad.add_ints(static_cast<int64_t>(i));
     }
   }
 
-  // Reorder the user outputs.
-  training_graph_info_.user_output_names.clear();
-  for (const auto& name : user_output_names_require_grad) {
-    training_graph_info_.user_output_names.emplace_back(name);
-  }
-
-  for (const auto& name : user_output_names_no_grad) {
-    training_graph_info_.user_output_names.emplace_back(name);
-  }
-
   std::vector<NodeArg*> yield_input_node_args;
   std::vector<NodeArg*> yield_output_node_args;
   for (const auto& name : training_graph_info_.user_output_names) {
     yield_input_node_args.emplace_back(gradient_graph.GetNodeArg(name));
   }
 
-  for (const auto& name : training_graph_info_.backward_output_grad_names) {
-    yield_output_node_args.emplace_back(gradient_graph.GetNodeArg(name));
+  for (const auto& element : training_graph_info_.backward_output_grad_names_map) {
+    yield_output_node_args.emplace_back(gradient_graph.GetNodeArg(element.first));
   }
 
-  gradient_graph.AddNode("YieldOp", "YieldOp", "Yield Op", yield_input_node_args, yield_output_node_args, {}, kMSDomain);
+  NodeAttributes attributes({{attribute_name, required_grad}});
+
+  gradient_graph.AddNode("YieldOp", "YieldOp", "Yield Op", yield_input_node_args, yield_output_node_args, &attributes, kMSDomain);
 }
 
 void ModuleGradientGraphBuilder::ReorderOutputs() {

orttraining/orttraining/core/framework/module_gradient_graph_builder.h

@@ -41,8 +41,9 @@ struct TrainingGraphInfo {
   std::vector<std::string> initializer_grad_names_to_train{};
   // The user outputs.
   std::vector<std::string> user_output_names{};
-  // The user output grad names that are actual required by the backward graph.
-  std::vector<std::string> backward_output_grad_names{};
+  // The user output grad names that are actually required by the backward graph 
+  // mapped to the index of the correspoinding output of inference graph.
+  std::unordered_map<std::string, size_t> backward_output_grad_names_map{};
 };
 
 class ModuleGradientGraphBuilder {

orttraining/orttraining/core/graph/training_op_defs.cc

@@ -2220,19 +2220,27 @@ Return true if all elements are true and false otherwise.
       .Output(0, "outputs_grad", "Gradient of outputs returned from pytorch.", "T", OpSchema::Variadic,
               /*is_homogeneous*/ false,
               /*min_arity*/ 1)
+      .Attr(
+        "required_grad",
+        "The indices of the outputs that require gradient outputs.",
+        AttributeProto::INTS)
       .TypeConstraint("T", OpSchema::all_tensor_types(), "Allow inputs and outputs to be any kind of tensor.")
       .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
-        // Assume the outputs and gradients are one-to-one matching
-        // TODO: The contrain is relaxed for now 
-        // ORT_ENFORCE(ctx.getNumInputs() == ctx.getNumOutputs(), "Yield op doesn't have the same number of inputs and output");
-
-        for (size_t i = 0; i < ctx.getNumOutputs(); ++i) {
-          propagateElemTypeFromInputToOutput(ctx, i, i);
-          auto typeProto = ctx.getInputType(i);
+        const std::string attribute_name = "required_grad";
+        auto required_grads = ctx.getAttribute(attribute_name);
+        if (nullptr == required_grads) {  // attribute not present
+          fail_type_inference("Value of attribute ", attribute_name, " not specified");
+        }
+        ORT_ENFORCE(ctx.getNumOutputs() == static_cast<size_t> (required_grads->ints_size()));
+        for (size_t i = 0, n = static_cast<size_t> (required_grads->ints_size()); i < n; ++i) {
+          size_t j = static_cast<size_t> (required_grads->ints(static_cast<int>(i)));
+          ORT_ENFORCE(ctx.getNumInputs() > j);
+          propagateElemTypeFromInputToOutput(ctx, j, i);
+          auto typeProto = ctx.getInputType(j);
           if (!hasShape(*typeProto)) {
             continue;
           }
-          propagateShapeFromInputToOutput(ctx, i, i);
+          propagateShapeFromInputToOutput(ctx, j, i);
         }
       });
 }

orttraining/orttraining/python/orttraining_pybind_state.cc

@@ -491,7 +491,7 @@ void addObjectMethodsForTraining(py::module& m) {
       .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)
-      .def_readwrite("backward_output_grad_names", &TrainingGraphInfo::backward_output_grad_names);
+      .def_readwrite("backward_output_grad_names_map", &TrainingGraphInfo::backward_output_grad_names_map);
 
   py::class_<ModuleGradientGraphBuilder> module_gradient_graph_builder(m, "ModuleGradientGraphBuilder");
   module_gradient_graph_builder.def(py::init([]() { return onnxruntime::make_unique<ModuleGradientGraphBuilder>(); }))

orttraining/orttraining/python/training/ortmodule.py

@@ -258,11 +258,14 @@ class ORTModule(torch.nn.Module):
                 # Use IO binding
                 # Push user output grads to ONNX backend.
                 backward_grad_output_ortvalue = []
-                for grad_output in grad_outputs[:len(self._onnx_graphs_info.backward_output_grad_names)]:
-                    # Force torch tensors to be contiguous before converting into OrtValue
+
+                # backward_output_grad_names_map only contains the subset of module outputs that need a gradient,
+                # we filter out the invalid entries in grad_outputs, accessing using the mapped index.
+
+                for _, i in self._onnx_graphs_info.backward_output_grad_names_map.items():
+                    grad_output = grad_outputs[i]
                     if not grad_output.is_contiguous():
                         grad_output = grad_output.contiguous()
-
                     backward_grad_output_ortvalue.append(onnxruntime.OrtValue.ortvalue_from_data_ptr(list(grad_output.size()), _utils.dtype_torch_to_numpy(
                         grad_output.dtype), grad_output.device.type, _utils.get_device_index(grad_output.device), grad_output.data_ptr()))
 

orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py

@@ -34,6 +34,38 @@ class NeuralNetSinglePositionalArgument(torch.nn.Module):
         out = self.fc2(out)
         return out
 
+class NeuralNetMultiplePositionalArgumentsMultipleOutputs0(torch.nn.Module):
+    def __init__(self, input_size, hidden_size, num_classes):
+        super(NeuralNetMultiplePositionalArgumentsMultipleOutputs0, self).__init__()
+
+        self.fc1 = torch.nn.Linear(input_size, hidden_size)
+        self.fc2 = torch.nn.Linear(input_size, hidden_size)
+        self.relu1 = torch.nn.ReLU()
+        self.relu2 = torch.nn.ReLU()
+
+    def forward(self, input1, input2):
+        model_input = input1 + input2
+        out1 = self.fc1(model_input)
+        out2 = self.fc2(model_input)
+        out1 = self.relu1(out1)
+        out2 = self.relu2(out2)
+        return out1, out2
+
+class NeuralNetMultiplePositionalArgumentsMultipleOutputs1(torch.nn.Module):
+    def __init__(self, input_size, hidden_size, num_classes):
+        super(NeuralNetMultiplePositionalArgumentsMultipleOutputs1, self).__init__()
+
+        self.fc1 = torch.nn.Linear(input_size, hidden_size)
+        self.relu = torch.nn.ReLU()
+        self.fc2 = torch.nn.Linear(hidden_size, num_classes)
+
+    def forward(self, input1, input2):
+        model_input = input1 + input2
+        out1 = self.fc1(model_input)
+        out1 = self.relu(out1)
+        out2 = self.fc2(out1)
+        return out1, out2
+
 class NeuralNetMultiplePositionalArguments(torch.nn.Module):
     def __init__(self, input_size, hidden_size, num_classes):
         super(NeuralNetMultiplePositionalArguments, self).__init__()
@@ -550,6 +582,71 @@ def test_changes_input_requires_grad_reinitializes_module_gradient_graph_builder
     model(x)
     assert module_gradient_graph_builder != model._module_gradient_graph_builder
 
+@pytest.mark.parametrize("device", ['cuda'])
+def test_input_requires_grad_backward_creates_input_grad_as_required0(device):
+    N, D_in, H, D_out = 32, 784, 500, 10
+    model = NeuralNetMultiplePositionalArgumentsMultipleOutputs0(D_in, H, D_out).to(device)
+    model = ORTModule(model)
+    x1 = torch.randn(N, D_in, device=device, requires_grad=True)
+    x2 = torch.randn(N, D_in, device=device, requires_grad=True)
+
+    y1, _ = model(x1, x2)
+    s1 = y1.sum()
+    s1.backward()
+    assert x1.grad is not None and x2.grad is not None
+
+    # named_params[0] and named_params[1] correspond to weight and bias for fc1, similarly
+    # named_params[2] and named_params[3] correspond to weight and bias for fc2.
+    named_params = list(model.named_parameters())
+    assert torch.count_nonzero(named_params[0][1].grad) > 0
+    assert torch.count_nonzero(named_params[1][1].grad) > 0
+    assert named_params[2][1].grad is None or torch.count_nonzero(named_params[2][1].grad) == 0
+    assert named_params[3][1].grad is None or torch.count_nonzero(named_params[3][1].grad) == 0
+
+    # Reset gradients
+    for param in named_params:
+        param[1].grad = None
+
+    _, y2 = model(x1,x2)
+    s2 = y2.sum()
+    s2.backward()
+    named_params = list(model.named_parameters())
+    assert named_params[0][1].grad is None or torch.count_nonzero(named_params[0][1].grad) == 0
+    assert named_params[1][1].grad is None or torch.count_nonzero(named_params[1][1].grad) == 0
+    assert torch.count_nonzero(named_params[2][1].grad) > 0
+    assert torch.count_nonzero(named_params[3][1].grad) > 0
+
+
+
+@pytest.mark.parametrize("x1_requires_grad, x2_requires_grad", [(True, True), (True, False), (False, False), (False, True)])
+def test_input_requires_grad_backward_creates_input_grad_as_required1(x1_requires_grad, x2_requires_grad):
+
+    def run_step(model, x1, x2):
+        y1, y2 = model(x1, x2)
+        s = y2.sum()
+        s.backward()
+
+    N, D_in, H, D_out = 32, 784, 500, 10
+    device = 'cuda'
+    pt_model = NeuralNetMultiplePositionalArgumentsMultipleOutputs1(D_in, H, D_out).to(device)
+    ort_model = ORTModule(pt_model)
+    pt_x1 = torch.randn(N, D_in, device=device, requires_grad=x1_requires_grad)
+    pt_x2 = torch.randn(N, D_in, device=device, requires_grad=x2_requires_grad)
+
+    ort_x1 = pt_x1.clone().detach()
+    ort_x2 = pt_x2.clone().detach()
+    ort_x1.requires_grad = x1_requires_grad
+    ort_x2.requires_grad = x2_requires_grad
+
+    run_step(pt_model, pt_x1, pt_x2)
+    run_step(ort_model, ort_x1, ort_x2)
+
+    assert not x1_requires_grad or ort_x1.grad is not None
+    assert not x2_requires_grad or ort_x2.grad is not None
+    assert not x1_requires_grad or torch.allclose(ort_x1.grad, pt_x1.grad)
+    assert not x2_requires_grad or torch.allclose(ort_x2.grad, pt_x2.grad)
+    _test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model)
+
 def test_gpu_reserved_memory_with_torch_no_grad():
     device = 'cuda'