[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 RecomputeUnit
 
 
 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 = [RecomputeUnit(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 与生产优化组合可联合训练。
 TestGlm52PretrainedEngine
     test_ep8_loss_curve_matches_reference: 预训练权重的 EP8 优化轨迹匹配数值基线。
     test_tilewise_fp8_loss_curve_matches_bf16: tilewise FP8 训练轨迹接近 BF16。
@@ -35,7 +35,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 RecomputeUnit, pad_to_max_length
 from xtuner.v1.utils.device import get_device
 from xtuner.v1.utils.test_utils import init_data_mesh
 
@@ -186,9 +186,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_sp_mtp_config()
+        model_cfg.recompute_cfg = [RecomputeUnit.SAVE_DSA_INDEXER]
         engine = TrainEngine(
             model_cfg=model_cfg,
             optim_cfg=AdamWConfig(lr=1e-3, foreach=False),

tests/model/test_glm52_moe.py

@@ -3,6 +3,7 @@
 TestGlm52Config
     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_region: GLM 只声明窄 DSA indexer 区间。
 TestGlm52CheckpointConversion
     test_tiny_model_round_trips_through_hf: tiny 主干与 MTP 参数可经公共 HF API 无损往返。
 TestGlm52RouterBias
@@ -32,6 +33,7 @@
 from xtuner.v1.model.moe.glm52 import DSAMLAConfig
 from xtuner.v1.module.mtp import MTPConfig
 from xtuner.v1.module.router.noaux_router import NoAuxRouterConfig
+from xtuner.v1.utils import RecomputeUnit
 from xtuner.v1.utils.test_utils import init_data_mesh
 
 
@@ -133,6 +135,30 @@ 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_region(self):
+        # GLM 的 attention 含原地 RMSNorm，不能沿用包住整个 attention 的通用区间；
+        # `True` 应保留新的 indexer unit，并排除宽 `save_attn`。
+        config = _tiny_glm52_config()
+        config.mtp_config = MTPConfig(num_layers=1, share_weights=True)
+        config.recompute_cfg = True
+
+        with torch.device("meta"):
+            model = config.build()
+
+        expected = [("dsa.indexer.begin", "dsa.indexer.end")]
+        assert model.default_recompute_cfg[RecomputeUnit.SAVE_DSA_INDEXER] == expected
+        assert RecomputeUnit.SAVE_ATTN not in model.default_recompute_cfg
+        assert expected[0] in model.recompute_intervals
+        assert model.layers["0"].self_attn.indexer.selective_checkpoint_topk
+        assert model.mtp_block.layers[0].decoder_layer.self_attn.indexer.selective_checkpoint_topk
+
+        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.layers["0"].self_attn.indexer.selective_checkpoint_topk
+        assert not default_model.mtp_block.layers[0].decoder_layer.self_attn.indexer.selective_checkpoint_topk
+
 
 @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 RecomputeUnit
 
 
 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 = [RecomputeUnit.SAVE_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

@@ -35,6 +35,8 @@
 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.glm52 import dsa_mla as glm52_dsa_mla
+from xtuner.v1.model.moe.glm52.glm52 import GLM52_RECOMPUTE_CFG
 from xtuner.v1.model.moe.moe import MOE_RECOMPUTE_CFG, MoE, MoEConfig, SequenceContext
 from xtuner.v1.model.utils import (
     RecomputeUnit,
@@ -169,7 +171,9 @@ def _run_once(self, ep_size: int, dispatcher: str, recompute_ratio: float):
         config = _build_moe_config(ep_size, dispatcher)
         with torch.device("meta"):
             model = MoE(config=config)._to_device_dtype(dtype=torch.bfloat16, skip_buffers_dtype=True)
-        model.fully_shard(fsdp_config=FSDPConfig(ep_size=ep_size, recompute_ratio=recompute_ratio, torch_compile=False))
+        model.fully_shard(
+            fsdp_config=FSDPConfig(ep_size=ep_size, recompute_ratio=recompute_ratio, torch_compile=False)
+        )
 
         torch.manual_seed(42)
         torch.cuda.manual_seed_all(42)
@@ -385,13 +389,15 @@ def _called_name(node: ast.Call) -> str | None:
 
 class TestMarkerVocabulary:
     @pytest.mark.parametrize(
-        "interval_map", [MOE_RECOMPUTE_CFG, DENSE_RECOMPUTE_CFG], ids=["moe", "dense"]
+        "interval_map",
+        [MOE_RECOMPUTE_CFG, DENSE_RECOMPUTE_CFG, GLM52_RECOMPUTE_CFG],
+        ids=["moe", "dense", "glm52"],
     )
     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 layer_module in (moe_decoder_layer, dense_decoder_layer, glm52_dsa_mla):
             for _, member in inspect.getmembers(layer_module, inspect.isclass):
                 if member.__module__ != layer_module.__name__:
                     continue

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,6 +24,7 @@
 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
@@ -41,6 +43,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():
@@ -244,6 +259,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)
+        decoder.self_attn.indexer.selective_checkpoint_topk = True
+        if compile_layer:
+            decoder = torch.compile(decoder, backend="eager", fullgraph=False)
+        source_block = apply_selective_checkpointing(
+            decoder,
+            [("dsa.indexer.begin", "dsa.indexer.end")],
+        )
+        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

@@ -22,6 +22,7 @@
     get_dsa_topk_indices,
     get_sparse_mla,
 )
+from xtuner.v1.utils import checkpoint_record
 
 from .dsa_topk_sharing import dsa_topk_source_layer
 
@@ -90,6 +91,29 @@ def __init__(
         self.k_norm = LayerNorm(index_head_dim, eps=1e-6)
         # weights_proj.weight: [index_n_heads, hidden_size]
         self.weights_proj = build_linear(hidden_size, index_n_heads, bias=False)
+        self.selective_checkpoint_topk = False
+
+    @torch.compiler.disable
+    def _select_topk_outside_checkpoint_replay(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        weights: torch.Tensor,
+        seq_ctx: SequenceContext,
+    ) -> torch.Tensor:
+        # Marker state is eager-only. Keep the graph break around the selection
+        # kernel itself so the index projections and SP gather stay compiled.
+        checkpoint_record("dsa.indexer.begin")
+        topk_ids = self.dsa_topk_indices_func(
+            q,
+            k,
+            weights,
+            seq_ctx,
+            index_head_dim=self.index_head_dim,
+            index_topk=self.index_topk,
+        )
+        checkpoint_record("dsa.indexer.end")
+        return topk_ids
 
     @torch.no_grad()
     def forward(
@@ -155,6 +179,8 @@ 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]
+        if self.selective_checkpoint_topk:
+            return self._select_topk_outside_checkpoint_replay(q, k, weights, seq_ctx)
         return self.dsa_topk_indices_func(
             q,
             k,