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

Pull Request resolved: pytorch#130402
Approved by: https://github.com/eellison
ghstack dependencies: pytorch#130391, pytorch#130392, pytorch#130503, pytorch#130393

Author: mlazos

test/inductor/test_cudagraph_trees.py

@@ -1825,7 +1825,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):
+        def _module_test(self, mod, name="weight", param_wrapping=True):
             with torch.device("cuda"):
 
                 def fn(x, mod):
@@ -1848,11 +1848,14 @@ def run_test():
                         self.assertEqual(exp_grad, compiled_grad)
 
                 run_test()
-                old = mod.weight.data
-                mod.weight.data = torch.rand_like(mod.weight.data)
+                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)
                 run_test()
                 # Run original version to verify we reuse the other recording
-                mod.weight.data = old
+                setattr(mod, name, old_attr)
                 run_test()
 
                 # Fwd + bwd graphs for each version of the function => 4 graphs
@@ -1877,6 +1880,18 @@ 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):
@@ -1894,6 +1909,30 @@ 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):