Disable compiled autograd for re-entrant autograd (#135795)

majian4work · pytorchmergebot · commit 0b667c073e71 · 2024-09-24T15:09:16.000Z
Fixes #135298 Pull Request resolved: #135795 Approved by: https://github.com/xmfan
diff --git a/test/inductor/test_compiled_autograd.py b/test/inductor/test_compiled_autograd.py
@@ -2686,10 +2686,7 @@ def fn(x):
 
         inp = torch.rand(10, 10, requires_grad=True)
         out = torch.utils.checkpoint.checkpoint(fn, inp, use_reentrant=True)
-        with self.assertRaisesRegex(
-            RuntimeError,
-            r"\(e.g. reentrant checkpointing\), this is not supported yet\.",
-        ), torch._dynamo.compiled_autograd.enable(torch.compile):
+        with torch._dynamo.compiled_autograd.enable(torch.compile):
             out.backward()
 
 
diff --git a/torch/_dynamo/compiled_autograd.py b/torch/_dynamo/compiled_autograd.py
@@ -337,7 +337,8 @@ def runtime_wrapper(compiled_fn, inputs, sizes, scalars, hooks):
                 for i in runtime_inputs_to_move:
                     inputs[i] = inputs[i].pin_memory().cuda(non_blocking=True)
 
-                return compiled_fn(inputs, sizes, scalars, hooks)
+                with disable():
+                    return compiled_fn(inputs, sizes, scalars, hooks)
             finally:
                 in_compiled_autograd_region = False
 
diff --git a/torch/_dynamo/variables/misc.py b/torch/_dynamo/variables/misc.py
@@ -946,7 +946,7 @@ def call_method(
         kwargs: "Dict[str, VariableTracker]",
     ) -> "VariableTracker":
         if name == "queue_callback":
-            if torch._dynamo.compiled_autograd.compiled_autograd_enabled:
+            if torch._dynamo.compiled_autograd.in_compiled_autograd_region:
                 assert (
                     tx.one_graph
                 ), "queue_callback() is only supported when Compiled Autograd is enabled with fullgraph=True"
