[Feature] Keep DSA top-k out of checkpoint replay

examples/v1/config/sft_glm5p2.py

@@ -8,6 +8,7 @@
 from xtuner.v1.model import get_model_config_from_hf
 from xtuner.v1.train import TrainerConfig
 from xtuner.v1.train.trainer import LoadCheckpointConfig
+from xtuner.v1.utils import RecomputeConfig, SaveUnit
 
 
 def _get_bool_env(name: str, default: bool = False) -> bool:
@@ -52,6 +53,8 @@ def _get_float8_config() -> Float8Config | None:
 model_cfg.compile_cfg = _get_bool_env("MODEL_COMPILE", False)
 model_cfg.float8_cfg = _get_float8_config()
 model_cfg.lm_loss_cfg = loss_cfg
+if recompute_units := os.environ.get("RECOMPUTE_CFG"):
+    model_cfg.recompute_cfg = RecomputeConfig(save=[SaveUnit(unit.strip()) for unit in recompute_units.split(",")])
 if hasattr(model_cfg.attention, "sparse_mla_backend"):
     model_cfg.attention.sparse_mla_backend = os.environ.get("SPARSE_MLA_BACKEND", "tilelang")
 

tests/engine/test_glm52_moe_train_engine.py

@@ -1,7 +1,7 @@
 """GLM-5.2 TrainEngine 的训练、优化组合与 DCP 持久化行为测试。
 
 TestGlm52OptimizedEngine
-    test_sp2_ep4_micro2_compile_offload_train_step: SP2、EP4、micro2、compile 与双 offload 可联合训练。
+    test_sp2_ep4_micro2_compile_offload_train_step: selective checkpoint 与生产优化组合可联合训练。
 TestGlm52ParallelHFCheckpoint
     test_fsdp2_ep4_mtp_hf_round_trip_preserves_weights: FSDP2、EP4、MTP 权重可经 HF 无损往返。
 TestGlm52PretrainedEngine
@@ -39,7 +39,7 @@
 from xtuner.v1.model.moe.glm52 import DSAMLAConfig, Glm52MoEConfig
 from xtuner.v1.module.mtp import MTPConfig
 from xtuner.v1.module.router.noaux_router import NoAuxRouter, NoAuxRouterConfig
-from xtuner.v1.utils import pad_to_max_length
+from xtuner.v1.utils import RecomputeConfig, SaveUnit, pad_to_max_length
 from xtuner.v1.utils.device import get_device
 from xtuner.v1.utils.test_utils import init_data_mesh
 
@@ -190,9 +190,11 @@ def _run_loss_curve(
 @unittest.skipUnless(torch.cuda.device_count() >= 8, "requires 8 CUDA devices")
 class TestGlm52OptimizedEngine(DeterministicDDPTestCase):
     def test_sp2_ep4_micro2_compile_offload_train_step(self):
-        # 验证生产优化组合经两次梯度累积后 loss、梯度与优化器状态均有效。
+        # 验证 DSA selective checkpoint 与 SP2、EP4、micro2、compile、双 offload
+        # 联合执行后，loss、梯度与优化器状态均有效。
         self.create_pg("cuda")
         model_cfg = _tiny_ep4_mtp_config()
+        model_cfg.recompute_cfg = RecomputeConfig(save=[SaveUnit.DSA_INDEXER])
         engine = TrainEngine(
             model_cfg=model_cfg,
             optim_cfg=AdamWConfig(lr=1e-3, foreach=False),

tests/model/test_glm52_moe.py

@@ -4,6 +4,7 @@
     test_save_hf_matches_transformers_and_engine_contracts: HF round-trip 与推理引擎字段契约一致。
     test_from_hf_preserves_glm_specific_behavior: HF 配置转换保留 DSA、router 与 MTP 语义。
     test_rejects_shared_physical_mtp_indexer: 非法的 physical MTP indexer 计划会被拒绝。
+    test_recompute_cfg_exposes_only_narrow_dsa_indexer_callable: GLM 只声明窄 DSA indexer callable。
 TestGlm52CheckpointConversion
     test_tiny_model_round_trips_through_hf: tiny 主干与 MTP 参数可经公共 HF API 无损往返。
 TestGlm52RouterBias
@@ -32,8 +33,10 @@
 from xtuner.v1.loss.ce_loss import CELossConfig
 from xtuner.v1.model import Glm52MoEConfig, get_model_config, get_model_config_from_hf
 from xtuner.v1.model.moe.glm52 import DSAMLAConfig
+from xtuner.v1.model.utils import KeptCallables
 from xtuner.v1.module.mtp import MTPConfig
 from xtuner.v1.module.router.noaux_router import NoAuxRouterConfig
+from xtuner.v1.utils import RecomputeConfig, SaveUnit
 from xtuner.v1.utils.test_utils import init_data_mesh
 
 
@@ -184,6 +187,27 @@ def test_rejects_shared_physical_mtp_indexer(self):
         with pytest.raises(ValueError, match="physical MTP indexer_types"):
             config.build()
 
+    def test_recompute_cfg_exposes_only_narrow_dsa_indexer_callable(self):
+        # GLM 不走通用 flash-attention op；`True` 应改为保留窄 indexer callable，
+        # 而不是声明一个实际不会命中的 ATTN unit。
+        config = _tiny_glm52_config()
+        config.mtp_config = MTPConfig(num_layers=1, share_weights=True)
+        config.recompute_cfg = RecomputeConfig(save=True)
+
+        with torch.device("meta"):
+            model = config.build()
+
+        expected = KeptCallables("xtuner.v1.model.moe.glm52.dsa_mla.DSAIndexer._select_topk")
+        assert model.default_recompute_cfg[SaveUnit.DSA_INDEXER] == expected
+        assert SaveUnit.ATTN not in model.default_recompute_cfg
+        assert model.keeps_any_recompute_unit
+
+        default_config = _tiny_glm52_config()
+        default_config.mtp_config = MTPConfig(num_layers=1, share_weights=True)
+        with torch.device("meta"):
+            default_model = default_config.build()
+        assert not default_model.keeps_any_recompute_unit
+
 
 @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
 class TestGlm52CheckpointConversion(DeterministicDDPTestCase):

tests/model/test_glm52_mtp_checkpoint_repro.py

@@ -1,7 +1,7 @@
 """GLM-5.2 MTP checkpoint 的真实训练回归测试。
 
 TestGlm52CompiledMTPCheckpoint
