diff --git a/examples/v1/config/sft_glm5p2.py b/examples/v1/config/sft_glm5p2.py index cab158ec2f..d8a6cb1c17 100644 --- a/examples/v1/config/sft_glm5p2.py +++ b/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") diff --git a/tests/loss/test_mtp_e2e_tv_loss.py b/tests/loss/test_mtp_e2e_tv_loss.py new file mode 100644 index 0000000000..f0f7c043f6 --- /dev/null +++ b/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) diff --git a/xtuner/v1/loss/__init__.py b/xtuner/v1/loss/__init__.py index d2f20b3a16..d02cbc6353 100644 --- a/xtuner/v1/loss/__init__.py +++ b/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", ] diff --git a/xtuner/v1/loss/mtp_loss.py b/xtuner/v1/loss/mtp_loss.py index a5aebc3b19..ad1aff5f18 100644 --- a/xtuner/v1/loss/mtp_loss.py +++ b/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, {}) diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py index 15dc0d628e..27c1dd891f 100644 --- a/xtuner/v1/model/moe/moe.py +++ b/xtuner/v1/model/moe/moe.py @@ -32,11 +32,12 @@ BalancingLossContext, BaseLossContext, LMHeadLossContext, + MTPE2ETVLossContext, MTPLossContext, ZLossConfig, ZLossContext, ) -from xtuner.v1.loss.mtp_loss import MTPLossConfig +from xtuner.v1.loss.mtp_loss import MTPE2ETVLossConfig, MTPLossConfig from xtuner.v1.model.base import ( DEFAULT_FLOAT8_CFG, BaseModel, @@ -146,6 +147,7 @@ class MoELossContextDict(TypedDict): balancing: BalancingLossContext | None z_loss: ZLossContext | None mtp: list[BaseLossContext] | None + mtp_e2e_tv: BaseLossContext | None class MoEConfig(TransformerConfig): @@ -454,7 +456,8 @@ def build_loss_ctx_batch( # type: ignore[override] - "lm": LM loss context - "balancing": Balancing loss context (if configured) - "z_loss": Z-loss context (if configured) - - "mtp": MTP loss contexts (if configured) + - "mtp": Per-depth MTP loss contexts (if configured) + - "mtp_e2e_tv": Joint multi-depth MTP TV loss context (if configured) Note: Auxiliary loss contexts are built without parameters. @@ -471,33 +474,59 @@ def build_loss_ctx_batch( # type: ignore[override] self._add_auxiliary_loss("balancing", self.config.balancing_loss_cfg, _data_batch, res) self._add_auxiliary_loss("z_loss", self.config.z_loss_cfg, _data_batch, res) - # Add MTP loss contexts if MTP is enabled + # Add MTP loss contexts if MTP is enabled. if self.config.mtp_config is not None: - for mtp_idx in range(self.config.mtp_config.num_layers): - mtp_loss_cfg = MTPLossConfig( - **self.config.lm_loss_cfg.model_dump(), - mtp_depth=mtp_idx + 1, + mtp_config = self.config.mtp_config + if mtp_config.loss_type == "e2e_tv": + tv_loss_cfg = MTPE2ETVLossConfig( + ignore_idx=self.config.lm_loss_cfg.ignore_idx, + mode="chunk", + chunk_size=mtp_config.tv_loss_chunk_size, + loss_reduction=self.config.lm_loss_cfg.loss_reduction, + num_steps=mtp_config.num_layers, detach_mtp_lm_head_weight=self.config.mtp_config.detach_mtp_lm_head_weight, ) - mtp_loss_ctx_list = self._build_loss_ctx(mtp_loss_cfg, _data_batch, sp_mesh) - if mtp_loss_ctx_list is not None: - mtp_loss_ctx_list = MTPLossContext.build_batches( # type: ignore[assignment] - cast(list[MTPLossContext], mtp_loss_ctx_list), # type: ignore[arg-type] + tv_loss_ctx_list = self._build_loss_ctx(tv_loss_cfg, _data_batch, sp_mesh) + if tv_loss_ctx_list is not None: + tv_loss_ctx_list = MTPE2ETVLossContext.build_batches( # type: ignore[assignment] + cast(list[MTPE2ETVLossContext], tv_loss_ctx_list), # type: ignore[arg-type] cu_seq_lens_list=cu_seq_lens_list, sp_mesh=sp_mesh, ) - for i, mtp_loss_ctx in enumerate(mtp_loss_ctx_list): - if "mtp" not in res[i]: - res[i]["mtp"] = [] - res[i]["mtp"].append(mtp_loss_ctx) # type: ignore[union-attr] - - # Ensure all microbatches have mtp key - for loss_ctx_dict in res: - if "mtp" not in loss_ctx_dict: + for i, tv_loss_ctx in enumerate(tv_loss_ctx_list): + res[i]["mtp_e2e_tv"] = tv_loss_ctx + for loss_ctx_dict in res: loss_ctx_dict["mtp"] = None + if "mtp_e2e_tv" not in loss_ctx_dict: + loss_ctx_dict["mtp_e2e_tv"] = None + else: + for mtp_idx in range(mtp_config.num_layers): + mtp_loss_cfg = MTPLossConfig( + **self.config.lm_loss_cfg.model_dump(), + mtp_depth=mtp_idx + 1, + detach_mtp_lm_head_weight=mtp_config.detach_mtp_lm_head_weight, + ) + mtp_loss_ctx_list = self._build_loss_ctx(mtp_loss_cfg, _data_batch, sp_mesh) + if mtp_loss_ctx_list is not None: + mtp_loss_ctx_list = MTPLossContext.build_batches( # type: ignore[assignment] + cast(list[MTPLossContext], mtp_loss_ctx_list), # type: ignore[arg-type] + cu_seq_lens_list=cu_seq_lens_list, + sp_mesh=sp_mesh, + ) + for i, mtp_loss_ctx in enumerate(mtp_loss_ctx_list): + if "mtp" not in res[i]: + res[i]["mtp"] = [] + res[i]["mtp"].append(mtp_loss_ctx) # type: ignore[union-attr] + + # Ensure all microbatches have both MTP keys. + for loss_ctx_dict in res: + if "mtp" not in loss_ctx_dict: + loss_ctx_dict["mtp"] = None + loss_ctx_dict["mtp_e2e_tv"] = None else: for loss_ctx_dict in res: loss_ctx_dict["mtp"] = None + loss_ctx_dict["mtp_e2e_tv"] = None return res # type: ignore[return-value] @@ -662,19 +691,31 @@ def _micro_batch_forward( has_mtp_loss = False for micro_batch_idx, (loss_ctx_dict, mtp_outputs) in enumerate(zip(loss_ctx_list, mtp_outputs_per_mb)): mtp_loss_ctx_list = loss_ctx_dict.get("mtp") - if mtp_loss_ctx_list is None: - continue - - micro_batch_mtp_losses = torch.tensor(0.0, device=DEVICE) - for mtp_idx, (mtp_hidden, mtp_ctx) in enumerate(zip(mtp_outputs, mtp_loss_ctx_list)): - mtp_loss, _ = self.lm_head(mtp_hidden["hidden_states"], cast(MTPLossContext, mtp_ctx)) - micro_batch_mtp_losses += mtp_loss - + mtp_tv_loss_ctx = loss_ctx_dict.get("mtp_e2e_tv") + for mtp_idx, mtp_hidden in enumerate(mtp_outputs): if keep_router: router_logits_list[micro_batch_idx][f"mtp_layer{mtp_idx}"] = mtp_hidden["router_logits"] - mtp_losses += micro_batch_mtp_losses / len(mtp_loss_ctx_list) - has_mtp_loss = True + if mtp_tv_loss_ctx is not None: + # The target branch is detached by the TV loss context. Keep the + # original, unnormalized hidden_states_list untouched so the main + # loss can still carry the MTP auxiliary loss below. + target_hidden_states = self.norm(hidden_states_list[micro_batch_idx]) + draft_hidden_states = [mtp_hidden["hidden_states"] for mtp_hidden in mtp_outputs] + mtp_loss, _ = self.lm_head( + (target_hidden_states, draft_hidden_states), + cast(MTPE2ETVLossContext, mtp_tv_loss_ctx), + ) + mtp_losses += mtp_loss + has_mtp_loss = True + elif mtp_loss_ctx_list is not None: + micro_batch_mtp_losses = torch.tensor(0.0, device=DEVICE) + for mtp_hidden, mtp_ctx in zip(mtp_outputs, mtp_loss_ctx_list): + mtp_loss, _ = self.lm_head(mtp_hidden["hidden_states"], cast(MTPLossContext, mtp_ctx)) + micro_batch_mtp_losses += mtp_loss + + mtp_losses += micro_batch_mtp_losses / len(mtp_loss_ctx_list) + has_mtp_loss = True if has_mtp_loss: # MTP routed experts feed the same balancing / z aux loss as the main MoE layers @@ -904,10 +945,12 @@ def _forward( output["extra_info"] = extra_info # MTP forward pass and loss computation + mtp_loss_ctx_list = loss_ctx.get("mtp") if loss_ctx is not None else None + mtp_tv_loss_ctx = loss_ctx.get("mtp_e2e_tv") if loss_ctx is not None else None if ( self.mtp_block is not None and loss_ctx is not None - and (mtp_loss_ctx_list := loss_ctx.get("mtp")) is not None + and (mtp_loss_ctx_list is not None or mtp_tv_loss_ctx is not None) ): mtp_seq_ctx = seq_ctx.copy( input_ids=input_ids.clone() if input_ids is not None else None, @@ -929,9 +972,9 @@ def _forward( seq_ctx=mtp_seq_ctx, ) - # Compute MTP losses for each depth - mtp_losses = torch.tensor(0.0, device=DEVICE) - for idx, (mtp_hidden, mtp_ctx) in enumerate(zip(mtp_outputs, mtp_loss_ctx_list)): + # Accumulate MTP auxiliary losses and collect every recurrent step. + draft_hidden_states: list[torch.Tensor] = [] + for idx, mtp_hidden in enumerate(mtp_outputs): mtp_hidden_states = mtp_hidden["hidden_states"] mtp_router_logits = mtp_hidden["router_logits"] mtp_router_weights = mtp_hidden["router_weights"] @@ -955,11 +998,24 @@ def _forward( num_tokens_global=mtp_num_tokens_global, world_size=mtp_z_world_size, ) - mtp_loss, _ = self.lm_head(mtp_hidden_states, cast(MTPLossContext, mtp_ctx)) - mtp_losses += mtp_loss - - # Average MTP losses across depths and scale - mtp_losses = mtp_losses / len(mtp_loss_ctx_list) + draft_hidden_states.append(mtp_hidden_states) + + if mtp_tv_loss_ctx is not None: + # One joint call is required: the acceptance products in Eq. 13 + # couple all recurrent MTP steps. + mtp_losses, _ = self.lm_head( + (hidden_states, draft_hidden_states), + cast(MTPE2ETVLossContext, mtp_tv_loss_ctx), + ) + else: + assert mtp_loss_ctx_list is not None + mtp_losses = torch.tensor(0.0, device=DEVICE) + for mtp_hidden_states, mtp_ctx in zip(draft_hidden_states, mtp_loss_ctx_list): + mtp_loss, _ = self.lm_head(mtp_hidden_states, cast(MTPLossContext, mtp_ctx)) + mtp_losses += mtp_loss + mtp_losses = mtp_losses / len(mtp_loss_ctx_list) + + # Both objectives are normalized over MTP depth internally; scale once. scaled_mtp_loss = mtp_losses * self.config.mtp_config.loss_scaling_factor # type: ignore # Add to total loss diff --git a/xtuner/v1/module/lm_head/lm_head.py b/xtuner/v1/module/lm_head/lm_head.py index 67e912d47c..1c0946bb5a 100644 --- a/xtuner/v1/module/lm_head/lm_head.py +++ b/xtuner/v1/module/lm_head/lm_head.py @@ -1,4 +1,4 @@ -from typing import Any, TypeAlias +from typing import Any, Sequence, TypeAlias import torch import torch.nn as nn @@ -14,6 +14,7 @@ Weight: TypeAlias = torch.Tensor | DTensor Bias: TypeAlias = torch.Tensor | DTensor | None HiddenStates: TypeAlias = torch.Tensor +LossHiddenStates: TypeAlias = torch.Tensor | tuple[torch.Tensor, Sequence[torch.Tensor]] Labels: TypeAlias = torch.Tensor @@ -25,11 +26,11 @@ def forward( @overload # type: ignore[override] def forward( - self, hidden_states: HiddenStates, loss_ctx: LMHeadLossContext + self, hidden_states: LossHiddenStates, loss_ctx: LMHeadLossContext ) -> tuple[Loss, tuple[Logits | None, dict[str, Any]]]: ... def forward( # type: ignore[override] - self, hidden_states: torch.Tensor, loss_ctx: LMHeadLossContext | None = None + self, hidden_states: LossHiddenStates, loss_ctx: LMHeadLossContext | None = None ) -> tuple[Loss | None, tuple[Logits | None, dict[str, Any]]]: """Forward pass of the language model head.""" if isinstance(self.weight, DTensor): @@ -43,10 +44,11 @@ def forward( # type: ignore[override] w = self.weight b = self.bias if loss_ctx is None: + assert isinstance(hidden_states, torch.Tensor) logits = F.linear(hidden_states, w, b) return None, (logits.float(), {}) else: - return loss_ctx.forward(hidden_states, w, b) + return loss_ctx.forward(hidden_states, w, b) # type: ignore[arg-type] @overload # type: ignore def __call__( @@ -55,7 +57,7 @@ def __call__( @overload # type: ignore def __call__( - self, hidden_states: HiddenStates, loss_ctx: LMHeadLossContext + self, hidden_states: LossHiddenStates, loss_ctx: LMHeadLossContext ) -> tuple[Loss, tuple[Logits | None, dict[str, Any]]]: ... __call__ = nn.Module.__call__ diff --git a/xtuner/v1/module/mtp/config.py b/xtuner/v1/module/mtp/config.py index d6bcf7a9b5..84b390d6b9 100644 --- a/xtuner/v1/module/mtp/config.py +++ b/xtuner/v1/module/mtp/config.py @@ -1,6 +1,6 @@ """Configuration for Multi-Token Prediction (MTP).""" -from typing import Annotated +from typing import Annotated, Literal from cyclopts import Parameter from pydantic import BaseModel, ConfigDict @@ -27,9 +27,15 @@ class MTPConfig(BaseModel): This is used in RL training. Default is False. detach_mtp_inputs (bool): Whether to detach the input embeddings and hidden states. This is used in RL training. Default is False. - loss_scaling_factor (float): Scaling factor for MTP loss. The total MTP loss - is computed as the average of losses across all depths, multiplied by - this factor. Default: 0.1. + loss_scaling_factor (float): Scaling factor for MTP loss. Per-depth CE is + averaged across depths; e2e TV is already normalized by the configured + speculative horizon. The selected objective is multiplied by this + factor. Default: 0.1. + loss_type (str): MTP training objective. ``"ce"`` preserves the existing + per-depth cross-entropy objective. ``"e2e_tv"`` directly optimizes the + expected multi-step rejection-sampling acceptance length. Default: ``"ce"``. + tv_loss_chunk_size (int): Number of token positions processed together when + computing the exact full-vocabulary TV overlap. Default: 128. Example: >>> # In model config @@ -50,3 +56,5 @@ class MTPConfig(BaseModel): detach_mtp_lm_head_weight: Annotated[bool, Parameter(group="model")] = False detach_mtp_inputs: Annotated[bool, Parameter(group="model")] = False loss_scaling_factor: Annotated[float, Parameter(group="model")] = 0.1 + loss_type: Annotated[Literal["ce", "e2e_tv"], Parameter(group="model")] = "ce" + tv_loss_chunk_size: Annotated[int, Parameter(group="model")] = 128