[Refactor] Switch gradient checkpointing to the non-reentrant implementation (1/3) - #1979
[Refactor] Switch gradient checkpointing to the non-reentrant implementation (1/3)#1979HAOCHENYE wants to merge 5 commits into
Conversation
There was a problem hiding this comment.
Please move the modifications of this file to the selective checkpointing branch as appropriate. This branch should only be responsible for non-reentrant changes and does not need to involve selective checkpointing
There was a problem hiding this comment.
Done — the whole file moved out of this PR. The contract is now the first commit of #1980 (fc688680), and this branch touches no selective-checkpointing file at all; git diff --name-only main..HEAD | grep selective is empty.
7410d37 to
d6dfeca
Compare
ErenAta16
left a comment
ErenAta16
left a comment
There was a problem hiding this comment.
Reviewed as the base of the stack. The direction is clearly right and one part of it is done unusually well, so I want to name that before the one gap I found.
Retiring the signature checker is properly justified
main carries _check_signature_of_forward, and tests/utils/test_checkpoint_wrapper_checker.py pins eight cases against it — five shapes it must reject (missing type hints, no raw tensor among the args, missing return annotation, no raw tensor in the return, return not a tuple) and three it must accept. Both are deleted here.
Deleting a guard plus its tests in the same commit is usually where I'd push back, but this one earns it. _KeywordOnlyBlock in the new test_recompute.py is exactly a shape the old checker rejected:
def forward(self, inputs: dict[str, torch.Tensor], *, scale: float) -> dict[str, torch.Tensor]:and its docstring says so plainly — "a forward shape only the non-reentrant implementation supports". So the constraint isn't being dropped, it's being demonstrated gone, which is the right way to retire a validator. The reentrant implementation needed raw tensors on both sides for autograd to reconnect the graph; non-reentrant uses saved-tensor hooks and doesn't. That reasoning is worth a line in the PR body, because from the diff alone the deletion reads as scope creep.
The gap: a tensor used on both sides of the checkpoint boundary
tests/utils/test_pytree_reentrant_checkpoint.py is also deleted, and I don't think its scenario comes back. What it pinned:
direct = direct_source * 2
nested = nested_source * 3
loss = block(direct, nested=[nested]).sum() + nested.square().sum()
loss.backward()
assert direct_source.grad == 30.0
assert nested_source.grad == 102.0nested feeds the checkpointed block and a term outside it, so its gradient has to accumulate down two paths without the checkpoint re-traversing the outer graph. Absolute values, not a comparison.
The three new tests don't reach that shape:
| test | what it covers |
|---|---|
test_wrapper_is_transparent_to_state_dict_and_attributes |
param names, state_dict, attribute passthrough |
test_non_tensor_signature_preserves_gradients |
wrapped vs plain, but x is consumed only inside the block |
test_recompute_matches_baseline_under_domino_ep |
recompute vs baseline under distributed EP |
The middle one is the closest, and its x never appears outside the wrapped module, so a double-accumulation bug wouldn't show up. The third is a DeterministicDDPTestCase — if that needs more than one device it won't be the safety net on a normal CI run either.
There's also a category difference worth being deliberate about: the deleted test asserted absolute gradients, the new ones assert equality with a non-recomputed baseline. Differential is the better choice for a refactor, since it catches recompute-specific divergence — but it passes if both paths are wrong the same way, which absolute pinning wouldn't.
Non-reentrant handles shared activations correctly by construction, so I'm not claiming a bug. I'm saying the one scenario that historically broke checkpointing implementations no longer has a test, and adding nested.square().sum() to the existing keyword-only test would restore it for roughly two lines.
Smaller
checkpointing.py goes 120 → 136 lines while dropping three public names (checkpoint_wrapper, pytree_reentrant_checkpoint, _check_signature_of_forward) in favour of apply_gradient_checkpointing. Anything outside this repo importing the old names breaks. If xtuner.v1.model.utils is treated as public API, a deprecation shim would be cheap; if it isn't, saying so in the PR body would settle it.
Since this is 1/3 and #1980 / #1981 build on it: both currently run the same four checks as this PR (detect_changes, lint, unit_test, build-n-publish), so the upper layers aren't under-tested relative to the base. Worth keeping an eye on if the stack grows.
7ba1b7e to
ff899c2
Compare
ErenAta16
left a comment
ErenAta16
left a comment
There was a problem hiding this comment.
Re-read the file at ff899c2a. The keyword-only test is meaningfully stronger than when I first looked, and one half of what I raised is now covered.
What closed. It compares the input gradient against a baseline rather than only the weight gradient:
plain({"x": x}, scale=2.0)["out"].square().sum().backward()
baseline_input_grad, x.grad = x.grad.clone(), None
wrapped({"x": x}, scale=2.0)["out"].square().sum().backward()
assert torch.equal(x.grad, baseline_input_grad)Clearing x.grad between the two runs is the detail that makes it work; without it the second backward accumulates onto the first and the assertion compares a sum against a single contribution.
What is still open: x has exactly one path to the loss. .square() is applied to the block's output, not to x outside the block, so the shape from the deleted test_pytree_reentrant_checkpoint.py (a tensor feeding the checkpointed region and a term outside it) is not reproduced.
Measured on torch 2.11.0, with f standing in for the checkpointed block:
single path, f(x).square().sum() -> x.grad = 36.0
double path, f(x).square().sum() + x.square().sum() -> x.grad = 40.0
difference = 4.0 ( = d/dx x^2 )
An implementation that lost the outside-path contribution produces 36.0, and since the current assertion only checks wrapped == plain, both sides would produce 36.0 and agree. Differential testing is the right instinct for a refactor, but it cannot see an error that both paths make identically, and dropping a second accumulation into the same leaf is exactly that kind of error.
Two lines on the existing test restore it:
plain_out = plain({"x": x}, scale=2.0)["out"]
(plain_out.square().sum() + x.square().sum()).backward()and the same on the wrapped side. x.grad then carries both contributions and the comparison has something to disagree about.
To be clear about what I am and am not claiming: non-reentrant checkpointing handles shared activations correctly by construction, so I am not reporting a bug. The point is that the scenario which historically broke checkpointing implementations is the one that lost its test in this PR, and it costs two lines to keep.
Unrelated, and worth saying because it is easy to miss. The two protocol tests added since my first pass are a good catch, and the comment explains a failure mode I would not have predicted:
特殊方法在类型上查找，
__getattr__看不到，必须逐个转发；只转发__getitem__时iter()会退化到序列协议，报出与真实原因无关的 "not subscriptable"。
Dunder lookup bypassing __getattr__, and a partial forward making iter() fall back to the sequence protocol and raise something that points at the wrong thing, is precisely the sort of wrapper bug that costs an afternoon. test_wrapper_does_not_claim_protocols_the_module_lacks is the half people forget.
| # Special methods are looked up on the type, so `__getattr__` never sees them and each protocol | ||
| # has to be forwarded explicitly. Which ones exist is decided per wrapped type, because defining | ||
| # them unconditionally changes what the wrapper *is* rather than what it forwards: a `__len__` | ||
| # that always exists makes `bool(wrapper)` call it, so `module or default` -- which every | ||
| # `nn.Module` answers True for -- would raise for any module that is not sized. | ||
| def __new__(cls, module: nn.Module, checkpointed_call: Callable[..., Any]) -> "CheckpointWrapper": | ||
| if cls is CheckpointWrapper: | ||
| cls = _wrapper_class_for(type(module)) | ||
| return super().__new__(cls) |
There was a problem hiding this comment.
This implementation is a bit hacky. Could you refer to FSDPModule? Through the approach of making a new subclass
There was a problem hiding this comment.
Done — switched to the fully_shard pattern: install_checkpointing builds type(f"Checkpoint{cls.__name__}", (CheckpointModule, cls), {}) and assigns module.__class__, so the module stays itself instead of being wrapped.
That removed more than the hacky part: __getattr__ forwarding, the per-type protocol table and __new__, the state-dict prefix hooks, the named_parameters override, and the four places outside this file that stripped _checkpoint_wrapped_module. from names (train_engine.py, internal_metrics.py, misc.py, base.py) — the prefix no longer exists.
One trap worth recording: the mixin overrides __call__, not forward. Hooks run inside nn.Module.__call__, so replacing forward would put them outside the checkpointed region — which is how the DSA top-k cache lifecycle broke once before. Verified the event order is enter-region → pre-hook → forward → post-hook → exit-region.
Also fixed a gap this exposed: the SAC path had its own checkpoint(...) call that did not flatten inputs, so selecting a unit re-broke activation offload. Both paths now go through checkpoint_flattened.
Verified on 30B: eager+SAC, domino+offload, and compile+domino+deepep all run, with memory matching the pre-rewrite numbers.
…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
…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.
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`. Two GLM-5.2 tests are left failing rather than skipped or marked `xfail`, so the gap is loud rather than quietly green: - 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. Their bodies change only where the removed API forced it: `checkpoint_wrapper` with an explicit `CheckpointImpl.REENTRANT` no longer exists, so the call site and the test name that referenced it are updated. MTP itself is healthy on the non-reentrant path -- 19 non-zero parameter gradients and 5 `None`, identical to base -- and 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.
…ed-tensor hooks Inputs are flattened before being handed to `checkpoint` and reassembled inside. Not for gradient correctness -- non-reentrant checkpointing gets nested inputs right either way -- but for who else gets to see them. `_CheckpointFrame.save_inputs` wraps only *top-level* tensor arguments into a `SavedVariable`, and constructing one is what fires the ambient `saved_tensors_hooks`. It runs just before the checkpoint installs its own hooks, so those ambient hooks are still the caller's. That is how activation offloading gets hold of a layer's inputs. Converting the decoder layers to `TypedDict` I/O turned the call into `decoder_layer(hidden_states_list, ...)`, so the only positional argument is a list. Its tensors are stored as plain references, reach no hook, and are never offloaded -- silently: gradients stay correct and nothing fails. Measured on Qwen3-MoE-30BA3, `ep_size=4`, `intra_layer_micro_batch=2`, 8k, eager, all2all, peak allocated at the last step: | | peak allocated | peak reserved | |---|---|---| | offload off | 97.73 GB | 122.40 GB | | offload on, before | 98.07 GB | 122.75 GB | | offload on, after | 96.09 GB | 119.29 GB | Enabling offload used to cost 0.34 GB and save nothing. Step-1 loss is bit-identical across all three. The regression test asserts by `data_ptr`, not by shape: a checkpointed region's output is easily the same shape as its input, and an earlier version of this test passed against the unfixed code for exactly that reason.
89dcd9c to
13af713
Compare
ErenAta16
left a comment
ErenAta16
left a comment
There was a problem hiding this comment.
Read the rewrite at 13af713a. The move to the fully_shard pattern is a clear improvement over the wrapper, and the saved-tensor-hooks commit is the most valuable thing in this PR. I checked its premise independently rather than from the commit message, and I have one correction about what that commit actually contains.
The hooks premise holds, derived from torch alone
Your reasoning is that _CheckpointFrame.save_inputs wraps only top-level tensor arguments into a SavedVariable, that constructing one is what fires the ambient hooks, and that a tensor nested in a container therefore reaches none of them. I tested that against torch.utils.checkpoint directly, with nothing from this repo involved, so the two sides are independent:
torch 2.11.0+cu128
A. nested tensor passed straight to checkpoint (no flattening)
list [x] input seen by ambient hook -> False
dict {'x': x} input seen by ambient hook -> False
B. top-level positional tensor (control)
positional tensor input seen by ambient hook -> True
C. flattened first, reassembled inside (what this PR does)
list [x] input seen by ambient hook -> True
dict {'x': x} input seen by ambient hook -> True
D. grad(no flatten) == grad(flatten) -> True
B is what makes A meaningful: the hook machinery is working, the tensor just never reaches it. D is the part that matters most, and it is why this could sit in the tree indefinitely. The gradients are bit-identical, so offload silently transfers nothing while every correctness signal stays green. Your measured numbers say the same thing from the other end, offload costing 0.34 GB and saving nothing.
Asserting on data_ptr rather than shape is the right call, and I would keep it even though _FlexibleBlock now uses nn.Linear(4, 6). The shape difference makes the current fixture safe; data_ptr is what keeps it safe when someone later writes a fixture whose output happens to match its input, which is the version that passed against unfixed code.
One correction: 13af713a is a test, not a fix
The commit is titled [Fix] and reads as a behaviour change, but the diff is 47 lines in tests/model/test_recompute.py and nothing else. The flattening was already on the branch: I ran the three shapes against checkpointing.py at the parent 356868f4 and at 13af713a, and both give True for nested-in-list, nested-in-dict and passed-by-keyword.
That is not a problem with the change, the test is exactly the right thing to add. It is worth the title saying so, because anyone bisecting a future offload regression will land on this commit, see [Fix], and reasonably conclude the behaviour changed here when it changed earlier in the stack.
The point I raised on 7 August is still open, and it is still not a bug
test_non_tensor_signature_preserves_gradients at line 113 has the same single-path shape:
plain({"x": x}, scale=2.0)["out"].square().sum().backward().square() applies to the block's output, so x reaches the loss only through the checkpointed region. Clearing x.grad between the runs is there now and that was the half that was actually wrong, so this is coverage rather than correctness.
To be concrete about the non-bug half, measured on 13af713a:
single path plain 9.585812 wrapped 9.585812 equal -> True
double path plain 15.417177 wrapped 15.417177 equal -> True
outside-path contribution 5.831363 ( = 2*sum(x) = 5.831363 )
The outside-path term matches the analytic derivative independently, so the double-path case genuinely passes; non-reentrant checkpointing handles a shared activation correctly, as I said before. The argument for adding it is only that the scenario which historically broke checkpointing implementations is the one that lost its test in this PR, and differential testing cannot see an error both sides make identically. Two lines on the existing test:
plain_out = plain({"x": x}, scale=2.0)["out"]
(plain_out.square().sum() + x.square().sum()).backward()Two things I checked and found fine
_CHECKPOINT_CLASSES does what its comment says. Two layers of the same type share one generated class and isinstance against the original still holds, so checkpointing N layers builds one class rather than N. That was the first thing I went looking for and it is already handled.
Double installation raises TypeError: Cannot create a consistent method resolution order (MRO) for bases CheckpointModule, CheckpointB, since the second call sees the generated class as the base. It is not reachable today: both production call sites are in fully_shard, iterating self.layers and self.mtp_block.layers, which are disjoint and visited once. So there is nothing to fix. Recording it only because that message names neither the module nor the cause, and if isinstance(module, CheckpointModule): return module would make the call idempotent for the cost of one line if fully_shard ever runs twice on a model.
Nothing blocking from me.
2 participants
Stack (bottom to top):
feat/sac-nonlegacy-base→main← you are herefeat/sac-engine→feat/sac-nonlegacy-baserecompute_cfgand per-model recompute declarations (3/3) #1981feat/sac-recompute-cfg-v2→feat/sac-engineSummary
Moves gradient checkpointing off the legacy reentrant implementation and converts decoder-layer I/O to
TypedDict. Nothing selective-checkpointing-specific lands here: the vocabulary and the engine both start at #1980.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.uses_dsa_topk_lifecyclecapability predicate that existed only to route to it.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: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 that was still reentrant. MTP block params went from 19 non-zero grads to 0 non-zero / 24Nonewhilemtp_lossstayed finite.CrossLayerTopKSharingRuntimedetected the checkpoint's original pass vianot torch.is_grad_enabled(), a reentrant-only proxy; under non-reentrant both passes have grad enabled.TestGlm52OptimizedEngine::test_sp2_ep4_micro2_compile_offload_train_stepnow passes, matching base.Test plan
tests/model/test_recompute.py,tests/model/test_glm52_mtp_checkpoint_repro.py,tests/module/attention/test_dsa_mla.py,tests/model/test_moe.py.tests/utils/test_checkpoint_wrapper_checker.pydeleted — it only exercised the removed checker.pre-commit run --from-ref <base> --to-ref <tip>: all hooks pass including mypy.Pre-existing failures confirmed identical on base
5c2275c7and not addressed here:test_internal_metrics.py, atest_moe.pyDDP-harness timeout that reproduces identically on base when run in isolation, and threeTestGlm52PretrainedEngineloss-curve tests (hard-coded reference curve appears to predate the available checkpoint).Activation offload
Converting the decoder layers to
TypedDictI/O turned the call intodecoder_layer(hidden_states_list, ...), and that silently disabled activation offload on themicro-batch path.
_CheckpointFrame.save_inputswraps only top-level tensor arguments into aSavedVariable, and constructing one is what fires the ambientsaved_tensors_hooks-- which ishow offload gets hold of a layer's inputs. Tensors inside a list reach no hook at all.
Inputs are therefore flattened before being handed to
checkpointand reassembled inside. Measuredat
ep_size=4,intra_layer_micro_batch=2, 8k, eager, peak allocated at the last step:Enabling offload used to cost 0.34 GB and save nothing. Step-1 loss is bit-identical across all
three, so this changes memory only.
Review comments addressed
selective_checkpointing.pymoved out of this PR entirely; the contract is now the first commit of [Feature] Add the selective checkpointing engine (2/3) #1980, and this PR touches no SAC file.xfail, so the gap stays loud. Their bodies change only where the removed API forced it.TestDominoEPRecomputeremoved — a checkpoint-wrapper unit test should not stand up a whole MoE model; that coverage belongs in the integration tests.torch.equal, notassert_close: non-reentrant recompute is bitwise exact, so the assertion should be too.CheckpointWrappernow forwards__len__/__iter__/__contains__alongside__getitem__. Special methods are looked up on the type, so__getattr__never saw them; forwarding only__getitem__madeiter(wrapper)fall back to the sequence protocol and report "not subscriptable" for a module whose real problem was that it is not iterable. Covered by two new tests, both verified red before green.