-    test_shared_mtp_depths_train_with_compile_and_topk_offload: 共享 MTP 深度可在 compile/offload 下训练。
+    test_shared_mtp_depths_train_with_selective_checkpoint_fp8_compile: selective checkpoint 与 FP8/MTP 兼容。
 TestGlm52MicroBatchMTPCheckpoint
     test_nested_micro_batch_inputs_preserve_gradients: EP2 micro2 的嵌套 embedding 梯度可正确反传。
 """
@@ -17,17 +17,19 @@
 from xtuner.v1.config import AdamWConfig, FSDPConfig
 from xtuner.v1.data_proto import SequenceContext
 from xtuner.v1.engine.train_engine import TrainEngine
+from xtuner.v1.float8.config import Float8Config, ScalingGranularity
 from xtuner.v1.loss.ce_loss import CELossConfig
 from xtuner.v1.model.base import ModelItem
 from xtuner.v1.model.moe.glm52 import DSAMLAConfig, Glm52MoEConfig
 from xtuner.v1.module.mtp import MTPConfig
 from xtuner.v1.module.router.noaux_router import NoAuxRouterConfig
+from xtuner.v1.utils import RecomputeConfig, SaveUnit
 
 
 def _tiny_mtp_config(ep_size: int, mtp_num_layers: int, compile_model: bool) -> Glm52MoEConfig:
     return Glm52MoEConfig(
         vocab_size=32,
-        max_position_embeddings=64,
+        max_position_embeddings=128,
         pad_token_id=0,
         eos_token_id=1,
         hf_eos_token_id=[1],
@@ -76,14 +78,32 @@ def _build_engine(
     ep_size: int,
     mtp_num_layers: int,
     compile_model: bool,
+    selective_indexer: bool = False,
+    float8: bool = False,
 ) -> TrainEngine:
+    model_cfg = _tiny_mtp_config(ep_size, mtp_num_layers, compile_model)
+    if selective_indexer:
+        model_cfg.recompute_cfg = RecomputeConfig(save=[SaveUnit.DSA_INDEXER])
+    if float8:
+        # Tile-wise FP8 requires every GEMM input dimension to be 128-aligned.
+        model_cfg.attention.q_lora_rank = 128
+        model_cfg.attention.kv_lora_rank = 128
+        model_cfg.attention.head_dim = 64
+        model_cfg.attention.qk_nope_head_dim = 64
+        model_cfg.attention.qk_rope_head_dim = 64
+        model_cfg.attention.v_head_dim = 64
+        model_cfg.attention.index_head_dim = 128
+        model_cfg.float8_cfg = Float8Config(
+            scaling_granularity_gemm=ScalingGranularity.TILEWISE,
+            scaling_granularity_grouped_gemm=ScalingGranularity.TILEWISE,
+        )
     engine = TrainEngine(
-        model_cfg=_tiny_mtp_config(ep_size, mtp_num_layers, compile_model),
+        model_cfg=model_cfg,
         optim_cfg=AdamWConfig(lr=1e-3, foreach=False),
         fsdp_cfg=FSDPConfig(
             ep_size=ep_size,
             cpu_offload=False,
-            recompute_ratio=0.0,
+            recompute_ratio=1.0 if selective_indexer else 0.0,
             torch_compile=compile_model,
         ),
         intra_layer_micro_batch=intra_layer_micro_batch,
@@ -92,8 +112,8 @@ def _build_engine(
     return engine
 
 
-def _model_item(engine: TrainEngine, start: int) -> ModelItem:
-    input_ids = torch.arange(start, start + 6).view(1, -1) % engine.model_cfg.vocab_size
+def _model_item(engine: TrainEngine, start: int, num_tokens: int = 5) -> ModelItem:
+    input_ids = torch.arange(start, start + num_tokens + 1).view(1, -1) % engine.model_cfg.vocab_size
     seq_ctx = SequenceContext.from_input_ids((input_ids[:, :-1],), device="cuda")
     data = {"seq_ctx": seq_ctx, "shifted_labels": input_ids[:, 1:]}
     loss_ctx = engine.model.build_loss_ctx_batch([data], sp_mesh=None)[0]
@@ -102,21 +122,24 @@ def _model_item(engine: TrainEngine, start: int) -> ModelItem:
 
 @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
 class TestGlm52CompiledMTPCheckpoint(DeterministicDDPTestCase):
-    def test_shared_mtp_depths_train_with_compile_and_topk_offload(self):
-        # 验证共享 MTP 深度可在 compile/offload 下训练且 loss 有限。
+    def test_shared_mtp_depths_train_with_selective_checkpoint_fp8_compile(self):
+        # 复现真实 SFT 的 main selective checkpoint + MTP checkpoint + FP8/compile
+        # 组合，并验证共享 MTP 深度可完成训练。
         self.create_pg("cuda")
         engine = _build_engine(
             intra_layer_micro_batch=1,
             ep_size=1,
             mtp_num_layers=2,
             compile_model=True,
+            selective_indexer=True,
+            float8=True,
         )
         try:
             with mock.patch.dict(
                 os.environ,
                 {"XTUNER_ACTIVATION_OFFLOAD": "0", "XTUNER_DSA_TOPK_OFFLOAD": "1"},
             ):
-                step_info = engine.train_step([_model_item(engine, 2)])
+                step_info = engine.train_step([_model_item(engine, 2, num_tokens=128)])
 
             assert math.isfinite(step_info["total_loss"])
             assert math.isfinite(step_info["logs_info"]["reduced_mtp_loss"])

tests/model/test_recompute.py

@@ -29,7 +29,6 @@
     test_input_tensors_reach_the_ambient_saved_tensor_hooks: 嵌套/关键字传入的输入也能进外层 hook。
 """
 
