[Feature] Add region-level recompute_cfg and per-model marker 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 region-level selective checkpointing engine (2/3) #1980 feat/sac-engine → feat/sac-nonlegacy-base
3. [Feature] Add region-level recompute_cfg and per-model marker 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, per-model default_recompute_cfg mapping semantic units to marker intervals, and checkpoint_record markers in the decoder-layer forwards. This is what makes the engine in PR2 reachable from a config.

User surface

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

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

None deliberately does not mean "model default", diverging from compile_cfg: keeping regions 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 which regions inside a selected layer are kept.

RecomputeUnit is a StrEnum, so configs round-trip readably (["save_attn", "save_mlp"]) and trainer resume preserves the setting — no Field(exclude=True).

Layering

checkpoint_record is called from xtuner/v1/module/decoder_layer/*.py, so the contract cannot live under model/ — importing it from there pulls xtuner/v1/model/__init__.py, which imports the decoder layers back:

ImportError: cannot import name 'MoEActFnConfig' from partially initialized module
'xtuner.v1.module.decoder_layer.moe_decoder_layer' (most likely due to a circular import)

The contract and marker session move to xtuner/v1/utils/selective_checkpointing.py; the engine stays at the model layer, where it may import uses_dsa_topk_lifecycle. model/utils/__init__.py re-exports, so no import site outside these files changes. A regression test pins the standalone import in a subprocess — in-process the check is vacuous, since xtuner.v1.model is already imported by then.

Markers

Explicit .begin / .end pairs rather than "end = the next region's start". _micro_batch_forward splits the dispatch/combine chain across four stage loops, so "the next region" is a different region there than in _forward; under the implicit convention the same interval would mean two different things in the two paths. Both paths are instrumented, and a static guard checks every marker name referenced by a default_recompute_cfg is actually recorded.

What each unit actually keeps — measured, not asserted

Eager, ep_size=2, kept-op identities from the real policy:

unit	kept	ops	collectives
save_attn	624	view/transpose/t/mm	{}
save_moe_gate	152	_to_copy/view/add/sum/div	{}
save_moe_dispatch	152	view/sum/permute/unpermute	{}
save_mlp	202	view/t/mm/silu	{}

Loss bit-identical to full recompute in every row, 48/48 live grads. save_moe_dispatch keeps the permutation/padding buffers either side of the all-to-all; the collective itself is always recomputed (PR2's invariant), and the docstring says so.

Under EP + torch.compile the effective set is save_moe_dispatch alone. A region is addressable only if its contents run in eager: save_mlp's markers do run in eager, but its region encloses the EP-compiled _shared_experts_forward, so nothing is kept. Documented per unit with the measured numbers, and PR2's diagnostics report it at runtime.

DSA models

DSA-lifecycle models route to reentrant (PR1) and cannot keep regions, so default_recompute_cfg declares no units for them — otherwise the config validates and the trainer logs "keeps these regions resident: [6 intervals]" while nothing is kept. Verified: non-DSA units=4 intervals=6, DSA units=0 intervals=0. PR2's wrap-site fallback remains as an independent guard for a direct caller. The gate resolves lazily on first access because it must read the built layers, which subclasses create after BaseModel.__init__ returns — an eager gate would inspect a layerless model, never fire, and look correct.

Test plan

Config tri-state resolution including recursive False propagation into nested compose configs (verified on a real Qwen3VLMoE30BA3Config, not a synthetic probe); enum round-trip through model_dump_json / model_validate; marker coverage across both MoE paths; the subprocess layering test, verified red by reintroducing an upward import. 33 passed / 4 skipped at the tip of the stack; every commit in this PR builds and passes independently (bisectable); pre-commit clean including mypy.