[Feature] Add the selective checkpointing engine (2/3)

[HAOCHENYE]

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 ← you are here
3. [Feature] Add recompute_cfg and per-model recompute declarations (3/3) #1981 feat/sac-recompute-cfg-v2 → feat/sac-engine

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

---

Summary

The contract and the engine. Per review on #1979, the shared vocabulary moved down here — xtuner/v1/model/utils/selective_checkpointing.py is the first commit of this PR, so #1979 is purely the non-reentrant switch.

RecomputeUnit is a StrEnum, so it round-trips readably through saved configs; KeptOps / KeptCallables are the two things a unit can resolve to; RecomputeTargetMap is the per-model table binding them.

The engine: a per-op policy and one seam the FSDP sharding paths call. The whole DecoderLayer gets one checkpoint(..., use_reentrant=False, context_fn=...); create_selective_checkpoint_contexts(policy_fn) then decides per op whether to keep or recompute.

The seam

def apply_selective_checkpointing(
    module, kept_ops=frozenset(), *, keeps_any_unit=False, preserve_rng_state=True
) -> nn.Module

One unconditional call at all five wrap sites (MoE layers, MTP, dense, both vision towers). With nothing selected it degenerates to plain full recompute — the same mechanism, not a separate path — so this PR alone is behaviour-preserving.

keeps_any_unit=False skips installing the policy entirely. Torch's own default already recomputes everything; running a dispatch mode to reach the same answer would only put itself in the way of every op.

How the policy knows an op belongs to a kept unit

Two routes, and the policy answers the same question either way:

• resolve_kept_ops — recognised from the op itself. Nothing installed, works identically inside and outside compiled code.
• in_recompute_unit — recognised from a ContextVar the wrapper sets around a callable. Only works because the model layer also takes those callables out of the compiled set; a ContextVar set inside compiled code is neither written nor readable.

resolve_kept_ops skips names this build does not register, so a model may list several backends' spellings of the same kernel (flash-attention v2 and v3) and only the one that exists resolves.

Policy invariants

Two classes of op are never kept, whatever unit they fall in:

• Collectives (c10d, _c10d_functional). Keeping one would elide it from the recompute pass, which is only correct while the op that allocated its destination buffer is kept too. A unit boundary between the allocation and the collective would leave the recompute reading an uninitialised buffer — silently, and differently on each rank.
• Mutating ops. The recompute pass would replay them on top of their own results.

