[Feature] Keep DSA top-k out of checkpoint replay - #1989
Open
jayhenry wants to merge 23 commits into
Open
Conversation
… peaks (#1987) fix muon all2all padding
pass cluster_name as an explicit sandbox create arg
…ntation
The decoder layers pinned `CheckpointImpl.REENTRANT`, whose autograd.Function
only tracks gradients for top-level torch.Tensor arguments. That restriction
shaped the surrounding code: `_check_signature_of_forward` existed to fail early
on any other signature, decoder layers had to take hidden states as varargs and
return a flat positional tuple, and the domino EP path had to be re-derived from
tuple slices at every consumer.
Move `checkpoint_wrapper` onto `torch.utils.checkpoint.checkpoint` with
`use_reentrant=False`. Because it is built on saved-tensor hooks, gradients flow
through arbitrary forward signatures, so:
- decoder layers (dense and MoE), MTP layers and the MTP block now return
TypedDicts keyed by output name instead of positional tuples, and take
micro-batch inputs as a list rather than varargs;
- `_check_signature_of_forward` and its test are deleted.
`apply_gradient_checkpointing` takes a `context_fn` seam for the upcoming
selective checkpointing work. It is only passed through when set: dynamo lowers
`checkpoint` to a higher-order op that rejects an explicit `context_fn`, so a
compiled layer must not receive one.
The MTP layers keep the reentrant path behind the existing
`mtp_checkpoint_use_reentrant` switch, now spelled
`apply_legacy_reentrant_checkpointing`; DSA top-k cache sharing across MTP depths
still depends on the grad-free original pass.
Verified on torch 2.10 / 4x H200 against a no-recompute baseline, MoE ep_size=4,
intra_layer_micro_batch=2, all2all, 8 layers, seq 4096, bf16:
eager baseline loss 11.3384666443 grad_norm 4.5116462708 peak 2189.9 MiB
recompute 11.3384666443 4.5117201805 730.2 MiB
compile baseline 11.3386135101 4.5160999298 2006.2 MiB
recompute 11.3386135101 4.5163722038 682.8 MiB
…nting Introduces the vocabulary the selective-checkpointing (SAC) layers agree on: `RecomputeUnit`, the user-facing semantic units of activation that may stay resident; `MarkerInterval` / `RecomputeIntervalMap`, the per-model declaration of how each unit maps to marker intervals; and `checkpoint_record`, the imperative marker model authors place in forward. `checkpoint_record` is a documented no-op stub here. The contextvars session behind it, the per-op policy, and the config resolution that turns user selections into intervals land with the SAC engine and the config layer.
The marker session is backed by contextvars, which Dynamo cannot trace. Reading a ContextVar inside a `fullgraph=True` region is a hard compile error rather than a graph break, and xtuner compiles `_pre_moe_forward`, `MoEBlock.forward`, `_shared_experts_forward`, `_post_moe_forward` and (in the non-EP config) `MoEDecoderLayer.forward` that way -- so any marker placed in those methods would break compilation once the session lands. Return early on `torch.compiler.is_compiling()`. Dynamo constant-folds the check, so the body never enters the graph.
…e new checkpointing Two silent regressions surfaced by tests/engine/test_glm52_moe_train_engine.py, which needs GLM5_2_TINY_MOE_PATH and was therefore not covered before. Both produce a finite loss, so every existing assertion still passed. 1. Reentrant `CheckpointFunction` cannot carry gradients through a non-tensor return. Once `MTPLayer.forward` started returning a TypedDict, its outputs came back from the grad-free original pass without a `grad_fn` and the MTP subgraph was detached from the loss: every mtp_block parameter ended with `grad is None` (base: 19 with non-zero grads) while mtp_loss stayed finite. The pytree reentrant wrapper already flattened inputs so the autograd boundary could see tensors nested in containers; flatten the outputs the same way and rebuild the structure outside the checkpoint. 2. DSA cross-layer top-k sharing recognizes a checkpoint's original pass by grad being disabled, which only the reentrant implementation provides. Under the non-reentrant one both passes run with grad enabled, so `checkpoint_active` was never set, `after_recompute_release` never ran, and the shared top-k was never freed. Route decoder layers carrying the DSA lifecycle to the reentrant path, selected by the new `uses_dsa_topk_lifecycle` predicate rather than by model name. Both this and the MTP switch disappear once the cache tracks the original/replay phase explicitly. test_recompute.py gains a guard for the dict-return gradient path, which is the failure mode a finite-loss assertion cannot catch.
A real `context_fn` -- including the `functools.partial` over `create_selective_checkpoint_contexts` that selective checkpointing uses -- compiles fine. What Dynamo's checkpoint higher-order op rejects is torch's own `noop_context_fn` being forwarded as the "no policy" default, which fails with `NotImplementedError: ... LazyVariableTracker context_fn`. Behaviour is unchanged; the previous comment claimed `context_fn` was broken under compile in general, which would have wrongly ruled out selective checkpointing for every compiled layer.
Nothing needs the reentrant implementation any more except DSA cross-layer top-k sharing, whose `_is_checkpoint_original_forward` infers the checkpoint phase from `torch.is_grad_enabled()` -- a proxy that only ever held for reentrant. Rather than keep a whole checkpoint implementation alive for that one consumer, remove it; restoring DSA is part of the pending GLM-5.2 compatibility work and is left to it. Removed: `apply_legacy_reentrant_checkpointing`, `_pytree_reentrant_checkpoint` (and its flatten/unflatten of nested inputs and outputs, which existed solely to work around reentrant's inability to see tensors inside containers or to return non-tensors), `FSDPConfig.mtp_checkpoint_use_reentrant`, and the two branches in `fully_shard` that chose between the implementations. Both collapse to an unconditional `apply_gradient_checkpointing`. `uses_dsa_topk_lifecycle` is left in place: it is DSA code, and the selective checkpointing engine stacked on this branch still consumes it. Two GLM-5.2 tests are marked `xfail` with the reason rather than deleted, so the gap stays visible to whoever does the compatibility work: - the DSA top-k cache is no longer released after recompute; - shared MTP depths under compile trip torch's "Recomputed values have different metadata" check. MTP itself is healthy on the non-reentrant path -- 19 non-zero parameter gradients and 5 `None`, identical to base -- and the domino EP parity is unchanged: loss 11.3384666443 bit-identical with and without recompute, grad-norm 4.5116610527 vs 4.5117554665, peak 2189.9 vs 730.2 MiB.
`uses_dsa_topk_lifecycle` existed for one reason: to select which decoder layers had to stay on the reentrant checkpoint implementation, because DSA cross-layer top-k sharing infers the checkpoint phase from `torch.is_grad_enabled()`. The previous commit removed that implementation, so the predicate answers a question no caller can act on any more -- a grep across this branch and the two stacked on it finds only the definition. The cache machinery it guarded (`_is_checkpoint_original_forward`, `checkpoint_active`, `after_recompute_release`) stays: it is the subject of the pending GLM-5.2 compatibility work, and the two `xfail` tests keep that gap visible. Only the strategy-selection hook goes. Also drops "reentrant" from the lifecycle-hook comment: non-reentrant recompute re-invokes the decoder module and its hooks just the same, so the conclusion is unchanged but the named mechanism no longer exists.
Turn the whole-layer checkpoint into a per-op decision: a contextvars marker session records which `checkpoint_record` intervals are open, and the policy behind `create_selective_checkpoint_contexts` keeps the ops inside them while recomputing the rest. The sharding paths call one entry point per selected layer, which also owns the reentrant fallback for DSA layers. Two op classes are never kept, whatever the markers say: ops with a mutable schema, whose cached tensor can be overwritten before backward reads it, and collectives, whose destination buffer may be allocated outside the interval.
Every layer `recompute_ratio` selects now goes through one entry point, which picks the recompute strategy the layer supports: region-level selective checkpointing, or the legacy reentrant path for layers carrying the DSA top-k lifecycle. `BaseModel.recompute_intervals` is the whole surface the config layer needs; keeping nothing reproduces today's whole-layer recompute. Add the regression tests: kept regions reproduce full-recompute gradients bit for bit and keep them alive, unbalanced and overlapping intervals are safe, in-place writes inside a kept region are rejected rather than silently doubled, and a kept region matches full recompute under domino EP both eager and compiled.
The warnings exist so that a `recompute_cfg` which keeps nothing does not look like a silent no-op, and they are deduplicated because every layer of a model reaches the same diagnosis. Keyed globally on the interval or the layer name, though, the second model in a process -- an RL reference model, a compose model's other tower -- was silenced by the first, which is the same silent no-op wearing a different hat. Key on the diagnosis instead.
The compile-granularity rule was wrong: a region is addressable only if its contents run in eager, not merely its endpoints. `SAVE_MLP` under EP falsifies the old rule -- both markers fire in the uncompiled layer body while the region encloses the compiled shared-expert forward, so nothing is kept. Neither diagnostic could fire for it: the markers ran, and the layer is not compiled as one region. That is a silent no-op of the same shape as the three this stack already found. The session now learns from the policy whether anything was actually kept while an interval was open, which is a statement about contents rather than endpoints, and reports intervals that opened and kept nothing. Diagnostics also aggregate over the owning model instead of firing per layer. Interval maps legitimately span dense and MoE layers, so a marker missing from one layer type is normal, and warning per layer fired on every correctly configured model -- which teaches users to ignore the warnings that matter.
The legacy reentrant path no longer exists, so a layer carrying the DSA top-k lifecycle takes the ordinary selective checkpointing path like any other layer. The test asserted a fallback that was removed with the path itself.
The two comments named `_MarkerSession.report_unreached`, a method that never existed under that name: unreached markers are reported from `_report_pass`, once the owner's layers have all completed a pass.
… layer `checkpoint_record` is called from `xtuner/v1/module/decoder_layer/*.py`, and importing it from there pulls in `xtuner/v1/model/__init__.py`, which imports the decoder layers back. A module the `module/` layer depends on cannot live under `model/`. Split rather than move: the vocabulary and the marker session go to `xtuner/v1/utils`, below both packages, while the policy and the wrapping stay at the model layer, where knowing `nn.Module`, `CheckpointWrapper` and the DSA lifecycle predicate is legitimate. The session's API becomes public because it is now driven across a package boundary. `xtuner.v1.model.utils` keeps exporting the same names, so no import site outside these files changes.
…rations Lets a user keep named activation regions resident instead of recomputing them, inside the layers `fsdp_cfg.recompute_ratio` already selects. The two knobs stay orthogonal: `recompute_ratio` picks the layers, `recompute_cfg` picks the regions inside a selected layer. `XTunerBaseModelConfig.recompute_cfg` is tri-state. `None` and `False` keep today's behaviour of recomputing everything, `True` keeps every region the model declares, and an explicit list of `RecomputeUnit` keeps exactly those. Unlike `compile_cfg`, `None` is not "the model default": `default_recompute_cfg` declares what an architecture *can* keep, not what is worth keeping, so an unset config never changes a run's memory profile. `False` additionally propagates into nested sub-model configs, which is the walk `compile_cfg` already had, now shared between the two switches instead of duplicated. `MoEDecoderLayer` and `DenseDecoderLayer` record paired markers around the attention, router, dispatch, combine, shared-expert and MLP regions, in both the single-batch and the domino EP path. How much of that survives `torch.compile` depends on where the markers sit, which the per-model `default_recompute_cfg` docstrings spell out unit by unit.
…tion `test_micro_batch_path_records_the_same_markers` compared marker name sets, so it would have passed even if the domino path wrapped entirely different operations. It now compares which of the layer's own operations each region encloses, verified by mutation: moving `moe.dispatch.end` past the expert GEMMs fails the new assertion and passed the old one. Along the way the two paths are made to agree exactly -- `_forward` now reshapes outside the dispatch and combine regions, as the domino path already did. Intervals resolve in `BaseModel.__init__`, next to `compile_cfg` and by the same rule: `default_recompute_cfg` is answerable from the config alone, so there is nothing to wait for. A config naming a unit the model does not support now fails before the run spends anything on materializing and sharding weights. `recompute_cfg=False` reaching every nested sub-model config gains a test on a shipped compose config, which nests three of them; the previous probe nested one. Verified by mutation: stopping the walk at the outer config fails it. Record why regions use explicit `.begin` / `.end` marker pairs instead of ending each region at the next region's start marker, and warn instead of silently resolving to nothing when `recompute_cfg=True` meets a model that declares no units.
"Keep the expert dispatch / combine communication region" tells a user they are keeping the communication, which is the one thing the unit does not keep: the SAC policy forces every c10d collective to be recomputed, because keeping one would elide it from the recompute pass. What the unit trades memory for is the permutation, padding and unpermutation work on either side of the all-to-all.
Measured against the SAC policy at ep_size=2: only SAVE_MOE_DISPATCH keeps ops under compile (152), while SAVE_ATTN, SAVE_MOE_GATE and SAVE_MLP keep none. The docstring claimed the shared-expert half of SAVE_MLP stayed effective. The rule was stated as "markers must sit outside compiled regions", which is necessary but not sufficient. SAVE_MLP's markers do run in eager -- they sit in `_forward`, which EP does not compile -- but the region encloses `_shared_experts_forward`, which EP does compile, and ops inside a compiled region execute as fused kernels that never reach the per-op policy. A region is addressable only if its contents run in eager, not merely its endpoints. SAVE_MOE_DISPATCH survives because the dispatcher calls it wraps are the one part of the MoE path that stays uncompiled.
jayhenry
force-pushed
the
feat/dsa-topk-selective-checkpoint
branch
from
July 30, 2026 07:15
74cb1fd to
985442d
Compare
Collaborator
|
Heads-up: #1988 was rebased onto the rewritten SAC stack (new head Your commit is written against APIs that no longer exist:
A unit now declares what it resolves to, via
Two things stopped me from just writing that:
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. |
HAOCHENYE
approved these changes
Aug 12, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stack (bottom to top):
feat/sac-nonlegacy-base→mainfeat/sac-engine→feat/sac-nonlegacy-baserecompute_cfgand per-model recompute declarations (3/3) #1981feat/sac-recompute-cfg-v2→feat/sac-enginefeat/unify-offload→feat/sac-recompute-cfg-v2feat/dsa-topk-selective-checkpoint→feat/unify-offload← you are here (top of stack)Summary
This PR implements the selective-checkpoint follow-up from #1978 on top of the current non-reentrant checkpoint stack.
save_dsa_indexerselective checkpointing around only the mutation-free DSA top-k selection kernel.SAVE_DSA_INDEXERis selected; an unsetrecompute_cfgpreserves the original compiled path.torch.compile.RECOMPUTE_CFG=save_dsa_indexerin the GLM SFT example.Selective interval
RecomputeUnit.SAVE_DSA_INDEXERresolves to the half-open intervaldsa.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" .-> DValidation
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
2.9.1+cu128; CUDA 12.8;conda activate pt29_glm1.GLM-5.2-30B-MTP-new, 5 main decoder layers + 1 MTP layer, 32.797B trainable parameters.SAMPLE_MAX_LENGTH=4096;PACK_MAX_LENGTH=16384; GBS8.1e-6; activation offload, DSA top-k offload, optimizer swap, and explicit GC disabled to isolate selective checkpointing.control → selective → selective2 → control2, yielding 256 all-rank steady-state records and 32 paired rank-0 steps per variant.RECOMPUTE_CFG; selective usedRECOMPUTE_CFG=save_dsa_indexer. Tokens were identical rank-by-rank and step-by-step.Results
save_dsa_indexer, 2 runs>=1.8 s)The two order-balanced aggregate TGS comparisons were:
control → selectiveselective2 → control2Across 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
aten.topkexecution during backward replay before the interval existed.torch.compile(fullgraph=False)tests both observe zero backward top-k executions and finite gradients.[('dsa.indexer.begin', 'dsa.indexer.end')]and emit no missing/inert marker warning..scratch/dsa_indexer_sac_regression/results/16k_kernel_{control,selective,selective2,control2}/.Conclusion
save_dsa_indexernow 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.