[Refactor] Pass GLM DSA top-k IDs explicitly

[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 region-level selective checkpointing engine (2/3) #1980 feat/sac-engine → feat/sac-nonlegacy-base
3. [Feature] Add region-level recompute_cfg and per-model marker 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 ← you are here (top of stack)

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

---

Summary

This is the stack-compatible port of #1978. The explicit DSA top-k dataflow and offload design stay the same, while the integration is adapted to the non-reentrant and selective-checkpointing APIs introduced by #1979–#1981.

• Pass contiguous torch.int32 dsa_topk_ids explicitly through GLM decoder layers, intra-layer micro-batches, and MTP inputs/results, removing the mutable SequenceContext cache and manual cross-layer lifecycle runtime.
• Preserve the current keyed TypedDict decoder contract, list-based micro-batch API, and keyword checkpoint inputs instead of restoring the old reentrant-compatible flattened tuples.
• Use one saved-tensor offload window per decoder call for hidden states and the final shared-layer top-k IDs, matching tensors by storage so views share the same offload decision. MTP top-k offload remains excluded.
• Keep the sharing topology stateless: each main-stack source layer starts a new ID chain, shared layers consume the explicit tensor, and shared MTP logical depths forward the returned IDs directly.
• Add opt-in save_dsa_indexer selective checkpointing around only the mutation-free DSA top-k selection kernel, leaving index projections and the SP gather compiled and replayable.

Checkpoint-stack adaptation

• AttnOutputs, dense decoder outputs, and MoE decoder outputs expose dsa_topk_ids as an optional keyed field, which non-reentrant checkpointing can replay without input/output flattening.
• RecomputeUnit.SAVE_DSA_INDEXER resolves to the half-open dsa.indexer.begin → dsa.indexer.end interval. Its eager island is enabled only when that unit is selected, so an unset recompute_cfg preserves the original compiled path.
• Main decoder layers and GLM MTP decoder layers use the same selected DSA interval. This also avoids mixing selective main-layer replay with plain MTP replay in the FP8/compile path.
• Sparse MLA backends consume the same contiguous int32 ID format across PyTorch, TileLang, cuDNN, and sequence-parallel paths.

Validation

• Full pre-commit suite passed: merge-conflict check, codespell, docformatter, pyupgrade, Ruff, Ruff format, mypy, and pydantic-extra-check.
• Original explicit-dataflow coverage: 40 focused tests covering selective checkpointing, recompute configuration, DSA backends/SP, GLM forward/backward, MTP, and the 8-GPU optimized engine path.
• New selective-indexer coverage:
  • 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.

Performance evidence from #1978

These measurements are from the original isolated implementation in #1978; this stacked port reran the functional coverage above.

• 16K, 8×GPU AdamW: rank-0 effective TGS 10,893.7 vs 10,957.3 baseline (-0.58%; below the 11K absolute target); peak allocated memory 104.69 vs 105.26 GB.
• 32K, 8×GPU Muon SP2+EP2+micro2: full-rank sequence TGS 1,769.0 vs 1,599.3 baseline (+10.61%); peak allocated memory 101.87 vs 102.29 GB; VmPin 42.87 GiB in both runs.

Follow-ups

• Evaluate fullgraph=True separately.
• Investigate the opt-in eager-island overhead before recommending save_dsa_indexer as a default.

---

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.

flowchart LR
    A["Compiled index projections"] --> 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

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.