[Feature] Add region-level recompute_cfg and per-model marker declarations (3/3)

tests/model/test_recompute.py

@@ -5,9 +5,24 @@
     test_non_tensor_signature_preserves_gradients: 关键字参数 + dict 返回值下梯度与不重算一致。
 TestDominoEPRecompute
     test_recompute_matches_baseline_under_domino_ep: domino EP 下重算与不重算的 loss/梯度一致。
+TestRecomputeCfgResolution
+    test_unset_cfg_keeps_full_recompute: `None` 不改变显存行为，解析为空区间。
+    test_true_selects_every_supported_unit: `True` 选中模型声明的全部 unit。
+    test_explicit_units_select_only_their_intervals: 显式 list 只解析出对应区间。
+    test_string_units_are_accepted: 配置文件里的字符串能解析成 RecomputeUnit。
+    test_unsupported_unit_is_rejected: 模型不支持的 unit 在构造时报错并列出支持项。
+    test_disable_propagates_into_nested_configs: `False` 递归关闭嵌套子模型配置。
+    test_disable_reaches_every_sub_model_of_a_real_compose_config: 真实 compose 配置的三个子配置都被关闭。
+    test_units_round_trip_through_json: enum 序列化成可读字符串并能读回。
+TestMarkerVocabulary
+    test_declared_intervals_have_markers: default_recompute_cfg 引用的 marker 都真实埋点。
+    test_micro_batch_path_covers_the_same_operations: 单路与 domino 路每个区间覆盖的算子一致。
 """
 
+import ast
+import inspect
 import os
+import textwrap
 
 import pytest
 import torch
@@ -17,9 +32,16 @@
 from xtuner._testing import DeterministicDDPTestCase
 from xtuner.v1.config import FSDPConfig
 from xtuner.v1.loss.ce_loss import CELossConfig
-from xtuner.v1.model.moe.moe import MoE, MoEConfig, SequenceContext
-from xtuner.v1.model.utils import apply_gradient_checkpointing
+from xtuner.v1.model.base import BaseModel, XTunerBaseModelConfig, _disable_nested_switch
+from xtuner.v1.model.compose.qwen3_vl import Qwen3VLMoE30BA3Config
+from xtuner.v1.model.dense.dense import DENSE_RECOMPUTE_CFG
+from xtuner.v1.model.moe.moe import MOE_RECOMPUTE_CFG, MoE, MoEConfig, SequenceContext
+from xtuner.v1.model.utils import (
+    RecomputeUnit,
+    apply_gradient_checkpointing,
+)
 from xtuner.v1.module.attention import MHAConfig
+from xtuner.v1.module.decoder_layer import dense_decoder_layer, moe_decoder_layer
 from xtuner.v1.module.router import NoAuxRouterConfig
 
 
@@ -185,3 +207,205 @@ def _run_once(self, ep_size: int, dispatcher: str, recompute_ratio: float):
         del model, out, loss
         torch.cuda.empty_cache()
         return loss_val, grad_norm, all_finite
+
+
+def _build_tiny_moe_config(**overrides) -> MoEConfig:
+    """A MoE small enough to instantiate on the meta device in a config-only test."""
+    router_config = NoAuxRouterConfig(
+        scoring_func="sigmoid",
+        router_scaling_factor=1.0,
+        n_group=1,
+        topk_group=1,
+        norm_topk_prob=True,
+    )
+    return MoEConfig(
+        vocab_size=256,
+        max_position_embeddings=128,
+        pad_token_id=0,
+        eos_token_id=0,
+        num_hidden_layers=2,
+        hidden_size=64,
+        intermediate_size=128,
+        rms_norm_eps=1e-6,
+        rope_theta=1e6,
+        hidden_act="silu",
+        attention=MHAConfig(num_attention_heads=4, num_key_value_heads=4, head_dim=16),
+        tie_word_embeddings=False,
+        n_routed_experts=4,
+        n_shared_experts=1,
+        num_experts_per_tok=2,
+        first_k_dense_replace=1,
+        hidden_factor=1.0,
+        moe_intermediate_size=64,
+        router=router_config,
+        compile_cfg=False,
+        **overrides,
+    )
+
+
+def _resolve_intervals(**overrides) -> list[tuple[str, str]]:
+    with torch.device("meta"):
+        return MoE(config=_build_tiny_moe_config(**overrides)).recompute_intervals
+
+
+class _NestedProbeConfig(XTunerBaseModelConfig):
+    """Stand-in for a sub-model config, as a compose model nests one."""
+
+
+class _ProbeConfig(XTunerBaseModelConfig):
+    text_config: XTunerBaseModelConfig
+
+
+class _ProbeModel(BaseModel):
+    """A model that contributes nothing but ``BaseModel.__init__``'s config resolution."""
+
+    config: _ProbeConfig
+
+
+class TestRecomputeCfgResolution:
+    def test_unset_cfg_keeps_full_recompute(self):
+        # `None` must not change the memory profile of an existing training run: unlike `compile_cfg`,
+        # it resolves to "retain nothing" rather than to the model's declared units.
+        assert _resolve_intervals(recompute_cfg=None) == []
+
+    def test_true_selects_every_supported_unit(self):
+        intervals = _resolve_intervals(recompute_cfg=True)
+
+        expected = [interval for unit_intervals in MOE_RECOMPUTE_CFG.values() for interval in unit_intervals]
+        assert intervals == expected
+
+    def test_explicit_units_select_only_their_intervals(self):
+        intervals = _resolve_intervals(recompute_cfg=[RecomputeUnit.SAVE_MOE_DISPATCH])
+
+        assert intervals == MOE_RECOMPUTE_CFG[RecomputeUnit.SAVE_MOE_DISPATCH]
+
+    def test_string_units_are_accepted(self):
+        # Configs arrive as JSON/py files where units are written as plain strings.
+        assert _resolve_intervals(recompute_cfg=["save_attn"]) == MOE_RECOMPUTE_CFG[RecomputeUnit.SAVE_ATTN]
+
+    def test_unsupported_unit_is_rejected(self):
+        # A model declaring no units cannot honour any selection, so this is a user configuration
+        # error rather than something to silently drop. It surfaces at construction, before the run
+        # spends anything on materializing and sharding weights.
+        with pytest.raises(ValueError, match="does not support"):
+            _ProbeModel(_ProbeConfig(text_config=_NestedProbeConfig(), recompute_cfg=[RecomputeUnit.SAVE_ATTN]))
+
+    def test_disable_propagates_into_nested_configs(self):
+        # A sub-model resolves its own switch, so `False` on the outer config only means something
+        # if it reaches the nested ones. `compile_cfg` must stay untouched: the walk is per switch.
+        config = _ProbeConfig(text_config=_NestedProbeConfig(), recompute_cfg=False)
+
+        model = _ProbeModel(config)
+
+        assert model.recompute_intervals == []
+        assert config.text_config.recompute_cfg is False
+        assert config.text_config.compile_cfg is None
+
+    def test_disable_reaches_every_sub_model_of_a_real_compose_config(self):
+        # The probe above has one nested config; a shipped compose config has three, one of them a
+        # further-derived MoE config. Exercised on the config walk rather than through the model,
+        # because constructing a 30B compose model is the expensive part and contributes nothing:
+        # what can regress here is which nested configs the walk reaches.
+        config = Qwen3VLMoE30BA3Config(recompute_cfg=False)
+
+        _disable_nested_switch(config, "recompute_cfg")
+
+        for sub_config in (config.vision_config, config.projector_config, config.text_config):
+            assert sub_config.recompute_cfg is False
+
+    def test_units_round_trip_through_json(self):
+        # Trainer resume reads the config back, and serialized runs are read by humans, so units
+        # must survive as their readable names.
+        config = _build_tiny_moe_config(recompute_cfg=[RecomputeUnit.SAVE_ATTN, RecomputeUnit.SAVE_MLP])
+
+        dumped = config.model_dump(mode="json")["recompute_cfg"]
+        assert dumped == ["save_attn", "save_mlp"]
+
+        restored = _build_tiny_moe_config(recompute_cfg=dumped)
+        assert restored.recompute_cfg == [RecomputeUnit.SAVE_ATTN, RecomputeUnit.SAVE_MLP]
+
+
+def _recorded_markers(func) -> set[str]:
+    """Marker names a function passes to ``checkpoint_record``."""
+    tree = ast.parse(textwrap.dedent(inspect.getsource(func)))
+    return {
+        node.args[0].value
+        for node in ast.walk(tree)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Name)
+        and node.func.id == "checkpoint_record"
+        and node.args
+        and isinstance(node.args[0], ast.Constant)
+    }
+
+
+def _region_coverage(func) -> dict[str, set[str]]:
+    """Which of the layer's own operations each marker region encloses.
+
+    Regions are keyed by the shared prefix of a ``<name>.begin`` / ``<name>.end`` pair, and an operation is a call on
+    ``self`` -- ``self.experts(...)``, ``self.dispatcher.dispatch(...)``. Calls on tensors are ignored: a region is
+    defined by the sub-modules it covers, not by the reshapes threaded between them.
+    """
+    tree = ast.parse(textwrap.dedent(inspect.getsource(func)))
+    # `ast.walk` is breadth-first; the marker state machine needs source order.
+    nodes = sorted(
+        (node for node in ast.walk(tree) if isinstance(node, ast.Call)),
+        key=lambda node: (node.lineno, node.col_offset),
+    )
+
+    coverage: dict[str, set[str]] = {}
+    active: set[str] = set()
+    for node in nodes:
+        name = _called_name(node)
+        if name == "checkpoint_record" and node.args and isinstance(node.args[0], ast.Constant):
+            region, _, edge = node.args[0].value.rpartition(".")
+            if edge == "begin":
+                active.add(region)
+                coverage.setdefault(region, set())
+            elif edge == "end":
+                active.discard(region)
+        elif name is not None and name.startswith("self."):
+            for region in active:
+                coverage[region].add(name.removeprefix("self."))
+    return coverage
+
+
+def _called_name(node: ast.Call) -> str | None:
+    """Dotted name of a call target, e.g. ``self.dispatcher.dispatch``."""
+    parts: list[str] = []
+    target: ast.expr = node.func
+    while isinstance(target, ast.Attribute):
+        parts.append(target.attr)
+        target = target.value
+    if not isinstance(target, ast.Name):
+        return None
+    parts.append(target.id)
+    return ".".join(reversed(parts))
+
+
+class TestMarkerVocabulary:
+    @pytest.mark.parametrize(
+        "interval_map", [MOE_RECOMPUTE_CFG, DENSE_RECOMPUTE_CFG], ids=["moe", "dense"]
+    )
+    def test_declared_intervals_have_markers(self, interval_map):
+        # A renamed marker would not fail anywhere at runtime: the interval would simply never open
+        # and the region would stay recomputed, silently costing the memory the user asked to keep.
+        recorded: set[str] = set()
+        for layer_module in (moe_decoder_layer, dense_decoder_layer):
+            for _, member in inspect.getmembers(layer_module, inspect.isclass):
+                if member.__module__ != layer_module.__name__:
+                    continue
+                for _, method in inspect.getmembers(member, inspect.isfunction):
+                    recorded |= _recorded_markers(method)
+
+        declared = {name for intervals in interval_map.values() for interval in intervals for name in interval}
+        assert declared <= recorded, f"markers referenced but never recorded: {sorted(declared - recorded)}"
+
+    def test_micro_batch_path_covers_the_same_operations(self):
+        # `_micro_batch_forward` re-implements the dispatch/combine chain across four stage loops.
+        # Comparing marker names alone would pass even if the domino path wrapped entirely different
+        # operations, so compare what each region actually encloses.
+        single = _region_coverage(moe_decoder_layer.MoEDecoderLayer._forward)
+        micro_batch = _region_coverage(moe_decoder_layer.MoEDecoderLayer._micro_batch_forward)
+
+        assert single == micro_batch

