[faresobeid]

Status on main

• revert(trainer): compile around the checkpoint again for MoE blocks with an fp32 router #4037 reverted perf(trainer): avoid eager fallback of MoE blocks with an fp32 router and activation checkpointing #3925, so main compiles around the checkpoint again (compile(AC(block))). The recompute then runs inside the graph and _checkpoint_contexts takes its torch.compiler.is_compiling() early return: the activation_checkpointing.py hunk is a no-op on main today. It becomes active when perf(trainer): avoid eager fallback of MoE blocks with an fp32 router and activation checkpointing #3925 re-lands.
• The model.py hunk (ring-CP DYNAMIC_INT_ARGS) is active now: ring CP reads these ints inside the compiled block, so without it every micro batch recompiles.
• The crash and the numbers below were measured with perf(trainer): avoid eager fallback of MoE blocks with an fp32 router and activation checkpointing #3925 on main.

Problem

MoE RL with context parallelism, activation checkpointing and compile crashes in the first backward:
RuntimeError: aten.empty_like.default encountered during backward but not found in storage.
Repro: samsja/mini-glm-moe, reverse-text, --trainer.model.cp 2 on 2 trainer GPUs, default ac/compile/moe_router_dtype (fp32 in RL). Ring CP crashes with empty_like from the quack RMSNorm. Ulysses CP crashes with aten.empty.memory_format from inductor autotuning, which runs during the recompute.

Root cause

• With the fp32 router, perf(trainer): avoid eager fallback of MoE blocks with an fp32 router and activation checkpointing #3925 (since reverted by revert(trainer): compile around the checkpoint again for MoE blocks with an fp32 router #4037) kept the checkpoint eager around the compiled MoE block. Backward therefore recomputes the block through its Dynamo frames, on the autograd device thread.
• The recompute never reuses the forward's cache entries. The forward input is a view, and its guards read hidden_states._base. The recompute input is not a view. So the recompute compiles its own entries on the autograd thread.
• Since torch 2.13, Dynamo settings are thread-local (ContextVar). The trainer's recompile_limit, capture_scalar_outputs and dynamic_sources (including max_seqlen) are set on the main thread only. The recompute compiles with the defaults: recompile_limit 8, ints specialized, .item() breaks the graph.
• Ring CP reads max_seqlen_q/k and local_k_slice from ring_flash_attn's DATA_PARAMS inside the compiled block. These change every micro batch, so each micro batch recompiles in both threads (112 of 137 recompile reasons on fake data).
• When the autograd thread hits recompile_limit (8), Dynamo runs the block eagerly. The SAC storage only holds the compiled region from the forward, so the replay fails. With ulysses, the recompute compiles a different graph than the forward. Its autotuning allocations run under the SAC replay mode, and the replay fails.

Fix

• activation_checkpointing.py: the recompute context applies the forward thread's recompile_limit, capture_scalar_outputs and dynamic_sources. It is skipped when the checkpoint itself is compiled; the recompute then runs inside the graph.
• model.py: max_seqlen_q, max_seqlen_k and local_k_slice are added to DYNAMIC_INT_ARGS, so ring CP stops recompiling per micro batch.
• Either change alone is not enough. With only the dynamic ints, the run still crashes, because the autograd thread does not see them.

Checks

• GPU, 1 node, released. Same rl config as the failing run (mini-glm-moe, reverse-text, 2 trainer GPUs, cp 2):

run	main	this PR
ring cp2, 20 steps	crash in step 1 (hit config.recompile_limit (8) on the autograd thread)	20/20 steps; 38 main-thread + 4 autograd-thread recompiles
ulysses cp2, 8 steps	crash in step 1 (aten.empty.memory_format ... not found in storage)	8/8 steps
ring cp2, fake data, step 3	12.8 s/step, 54 recompiles	2.5 s/step, 14 recompiles

• Fake-data trainer on main: cp2 with compile off, bf16 router, cp1/ep2 and selective AC all pass. Real packed data is needed to trigger the crash, because hidden-state shapes vary per micro batch.
• pytest -m "not gpu" tests/unit/orchestrator tests/unit/test_configs.py tests/unit/train/models/test_checkpointing.py: 302 passed.
• ruff check / ruff format --check: clean.
• No new test. The failure needs AC recompute on a CUDA autograd thread.

---

Note

Medium Risk
Touches activation-checkpoint recompute and torch.compile dynamic-source configuration on the training hot path; incorrect settings could affect compile stability or performance, but the change is narrowly scoped to AC + compile + CP failure modes.

Overview
Fixes first-backward crashes when activation checkpointing, torch.compile, and context parallelism run together (e.g. MoE RL with ring/Ulysses CP).

Activation checkpointing now routes selective checkpoint contexts through _checkpoint_contexts, which wraps the recompute path with the forward thread’s Dynamo/compiler settings (recompile_limit, capture_scalar_outputs, dynamic_sources) so autograd-thread recomputation of compiled blocks matches the forward compile instead of hitting default limits and eager fallback that breaks SAC replay. That patch is skipped when the checkpoint itself is compiled.

Compile setup extends DYNAMIC_INT_ARGS with ring CP sequence/slice keys (max_seqlen_q, max_seqlen_k, local_k_slice) so those ints are symbolic from the first compile and stop per–micro-batch recompiles inside compiled blocks.

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

Recipe confirmation: GLM-4.5-Air, cp=4 ulysses + full AC + ac_offloading, 2 nodes (1 train + 1 inference): main crashes in step 2 backward (aten.empty.memory_format ... not found in storage); with this PR, 8/8 steps, mismatch KL 0.0017-0.0053, peak memory 54-61 GiB.

[samsja]

not relevnat anymore since we removed ring cp ?