refactor(trainer): make activation checkpointing operator-based

[S1ro1]

Summary

• Replace method- and attribute-based selective activation checkpointing with one non-reentrant checkpoint region around each selected transformer block.
• Keep mode = "full" as the default. Full mode retains only mandatory non-replayable operations and recomputes everything else; selective mode additionally retains configured expensive operations through the same Prime-owned policy.
• Allow selective mode to replace its public default targets with explicit operator names or namespaces. Correctness-required targets remain active for every selective configuration.
• Register real runtime boundaries for Torch EP/CP collectives, DeepEP combine, ring attention, sparse MLA, dense/grouped DeepGEMM, MXFP8 transport, and MoE routing telemetry so checkpointing, autograd, and compile observe the same operations.
• Keep sparse MLA dispatcher, fake, and autograd registration beside the vendored forward/backward kernel entrypoints rather than in GLM model code.
• Collapse ring attention into one backend-neutral forward/backward custom-op pair. Its adapters call only the pinned FA2, FA3, and FA4 APIs explicitly, without signature probing or generated per-backend registrations.
• Remove checkpoint-only attention helpers and the model-specific checkpointing module. Generic attention, AFMoE, Qwen3.5, and Laguna keep their projection/output logic directly in forward; GLM DSA keeps only its substantive mla_latents and mla_up_proj helpers.
• Rebase onto the Torch 2.13, CUDA 13, and vLLM 0.28 runtime from current main, and update FlashAttention 4 to 0251105 with the shared TVM FFI version required by vLLM.
• Ignore FSDP replay-only profiler markers when selective checkpointing validates the saved forward operation sequence.
• Prefetch the final MoE block, experts, and separately wrapped router before backward replay so FSDP collectives occur in the same order as the saved forward.

Runtime contract

• Model implementations do not expose checkpoint target names or checkpoint-only methods.
• freq selects whole transformer blocks; the dispatched operator and selected target set determine whether selective AC saves or recomputes its result.
• Both modes retain mandatory MoE routing, telemetry, non-replayable DeepEP, and CUDA-to-CPU metadata operations so they execute once. Full mode recomputes every other operation; selective mode additionally retains its configured targets.
• The default selective targets retain supported attention, matrix multiplication, grouped GEMM, and distributed communication operators. Elementwise operations, norms, reshapes, and RoPE replay.
• Custom targets replace the selective defaults and match either an exact operator name or an operator namespace. They cannot exclude mandatory routing, telemetry, non-replayable communication, or CUDA-to-CPU metadata operations.
• Unsupported or missing backend operators fail at their normal call site; there is no checkpoint fallback layer.
• MXFP8 transport resolves mxfp8_moe inside its custom-op implementation. MXFP8 grouped compute still resolves it during runtime setup because the exported grouped GEMM is already an operator boundary.

Configuration

[trainer.model.ac]
mode = "selective" # full | selective
freq = 1
# targets = ["aten::mm", "prime_rl_collectives"]

Omitting targets uses the public Prime defaults. Any explicit list replaces those defaults; an empty list retains only the mandatory operations.

Validation

• Ruff, Ruff formatting, uv lock --check, and git diff --check passed.
• Focused config, checkpointing, and Qwen3.5 CP tests on the rebased dependency stack: 138 passed.
• A6000 CI regression suite passed all eight checks: four-GPU peak memory was 6.2467 GiB against the 6.2988 GiB baseline (down from 6.6185 GiB before the fix); one-GPU peak memory was 10.6738 GiB.
• Compiled selective AC with DeepEP completed a two-GPU, two-step samsja/mini-glm-moe training smoke after rebasing onto the merged MoE runtime.
• The pinned FA2, FA3, and FA4 varlen backends all use the [heads, total_tokens] LSE layout; the ring fake and backward paths enforce that single contract.
• Two-H200 compiled ring validation matched each backend's direct attention reference exactly for FA2, FA3, and FA4 output and dq, dk, and dv.
• Full H200 unit suite: 574 passed, 9 skipped.
• Eager and compiled output/gradient/replay validation passed for local MoE, Torch EP, DeepEP with and without chunking, CP, FA2/FA3 ring attention, sparse MLA, dense DeepGEMM, grouped DeepGEMM, and the production MXFP8 transport boundary with a fake numerical kernel.
• FSDP parameter/state-dict parity and selective AC plus activation offloading passed.
• Whole-model parity passed for Llama, Nemotron-H, GPT-OSS, AFMoE, Laguna, and Qwen3.5 full-attention layers. Sparse MLA parity also passed after restoring its substantive helper structure.
• On the stacked Qwen3.8 Flash Next implementation, selective and full activation checkpointing each completed 20 RL steps on 24 H200s with compile, FSDP, expert parallelism, SignSGD, NIXL policy reload, and CPU optimizer-state offload. Steps 2-20 averaged 8,008 tokens/s for selective and 7,211 tokens/s for full; both peaked at 130.4 GiB and held mismatch KL at 0.0014-0.0021 with a final value of 0.0016.

Hardware boundary

• Numerical MXFP8 grouped GEMM and transport still require SM100 hardware. The H200 validation exercised the production custom-op/autograd/compile boundary with a BF16 fake transport kernel.
• FA4 forward/backward numerical validation passed on H200 (SM90); Blackwell remains untested in this PR.
• Qwen3.5 GatedDeltaNet whole-model parity did not reach execution because the no-AC reference forward remained in FLA kernel compilation until the diagnostic jobs were cancelled. Qwen3.5 full-attention parity passed, and the GatedDeltaNet runtime code itself is unchanged.

---

Note

High Risk
Touches core training memory/throughput (activation checkpointing), distributed MoE/DeepEP/ring attention autograd, and compile boundaries—incorrect replay or double-counted routing would show up as silent wrong grads or NCCL deadlocks.

Overview
Replaces layer-specific selective activation checkpointing with whole-block non-reentrant checkpointing driven by a PyTorch selective operator policy (activation_checkpointing.py). Semantic targets like norm / attn_proj are gone; model.ac.targets is now optional operator names or namespaces that replace Prime’s defaults, while mandatory saves still cover MoE routing/top-k, DeepEP, and CUDA→CPU copies.

Registers distributed and kernel paths as torch.library custom ops so checkpoint replay, autograd, and compile see the same boundaries: new prime_rl_collectives for EP/CP all-to-all and gather/reduce-scatter, deepep::combine, unified prime_rl_ring ring attention (FA2/FA3/FA4), prime_rl::sparse_mla, FP8/grouped FP8 GEMMs, and record_moe_routing_statistics so routing counters are not double-applied on recompute. Inline autograd.Function wrappers in token dispatch, Mamba CP, and ring attention are removed in favor of these ops.

Simplifies trainer wiring and models: apply_ac always wraps whole transformer blocks; the old layers/checkpointing.py hook layer is deleted and attention modules no longer expose checkpoint-only attn_projections / output_proj splits. FSDP backward prefetch now includes the last block’s MoE experts/router. Docs, examples, and benchmarks drop legacy selective target lists; flash-attn-4 is bumped to 0251105 with a pinned apache-tvm-ffi override.

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

[cursor]

Cursor Bugbot has reviewed your changes and found 3 potential issues.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 1fe9da7. Configure here.}