Add custom autograd function to prevent input passthrough on ORTModule (#7694)

* Changes for investigation

* Gradient for Identity

* Keep Identity betwen YieldOp and GraphOutput

* Revert debugging changes

* Add custom autograd fn to prevent input passthrough on ORTModule

* Add comment

Co-authored-by: Sherlock Huang <bahuang@OrtTrainingDev3.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>

onnxruntime/core/optimizer/identity_elimination.cc

@@ -55,7 +55,7 @@ Status EliminateIdentity::Apply(Graph& graph, Node& node, RewriteRuleEffect& rul
     graph.RemoveNode(node.Index());
     // update input node's output def to the graph output
     input_node.MutableOutputDefs()[output_idx] = output;
-    rule_effect = RewriteRuleEffect::kRemovedCurrentNode; 
+    rule_effect = RewriteRuleEffect::kRemovedCurrentNode;
   }
   return Status::OK();
 }
@@ -65,17 +65,25 @@ bool EliminateIdentity::SatisfyCondition(const Graph& graph, const Node& node, c
     return true;
   }
 
+  bool node_output_is_graph_output = !graph.GetNodeOutputsInGraphOutputs(node).empty();
+
   // relax the condition if Identity is connecting to graph output
-  if (node.GetOutputEdgesCount() != 0 || node.OutputDefs().size() != 1 ||
-    graph.GetNodeOutputsInGraphOutputs(node).empty())
+  if (node.GetOutputEdgesCount() != 0 ||
+      node.OutputDefs().size() != 1 ||
+      !node_output_is_graph_output) {
     return false;
+  }
 
   const Node* p_input_node = graph_utils::GetInputNode(node, 0);
-  if (p_input_node == nullptr)
+  if (p_input_node == nullptr) {
     return false;
-    
+  }
+  if (p_input_node->OpType() == "YieldOp" && node_output_is_graph_output) {
+    return false;
+  }
+
   // skip if the src arg is also a graph output
-  int src_arg_index = graph_utils::GetNodeOutputIndexFromOutputName(*p_input_node, node.InputDefs()[0]->Name());  
+  int src_arg_index = graph_utils::GetNodeOutputIndexFromOutputName(*p_input_node, node.InputDefs()[0]->Name());
   if (graph.IsOutput(p_input_node->OutputDefs()[src_arg_index]))
     return false;
 
@@ -87,7 +95,7 @@ bool EliminateIdentity::SatisfyCondition(const Graph& graph, const Node& node, c
     }
   }
   // condition not met if there are more than 1 consumer for the same src arg
-  if (count > 1) 
+  if (count > 1)
     return false;
 
   return true;

orttraining/orttraining/core/framework/gradient_graph_builder.cc

@@ -56,17 +56,31 @@ GradientGraphBuilder::GradientGraphBuilder(Graph* graph,
     }
 
     const Node* node = graph_->GetProducerNode(name);
-    if (!node) {
-      ORT_THROW(name, " couldn't find the producer node.");
-    }
-
-    if (forward_reachable_nodes.find(node) == forward_reachable_nodes.end()) {
-      non_differentiable_y_node_arg_names_.insert(name);
-      LOGS(logger_, INFO) << "The model weights and inputs are non-differentiable from " << name << ". "
-                          << "ORT will assume no gradient will be provided for " << name << ".";
+    if (node) {
+      if (forward_reachable_nodes.find(node) == forward_reachable_nodes.end()) {
+        non_differentiable_y_node_arg_names_.insert(name);
+        LOGS(logger_, INFO) << "The model weights and inputs are non-differentiable from " << name << ". "
+                            << "ORT will assume no gradient will be provided for " << name << ".";
+      } else {
+        y_node_args_.insert(node_arg);
+        y_nodes_.insert(node);
+      }
     } else {
-      y_node_args_.insert(node_arg);
-      y_nodes_.insert(node);
+      const std::vector<const NodeArg*>& graph_inputs = graph_->GetInputs();
+      bool is_graph_input = false;
+      for (const auto input : graph_inputs) {
+        if (input->Name() == name) {
+          is_graph_input = true;
+          break;
+        }
+      }
+
+      if (is_graph_input) {
+        LOGS(logger_, INFO) << "NodeArg " << name << " cannot find a producer node, "
+                                                     "but it's a graph input.";
+      } else {
+        ORT_THROW(name, ": couldn't find the producer node, and it's not a graph input.");
+      }
     }
   }
 

orttraining/orttraining/core/graph/gradient_builder.cc

@@ -1434,6 +1434,11 @@ IMPLEMENT_GRADIENT_BUILDER(GetExpGradient) {
               {GI(0)})};
 }
 
+IMPLEMENT_GRADIENT_BUILDER(GetIdentityGradient) {
+  return std::vector<NodeDef>{
+      NodeDef("Identity", {GO(0)}, {GI(0)})};
+}
+
 IMPLEMENT_GRADIENT_BUILDER(GetFlattenGradient) {
   return std::vector<NodeDef>{
       NodeDef("Shape", {I(0)}, {IA("input_shape")}),

orttraining/orttraining/core/graph/gradient_builder.h

@@ -73,6 +73,7 @@ DECLARE_GRADIENT_BUILDER(GetAbsGradient)
 DECLARE_GRADIENT_BUILDER(GetMinMaxGradient)
 DECLARE_GRADIENT_BUILDER(GetTileGradient)
 DECLARE_GRADIENT_BUILDER(GetATenOpGradient)
+DECLARE_GRADIENT_BUILDER(GetIdentityGradient)
 
 }  // namespace training
 }  // namespace onnxruntime

orttraining/orttraining/core/graph/gradient_builder_registry.cc