tests/model/test_selective_checkpointing.py

@@ -5,6 +5,8 @@
     test_unbalanced_interval_is_safe: end marker 不执行时只多留驻，不影响梯度。
     test_overlapping_intervals_keep_the_union: 区间重叠时按并集留驻。
     test_marker_outside_session_is_noop: 不在会话内时埋点不做任何事。
+TestContractLayering
+    test_module_layer_imports_the_contract_without_the_model_layer: 契约模块必须能被 module/ 层单独导入。
 TestUnsupportedRegions
     test_a_second_model_still_gets_its_own_diagnosis: 诊断去重不跨模型，第二个模型仍会告警。
     test_in_place_op_in_kept_region_is_rejected: 留驻区间内的 in-place 写会明确报错而不是静默改梯度。
@@ -14,6 +16,8 @@
 """
 
 import os
+import subprocess
+import sys
 
 import pytest
 import torch
@@ -25,10 +29,10 @@
 from xtuner.v1.config import FSDPConfig
 from xtuner.v1.loss.ce_loss import CELossConfig
 from xtuner.v1.model.moe.moe import MoE, MoEConfig, SequenceContext
-from xtuner.v1.model.utils import apply_selective_checkpointing, checkpoint_record
+from xtuner.v1.model.utils import MarkerInterval, apply_selective_checkpointing, checkpoint_record
 from xtuner.v1.module.attention import MHAConfig
 from xtuner.v1.module.router import NoAuxRouterConfig
-from xtuner.v1.model.utils import selective_checkpointing as contract
+from xtuner.v1.utils import selective_checkpointing as contract
 
 
 class _MarkedBlock(nn.Module):
@@ -52,22 +56,6 @@ class _OtherMarkedBlock(_MarkedBlock):
     """A second layer class, standing in for another model living in the same process."""
 
 
-class _InPlaceBlock(nn.Module):
-    """A region whose body accumulates in place, which selective checkpointing cannot keep."""
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.linear = nn.Linear(4, 4)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        checkpoint_record("region.start")
-        hidden = self.linear(x)
-        accumulator = torch.zeros_like(hidden)
-        accumulator.add_(hidden)
-        checkpoint_record("region.end")
-        return accumulator
-
-
 class _EmptyRegionBlock(nn.Module):
     """A layer whose declared region encloses no op the policy can keep.
 
