fix(trainer): preserve checkpoint replay across autotuning and compile fallback

src/prime_rl/trainer/compile_compat.py

@@ -0,0 +1,47 @@
+"""PyTorch 2.13 compatibility for runtime autotuning under checkpointing.
+
+Modes opting out of compiler internals must not observe runtime
+autotuning allocations or benchmark executions during checkpoint recomputation.
+"""
+
+from contextlib import contextmanager
+
+import torch
+from torch.utils._python_dispatch import _get_current_dispatch_mode_stack, _pop_mode, _push_mode
+
+
+@contextmanager
+def _disable_modes_for_compile_internals():
+    modes = _get_current_dispatch_mode_stack()
+    retained = [mode for mode in modes if not mode.ignore_compile_internals()]
+    if len(retained) == len(modes):
+        yield
+        return
+
+    for _ in modes:
+        _pop_mode()
+    for mode in retained:
+        _push_mode(mode)
+    try:
+        yield
+    finally:
+        for _ in retained:
+            _pop_mode()
+        for mode in modes:
+            _push_mode(mode)
+
+
+def apply_autotune_checkpoint_patch() -> None:
+    """Apply the tested runtime backport once; refuse unvalidated Torch versions."""
+    if torch.__version__.split("+")[0] != "2.13.0":
+        raise RuntimeError(f"Autotuning checkpoint backport requires Torch 2.13.0; got {torch.__version__}")
+
+    from torch._inductor.runtime.triton_heuristics import CachingAutotuner
+
+    for name in ("bench", "autotune_to_one_config", "coordinate_descent_tuning"):
+        original = getattr(CachingAutotuner, name)
+        if getattr(original, "_prime_rl_checkpoint_backport", False):
+            continue
+        wrapped = _disable_modes_for_compile_internals()(original)
+        wrapped._prime_rl_checkpoint_backport = True
+        setattr(CachingAutotuner, name, wrapped)

src/prime_rl/trainer/model.py

@@ -16,7 +16,6 @@
 from huggingface_hub import snapshot_download
 from jaxtyping import Int
 from torch import Tensor
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointWrapper
 from torch.distributed.checkpoint.hf_storage import HuggingFaceStorageReader
 from torch.distributed.checkpoint.state_dict_loader import load as dcp_load
 from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, OffloadPolicy, fully_shard
@@ -842,18 +841,17 @@ def mark_dynamic_int_args() -> None:
 
 
 def apply_compile(model: nn.Module, compile_config: CompileConfig):
+    from prime_rl.trainer.compile_compat import apply_autotune_checkpoint_patch
+
+    apply_autotune_checkpoint_patch()
+    get_logger().info(f"Autotuning checkpoint dispatch backport active (torch={torch.__version__})")
     torch._dynamo.config.capture_scalar_outputs = True
     mark_dynamic_int_args()
     language_model = get_language_model(model)
     for layer_id in range(len(language_model.layers)):
         layer = language_model.layers[layer_id]
-        if isinstance(layer, CheckpointWrapper) and any(
-            isinstance(module, FSDPModule) for module in layer._checkpoint_wrapped_module.modules()
-        ):
-            # A nested FSDP unit's hooks always break the graph, and a break inside a compiled checkpoint
-            # sends the block to eager, so keep AC eager around the compiled block. pytorch/pytorch#196626
-            # removes the fp32 router's nested unit by letting it join the block's FSDP unit.
-            layer = layer._checkpoint_wrapped_module
+        # Compile the checkpoint wrapper as one region: an eager fallback during
+        # recomputation must not change the operations observed by an eager SAC context.
         # Doing it in-place avoids mangled fqn which can break checkpoint loading
         layer.compile(fullgraph=compile_config.fullgraph, mode=compile_config.mode)
     get_logger().info(
@@ -1016,7 +1014,7 @@ def setup_model(
             override_attr=config.vlm.vision_encoder_attr if config.vlm is not None else None,
         )
 
-    # the right order is AC -> FSDP -> Compile: compile needs to see which blocks hold nested FSDP units
+    # Keep the checkpoint wrapper inside the compiled region and apply FSDP before tracing.
     if config.ac is not None:
         apply_ac(model, config.ac)