@@ -105,6 +105,7 @@ void GradientBuilderRegistry::RegisterGradientBuilders() {
   REGISTER_GRADIENT_BUILDER("Max", GetMinMaxGradient);
   REGISTER_GRADIENT_BUILDER("Tile", GetTileGradient);
   REGISTER_GRADIENT_BUILDER("ATenOp", GetATenOpGradient);
+  REGISTER_GRADIENT_BUILDER("Identity", GetIdentityGradient);
 };
 
 }  // namespace training

orttraining/orttraining/python/training/ortmodule/_io.py

@@ -9,6 +9,54 @@ import inspect
 import torch
 import warnings
 
+
+class _OutputIdentityOp(torch.autograd.Function):
+    '''Internal class used to prepend Identity ops in model's outputs
+
+    This class is required to support ONNX models which passthrough [some of] the models's inputs
+    directly to the graph output. This is an issue because ONNX Runtime cannot build proper
+    gradient graph based on this pattern.
+
+    Adding a direct Identity Op to the user model doesn't work as the ONNX exporter would optimize it away,
+    resulting in the same issue.
+
+    Therefore a custom Autograd function was introduced to add an Identity right before the output
+    in a way the ONNX exporter will not optimize it away.
+
+    Given the model below
+
+    .. code-block:: python
+
+        class PassthroughNet(torch.nn.Module):
+            def __init__(self, input_size, hidden_size, num_classes):
+                super(PassthroughNet, self).__init__()
+                self.fc1_1 = torch.nn.Linear(input_size, hidden_size)
+                self.relu1 = torch.nn.ReLU()
+                self.fc1_2 = torch.nn.Linear(hidden_size, num_classes)
+            def forward(self, input1, passthrough_input):
+                out1 = self.fc1_2(self.relu1(self.fc1_1(input1)))
+                # use shape from passthrough_input
+                out1 = out1.view(passthrough_input.size()[0], -1)
+                return out1, passthrough_input
+
+    We can see `passthrough_input` is part of both model input and output and the resulting
+    ONNX subgraph would contain something like `output2 -> output2`.
+
+    By prepending each model output to an :class:`_OutputIdentityOp` op, the resulting
+    onnx subgraph for this example would be  `passthrough_input -> Identity -> output2`.
+
+    TODO: Remove once PyTorch 1.8.2 or newer is released
+    '''
+    @staticmethod
+    def forward(ctx, input):
+        return torch.nn.Identity()(input)
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output
+    @staticmethod
+    def symbolic(g, self):
+        return g.op("Identity", self)
+
 class _PrimitiveType(object):
     _primitive_types = {int, bool, float}
     @staticmethod
@@ -301,7 +349,8 @@ def _transform_output_to_flat_tuple(data):
         if data is None:
             return
         elif isinstance(data, torch.Tensor):
-            flat_data.append(data)
+            identity = _OutputIdentityOp.apply
+            flat_data.append(identity(data))
         elif isinstance(data, abc.Sequence):
             for value in data:
                 _flatten_data(value, flat_data)

orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py

@@ -1038,6 +1038,46 @@ def test_input_requires_grad_backward_creates_input_grad_as_required1(x1_require
     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)
 
+
+@pytest.mark.parametrize("device", ['cuda'])
+def test_model_with_bypass_input(device):
+    class NeuralNetWithBypassInput(torch.nn.Module):
+        def __init__(self, input_size, hidden_size, num_classes):
+            super(NeuralNetWithBypassInput, self).__init__()
+
+            self.fc1_1 = torch.nn.Linear(input_size, hidden_size)
+            self.relu1 = torch.nn.ReLU()
+            self.fc1_2 = torch.nn.Linear(hidden_size, num_classes)
+
+        def forward(self, input1, bypass_input):
+            out1 = self.fc1_2(self.relu1(self.fc1_1(input1)))
+            # use shape from bypass_input
+            out1 = out1.view(bypass_input.size()[0], -1)
+            return out1, bypass_input
+
+    def run_step(model, x1, x2):
+        y1, y2 = model(x1, x2)
+        loss = y1.sum() + y2.sum()
+        loss.backward()
+        return y1, y2
+
+    N, D_in, H, D_out = 32, 784, 500, 10
+    pt_model = NeuralNetWithBypassInput(D_in, H, D_out).to(device)
+    ort_model = ORTModule(copy.deepcopy(pt_model))
+
+    pt_x1 = torch.randn(N, D_in, device=device, requires_grad=True)
+    pt_x2 = torch.randn(N, D_in, device=device, requires_grad=True)
+    ort_x1 = pt_x1.clone()
+    ort_x2 = pt_x2.clone()
+
+    # Both y1 and y2's gradients are not None.
+    pt_y1, pt_y2 = run_step(pt_model, pt_x1, pt_x2)
+    ort_y1, ort_y2 = run_step(ort_model, ort_x1, ort_x2)
+
+    _test_helpers.assert_values_are_close(pt_y1, ort_y1, atol=1e-06)
+    _test_helpers.assert_values_are_close(pt_y2, ort_y2, atol=1e-06)
+    _test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model)
+
 def test_gpu_reserved_memory_with_torch_no_grad():
     device = 'cuda'
 

orttraining/orttraining/test/python/orttraining_test_ortmodule_bert_classifier.py

@@ -378,9 +378,9 @@ def main():
     if not args.pytorch_only:
         model = ORTModule(model)
 
-    # TODO: change it to False to stop saving ONNX models
-    model._save_onnx = True
-    model._save_onnx_prefix = 'BertForSequenceClassification'
+    # Just for future debugging
+    model._execution_manager(model._is_training())._save_onnx = False
+    model._execution_manager(model._is_training())._save_onnx_prefix = 'BertForSequenceClassification'
 
     # Tell pytorch to run this model on the GPU.
     if torch.cuda.is_available() and not args.no_cuda: