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

Base is PR2's branch. Review only this PR's own diff.

---

Summary

The user-facing half: a tri-state recompute_cfg, a per-model RecomputeTargetMap declaring what each unit resolves to, and the resolution that installs them at model construction — mirroring how compile_cfg is declared and resolved.

Design notes, including the routes measured and rejected: docs/design/selective_checkpointing_and_compile.md.

User surface

recompute_cfg: list[RecomputeUnit] | bool | None = None

• None / False — keep nothing (every selected layer recomputed whole)
• True — the model's declared units
• explicit list — exactly those units

None deliberately does not mean "model default", diverging from compile_cfg: keeping units resident trades memory for speed, so an unset config must not silently change an existing run's peak memory.

Orthogonal to recompute_ratio, which continues to decide which layers are checkpointed; recompute_cfg decides what inside a selected layer is kept.

Per-model declaration

MOE_RECOMPUTE_CFG: RecomputeTargetMap = {
    RecomputeUnit.SAVE_ATTN: KeptOps("flash_attn::_flash_attn_varlen_forward_v3", ...),
    RecomputeUnit.SAVE_MOE_GATE: KeptCallables(f"{_MOE_GATE}.project"),
    RecomputeUnit.SAVE_MOE_DISPATCH: KeptCallables(...),
}

Which of the two resolutions a unit gets is not a style choice. KeptOps works only when the op identifies the unit on its own — flash_attn::_flash_attn_varlen_forward appears nowhere else — and costs nothing. KeptCallables expresses anything, at the price of torch._dynamo.disable on those callables, so selecting one also relaxes fullgraph on the entries that would otherwise inline them.

That price is why attention moved to op identity: resolving save_attn by withdrawing _pre_moe_forward measured −9.9%, because that one method accounts for 128 of the layer's 154 captured calls. KeptCallables should always name the smallest callable that covers the unit.

MoEGate.forward is split

The routing projection moves into MoEGate.project so the marker encloses only it. Top-k selection normalises its weights in place (greedy.py, topk_weights /= ...), and an in-place op inside a kept unit is recomputed rather than kept — which would have made save_moe_gate keep nothing.

Measured

Qwen3-MoE-30BA3, ep_size=4 on 8 GPUs, deepep, torch 2.10, 16k domino (pack_max_length=8192,
intra_layer_micro_batch=2), mean of steps 5-8, three independent runs per setting:

unit	tgs (3 runs)	mean	peak allocated
none	10206.0 / 9990.6 / 10205.9	10134.2	84.4-85.3 GB
save_attn	10165.6 / 10204.6 / 10071.2	10147.2	90.4-91.2 GB
save_moe_gate	10095.2 / 10179.5 / 10144.4	10139.7	93.7-94.5 GB
save_moe_dispatch	OOM	—	—

Throughput is unchanged: the three means are within 0.1% of each other while the baseline's own
runs span 2.1%. A single run per setting reads as ±2% in whichever direction the variance fell, so
it is reported as three. Memory is the reproducible effect — +6.0 GB and +9.2 GB, each within
±0.4 GB across runs.

On this model a unit therefore costs memory and buys nothing. That is a property of where an MoE
layer's recompute actually goes, not of the mechanism: a resident-activation census finds 61% of
the resident set is the loss's fp32 logits, which no recompute unit addresses, and nothing inside
the MoE layers survives the whole-layer checkpoint at all.

save_moe_dispatch OOMs at every shape tried — with and without domino, and at
pack_max_length=4096 where it peaks at 111.9 GB against the baseline's 85.9 GB. It is verified
numerically for one step and declared for smaller models, not as a default here.

Numerical correctness

Checkpointing changes only the backward pass, so step-1 loss must be bit-identical whatever is
kept. It is, across every setting and both shapes — 2.46262765 at 8k, 2.38410378 at 4k,
including the save_moe_dispatch run that OOMs immediately after.

Gradients need a noise floor to interpret, so the baseline was run twice under identical settings:

	step-1 loss	step-1 grad_norm	vs baseline
baseline	2.46262765	24.39401245	—
baseline, second run	2.46262765	24.39329147	3.0e-5
save_attn	2.46262765	24.39325905	3.1e-5
save_moe_gate	2.46262765	24.39400864	1.6e-8

Every unit sits at or below the floor two identical runs produce. Measured eager
(torch_compile=False, all2all, no domino) so compilation cannot be the source of a difference.

Also in this PR

DeepEPDispatcher builds its low-latency buffer in __init__ instead of on first dispatch. The buffer is a process-wide singleton whose constructor all_gather_objects device ids and IPC handles, and that swaps a tensor's storage through aten.set_. Left lazy, the one-time setup lands inside the first checkpointed forward and fails the step with "Tensor cached during selective activation checkpoint has been mutated" — on the first step only, from a call site unrelated to the model code. hidden_size is threaded through build_dispatcher because the buffer is sized from it.

Test plan

tests/model/test_recompute.py and tests/model/test_selective_checkpointing.py, 27 tests. pre-commit passes including mypy. End-to-end numbers above.

[ErenAta16]

Reviewed the whole stack in order, so this is mostly about what changes for the finding I left on #1980.

The degenerate-interval gap moves onto a public surface

MarkerSession.record and .finish are byte-identical to the _MarkerSession versions in #1980. What changed is where they live and who can reach them:

#1980:  xtuner/v1/model/utils/selective_checkpointing.py   _MarkerSession   (private)
#1981:  xtuner/v1/utils/selective_checkpointing.py          MarkerSession    (public)

xtuner/v1/utils/__all__       exports MarkerInterval, RecomputeIntervalMap,
                              RecomputeUnit, checkpoint_record
xtuner/v1/model/utils/__all__ re-exports the same four

So the two intervals that never close —

(("mlp","mlp"),)         start == end, the elif is unreachable, opens and never closes
(("s","e"),) seq: e, s   discard is a no-op, then start opens it, nothing closes it

— are now reachable from a RecomputeIntervalMap written outside this repo, and MarkerInterval being tuple[str, str] with no validation is now part of the exported vocabulary rather than an internal detail.

The user surface itself is well guarded: recompute_cfg takes RecomputeUnit members and _resolve_recompute_cfg checks them against default_recompute_cfg, warning when a model declares no units. Users never write marker strings, which is the right split. The exposure is to whoever writes default_recompute_cfg for a model — and with the helpers exported from xtuner/v1/utils, that is now third-party model code, not only this repo's.

That makes the finish() check I suggested on #1980 cheaper to justify, not more expensive: it is the only place that sees an interval still open at the end of a completed pass, and it is now the boundary between an architecture's declaration and silently keeping every activation after the start marker.

def finish(self) -> None:
    if self._open:
        log_rank0.warning(
            f"Selective checkpointing: intervals {sorted(self._open)} were still open when the "
            f"region ended, so every op after their start marker was kept resident."
        )
    _report_pass(self._owner, self._intervals, self._recorded, self._kept)

A start == end pair could also be rejected where a RecomputeIntervalMap is declared — under the current record() semantics it can never mean anything useful, and catching it at declaration is better than at the end of a pass.

The rest of the user surface reads well

_resolve_recompute_cfg warning that recompute_cfg=True has no effect when a model declares no units is exactly the failure a user would otherwise attribute to the engine rather than to their model choice. Tri-state rather than a bool is also the right call — "keep everything the model declares" and "keep these specific units" are genuinely different intents and collapsing them into one flag would have forced a sentinel later.

Splitting the vocabulary (RecomputeUnit, MarkerInterval, checkpoint_record) into xtuner/v1/utils while the engine stays in model/utils is the right seam: model authors import the vocabulary, the sharding paths import the engine, and neither pulls the other in.

One documentation note: default_recompute_cfg is described in a comment as "the vocabulary of what an architecture can keep resident". That framing is worth promoting into the docstring on base.py:1033 — it is the sentence that explains why an empty map is a legitimate answer rather than a missing implementation, and it is the thing a model author reads first.

[ErenAta16]

Re-reviewed from scratch at 28ffa316. My earlier comments here were written against the MarkerSession / interval design and are obsolete; this replaces them.

The gap I raised on the old design was fixed, in a form that fits the new one. My original point was that the diagnostics only detected keeping too little and were silent when a selection turned out inert. That case is now explicit:

if selected is True and not supported:
    log_rank0.warning(
        f"`recompute_cfg=True` has no effect: {type(self).__name__} declares no recompute units, so every "
        "region is recomputed."
    )

and the neighbouring ValueError for an unsupported unit is the best error message in the stack, because it answers the question the user is about to ask:

Note that a compose model declares no units of its own -- set recompute_cfg on the sub-model config that owns the layers instead.

But one route to the same silence survives, one level further down. The warning above fires on an empty vocabulary. It cannot fire when the vocabulary is fine and the resolution comes back empty:

target = supported[unit]
if isinstance(target, KeptOps):
    op_names.extend(target.names)
...
self._selected_recompute_units.add(unit)
return resolve_kept_ops(tuple(op_names)), callables

_selected_recompute_units is added to before anything is resolved, and resolve_kept_ops drops unregistered names at debug level:

if packet is None:
    log_rank0.debug(f"Selective checkpointing: {name} is not registered in this build, skipping it.")
    continue

So for a user who explicitly selects a unit whose target is KeptOps, where none of the listed spellings exist in this build:

_selected_recompute_units  non-empty  ->  keeps_any_recompute_unit = True
_kept_ops                  empty
_kept_callables            empty      (the unit was KeptOps, not KeptCallables)

apply_selective_checkpointing then takes the keeps_any_unit=True branch, installs the selective-checkpoint policy, and the policy keeps nothing. The user gets full recompute plus a per-op dispatch mode — which is precisely the cost keeps_any_unit exists to avoid, per its own docstring:

torch's own default already recomputes everything, so running a policy to reach the same answer would only put a dispatch mode in the way of every op.

Nothing at info or above says so. This is not hypothetical for the case the skip was written for: the docstring gives flash-attention v2 versus v3 as the motivating example, and a build with neither is exactly how you end up here.

The skipping itself is right and I am not asking for it to raise — that would break the multi-spelling pattern it exists to support. What is missing is the distinction between some names resolving and none of them resolving. Something like: if a selected unit contributed op names and none of them resolved, warn naming the unit and the names it tried. Per unit rather than per name, so a v2/v3 pair still resolves quietly.

The KeptCallables path already has this property, which is what makes the asymmetry visible:

log_rank0.info(
    f"Keeping {sorted(self._selected_recompute_units)} resident excludes {sorted(self._kept_callables)} "
    f"from torch.compile, so the callers that reach them are compiled with `fullgraph=False`."
)

A user who selects a callable unit gets told what happened at info. A user who selects an op unit that resolved to nothing gets debug.

On _without_compiled_selected_regions. The reasoning is the part I would have got wrong:

Which callers those are is not knowable from the config: Dynamo inlines across call sites, so the break can surface in any compiled method upstream of the unit. Relaxing all of them is the only rule that is correct without tracing.

Relaxing every compiled caller rather than trying to compute the affected set looks over-broad until you notice that inlining makes the set undecidable statically. Worth keeping that sentence, because the obvious "optimisation" here is to narrow it, and the comment is the only thing that would stop someone.

The info log naming both the units and the excluded callables is the right level for it, since silently dropping fullgraph=True is something a user tuning compile behaviour needs to see.

One question on recompute_cfg=True semantics. units = list(supported) if selected is True else selected means True tracks whatever the model declares, so a model that gains a unit in a later release changes behaviour for a config that did not change. That is probably the intent, and it mirrors how compile_cfg works, but it is worth one line in the design doc: True is "whatever this version offers", not a stable set. Anyone pinning memory behaviour across upgrades needs the explicit list instead.

Design doc being in the PR is unusual and welcome, particularly the routes measured and rejected. That is normally the part that gets lost.

[HAOCHENYE]

missing docstring

[HAOCHENYE]

Done.

[HAOCHENYE]

