[Refactor] Unify SAC checkpointing and GLM DSA dataflow

tests/engine/test_glm52_moe_train_engine.py

@@ -36,8 +36,7 @@
 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
@@ -208,7 +207,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):
@@ -218,7 +216,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(
@@ -234,9 +231,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/engine/test_moe_train_engine_float8.py

@@ -20,7 +20,7 @@
 from xtuner.v1.utils.device import get_device
 from xtuner.v1.model.base import ModelItem
 from xtuner.v1.loss.ce_loss import CELossConfig
-from xtuner.v1.model.moe.moe import BalancingLossConfig
+from xtuner.v1.model.moe.moe import MOE_BLOCK_FORWARD, BalancingLossConfig
 
 
 
@@ -35,9 +35,9 @@ class TestMoEEngineFloat8(DeterministicDDPTestCase):
         "device,ep_size,hsdp_sharding_size,sim_tol,rtol",
         [
             ("cuda", 1, int(os.getenv("XTUNER_TEST_WORLD_SIZE", "8")), 0.01, 0.01),
-            # ep8 is a smoke/trend coverage for the FSDP shard-mesh-size-1 FP8 path.
-            # It shares the ep1 reference below, but is not expected to align step-by-step
-            # because EP changes routing/collective order and accumulates FP8 numeric drift.
+            # EP8 covers checkpoint replay across layer-varying routed-token shapes while MoEBlock
+            # remains fullgraph-compiled. It shares the EP1 reference below, but EP changes routing
+            # and collective order, so the two loss curves need not align step-by-step.
             # Observed 10-step loss:
             # [2.4714, 2.4714, 1.8044, 1.5210, 0.9570, 0.6952, 0.4370, 0.3123, 0.1714, 0.1100]
             ("cuda", 8, int(os.getenv("XTUNER_TEST_WORLD_SIZE", "8")), 0.01, 0.15),
@@ -66,6 +66,9 @@ def test_tile_wise_fp8(self, device, ep_size, hsdp_sharding_size, sim_tol, rtol)
             optim_cfg=optim_cfg,
             fsdp_cfg=fsdp_cfg,
         )
+        if ep_size > 1:
+            # Checkpoint replay must remain correct while the EP expert block stays fullgraph.
+            self.assertEqual(engine.model.compile_cfg.get(MOE_BLOCK_FORWARD), {"fullgraph": True})
         engine.from_hf(hf_path=QWEN3_MOE_PATH)
 
         loss_cfg = CELossConfig()

tests/model/test_fsdp_checkpoint.py

@@ -0,0 +1,174 @@
+import torch
+
+from xtuner._testing import DeterministicDDPTestCase
+from xtuner.v1.config import FSDPConfig
+from xtuner.v1.data_proto import SequenceContext
+from xtuner.v1.model.dense.qwen3 import Qwen3DenseConfig
+from xtuner.v1.module.attention import MHAConfig
+from xtuner.v1.utils.compile import is_compiled_function
+
+
+class TestFSDPCheckpoint(DeterministicDDPTestCase):
+    @property
+    def world_size(self) -> int:
+        return 2
+
+    def test_reentrant_checkpoint_keeps_fsdp_outside_recompute(self):
+        self.create_pg("cuda")
+        config = Qwen3DenseConfig(
+            vocab_size=64,
+            max_position_embeddings=64,
+            eos_token_id=2,
+            bos_token_id=1,
+            num_hidden_layers=2,
+            hidden_size=32,
+            intermediate_size=64,
+            rms_norm_eps=1e-6,
+            hidden_act="silu",
+            attention=MHAConfig(
+                num_attention_heads=4,
+                num_key_value_heads=2,
+                head_dim=8,
+                qk_norm=True,
+            ),
+            compile_cfg=False,
+        )
+        model = config.build().cuda()
+        grad_modes: list[bool] = []
+        original_layer = model.layers["0"]
+        original_layer.register_forward_pre_hook(lambda _module, _inputs: grad_modes.append(torch.is_grad_enabled()))
+
+        model.fully_shard(
+            FSDPConfig(
+                param_dtype=torch.bfloat16,
+                reduce_dtype=torch.bfloat16,
+                torch_compile=False,
+            )
+        )
+        checkpoint_calls = 0
+
+        def record_checkpoint_call(_module, _inputs, _output):
+            nonlocal checkpoint_calls
+            checkpoint_calls += 1
+
+        model.layers["0"].register_forward_hook(record_checkpoint_call)
+        input_ids = torch.randint(0, config.vocab_size, (1, 8), device="cuda")
+        output = model(SequenceContext.from_input_ids((input_ids,)))
+        assert output.logits is not None
+        output.logits.sum().backward()
+
+        # The original layer and its lifecycle hooks must be replayed, while
+        # the outer FSDP/checkpoint boundary is one logical forward only.
+        assert grad_modes == [False, True]
+        assert checkpoint_calls == 1
+
+    def test_qwen3_vl_checkpoint_compile_allows_pytree_boundary(self):
+        self.create_pg("cuda")
+        from xtuner.v1.model.compose.qwen3_vl.qwen3_vl_config import Qwen3VLVisionConfig
+
+        compile_target = "xtuner.v1.model.compose.qwen3_vl.modeling_vision.Qwen3VLVisionLayer.forward"
+        config = Qwen3VLVisionConfig(
+            depth=1,
+            hidden_size=32,
+            intermediate_size=64,
+            num_attention_heads=4,
+            patch_size=2,
+            temporal_patch_size=1,
+            spatial_merge_size=1,
+            num_position_embeddings=4,
+            deepstack_visual_indexes=[],
+            attn_impl="eager_attention",
+            compile_cfg={compile_target: {"fullgraph": True}},
+        )
+        model = config.build().cuda()
+        model.fully_shard(FSDPConfig(vision_recompute_ratio=1.0))
+        assert is_compiled_function(model.blocks[0].forward)
+
+        hidden_states = torch.randn(4, 32, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+        cu_seqlens = torch.tensor([0, 4], device="cuda", dtype=torch.int32)
+        cos = torch.ones(4, 8, device="cuda", dtype=torch.bfloat16)
+        sin = torch.zeros_like(cos)
+        output = model.blocks[0](hidden_states, cu_seqlens, 4, (cos, sin))
+        output.square().sum().backward()
+
+        assert hidden_states.grad is not None
+        assert torch.isfinite(hidden_states.grad).all()
+
+    def test_intern_s1_checkpoint_compile_allows_pytree_boundary(self):
+        self.create_pg("cuda")
+        from xtuner.v1.model.compose.intern_s1.intern_s1_config import InternS1VisionConfig
+
+        compile_target = "xtuner.v1.model.compose.intern_s1.modeling_vision.InternS1VisionLayer.forward"
+        config = InternS1VisionConfig(
+            image_size=(4, 4),
+            patch_size=(2, 2),
+            num_hidden_layers=1,
+            hidden_size=32,
+            intermediate_size=64,
+            num_attention_heads=4,
+            attn_impl="eager_attention",
+            compile_cfg={compile_target: {"fullgraph": True}},
+        )
+        model = config.build().cuda()
+        model.fully_shard(FSDPConfig(vision_recompute_ratio=1.0))
+        assert is_compiled_function(model.encoder.layer[0].forward)
+
+        hidden_states = torch.randn(1, 4, 32, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+        output = model.encoder.layer[0](hidden_states)
+        output.square().sum().backward()
+
+        assert hidden_states.grad is not None
+        assert torch.isfinite(hidden_states.grad).all()
+
+    def test_mixed_dense_checkpoint_compile_allows_pytree_boundary(self):
+        self.create_pg("cuda")
+        from xtuner.v1.loss.ce_loss import CELossConfig
+        from xtuner.v1.model.dense.qwen3_5_text import Qwen3_5_VLTextDenseConfig
+        from xtuner.v1.module.attention import GatedDeltaNetConfig
+
+        config = Qwen3_5_VLTextDenseConfig(
+            vocab_size=64,
+            max_position_embeddings=64,
+            eos_token_id=2,
+            num_hidden_layers=4,
+            hidden_size=128,
+            intermediate_size=256,
+            rms_norm_eps=1e-6,
+            hidden_act="silu",
+            attention=MHAConfig(
+                with_gate=True,
+                num_attention_heads=4,
+                num_key_value_heads=2,
+                head_dim=32,
+                qk_norm=True,
+                rms_norm_eps=1e-6,
+                rms_norm_type="zero_centered",
+            ),
+            linear_attention=GatedDeltaNetConfig(
+                num_value_heads=4,
+                num_key_heads=4,
+                key_head_dim=16,
+                value_head_dim=16,
+                conv_kernel_dim=4,
+                hidden_act="silu",
+                rms_norm_eps=1e-6,
+            ),
+        )
+        model = config.build().cuda()
+        model.fully_shard(FSDPConfig(recompute_ratio=1.0, torch_compile=True))
+        assert all(is_compiled_function(layer.forward) for layer in model.layers.values())
+
+        input_ids = torch.randint(0, config.vocab_size, (1, 16), device="cuda")
+        seq_ctx = SequenceContext.from_input_ids((input_ids[:, :-1],))
+        loss_config = CELossConfig(mode="eager")
+        loss_ctx = loss_config.build(
+            data={"shifted_labels": input_ids[:, 1:]},
+            sp_mesh=None,
+        )
+        loss_ctx = loss_config.loss_ctx_cls.build_batches([loss_ctx])[0]
+        output = model(seq_ctx, {"lm": loss_ctx})
+        assert output.loss is not None
+        output.loss.backward()
+
+        assert torch.isfinite(output.loss)
+        assert any(parameter.grad is not None for parameter in model.parameters())

tests/model/test_glm52_moe.py

@@ -9,6 +9,8 @@
 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 与梯度匹配完整序列。
 """
@@ -29,7 +31,7 @@
 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.test_utils import init_data_mesh
@@ -273,6 +275,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):