pytorch
diff --git a/‎test/inductor/test_compiled_autograd.py
Lines changed: 1 addition & 12 deletions b/‎test/inductor/test_compiled_autograd.py
Lines changed: 1 addition & 12 deletions
diff --git a/‎torch/utils/checkpoint.py
Lines changed: 17 additions & 0 deletions b/‎torch/utils/checkpoint.py
Lines changed: 17 additions & 0 deletions
@@ -4282,6 +4282,7 @@ def wrap_test_class(orig_cls):
     "test_full_backward_hook_double_backward",  # _pack_with_none
     "test_grad_mode_restored_reentrant",  # assertTrue
     "test_multi_grad_any_hooks",  # register_multi_grad_hook
+    "test_saved_variable_packing_unpacking_did_not_save_original_with_hooks",  # register_hooks
 }
 
 test_contexts = {
@@ -4345,19 +4346,7 @@ def wrap_test_class(orig_cls):
         "test_return_duplicate",  # gradient batching rule not implemented for aten::sym_size.int
         "test_return_duplicate_inplace",  # gradient batching rule not implemented for aten::sym_size.int
         "test_setitem",  # CopySlices accuracy error
-        "test_save_on_cpu_and_checkpoint",  # https://github.com/pytorch/pytorch/issues/147565
-        "test_checkpoint_detects_non_determinism",  # different error
-        "test_checkpointing_non_reentrant_autocast_cpu",  # saved != recompute
-        "test_checkpointing_non_reentrant_autocast_gpu",  # saved != recompute
         "test_checkpointing_without_reentrant_saved_object_identity",  # same as https://github.com/pytorch/pytorch/issues/136193
-        "test_saved_variable_packing_unpacking_did_not_save_original_with_hooks",  # register_hooks multiple times
-        "test_saved_variable_saved_original_inplace_detach",  # RuntimeError not raised
-        "test_access_saved_tensor_twice_without_recomputation_works",  # saved != recompute
-        "test_checkpointing_without_reentrant_dataparallel",  # https://github.com/pytorch/pytorch/issues/127115
-        "test_checkpointing",  # takes very very long
-        "test_checkpointing_without_reentrant_input_requires_grad_False",  # takes very very long
-        "test_checkpointing_without_reentrant_input_requires_grad_True",  # takes very very long
-        "test_checkpointing_without_reentrant_memory_savings",  # takes very very long
         "test_dtensor_different_gradient_placement",  # Dynamo failed to run FX node with fake tensors
         "test_dtensor_noncontiguous_output",  # Dynamo failed to run FX node with fake tensors
         "test_dtensor_partial_placement_graph_output",  # Dynamo failed to run FX node with fake tensors
 
@@ -328,6 +328,7 @@ def backward(ctx, *args):
 def noop_context_fn():
     return contextlib.nullcontext(), contextlib.nullcontext()
 
+# Note: [torch.compile and checkpoint]
 # TorchDynamo does not step inside utils.checkpoint function.  The flow
 # looks likes this
 #  1) TorchDynamo tries to wrap utils.checkpoint in a HigherOrderOp by
@@ -1106,6 +1107,8 @@ def pack_hook(x):
                     frame.x_metadatas.append(frame.metadata_fn(x))
             return holder
 
+        # See Note: [compiled autograd and checkpoint unpack hook]
+        @torch._disable_dynamo
         def unpack_hook(holder):
             gid = torch._C._current_graph_task_id()
             if gid == -1:
@@ -1541,3 +1544,17 @@ def recompute_fn(*inputs):
         )
 
     return
+
+# Note: [compiled autograd and checkpoint unpack hook]
+# When tracing via compiled autograd, this hook will be visible to the
+# compiler if the forward of this checkpointed region ran in eager.
+# If the forward had ran under compile, it would have been wrapped in a
+# higher order op. See Note: [torch.compile and checkpoint].
+#
+# Since we run the recomputation hook under a enable_grad context,
+# AOTDispatch will trace a joint graph for this hook, and may
+# save different activations than in eager. This conflicts with the
+# strict activation count checks in `frame.check_recomputed_tensors_match`.
+# So, we disable this hook to force it to recompute eager checkpointed regions
+# in eager. This could be removed if we can disable the partitioner for this
+# graph segment.