The mutating case only warns; it does not refuse. A mutable schema says the op is suspicious, not that it is unsafe — refusing on the schema alone rejects library calls that touch nothing the unit kept (deep_ep's set_). The genuinely unsafe case is a tensor the unit did keep being mutated afterwards, and torch's version counter catches exactly that. The warning runs in the forward and names the op, which torch's check — raised in backward, naming only the tensor — cannot.

Test plan

tests/model/test_selective_checkpointing.py: kept and fully-recomputed units produce the same gradients, for both resolutions; an unregistered op name is skipped; an in-place op inside a kept unit is recomputed rather than refused; mutating a kept tensor raises. pre-commit passes including mypy.

[HAOCHENYE]

https://github.com/InternLM/xtuner/pull/1980/changes#diff-91441decf6cc57947a4b3a6995d4677108c25f1834232b062e2f196f6e491400R321

There is a major design flaw: compile cannot be compatible with selective checkpointing. This is completely unacceptable. We need to discuss again.

[ErenAta16]

Read this as 2/3 on top of #1979. The correctness reasoning is unusually well documented — the comments on why collectives and mutable ops are never kept, why context_fn has to be module-level for Dynamo, and why the early-stopped recompute pass is excluded from diagnostics are all things that would otherwise be rediscovered painfully. I checked the one thing I expected to be wrong and it isn't: _forward_with_marker_session builds a fresh _MarkerSession per call, and since use_reentrant=False re-invokes the wrapped function on recompute, the markers replay in the same order and the policy sees identical state in both passes.

One gap, and it's on the side the diagnostics don't watch.

The diagnostics only detect keeping too little

Both warnings are phrased around under-keeping:

the intervals ... have no effect and every selected layer is recomputed whole

this interval keeps nothing resident and its region is recomputed

and _report_pass skips any interval that kept something:

if interval in state.kept_intervals or interval in state.reported:
    continue

So an interval that keeps everything is silent by construction. Two ways to get one, both traced through record():

intervals=(("mlp","mlp"),)          seq: mlp, x, mlp, y
  keeping after each: True, True, True, True      <- never closes

intervals=(("s","e"),)              seq: e, s
  keeping after each: False, True                 <- never closes

The first is start == end: if name == start matches, so the elif name == end branch is unreachable for that interval and it can only ever open. MarkerInterval is tuple[str, str] with nothing checking the two differ.

The second is an end recorded before its start — discard on a closed interval is a no-op, then start opens it and nothing closes it. Reachable from a conditional forward that hits one marker on a branch the other doesn't cover, or from an interval declared in the wrong order in a RecomputeIntervalMap.

In both cases every op after that point takes MUST_SAVE, so the layer stops checkpointing and the whole activation set stays resident. The symptom is memory — an OOM, or a quiet increase — and the diagnostics say nothing, because the interval did keep something and gets skipped.

That's the worse failure of the two the module can have. Under-keeping costs recompute time and gets a warning; over-keeping costs memory and gets silence.

Both are detectable with what's already in hand

_MarkerSession.finish() already runs only for passes that completed, which is exactly the right place:

def finish(self) -> None:
    if self._open:
        log_rank0.warning(
            f"Selective checkpointing: intervals {sorted(self._open)} were still open at the end of "
            f"the layer, so every op after their start marker was kept resident. Check that each "
            f"interval's end marker runs on the same path as its start."
        )
    _report_pass(self._owner, self._intervals, self._recorded, self._kept)

That catches the reversed-order case, and the start == end case falls out of it too since such an interval is always still open. A cheaper static check for the second alone would be rejecting start == end wherever a RecomputeIntervalMap is declared — it can never mean anything useful under the current record() semantics.

Small

The "any interval open ⇒ keep" rule with no pairing asserted is a deliberate choice and the comment says so, which I'd keep. My note above isn't asking for pairing — it's asking for a signal when the set is non-empty at the end of a pass, which is compatible with unpaired semantics and is just as true for nested and overlapping intervals.

[ErenAta16]

Re-reviewed from scratch at ff5d378f, since the design changed substantially after my earlier pass and my old comments were written against the interval/MarkerSession shape that no longer exists. Treat this as replacing them.

The op-identity design is a better fit for the problem than intervals were. Intervals had to be opened and closed correctly by model code, and the failure mode was a marker that never closed. Here a unit is either a set of op overloads or a wrapped callable, and neither can be left half-open.

Three things I checked expecting to find a problem, and did not.

_checkpoint_policy refuses collectives before it refuses anything else:

if op.namespace in _NEVER_KEPT_NAMESPACES:
    return CheckpointPolicy.MUST_RECOMPUTE

with the reasoning that keeping a collective would elide it from the recompute pass, and a unit boundary between the destination allocation and the collective leaves the recompute reading an uninitialised buffer, per rank. That is the kind of thing that produces a silent wrong-gradient bug on eight GPUs and not on one, and it is handled before any other consideration.

context_fn is a module-level partial, with the comment saying why:

Dynamo's checkpoint higher-order op rejects anything else (lambdas, closures, bound methods)

in_recompute_unit deliberately skips functools.wraps, again with the reason recorded (__wrapped__ makes Dynamo resolve through to the unbound function and lose self). Both are the sort of constraint that gets "cleaned up" by someone who does not know why it is there, and both now carry their reason at the site.

And the ContextVar survives the recompute pass for the right reason: use_reentrant=False re-invokes the wrapped callable, so in_unit runs again and sets the var again in whichever thread the recompute lands on. The policy therefore sees the same answer in both passes rather than relying on state captured during forward.

The warn-versus-reject decision is right, and the docstring is the best part of the PR.

Whether that actually happens is not decidable from the op's schema, only from whether the tensor it writes to was cached. Torch decides exactly that, by version counter [...] Refusing here on the schema alone would reject units that are perfectly safe.

Reporting in forward where the op name is known, and leaving the actual decision to torch's version check in backward, is the correct division. The test that pins it (test_mutating_a_kept_tensor_is_caught_by_torch, asserting torch's "has been mutated") proves the delegation works rather than assuming it.

_writes_only_through_out is also correct rather than accidentally correct: all() over the write-aliasing arguments is not vacuous here, because this is only reached when op._schema.is_mutable is already true, which means at least one such argument exists.

Small: a comment names a function that no longer exists.

# Mutating ops that leave tensor *values* alone, so a kept unit may contain them. Anything else with
# a mutable schema goes through `_reject_non_replayable_op_in_kept_unit`.

The function is _warn_non_replayable_op_in_kept_unit. Left over from when it refused. Worth fixing precisely because reject-versus-warn is the distinction the docstring three lines below spends a paragraph establishing, so a reader who trusts the comment concludes the opposite of what the code does.

Small: _REPORTED_MUTATING_OPS is process-global and never cleared.

That matches the stated intent ("Reported once per op, not once per layer or per step"), and set.add under the GIL is fine for the threading involved. The consequence worth being deliberate about is that a process which builds a second model, or re-applies checkpointing with a different config, gets no warning the second time even though the second configuration is a different question. For a training script that is almost always right; for a notebook or a sweep it hides the signal. A comment saying "per process, deliberately" would settle it.

