From 3ec986b74544640e5813042373a7dabf91c46f44 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Thu, 30 Jul 2026 03:58:23 +0000 Subject: [PATCH 1/2] [Feature] Keep DSA top-k out of checkpoint replay --- examples/v1/config/sft_glm5p2.py | 3 + tests/engine/test_glm52_moe_train_engine.py | 8 ++- tests/model/test_glm52_moe.py | 24 ++++++++ .../model/test_glm52_mtp_checkpoint_repro.py | 41 ++++++++++--- tests/model/test_recompute.py | 16 ++++-- tests/module/attention/test_dsa_mla.py | 57 ++++++++++++++++--- xtuner/v1/model/moe/glm52/dsa_mla.py | 27 ++++++--- xtuner/v1/model/moe/glm52/glm52.py | 23 +++++++- xtuner/v1/model/moe/moe.py | 10 +++- xtuner/v1/utils/selective_checkpointing.py | 3 + 10 files changed, 175 insertions(+), 37 deletions(-) diff --git a/examples/v1/config/sft_glm5p2.py b/examples/v1/config/sft_glm5p2.py index cab158ec2..1556e3fff 100644 --- a/examples/v1/config/sft_glm5p2.py +++ b/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") diff --git a/tests/engine/test_glm52_moe_train_engine.py b/tests/engine/test_glm52_moe_train_engine.py index efb94f710..a0d95ed4f 100644 --- a/tests/engine/test_glm52_moe_train_engine.py +++ b/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), diff --git a/tests/model/test_glm52_moe.py b/tests/model/test_glm52_moe.py index b9b82800e..40f32cc39 100644 --- a/tests/model/test_glm52_moe.py +++ b/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): diff --git a/tests/model/test_glm52_mtp_checkpoint_repro.py b/tests/model/test_glm52_mtp_checkpoint_repro.py index fba6d5127..d6a12bc35 100644 --- a/tests/model/test_glm52_mtp_checkpoint_repro.py +++ b/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"]) diff --git a/tests/model/test_recompute.py b/tests/model/test_recompute.py index 07160861a..15e0f4421 100644 --- a/tests/model/test_recompute.py +++ b/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 diff --git a/tests/module/attention/test_dsa_mla.py b/tests/module/attention/test_dsa_mla.py index 02dc77c04..86626ef63 100644 --- a/tests/module/attention/test_dsa_mla.py +++ b/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( diff --git a/xtuner/v1/model/moe/glm52/dsa_mla.py b/xtuner/v1/model/moe/glm52/dsa_mla.py index ebb57f5d1..dafe5c091 100644 --- a/xtuner/v1/model/moe/glm52/dsa_mla.py +++ b/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): diff --git a/xtuner/v1/model/moe/glm52/glm52.py b/xtuner/v1/model/moe/glm52/glm52.py index c2c3b034a..9838e4aa1 100644 --- a/xtuner/v1/model/moe/glm52/glm52.py +++ b/xtuner/v1/model/moe/glm52/glm52.py @@ -7,6 +7,7 @@ from pydantic import Field, computed_field from typing_extensions import Self, override + try: from transformers.models.glm_moe_dsa import GlmMoeDsaConfig as HFGlmMoeDsaConfig except ImportError: @@ -14,7 +15,14 @@ from xtuner.v1.data_proto import SequenceContext from xtuner.v1.model.base import DEFAULT_FLOAT8_CFG, TorchCompileOption -from xtuner.v1.model.moe.moe import BalancingLossConfig, MoE, MoEConfig, ZLossConfig +from xtuner.v1.model.moe.moe import ( + MOE_RECOMPUTE_CFG, + BalancingLossConfig, + MoE, + MoEConfig, + ZLossConfig, +) +from xtuner.v1.model.utils import KeptCallables, RecomputeTargetMap, SaveUnit from xtuner.v1.module.decoder_layer.dense_decoder_layer import ( DenseDecoderLayerMicroBatchOutput, DenseDecoderLayerOutput, @@ -62,6 +70,14 @@ MOE_EP_COMPILE_CFG = MOE_NON_EP_COMPILE_CFG.copy() MOE_EP_COMPILE_CFG.pop("xtuner.v1.model.moe.glm52.decoder_layer.GLM52MoEDecoderLayer.forward") +# GLM uses sparse MLA rather than the generic flash-attention op named by +# ``SaveUnit.ATTN``. Keep the MoE units and expose the narrow no-grad top-k +# selection callable as the GLM-specific attention saving unit. +GLM52_RECOMPUTE_CFG: RecomputeTargetMap = { + unit: target for unit, target in MOE_RECOMPUTE_CFG.items() if unit is not SaveUnit.ATTN +} +GLM52_RECOMPUTE_CFG[SaveUnit.DSA_INDEXER] = KeptCallables("xtuner.v1.model.moe.glm52.dsa_mla.DSAIndexer._select_topk") + class Glm52MoE(MoE): dense_decoder_layer_cls = GLM52DenseDecoderLayer @@ -76,6 +92,11 @@ def default_compile_cfg(self) -> dict[str, TorchCompileOption]: return MOE_EP_COMPILE_CFG return MOE_NON_EP_COMPILE_CFG + @property + @override + def default_recompute_cfg(self) -> RecomputeTargetMap: + return GLM52_RECOMPUTE_CFG + @override def _configure_model_specific_layers(self) -> None: dsa_layers: list[DSAMultiLatentAttention] = [] diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index e2b39136c..6aae3d95c 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -1380,9 +1380,13 @@ def fully_shard( if self._should_recompute(None, mtp_idx=mtp_idx) or ( self.config.mtp_config is not None and self.config.mtp_config.share_weights ): # share mtp head must recompute - # Recompute units are declared against the decoder layer, so an MTP layer gets - # the same mechanism with nothing kept: recomputed whole, as before. - mtp_layer = apply_selective_checkpointing(mtp_layer) + # MTP contains the same decoder units as the main stack. In particular, GLM's + # no-grad DSA IDs must be kept here too or shared-depth replay reruns the indexer. + mtp_layer = apply_selective_checkpointing( + mtp_layer, + self.kept_ops, + keeps_any_unit=self.keeps_any_recompute_unit, + ) self.mtp_block.layers[mtp_idx] = mtp_layer reshard_after_forward = mtp_idx != len(self.mtp_block.layers) - 1 diff --git a/xtuner/v1/utils/selective_checkpointing.py b/xtuner/v1/utils/selective_checkpointing.py index 50d7a41e4..ef75dc065 100644 --- a/xtuner/v1/utils/selective_checkpointing.py +++ b/xtuner/v1/utils/selective_checkpointing.py @@ -79,6 +79,9 @@ class SaveUnit(StrEnum): is always visible to the checkpoint policy, compiled or not. """ + DSA_INDEXER = "dsa_indexer" + """Keep the no-grad DSA top-k selection result, without keeping its projections.""" + MOE_GATE = "moe_gate" """Keep the MoE router: gating projection, top-k selection, and routing weights.""" From 4511a97367937fc2447605d4bc7aeb51e2f1ffc2 Mon Sep 17 00:00:00 2001 From: zhaopenghao Date: Fri, 14 Aug 2026 12:44:47 +0000 Subject: [PATCH 2/2] [Fix] Remove stale DSA top-k cache preparation --- xtuner/v1/model/moe/moe.py | 10 ---------- xtuner/v1/model/utils/checkpointing.py | 3 +-- 2 files changed, 1 insertion(+), 12 deletions(-) diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 6aae3d95c..cc2b4ca5c 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -560,14 +560,6 @@ def post_micro_batch_forward(self, batch_outputs: Sequence[MoEModelOutputs]) -> moe_info = cast(MoEBatchForwardInfo, base_info) return moe_info - @staticmethod - def _prepare_seq_ctx_topk_cache(seq_ctx_list: Sequence[SequenceContext]) -> None: - # A slot owns both runtime offload state and pinned storage. TrainEngine - # finishes backward before the next accumulation group, so only contexts - # in one model call can be live concurrently and need distinct slots. - for offload_slot, seq_ctx in enumerate(seq_ctx_list): - seq_ctx.dsa_topk_cache.offload_slot = offload_slot - def _micro_batch_forward( self, seq_ctx_list: list[SequenceContext], @@ -579,7 +571,6 @@ def _micro_batch_forward( This method processes multiple micro-batches in parallel, similar to how MoEDecoderLayer handles micro-batching at the layer level. """ - self._prepare_seq_ctx_topk_cache(seq_ctx_list) if self.config.return_hidden_states: raise NotImplementedError @@ -854,7 +845,6 @@ def _forward( loss_ctx: MoELossContextDict | None, return_router_logits: bool = False, ) -> MoEModelOutputs: - self._prepare_seq_ctx_topk_cache([seq_ctx]) input_ids = seq_ctx.input_ids position_ids = seq_ctx.position_ids diff --git a/xtuner/v1/model/utils/checkpointing.py b/xtuner/v1/model/utils/checkpointing.py index 81b8a092c..1a4113034 100644 --- a/xtuner/v1/model/utils/checkpointing.py +++ b/xtuner/v1/model/utils/checkpointing.py @@ -129,8 +129,7 @@ class CheckpointModule: It overrides ``__call__`` rather than ``forward`` on purpose. Hooks registered on the module run inside ``nn.Module.__call__``, so replacing ``forward`` would leave them *outside* the - checkpointed region -- and hooks that maintain per-forward state, such as the DSA top-k cache - lifecycle, must see the recompute pass too. + checkpointed region and skip them during recomputation. """ _checkpointed_call: Callable[..., Any]