[Feature] Keep DSA top-k out of checkpoint replay

[jayhenry]

Stack (bottom to top):

1. [Refactor] Switch gradient checkpointing to the non-reentrant implementation (1/3) #1979 feat/sac-nonlegacy-base → main
2. [Feature] Add the selective checkpointing engine (2/3) #1980 feat/sac-engine → feat/sac-nonlegacy-base
3. [Feature] Add recompute_cfg and per-model recompute declarations (3/3) #1981 feat/sac-recompute-cfg-v2 → feat/sac-engine
4. [Refactor] Pass GLM DSA top-k IDs explicitly #1988 feat/unify-offload → feat/sac-recompute-cfg-v2
5. [Feature] Keep DSA top-k out of checkpoint replay #1989 feat/dsa-topk-selective-checkpoint → feat/unify-offload ← you are here (top of stack)

Base is PR4's branch. Review only this PR's own diff; GitHub shows it against PR4.

---

Summary

This PR implements the selective-checkpoint follow-up from #1978 on top of the current non-reentrant checkpoint stack.

• Add opt-in save_dsa_indexer selective checkpointing around only the mutation-free DSA top-k selection kernel.
• Keep index projection, RoPE, sequence-parallel gathering, and sparse attention outside the saved interval so they remain compiled and replayable.
• Enable the eager island only when SAVE_DSA_INDEXER is selected; an unset recompute_cfg preserves the original compiled path.
• Apply the same interval to GLM main and MTP decoder layers, avoiding a mixed selective/plain-checkpoint path under FP8 and torch.compile.
• Expose RECOMPUTE_CFG=save_dsa_indexer in the GLM SFT example.

Selective interval

RecomputeUnit.SAVE_DSA_INDEXER resolves to the half-open interval dsa.indexer.begin → dsa.indexer.end. The markers enclose only the actual top-k selection, which is mutation-free and safe for selective-checkpoint caching.

flowchart LR
    A["Compiled index projection and RoPE"] --> B["Compiled SP gather"]
    B --> C["dsa.indexer.begin"]
    C --> D["Eager mutation-free top-k selection"]
    D --> E["dsa.indexer.end"]
    E --> F["Compiled sparse MLA"]
    G["Checkpoint replay"] -. "SAC cache hit: skip top-k execution" .-> D

Loading

Validation

• Full pre-commit suite passed: merge-conflict check, codespell, docformatter, pyupgrade, Ruff, Ruff format, mypy, and pydantic-extra-check.
• 22 CPU tests passed, including eager and compiled replay counters that observe zero backward top-k executions.
• A real single-GPU FP8 + compile + selective main checkpoint + shared MTP checkpoint training step passed.
• An 8-GPU SP2 + EP4 + micro2 + compile + selective checkpoint + activation/top-k offload training step passed.
• Four 28-step, 8-GPU 16K SFT runs completed without NaN, OOM, CUDA, NCCL, or traceback failures.

---

DSA top-k selective-checkpoint regression (2026-07-29)

Summary

The narrow interval is functionally correct and bounded: checkpoint replay does not execute the top-k selection again, the default path remains unchanged, and the 16K production-style run needs less than 1 GB of additional peak allocated memory. On this workload the opt-in path costs about 3.15% aggregate TGS, so it remains an explicit memory-for-replay policy rather than a default.

Configuration and methodology

• Hardware/runtime: single node, 8×NVIDIA H200 141 GB; PyTorch 2.9.1+cu128; CUDA 12.8; conda activate pt29_glm1.
• Model: GLM-5.2-30B-MTP-new, 5 main decoder layers + 1 MTP layer, 32.797B trainable parameters.
• Data: Alpaca; SAMPLE_MAX_LENGTH=4096; PACK_MAX_LENGTH=16384; GBS8.
• Parallelism/operators: EP4, SP1, micro1, all-to-all dispatcher, cuDNN DSA sparse MLA, tile-wise FP8, model/torch compile, chunked CE loss with chunk size 2048.
• Optimizer: AdamW, LR 1e-6; activation offload, DSA top-k offload, optimizer swap, and explicit GC disabled to isolate selective checkpointing.
• Each run used 28 steps; statistics use steps 13–28. The order was control → selective → selective2 → control2, yielding 256 all-rank steady-state records and 32 paired rank-0 steps per variant.
• Control used unset RECOMPUTE_CFG; selective used RECOMPUTE_CFG=save_dsa_indexer. Tokens were identical rank-by-rank and step-by-step.

