[Fix] Run GLM-5.2 shared MTP for 7 logical steps

examples/v1/config/sft_glm5p2.py

@@ -1,11 +1,13 @@
 import os
+from typing import Literal, cast
 
 from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig, MuonConfig
 from xtuner.v1.datasets import OpenaiTokenizeFunctionConfig
 from xtuner.v1.datasets.config import DataloaderConfig, DatasetConfig
 from xtuner.v1.float8.config import Float8Config, ScalingGranularity
 from xtuner.v1.loss import CELossConfig
 from xtuner.v1.model import get_model_config_from_hf
+from xtuner.v1.module.mtp import MTPConfig
 from xtuner.v1.train import TrainerConfig
 from xtuner.v1.train.trainer import LoadCheckpointConfig
 
@@ -47,11 +49,33 @@ def _get_float8_config() -> Float8Config | None:
 )
 
 model_cfg = get_model_config_from_hf(GLM5_2_MODEL_PATH)
+# The released checkpoint contains one *physical* MTP layer
+# (``num_nextn_predict_layers=1``).  GLM-5.2 training reuses that layer for
+# seven logical prediction depths; ``MTPConfig.num_layers`` controls the
+# recurrent depth while ``share_weights=True`` keeps construction/checkpoint
+# loading at one physical layer.
+model_cfg.mtp_config = MTPConfig(
+    num_layers=7,
+    share_weights=True,
+    detach_mtp_lm_head_weight=False,
+    detach_mtp_inputs=False,
+    loss_scaling_factor=0.1,
+)
 model_cfg.dispatcher = _get_dispatcher()
 model_cfg.ep_size = ep_size
 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 model_cfg.mtp_config is not None:
+    # GLM-5.2 uses the joint multi-step TV objective with rejection-sampling
+    # verification: https://z.ai/blog/glm-5.2. CE remains the compatibility
+    # default; opt in with MTP_LOSS_TYPE=e2e_tv.
+    mtp_loss_type = os.environ.get("MTP_LOSS_TYPE", "ce")
+    if mtp_loss_type not in ("ce", "e2e_tv"):
+        raise ValueError(f"Unsupported MTP_LOSS_TYPE={mtp_loss_type!r}. Use ce or e2e_tv.")
+    model_cfg.mtp_config.loss_type = cast(Literal["ce", "e2e_tv"], mtp_loss_type)
+    model_cfg.mtp_config.loss_scaling_factor = float(os.environ.get("MTP_LOSS_WEIGHT", "0.1"))
+    model_cfg.mtp_config.tv_loss_chunk_size = int(os.environ.get("MTP_TV_LOSS_CHUNK_SIZE", "128"))
 if hasattr(model_cfg.attention, "sparse_mla_backend"):
     model_cfg.attention.sparse_mla_backend = os.environ.get("SPARSE_MLA_BACKEND", "tilelang")
 

tests/loss/test_mtp_e2e_tv_loss.py

