[Refactor] Switch gradient checkpointing to the non-reentrant implementation (1/3) - #1979
Draft
HAOCHENYE wants to merge 7 commits into
Draft
[Refactor] Switch gradient checkpointing to the non-reentrant implementation (1/3)#1979HAOCHENYE wants to merge 7 commits into
HAOCHENYE wants to merge 7 commits into
Conversation
…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.
This was referenced Jul 28, 2026
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.
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
1 participant
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→main← you are herefeat/sac-engine→feat/sac-nonlegacy-baserecompute_cfgand per-model marker declarations (3/3) #1981feat/sac-recompute-cfg-v2→feat/sac-engineSummary
Foundation for region-level selective checkpointing (design:
docs/design/selective_checkpointing.md). Moves gradient checkpointing off the legacy reentrant implementation, converts decoder-layer I/O toTypedDict, and lands the shared vocabulary the two PRs above build on.What changed
checkpoint_wrapper(legacyCheckpointImpl.REENTRANT) is replaced byapply_gradient_checkpointing(..., context_fn=...)built ontorch.utils.checkpoint(use_reentrant=False). Thecontext_fnseam is what PR2 uses for SAC.TypedDicts (MoEDecoderLayerOutputand friends) instead of positional tuples; micro-batching is selected by passing a list rather than varargs. Thelen == 4*nasserts in the MTP path are gone._check_signature_of_forwardis deleted — it existed only to fail fast on reentrant's "tensors must be top-level positional args" restriction, whichuse_reentrant=Falsedoes not have.pytree_reentrant_checkpointis no longer exported, but the function itself remains (private, as_pytree_reentrant_checkpoint) becauseapply_legacy_reentrant_checkpointingstill needs it for the two paths below.RecomputeUnit(aStrEnum, so it round-trips readably through saved configs),MarkerInterval,RecomputeIntervalMap,checkpoint_record.Why the reentrant pin existed
It was believed to be required by domino EP (
intra_layer_micro_batch > 1). It is not: it came in with the original FSDP-sharding refactor (a4a506f) and was never revisited.Verified on torch 2.10, 4×H200, MoE
ep_size=4,intra_layer_micro_batch=2, all2all, 8 layers:Loss bit-identical; grad-norm within ~5e-6 relative (bf16 + async EP reduction order); memory within 0.08% of the old reentrant path in eager (729.6 → 730.2 MiB) and 0.7% compiled (678.3 → 682.8 MiB). Same result with
torch.compileon, and on dense (Qwen3,fullgraph=True).Two reentrant paths deliberately remain
Both documented in place with their exit condition:
FSDPConfig.mtp_checkpoint_use_reentrant— shared MTP depths share one DSA top-k cache, and only a grad-free original pass advances it consistently.uses_dsa_topk_lifecycle(not a model-name check).CrossLayerTopKSharingRuntimedetects the checkpoint's original pass vianot torch.is_grad_enabled(), which is a reentrant-only proxy; under non-reentrant both passes have grad enabled and the top-k cache is never released.Both disappear once the DSA cache tracks original/replay phase explicitly instead of inferring it. That is scoped follow-up work, not in this stack.
Two silent regressions found and fixed during development
Both had finite loss and passed every existing assertion:
CheckpointFunctioncannot carry gradient through a non-tensor return, and theTypedDictchange hit the one path kept reentrant. MTP block params went from 19 non-zero grads to 0 non-zero / 24Nonewhilemtp_lossstayed finite. Fixed by flattening outputs as well as inputs in the pytree reentrant wrapper; grads now match base exactly (19/0/5).TestGlm52OptimizedEngine::test_sp2_ep4_micro2_compile_offload_train_stepnow passes, matching base.Test plan
tests/model/test_recompute.py(new), includingtest_legacy_reentrant_preserves_gradients_through_dict_return, verified red-before-green by reverting the output-flattening.tests/model/test_glm52_mtp_checkpoint_repro.py,tests/module/attention/test_dsa_mla.py,tests/utils/test_pytree_reentrant_checkpoint.py,tests/model/test_moe.py.tests/utils/test_checkpoint_wrapper_checker.pydeleted — it only exercised the removed checker.pre-commit run --files <changed>: all hooks pass including mypy.Pre-existing failures confirmed identical on base
5c2275c7and not addressed here:test_internal_metrics.py, and atest_moe.pyDDP-harness timeout that reproduces identically on base when the test is run in isolation (a full-file run shows a different count, which is why the isolation result is the meaningful one), plus threeTestGlm52PretrainedEngineloss-curve tests (hard-coded reference curve appears to predate the available checkpoint).