Results

Metric	Control, 2 runs	save_dsa_indexer, 2 runs	Delta / note
All-rank mean step time	1.523048 s	1.572652 s	+3.26%
Rank-0 median step time	1.513600 s	1.570650 s	+3.77%
Aggregate effective TGS	10,643.4	10,307.7	-3.15%
Aggregate sequence TGS	10,757.4	10,418.1	-3.15%
Effective-token utilization	98.940%	98.940%	identical input
Peak allocated range	102.40–104.50 GB	102.92–105.40 GB	max +0.90 GB
Peak reserved range	122.95–124.47 GB	122.41–125.03 GB	max +0.56 GB
Rank-0 slow steps (>=1.8 s)	0/32	0/32	none
Step-28 reduced LLM loss	7.89370728	7.89407897	+0.00037170
Step-28 reduced MTP loss	0.73700494	0.73717985	+0.00017491
Step-28 local loss	8.63071204	8.63125849	+0.00054646
Step-28 grad norm	17.89494896	17.90747166	+0.01252270
Training health	normal	normal	no NaN/OOM/CUDA/NCCL/traceback

The two order-balanced aggregate TGS comparisons were:

Order	Control TGS	Selective TGS	Selective vs control
control → selective	10,720.2	10,296.5	-3.95%
selective2 → control2	10,567.7	10,318.8	-2.36%

Across the 32 paired rank-0 steady steps, selective step time was +3.31% on average and +3.68% at the median; the approximate 95% interval was +2.53% to +4.09%. The repeat confirms that the slowdown is not an execution-order artifact.

Functional regression

• The RED test observed one real aten.topk execution during backward replay before the interval existed.
• The final eager and torch.compile(fullgraph=False) tests both observe zero backward top-k executions and finite gradients.
• The production-style selective logs declare only [('dsa.indexer.begin', 'dsa.indexer.end')] and emit no missing/inert marker warning.
• A minimized FP8/compile/MTP regression reproduced the original symbolic-shape failure when main layers used SAC but MTP used plain checkpointing; applying the same DSA interval to GLM MTP made the real training step pass.
• Regression logs are under .scratch/dsa_indexer_sac_regression/results/16k_kernel_{control,selective,selective2,control2}/.

Conclusion

save_dsa_indexer now provides the requested mutation-safe selective interval without changing the default compiled graph. It reliably removes the top-k selection from checkpoint replay and is compatible with non-reentrant checkpointing, MTP, FP8, compile, SP/EP, micro-batching, and offload. The measured 16K trade-off is +0.90 GB maximum peak allocation and -3.15% aggregate TGS, which is why the policy is opt-in.

[HAOCHENYE]

Heads-up: #1988 was rebased onto the rewritten SAC stack (new head 9ba1c5f3), so this branch needs a restack too. I deliberately did not do it for you — unlike #1988, this one is not a mechanical rebase.

Your commit is written against APIs that no longer exist:

• mtp_recompute_intervals, MarkerInterval, DSA_INDEXER_INTERVAL — the marker-interval model is gone;
• apply_selective_checkpointing(layer, intervals, owner=..., layer_compiled_as_one_region=...) — the signature is now (module, kept_ops, *, keeps_any_unit=..., preserve_rng_state=...);
• _compiles_whole_decoder_layer() — removed.

A unit now declares what it resolves to, via RecomputeTargetMap: either KeptOps (op names, recognised by identity, costs no compilation) or KeptCallables (callable names, recognised by a contextvar, at the price of torch._dynamo.disable on those callables).

SAVE_DSA_INDEXER looks like a natural KeptOps target, since sparse_mla::tilelang_dsa_topk_indices is a registered custom op — and if that holds it is strictly better than the marker version: no compilation given up, which would also make your _compiles_whole_decoder_layer() -> False override unnecessary.

Two things stopped me from just writing that:

1. Only the tilelang backend registers an op; the torch backend is plain Python, so op identity would not cover it.
2. The kernel runs under no_grad, and whether keeping it by op is equivalent to your marker region for a no-grad kernel is not something I can settle by reading — it needs a GLM-5.2 run with tilelang, and I have no checkpoint here.

Guessing would put code you did not write, and that I cannot verify, into your PR. Your call on the mapping; happy to do the mechanical part once you have decided.