[Feature] Add the region-level 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 region-level selective checkpointing engine (2/3) #1980 feat/sac-engine → feat/sac-nonlegacy-base ← you are here
3. [Feature] Add region-level recompute_cfg and per-model marker 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 selective-checkpointing engine: a marker session, a per-op policy, and one seam that the FSDP sharding paths call. 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. Default is MUST_RECOMPUTE; ops inside an active marker interval are MUST_SAVE.

The seam

def apply_selective_checkpointing(
    module, intervals=(), *, preserve_rng_state=True, layer_compiled_as_one_region=False
) -> nn.Module

One unconditional call at all five wrap sites (MoE layers, MTP, dense, both vision towers). Empty intervals == today's full-layer recompute, so this PR alone is behaviour-preserving. The strategy choice lives inside: DSA-lifecycle layers → legacy reentrant, empty intervals → plain checkpointing, otherwise SAC.

Nesting: FSDP outermost → CheckpointWrapper.forward → checkpoint boundary → marker session → module.__call__ → hooks → module.forward. The session is opened inside the region so the recompute pass rebuilds its own state, and module.__call__ (not .forward) keeps module hooks inside the region — an earlier in-place forward replacement moved them out and broke the DSA top-k cache lifecycle.

Policy invariants

Two rules apply regardless of marker state, because they are properties of the mechanism:

• Ops that write through anything other than an out argument are rejected inside a kept region. "Never keep mutating ops" is not sufficient: if a region keeps zeros_like and recomputes add_, recompute pops the forward's already-summed buffer and adds again — finite loss, no error, wrong gradients. Out-variants are exempt on the schema property alias_info.is_write, not on op name, because inductor's aten.mm.out is what the policy sees on every compiled layer.
• c10d / _c10d_functional collectives are always recomputed. Keeping a collective is only correct while its destination buffer's allocating op sits in the same interval; an interval boundary splitting them would return garbage silently.

Measured on a real 2-layer MoE, ep_size=2: keeping only the expert GEMMs gives MUST_SAVE 18 / MUST_RECOMPUTE 235; spanning the all2all gives 247 / 6 with forced recomputes {'c10d:collective': 2, '_c10d_functional:collective': 4}. These were obtained by injecting markers by hand — this PR adds no markers to model code, so the numbers are not reachable from a config until PR #1981 lands.

Compile behaviour

contextvars is untraceable by Dynamo, so both checkpoint_record and the session bind early-return under torch.compiler.is_compiling(). Consequences, all verified:

• Under compile, a region is addressable only if its contents run in eager — not merely its endpoints. Markers can sit outside a compiled region and still keep nothing if the region encloses one.
• Dense and both vision towers compile the whole layer, so they degrade to plain full recompute — exactly today's behaviour, no regression, and the layer is still checkpointed (1463.0 → 1173.8 MiB on a 4-layer probe).

Diagnostics

A configuration that silently does nothing is the failure mode this feature is most prone to, so three cases are reported, once per model (deduped on the whole diagnosis, held in a WeakKeyDictionary):

• markers that no layer recorded;
• layers compiled as one region, where intervals cannot apply;
• intervals that opened but kept nothing — a statement about contents, which is what catches the case above where markers fire but their region is compiled. Also covers a region containing only collectives.

Test plan

tests/model/test_selective_checkpointing.py (new): kept-region output and gradients bit-identical to full recompute (atol=rtol=0) plus gradient liveness; unbalanced interval safe; overlapping intervals union; marker outside a session is a no-op; in-place-in-kept-region rejected; DSA layer falls back to reentrant; two GPU cases (2 ranks, ep=2, intra_layer_micro_batch=2) eager and compiled. Diagnostic tests verified red-before-green — the first version passed vacuously because caplog never sees loguru output.

Plus tests/model/test_recompute.py and tests/module/attention/test_dsa_mla.py green; pre-commit clean 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.