[aotautograd] Fix inplace checks in autograd backward functions during functionalization

azahed98 · azahed98 · commit 6522e39cc25e · 2026-03-11T16:14:03.000-07:00
diff --git a/test/functorch/test_aotdispatch.py b/test/functorch/test_aotdispatch.py
@@ -9118,6 +9118,45 @@ def _inps():
         self.assertEqual(ref_inps_after_fw, inps_after_fw)
         self.assertEqual(ref_inps_after_bw, inps_after_bw)
 
+    def test_mutations_in_bw_requires_grad_input(self):
+        # Inplace mutation of a requires_grad=True forward input during the
+        # backward pass should not trigger the check_inplace error
+        class AF(torch.autograd.Function):
+            @staticmethod
+            def forward(ctx, dummy, inplace_tensor):
+                ctx.save_for_backward(inplace_tensor)
+                return dummy.clone()
+
+            @staticmethod
+            def backward(ctx, grad_output):
+                (inplace_tensor,) = ctx.saved_tensors
+                inplace_tensor.mul_(2.0)
+                return grad_output, None
+
+        def fn(dummy, inplace_tensor):
+            return AF.apply(dummy, inplace_tensor)
+
+        def _inps():
+            dummy = torch.zeros((2,), requires_grad=True)
+            # requires_grad=True is what triggered the bug
+            inplace_tensor = torch.ones((2,), requires_grad=True)
+            return dummy, inplace_tensor
+
+        inps = _inps()
+        out = fn(*inps)
+        ref_inps_after_fw = [x.clone().detach() for x in inps]
+        out.sum().backward()
+        ref_inps_after_bw = [x.clone().detach() for x in inps]
+
+        inps = _inps()
+        out = torch.compile(fn, backend="aot_eager", fullgraph=True)(*inps)
+        inps_after_fw = [x.clone().detach() for x in inps]
+        out.sum().backward()
+        inps_after_bw = [x.clone().detach() for x in inps]
+
+        self.assertEqual(ref_inps_after_fw, inps_after_fw)
+        self.assertEqual(ref_inps_after_bw, inps_after_bw)
+
     def test_mutation_of_input_in_fw_and_bw(self):
         class AF(torch.autograd.Function):
             @staticmethod
diff --git a/torch/_functorch/_aot_autograd/graph_capture_wrappers.py b/torch/_functorch/_aot_autograd/graph_capture_wrappers.py
@@ -969,7 +969,11 @@ def _post_forward(primals: Any) -> None:
                                 raise AssertionError(
                                     f"expected both before and after to be Tensors, got {type(before)} and {type(after)}"
                                 )
-                            before.copy_(after)
+                            # no_grad prevents the FakeTensor's requires_grad from
+                            # triggering check_inplace during tracing.  The
+                            # requires_grad case is checked at runtime instead
+                            with torch.no_grad():
+                                before.copy_(after)
                         meta.indices_of_inputs_that_requires_grad_with_mutations_in_bw.append(
                             idx
                         )

Original file line number	Diff line number	Diff line change
`@@ -969,7 +969,11 @@ def _post_forward(primals: Any) -> None:`
`969`	`969`	`raise AssertionError(`
`970`	`970`	`f"expected both before and after to be Tensors, got {type(before)} and {type(after)}"`
`971`	`971`	`)`
`972`		`- before.copy_(after)`
	`972`	`+ # no_grad prevents the FakeTensor's requires_grad from`
	`973`	`+ # triggering check_inplace during tracing. The`
	`974`	`+ # requires_grad case is checked at runtime instead`
	`975`	`+ with torch.no_grad():`
	`976`	`+ before.copy_(after)`
`973`	`977`	`meta.indices_of_inputs_that_requires_grad_with_mutations_in_bw.append(`
`974`	`978`	`idx`
`975`	`979`	`)`