preserve module signature multiple calls (#137999)

avikchaudhuri · facebook-github-bot · commit 2755f65abd77 · 2024-10-17T12:51:03.000-07:00
Summary: Pull Request resolved: #137999 Test Plan: fixed tests Reviewed By: tugsbayasgalan Differential Revision: D64406945
diff --git a/test/export/test_export.py b/test/export/test_export.py
@@ -6168,11 +6168,16 @@ def test_unflatten_no_unroll(self):
         inp = (torch.ones(1),)
 
         class N(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.const = torch.ones(1) * 4
+                self.buf = torch.nn.Buffer(torch.ones(1) * 4)
+
             def forward(self, x, b):
                 if b:
-                    return x + 1
+                    return x + self.const + 1
                 else:
-                    return x + 2
+                    return x + 2 * (self.buf + 1) - self.const
 
         class M(torch.nn.Module):
             def __init__(self):
@@ -6188,37 +6193,39 @@ def forward(self, x):
         m = M()
         eager_result = m(*inp)
 
-        if not is_retracebility_test(self._testMethodName):
-            ep = export(M(), inp, preserve_module_call_signature=("n",))
-            with self.assertRaisesRegex(
-                ValueError,
-                "Cannot unflatten multiple calls to module n while preserving its signature",
-            ):
-                torch.export.unflatten(ep)
+        def test(ep, swap):
+            epm = ep.module()
+            ufm = torch.export.unflatten(ep)
 
-        ep = export(M(), inp)
-        print(ep)
-        epm = ep.module()
-        ufm = torch.export.unflatten(ep)
+            exported_result = epm(*inp)
+            self.assertTrue(torch.allclose(exported_result, eager_result))
 
-        exported_result = epm(*inp)
-        self.assertTrue(torch.allclose(exported_result, eager_result))
+            unflattened_result = ufm(*inp)
+            self.assertTrue(torch.allclose(unflattened_result, eager_result))
 
-        unflattened_result = ufm(*inp)
-        self.assertTrue(torch.allclose(unflattened_result, eager_result))
+            for fqn, mod in swap.items():
+                ufm.set_submodule(fqn, mod)
+            unflattened_result = ufm(*inp)
+            self.assertTrue(torch.allclose(unflattened_result, eager_result))
+
+        if not is_retracebility_test(self._testMethodName):
+            test(
+                export(M(), inp, preserve_module_call_signature=("n",)),
+                swap={"n": N(), "n@1": N()},
+            )
 
         class _N(torch.nn.Module):
             def forward(self, x):
-                return x + 1
+                return x + 5
 
         class _N_1(torch.nn.Module):
             def forward(self, x):
-                return x + 2
+                return x + 6
 
-        ufm.set_submodule("n", _N())
-        ufm.set_submodule("n@1", _N_1())
-        unflattened_result = ufm(*inp)
-        self.assertTrue(torch.allclose(unflattened_result, eager_result))
+        test(
+            export(M(), inp),
+            swap={"n": _N(), "n@1": _N_1()},
+        )
 
     def test_preserve_module_call_signature_unflatten_specialization(self):
         class N(torch.nn.Module):
@@ -6257,7 +6264,7 @@ def forward(self, x):
             unflattened_result = ufm(*inp)
             self.assertTrue(torch.allclose(unflattened_result, eager_result))
 
-    def test_unflatten_multiple_graphs_preserve_signature_error(self):
+    def test_unflatten_multiple_graphs_preserve_signature_no_error(self):
         class N(torch.nn.Module):
             def forward(self, x, b):
                 if b:
@@ -6271,33 +6278,31 @@ def __init__(self):
                 self.n = N()
 
             def forward(self, x):
+                x = x + 3
                 x = self.n(x, True)
-                x = x + 1
+                x = x + 4
                 x = self.n(x, False)
-                x = x + 1
+                x = x + 5
                 return x
 
         inp = (torch.ones(1),)
         m = M()
         eager_result = m(*inp)
 
-        if not is_retracebility_test(self._testMethodName):
-            ep = export(M(), inp, preserve_module_call_signature=("n",))
-            with self.assertRaisesRegex(
-                ValueError,
-                "Cannot unflatten multiple calls to module n while preserving its signature",
-            ):
-                torch.export.unflatten(ep)
+        def test(ep):
+            epm = ep.module()
+            ufm = torch.export.unflatten(ep)
 
-        ep = export(M(), inp)
-        epm = ep.module()
-        ufm = torch.export.unflatten(ep)
+            exported_result = epm(*inp)
+            self.assertTrue(torch.allclose(exported_result, eager_result))
 
-        exported_result = epm(*inp)
-        self.assertTrue(torch.allclose(exported_result, eager_result))
+            unflattened_result = ufm(*inp)
+            self.assertTrue(torch.allclose(unflattened_result, eager_result))
 
-        unflattened_result = ufm(*inp)
-        self.assertTrue(torch.allclose(unflattened_result, eager_result))
+        if not is_retracebility_test(self._testMethodName):
+            test(export(M(), inp, preserve_module_call_signature=("n",)))
+
+        test(export(M(), inp))
 
     def test_unflatten_multiple_graphs_state(self):
         class N(torch.nn.Module):
@@ -6331,6 +6336,16 @@ def forward(self, x):
         m = M()
         eager_result = m(*inp)
 
+        if not is_retracebility_test(self._testMethodName):
+            with self.assertRaisesRegex(
+                ValueError,
+                r"Found multiple calls of module n that mutate buffer n.buf",
+            ):
+                # Unflattening while preserving signatures is NYI for this case.
+                torch.export.unflatten(
+                    export(M(), inp, preserve_module_call_signature=("n",))
+                )
+
         ep = export(M(), inp)
         epm = ep.module()
         ufm = torch.export.unflatten(ep)
diff --git a/torch/_export/passes/collect_tracepoints_pass.py b/torch/_export/passes/collect_tracepoints_pass.py
@@ -66,23 +66,38 @@ def get_arg_spec(arg):
                     node.meta["nn_module_stack"].popitem()
                 else:
                     nn_module_stack = None
+
+        def copy_sig(sig):
+            from torch.export.exported_program import ModuleCallSignature
+
+            return ModuleCallSignature(
+                inputs=[],
+                outputs=[],
+                in_spec=sig.in_spec,
+                out_spec=sig.out_spec,
+                forward_arg_names=None,
+            )
+
         for module in gm.modules():
             if not isinstance(module, torch.fx.GraphModule):
                 continue
             for node in module.graph.nodes:
                 if node.op != "call_function":
                     continue
                 if node.target == torch.ops.higher_order._export_tracepoint:
+                    path = node.kwargs["path"]
+                    module_key = next(reversed(node.meta["nn_module_stack"]))
+                    if "@" in module_key:
+                        call_path = f"{path}@{module_key.split('@')[-1]}"
+                        if call_path not in self.specs:
+                            self.specs[call_path] = copy_sig(self.specs[path])
+                        path = call_path
+                    kind = node.kwargs["kind"]
                     for i, arg in enumerate(node.args):
-                        kind = node.kwargs["kind"]
                         if kind == "module_call_inputs":
-                            self.specs[node.kwargs["path"]].inputs.append(
-                                get_arg_spec(arg)
-                            )
+                            self.specs[path].inputs.append(get_arg_spec(arg))
                         elif kind == "module_call_outputs":
-                            self.specs[node.kwargs["path"]].outputs.append(
-                                get_arg_spec(arg)
-                            )
+                            self.specs[path].outputs.append(get_arg_spec(arg))
                         else:
                             raise AssertionError(f"Unknown tracepoint kind: {kind}")
                         if isinstance(arg, torch.fx.Node):
diff --git a/torch/export/_trace.py b/torch/export/_trace.py
@@ -505,25 +505,29 @@ def _get_module_hierarchy(mod: torch.nn.Module) -> Dict[str, str]:
 
 
 def _make_module_call_graph(
-    module_hierarchy: Dict[str, str],
     in_spec: TreeSpec,
     out_spec: TreeSpec,
     module_call_signatures: Dict[str, ModuleCallSignature],
     forward_arg_names: Optional[List[str]] = None,
 ) -> List[ModuleCallEntry]:
-    ret = [
+    original = [
         ModuleCallEntry(fqn=fqn, signature=module_call_signatures.get(fqn))
-        for fqn in module_hierarchy
+        for fqn in _EXPORT_MODULE_HIERARCHY  # type: ignore[union-attr]
     ]
-    assert ret[0].fqn == ""
-    ret[0].signature = ModuleCallSignature(
+    assert original[0].fqn == ""
+    original[0].signature = ModuleCallSignature(
         inputs=[],
         outputs=[],
         in_spec=in_spec,
         out_spec=out_spec,
         forward_arg_names=forward_arg_names,
     )
-    return ret
+    additional = [
+        ModuleCallEntry(fqn=fqn, signature=signature)
+        for fqn, signature in module_call_signatures.items()
+        if fqn not in _EXPORT_MODULE_HIERARCHY  # type: ignore[operator]
+    ]
+    return [*original, *additional]
 
 
 def _export_to_torch_ir(
@@ -1103,7 +1107,6 @@ def _get_module_call_graph(
 
     assert _EXPORT_MODULE_HIERARCHY is not None
     module_call_graph = _make_module_call_graph(
-        _EXPORT_MODULE_HIERARCHY,
         original_in_spec,
         out_spec,
         module_call_signatures,
diff --git a/torch/export/unflatten.py b/torch/export/unflatten.py
@@ -219,6 +219,19 @@ def __init__(
         if export_module.graph_signature.backward_signature is not None:
             raise ValueError("Unflattening on JointExportModule NYI")
 
+        preserved_module_targets_with_multiple_calls = [
+            entry.fqn.split("@")[0]
+            for entry in export_module.module_call_graph
+            if "@" in entry.fqn
+        ]
+        for buf in export_module.graph_signature.buffers_to_mutate.values():
+            for fqn in preserved_module_targets_with_multiple_calls:
+                if buf.startswith(fqn + "."):
+                    raise ValueError(
+                        f"Found multiple calls of module {fqn} that mutate buffer {buf}. "
+                        "Unflattening while preserving signatures is NYI for this case."
+                    )
+
         fqn_list = [entry.fqn for entry in export_module.module_call_graph]
         assert fqn_list[0] == ""
         export_graph = deepcopy(export_module.graph)
@@ -766,6 +779,9 @@ def __init__(
         self.verbose = False
 
         self.fqn, num_calls = self.module_stack[-1]
+        # generate call name for self.fqn
+        self.child_fqn = _call_name(self.fqn, num_calls + 1)
+
         if module is not None:
             self.module = module
         else:
@@ -779,16 +795,7 @@ def __init__(
 
         self.parent_call_module: Optional[torch.fx.Node] = None
         if parent is not None:
-            if self.fqn in module_call_graph and num_calls == 1:
-                raise ValueError(
-                    f"Cannot unflatten multiple calls to module {self.fqn} while preserving its signature "
-                    "because each of these calls might have generated a different specialized graph. "
-                    f"If the reason you want to preserve the signature is to swap {self.fqn} with another module, "
-                    "consider using _swap_modules() directly on the exported program instead of unflattening it."
-                )
-            # generate call name for self.fqn
-            child_fqn = _call_name(self.fqn, num_calls + 1)
-            accessor = _compute_accessor(parent.fqn, child_fqn)
+            accessor = _compute_accessor(parent.fqn, self.child_fqn)
             _add_submodule(parent.module, accessor, self.module)
             self.parent_call_module = parent.graph.call_module(accessor)
             self.seen_modules[self.module_id].append(
@@ -802,7 +809,7 @@ def __init__(
                 )
             )
 
-        signature = module_call_graph.get(self.fqn)
+        signature = module_call_graph.get(self.child_fqn)
         if signature is not None and self.parent is not None:
             assert signature.in_spec.num_children == 2
             args_spec = signature.in_spec.children_specs[0]
@@ -949,7 +956,7 @@ def remap_input(self, x):
     def finalize_outputs(self):
         orig_outputs = []
 
-        signature = self.module_call_graph.get(self.fqn)
+        signature = self.module_call_graph.get(self.child_fqn)
         if signature is not None and self.parent is not None:
             for output in signature.outputs:
                 if isinstance(output, (TensorArgument, SymIntArgument)):
@@ -1294,7 +1301,7 @@ def _sink_params(
         state_name = None
         for sn in inputs_to_state[node.name]:
             sn_split = sn.split(".")
-            if sn_split[: len(scope)] == scope:
+            if sn_split[: len(scope)] == [x.split("@")[0] for x in scope]:
                 state_name = sn_split
                 break
 
