Update on "[cudagraphs] Fix issue in collecting static_input_idxs"

bdhirsh · bdhirsh · commit 8ccbcc72ac84 · 2025-04-29T15:47:19.000-07:00
related to #152275 cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx ipiszy chenyang78 kadeng muchulee8 amjames chauhang aakhundov [ghstack-poisoned]
diff --git a/test/inductor/test_inductor_freezing.py b/test/inductor/test_inductor_freezing.py
@@ -1,5 +1,6 @@
 # Owner(s): ["module: inductor"]
 import contextlib
+import copy
 import functools
 import importlib
 import itertools
@@ -376,19 +377,33 @@ def foo(mod, x):
             mod(x)
 
     def test_static_indices_cudagraph(self):
-        mod = torch.nn.Linear(2, 2).to(self.device)
+        if self.device != "cuda":
+            return
+
+        mod1 = torch.nn.Sequential(
+            torch.nn.Linear(2, 2).to(self.device), torch.nn.Linear(2, 2).to(self.device)
+        )
+        mod2 = copy.deepcopy(mod1)
 
-        def fn(x):
-            return mod(x) + x
+        def fn(x, y, mod):
+            x.add_(1)
+            getattr(mod, "0").bias.add_(2)
+            getattr(mod, "1").weight.add_(3)
+            return mod(x) + y
 
-        x = torch.randn(2, 2, device=self.device)
+        x1 = torch.randn(2, 2, device=self.device)
+        y1 = torch.randn(2, 2, device=self.device)
+        x2 = x1.clone()
+        y2 = y1.clone()
 
         opt_fn = torch.compile(fn, mode="reduce-overhead")
 
         with torch.no_grad():
-            ref = fn(x)
-            res = opt_fn(x)
+            ref = fn(x1, y1, mod1)
+            res = opt_fn(x2, y2, mod2)
         self.assertEqual(ref, res)
+        self.assertEqual(x1, x2)
+        self.assertEqual(y1, y2)
 
     def test_rng_op(self):
         @torch.compile()
diff --git a/torch/_functorch/aot_autograd.py b/torch/_functorch/aot_autograd.py
@@ -19,7 +19,6 @@
 from torch._dynamo.utils import (
     CompileEventLogger,
     dynamo_timed,
-    is_parameter_freezing,
     preserve_rng_state,
     set_feature_use,
 )
@@ -1038,10 +1037,7 @@ def _try_get_metadata_from_dynamo(
         seen_sources.add(source)
         aot_autograd_arg_pos_to_source.append(source)
 
-        # For freezing, the params are not lifted in the inductor Fx graph, so
-        # don't mark the params as static.
-        if not is_parameter_freezing():
-            static_input_indices.append(i)
+        static_input_indices.append(i)
 
     # Collect the dynamo graph inputs
     # TODO(mlazos): Revisit if this is still needed. With Dynamo install ID
@@ -1057,12 +1053,8 @@ def _try_get_metadata_from_dynamo(
         # input[i] in dynamo is now:
         # input[i + len(extra_params)] in AOT,
         # where extra_params are the params/buffers that dynamo baked into the
-        # OutputGraph. The special case is freezing, where the params are not
-        # lifted.
-        if is_parameter_freezing():
-            actual_pos = pos
-        else:
-            actual_pos = pos + len(param_keys)
+        # OutputGraph
+        actual_pos = pos + len(param_keys)
 
         if "tensor_dict" in node.meta and node.meta["tensor_dict"].get(
             "_dynamo_static_input_type", None
diff --git a/torch/_inductor/freezing.py b/torch/_inductor/freezing.py
@@ -52,14 +52,21 @@ def replace_params_with_constants(
         in (MutationType.MUTATED_IN_GRAPH, MutationType.MUTATED_OUT_GRAPH)
     ]
 
+    static_indices_new = []
+    static_indices_offset = 0
     for i, (real_input, node) in enumerate(zip(flat_params, fake_inp_nodes)):
         if i in mutated_inps or i in aliased_input_args:
             preserved_arg_indices.append(i)
-            continue
-        replace_node_with_constant(gm, node, real_input)
+            if i in fw_metadata.static_input_indices:
+                new_static_index = i - static_indices_offset
+                static_indices_new.append(new_static_index)
+        else:
+            replace_node_with_constant(gm, node, real_input)
+            static_indices_offset += 1
     # add on non param inputs
     preserved_arg_indices.extend(range(len(flat_params), len(params)))
     # is this necessary ?
+    fw_metadata.static_input_indices = static_indices_new
     gm.recompile()
     return preserved_arg_indices
 
