Revert "Add buffer static input tests to cudagraph trees (pytorch#130402)"

This reverts commit 80236dc.

Reverted pytorch#130402 on behalf of https://github.com/clee2000 due to broke lint for torch/_functorch/_aot_autograd/subclass_utils.py https://github.com/pytorch/pytorch/actions/runs/9948630877/job/27483551649 https://hud.pytorch.org/pytorch/pytorch/commit/80236dca90b0874cb2b6f9c9fa5f159c55726401 lint was green on PR, probably a landrace ([comment](pytorch#130393 (comment)))

Author: pytorchmergebot

test/inductor/test_cudagraph_trees.py

@@ -1855,7 +1855,7 @@ def run_static_input_param_test(self, fn_eager, num_graphs):
 
                 self.assertEqual(self.get_manager().new_graph_id().id, num_graphs)
 
-        def _module_test(self, mod, name="weight", param_wrapping=True):
+        def _module_test(self, mod):
             with torch.device("cuda"):
 
                 def fn(x, mod):
@@ -1878,14 +1878,11 @@ def run_test():
                         self.assertEqual(exp_grad, compiled_grad)
 
                 run_test()
-                old_attr = getattr(mod, name)
-                modified_attr = torch.rand_like(old_attr)
-                if param_wrapping:
-                    modified_attr = torch.nn.Parameter(modified_attr)
-                setattr(mod, name, modified_attr)
+                old = mod.weight.data
+                mod.weight.data = torch.rand_like(mod.weight.data)
                 run_test()
                 # Run original version to verify we reuse the other recording
-                setattr(mod, name, old_attr)
+                mod.weight.data = old
                 run_test()
 
                 # Fwd + bwd graphs for each version of the function => 4 graphs
@@ -1910,18 +1907,6 @@ def test_multi_dispatch_single_compile_builtin_module(self):
             # Note: Linear is a builtin module so we enable that config setting above
             self._module_test(torch.nn.Linear(2, 3, device="cuda"))
 
-        @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
-        def test_multi_dispatch_single_compile_builtin_module_buffers(self):
-            # Verify that we don't recompile when changing the buffer of a builtin module
-            # and that we record another cudagraph
-            self._module_test(
-                torch.nn.BatchNorm1d(2, device="cuda"),
-                name="running_mean",
-                param_wrapping=False,
-            )
-
-        @torch._inductor.config.patch("triton.cudagraphs", True)
         @torch._dynamo.config.patch("error_on_recompile", True)
         @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
         def test_multi_dispatch_custom_module(self):
@@ -1939,30 +1924,6 @@ def forward(self, x):
                 TestModule(torch.nn.Parameter(torch.rand([2, 2], device="cuda")))
             )
 
-        @torch._dynamo.config.patch("error_on_recompile", True)
-        @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
-        def test_multi_dispatch_custom_module_buffer(self):
-            # Test that we can correctly dispatch multiple graphs
-            # if buffers of a custom module change
-            class TestModule(torch.nn.Module):
-                def __init__(self, param, buf) -> None:
-                    super().__init__()
-                    self.weight = param
-                    self.register_buffer("buf", buf)
-
-                def forward(self, x):
-                    return x * self.weight + self.buf
-
-            self._module_test(
-                TestModule(
-                    torch.nn.Parameter(torch.rand([2, 2], device="cuda")),
-                    torch.rand([2, 2], device="cuda"),
-                ),
-                name="buf",
-                param_wrapping=False,
-            )
-
-        @torch._inductor.config.patch("triton.cudagraphs", True)
         @torch._dynamo.config.patch("error_on_recompile", True)
         @torch._dynamo.config.patch("inline_inbuilt_nn_modules", True)
         def test_multi_dispatch_child_node(self):