-
 import pydoc
 
 import pytest
@@ -41,6 +40,7 @@
 from xtuner.v1.model.base import BaseModel, TorchCompileOption, 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.glm52.glm52 import GLM52_RECOMPUTE_CFG
 from xtuner.v1.model.moe.moe import MOE_RECOMPUTE_CFG, MoE, MoEConfig
 from xtuner.v1.model.utils import (
     KeptCallables,
@@ -265,7 +265,9 @@ def test_unsupported_unit_is_rejected(self):
         # 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=RecomputeConfig(save=[SaveUnit.ATTN])))
+            _ProbeModel(
+                _ProbeConfig(text_config=_NestedProbeConfig(), recompute_cfg=RecomputeConfig(save=[SaveUnit.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
@@ -303,7 +305,11 @@ def test_units_round_trip_through_json(self):
 
 
 class TestDeclaredTargets:
-    @pytest.mark.parametrize("target_map", [MOE_RECOMPUTE_CFG, DENSE_RECOMPUTE_CFG], ids=["moe", "dense"])
+    @pytest.mark.parametrize(
+        "target_map",
+        [MOE_RECOMPUTE_CFG, DENSE_RECOMPUTE_CFG, GLM52_RECOMPUTE_CFG],
+        ids=["moe", "dense", "glm52"],
+    )
     def test_declared_targets_resolve(self, target_map):
         # 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
@@ -362,7 +368,9 @@ def test_a_callable_unit_keeps_its_callers_compiled(self):
 
     def test_no_unit_withdraws_the_method_that_holds_most_compilation(self):
         for unit in MOE_RECOMPUTE_CFG:
-            assert _PRE_MOE_FORWARD in self._compile_cfg(recompute_cfg=RecomputeConfig(save=[unit])), f"{unit} withdrew {_PRE_MOE_FORWARD}"
+            assert _PRE_MOE_FORWARD in self._compile_cfg(recompute_cfg=RecomputeConfig(save=[unit])), (
+                f"{unit} withdrew {_PRE_MOE_FORWARD}"
+            )
 
     def test_attention_is_kept_by_op_identity(self):
         # The attention kernel is a custom op, so it reaches the policy from inside a compiled

tests/module/attention/test_dsa_mla.py

@@ -6,6 +6,7 @@
     test_packed_inputs_respect_causal_boundaries_and_backward: packed attention 遵守分段因果边界并可反传。
     test_shared_layer_consumes_explicit_topk_ids: shared layer 复用显式 top-k IDs，漏传时立即报错。
     test_selective_checkpoint_preserves_explicit_topk_storage: 显式 IDs 穿过 non-reentrant selective checkpoint。
+    test_selective_checkpoint_keeps_indexer_out_of_replay: selective checkpoint 不在反向中重跑 indexer。
 TestAcceleratedSparseMLA
     test_tilelang_forward_backward_matches_torch: TileLang 前反向数值与 PyTorch 后端一致。
     test_compiled_cudnn_backward_matches_tilelang: 编译后的 cuDNN DSA 前反向与 TileLang 一致。
@@ -23,13 +24,15 @@
 import pytest
 import torch
 import torch.distributed as dist
+from torch.utils._python_dispatch import TorchDispatchMode
 
 from xtuner._testing import DeterministicDDPTestCase
 from xtuner.v1.data_proto import SequenceContext
 from xtuner.v1.model.moe.glm52 import DSAMLAConfig
 from xtuner.v1.model.moe.glm52.decoder_layer import GLM52DenseDecoderLayer
-from xtuner.v1.model.utils import apply_selective_checkpointing
+from xtuner.v1.model.utils import apply_selective_checkpointing, in_recompute_unit
 from xtuner.v1.ops.sparse_mla import dsa_topk_indices, sparse_mla
+from xtuner.v1.utils import SaveUnit
 from xtuner.v1.utils.test_utils import init_data_mesh
 
 
@@ -41,6 +44,19 @@
 CUDNN_DQ_RTOL = 5e-2
 
 
+class _TopKExecutionCounter(TorchDispatchMode):
+    """Count real top-k executions below the selective-checkpoint cache mode."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.count = 0
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        if func is torch.ops.aten.topk.default:
+            self.count += 1
+        return func(*args, **(kwargs or {}))
+
+
 @cache
 def _tilelang_sparse_mla_available() -> bool:
     if not torch.cuda.is_available():
@@ -213,14 +229,8 @@ def test_shared_layer_consumes_explicit_topk_ids(self):
     def test_selective_checkpoint_preserves_explicit_topk_storage(self):
         # 验证显式 IDs 可穿过 non-reentrant selective checkpoint 并完成反向。
         torch.manual_seed(0)
-        source_block = apply_selective_checkpointing(
-            _tiny_dsa_decoder(["full", "shared"], layer_idx=0),
-            [("mlp.begin", "mlp.end")],
-        )
-        shared_block = apply_selective_checkpointing(
-            _tiny_dsa_decoder(["full", "shared"], layer_idx=1),
-            [("mlp.begin", "mlp.end")],
-        )
+        source_block = apply_selective_checkpointing(_tiny_dsa_decoder(["full", "shared"], layer_idx=0))
+        shared_block = apply_selective_checkpointing(_tiny_dsa_decoder(["full", "shared"], layer_idx=1))
         hidden_states = torch.randn(1, 4, 4, requires_grad=True)
         position_embeddings = (torch.ones(1, 4, 2), torch.zeros(1, 4, 2))
         seq_ctx = SequenceContext.from_input_ids((torch.tensor([[1, 2, 3, 4]]),), device="cpu")
@@ -244,6 +254,35 @@ def test_selective_checkpoint_preserves_explicit_topk_storage(self):
         assert torch.isfinite(hidden_states.grad).all()
         assert source_ids.dtype == torch.int32
 
+    @pytest.mark.parametrize("compile_layer", [False, True], ids=["eager", "compile"])
+    def test_selective_checkpoint_keeps_indexer_out_of_replay(self, compile_layer):
+        # Counter 位于 SAC cache mode 之下：命中 cache 的 top-k 不会到达这里，
+        # 因而只统计 backward 中真正再次执行的 indexer。
+        torch.manual_seed(0)
+        decoder = _tiny_dsa_decoder(["full"], layer_idx=0)
+        indexer = decoder.self_attn.indexer
+        indexer._select_topk = torch._dynamo.disable(  # type: ignore[method-assign]
+            in_recompute_unit(SaveUnit.DSA_INDEXER, indexer._select_topk)
+        )
+        if compile_layer:
+            decoder = torch.compile(decoder, backend="eager", fullgraph=False)
+        source_block = apply_selective_checkpointing(decoder, keeps_any_unit=True)
+        hidden_states = torch.randn(1, 4, 4, requires_grad=True)
+        position_embeddings = (torch.ones(1, 4, 2), torch.zeros(1, 4, 2))
+        seq_ctx = SequenceContext.from_input_ids((torch.tensor([[1, 2, 3, 4]]),), device="cpu")
+
+        outputs = source_block(
+            hidden_states,
+            position_embeddings=position_embeddings,
+            seq_ctx=seq_ctx,
+        )
+        counter = _TopKExecutionCounter()
+        with counter:
+            outputs["hidden_states"].square().mean().backward()
+
+        assert counter.count == 0
+        assert torch.isfinite(hidden_states.grad).all()
+
 
 class TestAcceleratedSparseMLA:
     @pytest.mark.skipif(

xtuner/v1/model/moe/glm52/dsa_mla.py

@@ -94,6 +94,24 @@ def __init__(
         # parameters must not be registered with the training optimizer.
         self.requires_grad_(False)
 
+    def _select_topk(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        weights: torch.Tensor,
+        seq_ctx: SequenceContext,
+    ) -> torch.Tensor:
+        # This narrow callable is the SAC unit boundary: index projections and
+        # the SP gather remain recomputed, while only the discrete IDs are kept.
+        return self.dsa_topk_indices_func(
+            q,
+            k,
+            weights,
+            seq_ctx,
+            index_head_dim=self.index_head_dim,
+            index_topk=self.index_topk,
+        )
+
     @torch.no_grad()
     def forward(
         self,
@@ -158,14 +176,7 @@ def forward(
         # k: [bsz, S_g, Di]
         k = gather_for_sequence_parallel(k, dim=1, sp_mesh=seq_ctx.sequence_parallel_mesh)
         # returns topk_indices: [S, 1, K]
-        return self.dsa_topk_indices_func(
-            q,
-            k,
-            weights,
-            seq_ctx,
-            index_head_dim=self.index_head_dim,
-            index_topk=self.index_topk,
-        )
+        return self._select_topk(q, k, weights, seq_ctx)
 
 
 class DSAMLAConfig(MLAConfig):