@@ -0,0 +1,74 @@
+from types import SimpleNamespace
+
+import torch
+import torch.nn.functional as F
+
+from xtuner.v1.loss.mtp_loss import MTPE2ETVLossConfig, MTPE2ETVLossContext
+from xtuner.v1.utils.device import get_device
+
+
+DEVICE = torch.device(get_device())
+
+
+def _build_context(*, detach_head: bool = False) -> MTPE2ETVLossContext:
+    # Two packed samples with lengths 5 and 3. For gamma=2, only positions
+    # [0, 1, 2] and [5] have a complete two-token speculative horizon.
+    shifted_labels = torch.arange(8, device=DEVICE).unsqueeze(0)
+    seq_ctx = SimpleNamespace(cu_seq_lens_k=torch.tensor([0, 5, 8], dtype=torch.int32, device=DEVICE))
+    cfg = MTPE2ETVLossConfig(
+        mode="chunk",
+        chunk_size=2,
+        loss_reduction="token",
+        num_steps=2,
+        detach_mtp_lm_head_weight=detach_head,
+    )
+    ctx = cfg.build({"shifted_labels": shifted_labels, "seq_ctx": seq_ctx})
+    assert ctx is not None
+    return MTPE2ETVLossContext.build_batches(
+        [ctx],
+        cu_seq_lens_list=[seq_ctx.cu_seq_lens_k],
+    )[0]
+
+
+def test_e2e_tv_matches_expected_acceptance_length_and_detaches_teacher():
+    torch.manual_seed(0)
+    target = torch.randn(1, 8, 4, device=DEVICE, requires_grad=True)
+    drafts = [torch.randn(1, 8, 4, device=DEVICE, requires_grad=True) for _ in range(2)]
+    head_weight = torch.randn(7, 4, device=DEVICE, requires_grad=True)
+    ctx = _build_context()
+
+    actual, _ = ctx.forward((target, drafts), head_weight)
+
+    valid_positions = torch.tensor([0, 1, 2, 5], device=DEVICE)
+    target_positions = (
+        torch.tensor([1, 2, 3, 6], device=DEVICE),
+        torch.tensor([2, 3, 4, 7], device=DEVICE),
+    )
+    overlaps = []
+    for draft, positions in zip(drafts, target_positions):
+        p = F.softmax(F.linear(target[0, positions].detach(), head_weight.detach()).float(), dim=-1)
+        q = F.softmax(F.linear(draft[0, valid_positions], head_weight).float(), dim=-1)
+        overlaps.append(torch.minimum(p, q).sum(dim=-1))
+    alphas = torch.stack(overlaps, dim=-1)
+    expected = (1.0 - torch.cumprod(alphas, dim=-1).mean(dim=-1)).mean()
+
+    torch.testing.assert_close(actual, expected)
+    actual.backward()
+    assert target.grad is None
+    assert all(draft.grad is not None and torch.count_nonzero(draft.grad) > 0 for draft in drafts)
+    assert head_weight.grad is not None and torch.count_nonzero(head_weight.grad) > 0
+
+
+def test_e2e_tv_can_detach_shared_lm_head():
+    torch.manual_seed(1)
+    target = torch.randn(1, 8, 4, device=DEVICE, requires_grad=True)
+    drafts = [torch.randn(1, 8, 4, device=DEVICE, requires_grad=True) for _ in range(2)]
+    head_weight = torch.randn(7, 4, device=DEVICE, requires_grad=True)
+    ctx = _build_context(detach_head=True)
+
+    loss, _ = ctx.forward((target, drafts), head_weight)
+    loss.backward()
+
+    assert target.grad is None
+    assert head_weight.grad is None
+    assert all(draft.grad is not None and torch.count_nonzero(draft.grad) > 0 for draft in drafts)

tests/model/test_glm52_sft_mtp_config.py

