[Feature] Add end-to-end TV loss for MTP SFT

examples/v1/config/sft_glm5p2.py

@@ -1,4 +1,5 @@
 import os
+from typing import Literal, cast
 
 from xtuner.v1.config import AdamWConfig, FSDPConfig, LRConfig, MuonConfig
 from xtuner.v1.datasets import OpenaiTokenizeFunctionConfig
@@ -52,6 +53,16 @@ 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 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)

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",
 ]

xtuner/v1/loss/mtp_loss.py

@@ -1,9 +1,12 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from typing import Any
+from typing import Any, Sequence
 
 import torch
+import torch.distributed as dist
 import torch.nn.functional as F
 from torch.distributed.device_mesh import DeviceMesh
+from torch.distributed.nn.functional import all_reduce
+from torch.utils.checkpoint import checkpoint
 
 from xtuner.v1.loss.ce_loss import CELossConfig, CELossKwargs, LMHeadLossContext
 from xtuner.v1.loss.utils import sp_split
@@ -214,3 +217,219 @@ def _kl_loss_fn(
         )
 
         return kl_loss, (None, {})
+
+
+class MTPE2ETVLossKwargs(CELossKwargs):
+    """Inputs used to mask and globally normalize the joint MTP TV loss."""
+
+    sp_mesh: DeviceMesh | None = None
+
+    def sp_split(self, sp_mesh: DeviceMesh) -> "MTPE2ETVLossKwargs":
+        super().sp_split(sp_mesh)
+        self.sp_mesh = sp_mesh
+        return self
+
+
+class MTPE2ETVLossConfig(CELossConfig):
+    """End-to-end TV objective for multi-step rejection-sampling acceptance.
+
+    The objective follows Eq. 13 of https://arxiv.org/abs/2606.12370. Unlike
+    :class:`MTPLossConfig`, one context covers every logical MTP depth because
+    the prefix-acceptance products couple the per-depth distribution overlaps.
+    """
+
+    num_steps: int
+    detach_mtp_lm_head_weight: bool = False
+
+    @property
+    def loss_ctx_cls(self) -> type["MTPE2ETVLossContext"]:
+        return MTPE2ETVLossContext
+
+    @property
+    def _loss_kwargs_cls(self) -> type["MTPE2ETVLossKwargs"]:
+        return MTPE2ETVLossKwargs
+
+    def build(self, data: dict, sp_mesh: DeviceMesh | None = None) -> "MTPE2ETVLossContext | None":
+        from xtuner.v1.module.mtp.utils import roll_packed_tensor
+
+        if "shifted_labels" not in data:
+            return None
+
+        shifted_labels = data["shifted_labels"]
+        cu_seq_lens = data["seq_ctx"].cu_seq_lens_k
+        padded_len = int(cu_seq_lens[-1].item())
+        seq_len = shifted_labels.shape[-1]
+        if padded_len > seq_len:
+            pad = torch.full(
+                (*shifted_labels.shape[:-1], padded_len - seq_len),
+                fill_value=self.ignore_idx,
+                dtype=shifted_labels.dtype,
+                device=shifted_labels.device,
+            )
+            shifted_labels = torch.cat([shifted_labels, pad], dim=-1)
+
+        # Eq. 13 assumes a fixed gamma. Restrict training to starting positions
+        # for which every one of the configured MTP steps has a supervised target.
+        valid_mask = torch.ones_like(shifted_labels, dtype=torch.bool)
+        for depth in range(1, self.num_steps + 1):
+            rolled_labels = roll_packed_tensor(
+                shifted_labels,
+                cu_seq_lens,
+                shifts=-depth,
+                dim=-1,
+                fill_value=self.ignore_idx,
+            )
+            valid_mask.logical_and_(rolled_labels != self.ignore_idx)
+
+        # Reuse LMHeadLossContext.build_batches() for token/sample weighting and
+        # DP/SP/global-batch calibration. Label values are immaterial to TV loss;
+        # only ignore_idx positions are consumed by the weighting code.
+        mask_labels = torch.where(
+            valid_mask,
+            torch.zeros_like(shifted_labels),
+            torch.full_like(shifted_labels, self.ignore_idx),
+        )
+        loss_kwargs = MTPE2ETVLossKwargs(shifted_labels=mask_labels, sp_mesh=sp_mesh).to(DEVICE)
+        if sp_mesh is not None and sp_mesh.size() > 1:
+            loss_kwargs = loss_kwargs.sp_split(sp_mesh)
+        return MTPE2ETVLossContext(self, loss_kwargs)
+
+
+def _shift_target_hidden_states(
+    target_hidden_states: torch.Tensor,
+    num_steps: int,
+    sp_mesh: DeviceMesh | None,
+) -> list[torch.Tensor]:
+    """Left-shift target states with a small halo across sequence-parallel
+    ranks."""
+
+    target_hidden_states = target_hidden_states.detach()
+    batch_size, local_length, hidden_size = target_hidden_states.shape
+    if local_length == 0:
+        raise ValueError("MTP e2e TV loss requires a non-empty local sequence shard.")
+
+    sp_size = 1 if sp_mesh is None else sp_mesh.size()
+    if sp_size == 1:
+        halo = target_hidden_states.new_zeros((batch_size, num_steps, hidden_size))
+        extended = torch.cat((target_hidden_states, halo), dim=1)
+    elif local_length >= num_steps:
+        assert dist.is_initialized(), "Sequence parallelism requires torch.distributed to be initialized."
+        assert sp_mesh is not None
+        group = sp_mesh.get_group()
+        sp_rank = dist.get_rank(group)
+        prefix = target_hidden_states[:, :num_steps].contiguous()
+        prefixes = [torch.empty_like(prefix) for _ in range(sp_size)]
+        dist.all_gather(prefixes, prefix, group=group)
+        halo = prefixes[sp_rank + 1] if sp_rank + 1 < sp_size else torch.zeros_like(prefix)
+        extended = torch.cat((target_hidden_states, halo), dim=1)
+    else:
+        # This path is only relevant to tiny tests or extremely short shards.
+        # Gather complete shards because one horizon can cross several ranks.
+        assert dist.is_initialized(), "Sequence parallelism requires torch.distributed to be initialized."
+        assert sp_mesh is not None
+        group = sp_mesh.get_group()
+        sp_rank = dist.get_rank(group)
+        shards = [torch.empty_like(target_hidden_states) for _ in range(sp_size)]
+        dist.all_gather(shards, target_hidden_states.contiguous(), group=group)
+        global_hidden_states = torch.cat(shards, dim=1)
+        start = sp_rank * local_length
+        tail = global_hidden_states[:, start : start + local_length + num_steps]
+        if tail.shape[1] < local_length + num_steps:
+            tail = torch.cat(
+                (
+                    tail,
+                    target_hidden_states.new_zeros(
+                        (batch_size, local_length + num_steps - tail.shape[1], hidden_size)
+                    ),
+                ),
+                dim=1,
+            )
+        extended = tail
+
+    return [extended[:, depth : depth + local_length] for depth in range(1, num_steps + 1)]
+
+
+class MTPE2ETVLossContext(LMHeadLossContext):
+    """Compute exact full-vocabulary e2e TV loss over all MTP depths."""
+
+    loss_cfg: MTPE2ETVLossConfig
+    loss_kwargs: MTPE2ETVLossKwargs
+
+    @staticmethod
+    def _tv_overlap(
+        draft_hidden_states: torch.Tensor,
+        target_hidden_states: torch.Tensor,
+        head_weight: torch.Tensor,
+    ) -> torch.Tensor:
+        # The target distribution is a teacher for this objective. Gradients are
+        # defined only with respect to the draft distribution (Bebop Appendix F).
+        with torch.no_grad():
+            target_logits = F.linear(target_hidden_states, head_weight.detach()).float()
+            target_probs = F.softmax(target_logits, dim=-1)
+
+        draft_logits = F.linear(draft_hidden_states, head_weight).float()
+        draft_probs = F.softmax(draft_logits, dim=-1)
+        return torch.minimum(target_probs, draft_probs).sum(dim=-1).clamp_(0.0, 1.0)
+
+    def forward(  # type: ignore[override]
+        self,
+        hidden_states: tuple[torch.Tensor, Sequence[torch.Tensor]],
+        head_weight: torch.Tensor,
+        head_bias: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, tuple[None, dict[str, Any]]]:
+        if head_bias is not None:
+            raise NotImplementedError("MTP e2e TV loss does not support head_bias.")
+
+        target_hidden_states, draft_hidden_states = hidden_states
+        if len(draft_hidden_states) != self.loss_cfg.num_steps:
+            raise ValueError(
+                f"Expected {self.loss_cfg.num_steps} MTP outputs for e2e TV loss, got {len(draft_hidden_states)}."
+            )
+        if self.loss_cfg.detach_mtp_lm_head_weight:
+            head_weight = head_weight.detach()
+
+        shifted_targets = _shift_target_hidden_states(
+            target_hidden_states,
+            self.loss_cfg.num_steps,
+            self.loss_kwargs.sp_mesh,
+        )
+        loss_weight = self.loss_kwargs.loss_weight
+        assert loss_weight is not None, "loss_weight can not be None"
+        valid_mask = loss_weight != 0
+
+        if not valid_mask.any():
+            loss = head_weight.sum() * 0.0 + sum(draft.sum() * 0.0 for draft in draft_hidden_states)
+        else:
+            selected_weight = loss_weight[valid_mask]
+            alpha_per_step: list[torch.Tensor] = []
+            chunk_size = self.loss_cfg.chunk_size
+            assert chunk_size is not None and chunk_size > 0, "A positive chunk_size is required for e2e TV loss."
+
+            for draft_hidden, target_hidden in zip(draft_hidden_states, shifted_targets):
+                selected_draft = draft_hidden[valid_mask]
+                selected_target = target_hidden[valid_mask]
+                alpha_chunks: list[torch.Tensor] = []
+                for draft_chunk, target_chunk in zip(
+                    torch.split(selected_draft, chunk_size, dim=0),
+                    torch.split(selected_target, chunk_size, dim=0),
+                ):
+                    if torch.is_grad_enabled() and (draft_chunk.requires_grad or head_weight.requires_grad):
+                        alpha = checkpoint(
+                            self._tv_overlap,
+                            draft_chunk,
+                            target_chunk,
+                            head_weight,
+                            use_reentrant=False,
+                        )
+                    else:
+                        alpha = self._tv_overlap(draft_chunk, target_chunk, head_weight)
+                    alpha_chunks.append(alpha)
+                alpha_per_step.append(torch.cat(alpha_chunks, dim=0))
+
+            alphas = torch.stack(alpha_per_step, dim=-1)
+            normalized_expected_acceptance = torch.cumprod(alphas, dim=-1).mean(dim=-1)
+            loss = ((1.0 - normalized_expected_acceptance) * selected_weight).sum()
+
+        if dist.is_initialized():
+            loss = all_reduce(loss, op=dist.ReduceOp.SUM, group=dist.group.WORLD)
+        return loss, (None, {})