A defensive note on the public signature. apply_selective_checkpointing silently ignores a non-empty kept_ops when keeps_any_unit=False:

if not keeps_any_unit:
    return apply_gradient_checkpointing(module, preserve_rng_state=preserve_rng_state)

In this PR both arguments come from the same pair of model properties, so they cannot disagree, and in #1980 they are stubs. But the function is exported from xtuner.v1.model.utils, so an external caller passing ops without the flag gets full recompute and no indication the ops were dropped. Either an assert, or folding the flag into bool(kept_ops) or has_callable_units, would make the two impossible to desynchronise. Not blocking, and I would not restructure the signature for it.

The parametrised tests reading _assert_matches(selective, full_recompute) on exact equality rather than allclose is the right call for this: any difference at all means the save-list and the recompute disagree, and a tolerance would hide exactly the bug the engine could introduce.

[HAOCHENYE]

Please write docs and comments in chinese.

[HAOCHENYE]

Done — the docstrings and comments in this file are all Chinese now.

[HAOCHENYE]

I feel the overall testing logic is too indirect. Could we directly count the number of recomputes through a DispatchMode to observe whether SAC is effective? I understand the goals of the test as follows:

1. Accuracy is fully aligned with the non-recompute version
2. Recompute takes effect
3. Recompute + SAC takes effect

[HAOCHENYE]

Done — added TestRecomputeIsObservable, which counts op executions through a TorchDispatchMode instead of asserting on what the policy returned. Your three goals map onto three assertions, and aten.tanh makes them visible directly:

	tanh executions
no recompute	2
recompute	4
recompute + save attn	2

Goal 1 is a bitwise gradient comparison against the non-recomputed run; goal 2 is the 2 -> 4 doubling; goal 3 is 4 -> 2 for the kept op. The kept case also re-checks gradients, so "ran fewer ops" cannot pass as "computed the right thing".

[ErenAta16]

Withdrawing the finding I left on 1 August. It targeted a mechanism that no longer exists, and I would rather say so here than leave it sitting in the thread as if it still applied.

What it was. I reported that _MarkerSession could only ever detect keeping too little, and that two shapes produced an interval that opened and never closed, so a region kept every activation with no diagnostic: start == end, where the elif name == end branch is unreachable, and an end recorded before its start, where discard on a closed interval is a no-op.

Why it is void. Checked the current head rather than assuming from the diff. MarkerInterval, MarkerSession, RecomputeIntervalMap and checkpoint_record return nothing across the tree:

git grep -l "MarkerInterval\|MarkerSession\|RecomputeIntervalMap\|checkpoint_record" -- '*.py'
(no output)

Both failure shapes needed an author-written open/close pair that could be mis-ordered. The replacement has no pair to get wrong. A unit's scope is now a context manager over a ContextVar:

@contextmanager
def recompute_unit(unit: SaveUnit) -> Iterator[None]:
    token = _ACTIVE_UNIT.set(unit)
    try:
        yield
    finally:
        _ACTIVE_UNIT.reset(token)

installed by in_recompute_unit around the callable itself, so the close is structural rather than declared. There is no ordering for a model author to state, and no path, including an exception out of the wrapped callable, that leaves a unit active past its region. The class of bug I described cannot be expressed in this design. That is a better answer than the finish() warning I proposed, which was a detector for a hazard that has now been removed instead.

One thing that did survive the redesign, and is already handled. The old finding was really about a selection that silently does nothing. The new shape of that is resolve_kept_ops skipping an unregistered name at debug level: if every name a unit declares fails to resolve, the unit keeps nothing, the region stays fully recomputed, and the user's save=[...] is silently inert.

You have already covered it, and the comment on TestDeclaredTargets.test_declared_targets_resolve states the reasoning more precisely than my original review did:

A renamed method or op would not fail anywhere at runtime on its own: the unit would simply keep nothing and the region would stay recomputed, silently costing the memory the user asked to keep.

Asserting at declaration time rather than warning at runtime is the right place for it, and assert resolve_kept_ops(target.names) phrased as "some name must resolve" rather than "all must" is correct given a build registers one flash-attention version.

On the new test. TestRecomputeIsObservable counting aten.tanh through a TorchDispatchMode is a real improvement over asserting on policy return values, because 2 / 4 / 2 is a property of what ran rather than of what the policy said it intended. Pairing the kept case with a gradient comparison is the part that matters most: without it, "ran fewer ops" and "computed the right thing" are not distinguishable, and the failure this whole mechanism can produce is exactly one where the count looks right and the numbers are wrong.

Nothing blocking from me on this PR.