[hallerite]

Depends on #3951 and is stacked on fix/ac-recompute-dynamo-settings. The diff contains only the remaining autotuning and checkpoint-boundary protections; #3951 supplies consistent forward/recompute compiler settings and dynamic ring-attention integers.

Selective activation checkpointing can still fail when runtime autotuning executes under replay dispatch modes, or when a compiled forward is replayed eagerly. This change addresses those two cases:

• Install a PyTorch 2.13.0-only backport in shared apply_compile (RL and SFT). It filters dispatch modes opting out via ignore_compile_internals() around CachingAutotuner.bench, autotune_to_one_config, and coordinate_descent_tuning, restoring the stack afterward. Installation is idempotent and rejects other Torch versions because the implementation uses private runtime APIs.
• Keep CheckpointWrapper inside the compiled region rather than unwrapping it for nested FSDP blocks. This avoids an eager selective-checkpoint context spanning a compiled-to-eager replay transition while preserving AC → FSDP → compile setup order.

Performance limitation: nested FSDP graph breaks can cause the checkpointed block to run eagerly. This is a correctness safeguard, not a claim to preserve all compilation throughput. Router precision-group changes and parameter-alias activation-offloading work are outside this PR.

Validation on Torch 2.13.0+cu130:

• fix(trainer): recompile AC-recomputed compiled blocks with the forward's dynamo settings #3951 alone completes 12 varying-length reduced-GLM steps (CP2 Ulysses, EP2, FP32 router/master/reductions, full AC and activation offloading), including a cold-cache repeat. Its parent fails on the third backward with aten.as_strided ... not found in storage.

• A synthetic forced compiled-forward/eager-replay transition still fails with fix(trainer): recompile AC-recomputed compiled blocks with the forward's dynamo settings #3951 contexts alone (aten.empty_like ... not found in storage). Putting the checkpoint inside compilation passes 12 shapes with output/input-gradient checks.

• A controlled cold CUDA kernel launch during recomputation, using fix(trainer): recompile AC-recomputed compiled blocks with the forward's dynamo settings #3951 contexts, fails without the autotuning backport (aten.empty.memory_format ... not found in storage). With the backport it executes real autotuning (2 benchmark calls) and passes the analytical gradient check. This deliberately induced path establishes coverage, not that ordinary fix(trainer): recompile AC-recomputed compiled blocks with the forward's dynamo settings #3951 training necessarily hits it.

• Combined branch: 12 reduced-GLM optimizer steps and state-dict roundtrip pass on both GPUs. The compiler records no compiled blocks in this nested-FSDP case: it runs eagerly, confirming the performance limitation above.

• Combined branch: forced-fallback CUDA test passes 12 shapes with output/input-gradient checks; dispatch filtering, nesting, exception restoration and idempotency checks pass; existing checkpoint tests: 4 passed.

• Ruff check, Ruff format check and git diff --check pass. Signed head: c520a6156420512ab092dfe9ac5e9d6f177cae33.

Full-model validation of #3951 alone is queued as Slurm job 3809 (16 H200s, GLM-4.5-Air, 12 saved-batch replay steps); it is not counted as passed. No SFT end-to-end validation or full-model numerical-equivalence claim. No hosted-run changes.

---

Note

High Risk
Touches compiled training, SAC replay, and private Torch 2.13 Inductor/Dynamo APIs; misalignment here breaks backward or only surfaces on specific Torch versions.

Overview
Fixes backward failures when training with torch.compile + selective activation checkpointing (SAC) by keeping forward and recomputation aligned on which ATen ops run and which Dynamo settings apply.

Adds a Torch 2.13-only backport (compile_compat.py) that wraps Inductor CachingAutotuner bench/autotune paths so dispatch modes that opt out of compiler internals are temporarily removed. That stops Triton autotuning from allocating or benchmarking under SAC modes during checkpoint replay (the aten.empty.memory_format / missing-op mismatch class of bugs). The patch is applied once from shared apply_compile.

Activation checkpointing now routes SAC through _checkpoint_contexts: on the autograd recompute thread it reapplies the forward pass’s recompile_limit, capture_scalar_outputs, and dynamic_sources so a compiled block doesn’t hit Dynamo’s recompile cap and fall back to eager during replay.

Compile setup no longer unwraps CheckpointWrapper around nested FSDP blocks—layers are compiled in place with the wrapper inside the compiled region. DYNAMIC_INT_ARGS is extended for ring context-parallel (max_seqlen_q, max_seqlen_k, local_k_slice) so varying sequence metadata doesn’t force per-shape recompiles.

Tradeoff: restoring correctness may mean some checkpointed blocks run more eagerly when FSDP graph breaks remain; this is not a general fix for every compile/eager boundary.

^{Reviewed by Cursor Bugbot for commit c520a61. Bugbot is set up for automated code reviews on this repo. Configure here.}