[export] Downgrade captured buffers as normal constants. (pytorch#166777)

Summary:
make_fx() will register tensor constants as new buffers while tracing a shuffle graph for dynamo graph capture. This breaks the invariance that the resulting graph looks identical to the original eager model in terms of state dict.

So we need to de-register the buffers and set them as plain tensor constants.

Test Plan:
pytest test/export/test_experimental.py

Pull Request resolved: pytorch#166777
Approved by: https://github.com/tugsbayasgalan
ghstack dependencies: pytorch#166775, pytorch#166776

Author: zhxchen17

test/export/test_experimental.py

@@ -579,6 +579,26 @@ def forward(self, args_0):
         )
         self.assertEqual(gm(*test_inputs), foo(*test_inputs))
 
+    def test_dynamo_graph_capture_with_tensor_constant(self):
+        outer = torch.randn(2, 3)
+
+        class MyModel(torch.nn.Module):
+            def forward(self, x):
+                z = x + outer
+                return z
+
+        foo = MyModel()
+
+        def make_inputs():
+            return (torch.randn(2, 3),)
+
+        trace_inputs = make_inputs()
+        gm = dynamo_graph_capture_for_export(foo)(*trace_inputs)
+        test_inputs = make_inputs()
+        self.assertEqual(gm(*test_inputs), foo(*test_inputs))
+        self.assertEqual(len(list(gm.buffers())), len(list(foo.buffers())))
+        self.assertEqual(len(list(gm.parameters())), len(list(foo.parameters())))
+
     @unittest.skipIf(not TEST_CUDA, "CUDA not available")
     def test_dynamo_graph_capture_fx_graph_annotate_overlap_pass(self):
         class DummyOp(torch.autograd.Function):

torch/_dynamo/functional_export.py

@@ -450,6 +450,14 @@ def _suggest_or_raise_constraint_violation(
         raise constraint_violation_error
 
 
+def _normalize_shuffle_graph(shuffle_gm: torch.fx.GraphModule) -> None:
+    shuffle_gm.graph.eliminate_dead_code()
+    shuffle_gm.recompile()
+    for name, buffer in list(shuffle_gm.named_buffers()):
+        delattr(shuffle_gm, name)
+        setattr(shuffle_gm, name, buffer)
+
+
 @dataclass(frozen=True)
 class PyTreeifyOutput:
     graph_module: torch.fx.GraphModule
@@ -526,8 +534,7 @@ def backend_dummy(*example_inputs):
     in_shuffle_graph = make_fx(
         InShuffle(), tracing_mode="symbolic", proxy_module_inputs=True
     )(*flat_real_args)
-    in_shuffle_graph.graph.eliminate_dead_code()
-    in_shuffle_graph.recompile()
+    _normalize_shuffle_graph(in_shuffle_graph)
 
     output_node = next(iter(reversed(backend_input.graph_module.graph.nodes)))
 
@@ -575,8 +582,7 @@ def backend_dummy(*example_inputs):
     out_shuffle_graph = make_fx(
         out_shuffle, tracing_mode="symbolic", proxy_module_inputs=True
     )(*flat_out_shuffle_args)
-    out_shuffle_graph.graph.eliminate_dead_code()
-    out_shuffle_graph.recompile()
+    _normalize_shuffle_graph(out_shuffle_graph)
 
     assert out_shuffle.out_spec is not None
     return PyTreeifyOutput(