[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。
@@ -32,11 +32,10 @@
 from xtuner.v1.loss.ce_loss import CELossConfig
 from xtuner.v1.model import get_model_config_from_hf
 from xtuner.v1.model.base import ModelItem
-from xtuner.v1.model.moe.glm52 import Glm52MoEConfig
-from xtuner.v1.module.attention import DSAMLAConfig
+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
 
@@ -187,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),
@@ -204,7 +205,6 @@ def test_sp2_ep4_micro2_compile_offload_train_step(self):
         engine.init_model_weights()
         sp_mesh = init_data_mesh(str(DEVICE), sp_size=2)["sp"]
         data_batches = []
-        seq_ctx_list = []
 
         try:
             for micro_batch_idx in range(4):
@@ -214,7 +214,6 @@ def test_sp2_ep4_micro2_compile_offload_train_step(self):
                 data = {"seq_ctx": full_seq_ctx, "shifted_labels": input_ids[:, 1:]}
                 loss_ctx = engine.model.build_loss_ctx_batch([data], sp_mesh=sp_mesh)[0]
                 seq_ctx = full_seq_ctx.split(sp_mesh)
-                seq_ctx_list.append(seq_ctx)
                 data_batches.append(ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx))
 
             with mock.patch.dict(
@@ -230,9 +229,6 @@ def test_sp2_ep4_micro2_compile_offload_train_step(self):
             self.assertTrue(math.isfinite(step_info["logs_info"]["reduced_mtp_loss"]))
             self.assertTrue(math.isfinite(float(grad_norm)))
             self.assertTrue(engine.optimizer.state)
-            for seq_ctx in seq_ctx_list:
-                self.assertEqual(seq_ctx.dsa_topk_cache.indices, {})
-                self.assertEqual(seq_ctx.dsa_topk_cache.offloaded, {})
         finally:
             del engine
             torch.cuda.empty_cache()

tests/model/test_glm52_moe.py

@@ -3,11 +3,14 @@
 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
     test_scratch_init_zeroes_main_and_mtp_biases: 从头初始化清零主干与 MTP router bias。
     test_update_bias_handles_main_and_shared_mtp_loads: bias 更新覆盖主干并聚合共享 MTP 深度。
+TestGlm52ExplicitDsaDataflow
+    test_model_forward_backward_with_explicit_dsa_dataflow: 模型通过显式 IDs 完成前反向。
 TestGlm52SequenceParallel
     test_mtp_loss_and_gradients_match_full_sequence: SP2 的 MTP loss 与梯度匹配完整序列。
 """
@@ -27,9 +30,10 @@
 from xtuner.v1.data_proto import SequenceContext
 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.module.attention import DSAMLAConfig
+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
 
 
@@ -131,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):
@@ -222,6 +250,28 @@ def test_update_bias_handles_main_and_shared_mtp_loads(self):
         )
 
 
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+class TestGlm52ExplicitDsaDataflow:
+    def test_model_forward_backward_with_explicit_dsa_dataflow(self):
+        # 验证 GLM public forward/backward 经显式 DSA IDs 数据流产生有限 loss 和梯度。
+        config = _tiny_glm52_config()
+        config.mtp_config = None
+        model = config.build().to(device="cuda", dtype=torch.bfloat16)
+        model.init_weights()
+
+        input_ids = torch.tensor([[2, 3, 4, 5]], device="cuda")
+        shifted_labels = torch.tensor([[3, 4, 5, 6]], device="cuda")
+        seq_ctx = SequenceContext.from_input_ids((input_ids,), device="cuda")
+        data = {"seq_ctx": seq_ctx, "shifted_labels": shifted_labels}
+        loss_ctx = model.build_loss_ctx_batch([data], sp_mesh=None)[0]
+
+        output = model(seq_ctx=seq_ctx, loss_ctx=loss_ctx)
+        output["loss"].backward()
+
+        assert torch.isfinite(output["loss"])
+        assert any(parameter.grad is not None for parameter in model.parameters())
+
+
 @unittest.skipUnless(torch.cuda.device_count() >= 2, "requires 2 CUDA devices")
 class TestGlm52SequenceParallel(DeterministicDDPTestCase):
     def test_mtp_loss_and_gradients_match_full_sequence(self):

tests/model/test_glm52_mtp_checkpoint_repro.py

@@ -1,7 +1,7 @@
-"""GLM-5.2 MTP reentrant checkpoint 的真实训练回归测试。
+"""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,18 +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 Glm52MoEConfig
-from xtuner.v1.module.attention import DSAMLAConfig
+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],
@@ -77,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,
@@ -93,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]
@@ -103,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):
-        # 验证默认 reentrant checkpoint 可训练共享 MTP 深度且 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"])