The semantics here already seem a bit off. I suggest adjusting the design: change recompute_cfg into a dictionary, and migrate recompute_ratio from the fsdp cfg (BC needs to be considered). One field in recompute_cfg represents the region of sac (I trust you're better at naming than I am). Then, migrate all recompute-related settings into this recompute_cfg.

[HAOCHENYE]

Done. recompute_cfg is now a RecomputeConfig:

recompute_cfg=RecomputeConfig(ratio=0.25, save=["attn"])

ratio / vision_ratio decide which layers are recomputed, save decides what stays resident inside them — the pair that previously lived in two different configs.

On BC: fsdp_cfg.recompute_ratio and vision_recompute_ratio stay as deprecated aliases. Their default becomes None rather than 1.0, which is what makes "not set" distinguishable from "explicitly set to 1.0" — with 1.0 as the default, every run would look like it was requesting a migration and would overwrite whatever the user put in the new field. A set value is migrated with a warning; setting both the old and the new field raises rather than silently picking one. The exit condition is written on the fields.

In-repo configs (autotest/config/*_recompute.py) are moved to the new field, so CI exercises the new path rather than the shim. Three regression tests cover the migration, each verified red before green — including the None-sentinel one, which fails with a spurious "both set" error if the default goes back to 1.0.

Verified on 30B that the rename and the config change are behaviour-neutral: peak allocated 83.12 GB with no unit and 89.09 GB with save=["attn"], against 83.12 / 89.08 before.

[HAOCHENYE]

It’s the same here. This is actually not a recompute unit, but a save unit. The design needs to be updated and corrected

[HAOCHENYE]

Agreed — done. RecomputeUnit is now SaveUnit, and the members drop the redundant prefix since the class carries it: SAVE_ATTN -> ATTN, serialized as "attn". Selecting a member means "save this"; everything unselected is recomputed, which is the default and costs nothing.

[ErenAta16]

The finding I left here on 1 August is withdrawn along with its parent; I have written the reasoning out on #1980 rather than repeat it. Short version: it described MarkerSession intervals that open and never close, that mechanism is gone from the tree, and the replacement scopes a unit with a try/finally ContextVar reset installed on the callable, so there is no author-declared pairing left to get wrong.

SaveUnit is the right correction. "Selecting a member means save this, everything unselected is recomputed" is a rule you can hold in your head, which the old phrasing was not.

The rename stops one layer short, though. The enum changed name; the three exported functions that carry it did not:

xtuner/v1/utils/selective_checkpointing.py
  __all__ = [..., "active_recompute_unit", "recompute_unit", ...]
  def active_recompute_unit() -> SaveUnit | None
  def recompute_unit(unit: SaveUnit) -> Iterator[None]
  _ACTIVE_UNIT: ContextVar[SaveUnit | None] = ContextVar("xtuner_recompute_unit", default=None)

xtuner/v1/model/utils/selective_checkpointing.py
  def in_recompute_unit(unit: SaveUnit, func: Any) -> Any

So a reader meets recompute_unit(SaveUnit.ATTN) and in_recompute_unit(SaveUnit.ATTN, fn), which read as "recompute attn" while meaning the exact opposite. That is the same inversion the enum rename just removed, left standing on the functions that carry the enum. in_save_unit / save_unit / active_save_unit would finish it, and all four names including the ContextVar label are internal to this stack, so the change is mechanical. Optional, and not worth holding a merge for, but it is cheaper to do now than after the vocabulary is exported and in use.

On the exported surface. KeptOps, KeptCallables, SaveUnit, RecomputeTarget and RecomputeTargetMap are exported from both xtuner/v1/utils and xtuner/v1/model/utils, so a model outside this repo can declare a target map. TestDeclaredTargets.test_declared_targets_resolve catches a name that resolves to nothing, which is the failure that would otherwise be silent, and its comment says exactly why:

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.

That check runs over MOE_RECOMPUTE_CFG and DENSE_RECOMPUTE_CFG, so it covers the in-repo declarations and not out-of-repo ones, which get log_rank0.debug and an inert selection. Given the vocabulary is public, it is worth deciding whether that is intended. It is a reasonable answer either way, and I raise it because the decision is currently implicit rather than because I think it is wrong.

On the config change. Defaulting the deprecated fsdp_cfg.recompute_ratio to None rather than 1.0 is the detail that makes the migration safe, and it is the kind of thing that usually gets found the hard way: with 1.0 as the default, every run looks like it is requesting a migration and would overwrite the new field. Having a regression test that fails specifically if the sentinel goes back to 1.0 is the right guard, since that is a one-character regression that no other test would notice. Moving autotest/config/*_recompute.py onto the new field so CI exercises the real path instead of the shim is the other half of it.