diff --git a/examples/v1/config/sft_glm5p2.py b/examples/v1/config/sft_glm5p2.py index cab158ec2f..0bbd3a9de3 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 @@ -6,6 +7,7 @@ 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") 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/tests/model/test_glm52_sft_mtp_config.py b/tests/model/test_glm52_sft_mtp_config.py new file mode 100644 index 0000000000..1b07acf523 --- /dev/null +++ b/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() 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