@@ -163,6 +151,20 @@ def test_marker_outside_session_is_noop(self):
         assert inputs.grad is not None
 
 
+class TestContractLayering:
+    def test_module_layer_imports_the_contract_without_the_model_layer(self):
+        # 契约（含 marker session）之所以在 xtuner/v1/utils 而不是挨着 engine，就是因为
+        # `checkpoint_record` 的调用点在 xtuner/v1/module 的 forward 里：一旦契约里出现
+        # 指向 model/ 或 module/ 的 import，这条独立导入就会变成循环导入而失败。
+        # 必须用干净的解释器：同进程里 xtuner.v1.model 早就被导入了，测不出这个性质。
+        result = subprocess.run(
+            [sys.executable, "-c", "import xtuner.v1.module.decoder_layer.moe_decoder_layer"],
+            capture_output=True,
+            text=True,
+        )
+        assert result.returncode == 0, result.stderr
+
+
 class TestUnsupportedRegions:
     def test_in_place_op_in_kept_region_is_rejected(self):
         # 留驻的张量会原样交给重算那趟，read-modify-write 会在重算时二次累加：loss 有限、
@@ -263,8 +265,8 @@ class _RegionMoE(MoE):
     _REGION = ("moe.dispatch", "moe.combine.end")
 
     @property
-    def recompute_intervals(self) -> tuple[tuple[str, str], ...]:
-        return (self._REGION,)
+    def recompute_intervals(self) -> list[MarkerInterval]:
+        return [self._REGION]
 
     def fully_shard(self, *args, **kwargs):
         for layer in self.layers.values():