@@ -0,0 +1,108 @@
+"""Regression coverage for the canonical GLM-5.2 SFT recipe.
+
+The released checkpoint advertises one physical MTP layer through
+``num_nextn_predict_layers=1``.  The SFT recipe must keep that checkpoint
+layout while overriding XTuner's logical/recurrent depth to seven.
+"""
+
+import os
+import runpy
+import tempfile
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+from unittest import mock
+
+
+RECIPE_PATH = Path(__file__).resolve().parents[2] / "examples/v1/config/sft_glm5p2.py"
+
+
+class _CaptureTrainerConfig:
+    """Small stand-in that records recipe arguments without initializing a trainer."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        self.__dict__.update(kwargs)
+
+
+def _run_glm52_sft_recipe_regression() -> None:
+    """The SFT recipe should run one physical MTP block for seven depths."""
+    model_cfg = SimpleNamespace(
+        num_nextn_predict_layers=1,
+        mtp_config=SimpleNamespace(num_layers=1, share_weights=True),
+        attention=SimpleNamespace(),
+    )
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_path = Path(tmp_dir)
+        recipe_env = {
+            "GLM5_2_MODEL_PATH": str(tmp_path / "model"),
+            "ALPACA_PATH": str(tmp_path / "alpaca"),
+            "WORK_DIR": str(tmp_path / "work"),
+            "EP_SIZE": "1",
+            "INTRA_LAYER_MICRO_BATCH": "1",
+            "WORLD_SIZE": "1",
+            "GLOBAL_BATCH_SIZE": "1",
+            "SAMPLE_MAX_LENGTH": "4096",
+            "PACK_MAX_LENGTH": "16384",
+            "TOTAL_STEP": "10",
+            "LOSS_MODE": "chunk",
+            "LOSS_CHUNK_SIZE": "1024",
+            "FP8": "0",
+            "MODEL_COMPILE": "0",
+            "DATASET_TYPE": "alpaca",
+            "DATASET_SAMPLE_RATIO": "1.0",
+            "CACHE_TAG": "test",
+            "PACK_LEVEL": "soft",
+            "PACK_CHUNK_SIZE": "10000",
+            "DATALOADER_NUM_WORKERS": "0",
+            "PACK_WORKERS": "0",
+            "GLOBAL_PACK": "1",
+            "GROUP_BY_LENGTH": "1",
+            "LR": "1e-6",
+            "OPTIMIZER": "adamw",
+            "ADAMW_FOREACH": "0",
+            "SWAP_OPTIMIZER": "0",
+            "LR_TYPE": "cosine",
+            "WARMUP_RATIO": "0",
+            "CPU_OFFLOAD": "0",
+            "TORCH_COMPILE": "0",
+            "DISPATCHER": "none",
+            "SPARSE_MLA_BACKEND": "torch",
+            "SP_SIZE": "1",
+            "CHECKPOINT_INTERVAL": "200",
+            "CHECKPOINT_MAX_KEEP": "3",
+            "HF_INTERVAL": "200",
+            "HF_MAX_KEEP": "3",
+            "PROFILE_MEMORY": "0",
+            "PROFILE_TIME": "0",
+            "PROFILE_STEP": "2,3",
+            "DEBUG_SKIP_SAVE": "0",
+        }
+        # Deliberately omit these so they exercise the recipe's true defaults.
+        with mock.patch.dict(os.environ, recipe_env, clear=False):
+            os.environ.pop("STRICT_LOAD", None)
+            os.environ.pop("LOAD_CHECKPOINT_PATH", None)
+            with (
+                mock.patch("xtuner.v1.model.get_model_config_from_hf", return_value=model_cfg),
+                mock.patch("xtuner.v1.train.TrainerConfig", _CaptureTrainerConfig),
+            ):
+                namespace = runpy.run_path(str(RECIPE_PATH), run_name="__glm52_sft_recipe_test__")
+
+    trainer_cfg = namespace["trainer"]
+    recipe_model_cfg = namespace["model_cfg"]
+
+    assert recipe_model_cfg is model_cfg
+    assert trainer_cfg.model_cfg is model_cfg
+    assert model_cfg.num_nextn_predict_layers == 1
+    assert model_cfg.mtp_config.num_layers == 7
+    assert model_cfg.mtp_config.share_weights is True
+    assert model_cfg.mtp_config.detach_mtp_lm_head_weight is False
+    assert model_cfg.mtp_config.detach_mtp_inputs is False
+    assert model_cfg.mtp_config.loss_scaling_factor == 0.1
+    assert trainer_cfg.strict_load is True
+
+
+class TestGlm52SftMtpConfig:
+    def test_glm52_sft_recipe_uses_recurrent_mtp_depth_without_changing_checkpoint_depth(self) -> None:
+        """The SFT recipe should run one physical MTP block for seven depths."""
+        _run_glm52_sft_recipe_regression()

xtuner/v1/loss/__init__.py

@@ -10,7 +10,7 @@
     ZLossContext,
     ZLossKwargs,
 )
-from .mtp_loss import MTPLossContext
+from .mtp_loss import MTPE2ETVLossContext, MTPLossContext
 from .rl_loss import LogProbConfig, LogProbContext
 
 
@@ -31,6 +31,7 @@
     "BaseLossKwargs",
     "LMHeadLossContext",
     "MTPLossContext",
+    "MTPE2ETVLossContext",
     "LogProbConfig",
     "LogProbContext",
 ]