Fix bug in graph partitioner

Summary: Title

Test Plan: CI

Reviewed By: PaulZhang12

Differential Revision: D56688411

Author: tugsbayasgalan

test/export/test_export.py

@@ -3968,16 +3968,16 @@ def forward(self, b_pred, b_t, x, y):
             """\
 def forward(self, b_t, x, y):
     submod_3 = self.submod_1
-    add_1 = torch._higher_order_ops.wrap.wrap_with_set_grad_enabled(True, submod_3, b_t, x, y);  submod_3 = b_t = x = y = None
+    add_1 = torch._higher_order_ops.wrap.wrap_with_set_grad_enabled(True, submod_3, x, b_t, y);  submod_3 = x = b_t = y = None
     return (add_1,)""",
         )
 
         self.assertExpectedInline(
             str(exported_program.graph_module.true_graph_0.submod_1.code.strip()),
             """\
-def forward(self, b_t, x, y):
-    sub = torch.ops.aten.sub.Tensor(b_t, 1);  b_t = None
-    add = torch.ops.aten.add.Tensor(sub, x);  sub = x = None
+def forward(self, x, b_t, y):
+    sub = torch.ops.aten.sub.Tensor(x, 1);  x = None
+    add = torch.ops.aten.add.Tensor(sub, b_t);  sub = b_t = None
     add_1 = torch.ops.aten.add.Tensor(add, y);  add = y = None
     return add_1""",
         )

torch/_export/passes/replace_set_grad_with_hop_pass.py

@@ -160,7 +160,58 @@ def _maybe_inline_or_replace_with_hop(node: torch.fx.Node):
     return gm
 
 
-def replace_set_grad_with_hop_pass(gm: torch.fx.GraphModule):
+def replace_set_grad_with_hop_pass(gm: torch.fx.GraphModule, graph_signature):
+    original_output_nodes = ()
+    for node in gm.graph.nodes:
+        if node.op == "output":
+            original_output_nodes = node.args
+
+    gm = _replace_set_grad_with_hop_pass(gm)
+
+    updated_output_nodes = ()
+    for node in gm.graph.nodes:
+        if node.op == "output":
+            updated_output_nodes = node.args
+
+    # It could be updated because replace_set_grad_with_hop_pass adds new
+    # getitem nodes into the graph.
+    assert len(original_output_nodes) == len(updated_output_nodes)
+    old_output_name_to_new_output_name = {}
+    for k, v in zip(*original_output_nodes, *updated_output_nodes):
+        if isinstance(k, torch.fx.Node) and isinstance(k, torch.fx.Node):
+            old_output_name_to_new_output_name[k.name] = v.name
+        # If there is constant return in the end, it should be true for updated outputs too
+        else:
+            constant_types = (int, float, str, type(None))
+            assert k == v
+            assert isinstance(k, constant_types)
+
+    buffers_to_mutate_copy = graph_signature.buffers_to_mutate.copy()
+    user_inputs_to_mutate_copy = graph_signature.user_inputs_to_mutate.copy()
+    for k in old_output_name_to_new_output_name:
+        if k in graph_signature.buffers_to_mutate:
+            graph_signature.buffers_to_mutate[
+                old_output_name_to_new_output_name[k]
+            ] = buffers_to_mutate_copy[k]
+            if k not in old_output_name_to_new_output_name.values():
+                del graph_signature.buffers_to_mutate[k]
+
+        if k in graph_signature.user_inputs_to_mutate:
+            graph_signature.user_inputs_to_mutate[
+                old_output_name_to_new_output_name[k]
+            ] = user_inputs_to_mutate_copy[k]
+            if k not in old_output_name_to_new_output_name.values():
+                del graph_signature.user_inputs_to_mutate[k]
+
+    for i, k in enumerate(graph_signature.user_outputs):
+        if k in old_output_name_to_new_output_name:
+            new_k = old_output_name_to_new_output_name[k]
+            graph_signature.user_outputs[i] = new_k
+
+    return gm, graph_signature
+
+
+def _replace_set_grad_with_hop_pass(gm: torch.fx.GraphModule):
     new_gm = _sequential_split_and_maybe_inline_subgraphs(gm)
 
     # recursively call
@@ -169,7 +220,7 @@ def replace_set_grad_with_hop_pass(gm: torch.fx.GraphModule):
             subgm = getattr(new_gm, node.target)
             if not isinstance(subgm, torch.fx.GraphModule):
                 continue
-            new_subgm = replace_set_grad_with_hop_pass(subgm)
+            new_subgm = _replace_set_grad_with_hop_pass(subgm)
             setattr(new_gm, node.target, new_subgm)
 
     new_gm.recompile()

torch/export/_trace.py

@@ -539,7 +539,7 @@ def _compiling_state_context():
             replace_set_grad_with_hop_pass,
         )
 
-        gm = replace_set_grad_with_hop_pass(gm)
+        gm, graph_signature = replace_set_grad_with_hop_pass(gm, graph_signature)
 
     # Remove nn_module_stack, stack_trace metadata from all placeholders/inputs nodes.
     for _mod in gm.modules():

torch/fx/passes/split_module.py

@@ -457,6 +457,11 @@ def instantiate_node_partition_mapping(node):
     )
 
     already_constructed_attr_nodes = set()
+
+    # We actually need to insert the placeholder nodes in the original order
+    # otherwise graph signature will be wrong.
+    original_order = [node for node in m.graph.nodes if node.op == "placeholder"]
+
     for partition_name in construct_order_partitions:
         partition = partitions[partition_name]
 
@@ -475,8 +480,17 @@ def instantiate_node_partition_mapping(node):
         if keep_original_order:
             # first get the attr nodes required by this partition
             orig_mod_attr_nodes: List[Node] = [
-                orig_mod_env[key] for key in partition.inputs
+                orig_mod_env[key] for key in partition.inputs if key not in original_order
             ]
+
+            for node in original_order:
+                if node in already_constructed_attr_nodes:
+                    continue  # already added this attr to the base graph
+                base_mod_env, based_mod_attrs = construct_graph(
+                    node, base_mod_env, base_mod_attrs
+                )
+                already_constructed_attr_nodes.add(node)
+
             # Construct GraphModule for this partition
             for node in orig_mod_attr_nodes:  # type: ignore[attr-defined]
                 if node in already_constructed_attr_nodes: