From 5463195157b98b90d77d82cf9f799ecc47dfd6b6 Mon Sep 17 00:00:00 2001 From: llamafactory-mtp Date: Sat, 25 Jul 2026 10:23:23 +0800 Subject: [PATCH 1/4] feat(v1): add Multi-Token Prediction (MTP) layer support Adapt the FSDP2 MTP implementation from MindSpeed-LLM (mindspeed_llm/fsdp2/models/common/mtp.py) to the LlamaFactory v1 architecture, making it HuggingFace-transformers generic so it can be attached to any Llama/Qwen3/Mistral-style decoder-only causal LM. Components: * MultiTokenPredictionBlock / MultiTokenPredictionLayer (src/llamafactory/v1/plugins/model_plugins/mtp.py): K MTP heads that reuse the base model's decoder layer class, with shared enorm/hnorm norms, e_proj/h_proj projections and a final_layernorm, mirroring MindSpeed-LLM. Head k predicts token p+k+2 from position p. * MTPModelPlugin: grafts the block onto the model as `model.mtp` and patches `model.forward` to emit `mtp_logits` during training. * compute_mtp_loss: loss_weights-weighted per-head CE mean, averaged over heads; total loss = lm_loss + loss_scale * mtp_loss (loss_scale=0.3). * ModelArguments.mtp_config plugin field + ModelEngine wiring. * SFTTrainer MTP-aware compute_loss (non context-parallel path). * FSDP2: MTP inner decoder layers are sharded automatically by the existing prepare_model loop (wraps all modules of the decoder-layer class, including mtp.layers.*.layer). * Example YAML, unit tests, and docs. Verified on llamafactory/tiny-random-qwen2.5 (attach, forward shapes, loss + backward, head-offset convention). --- docs/v1_mtp.md | 64 +++ .../v1/train_full/train_full_mtp_fsdp2.yaml | 36 ++ src/llamafactory/v1/config/model_args.py | 11 + src/llamafactory/v1/core/base_trainer.py | 5 + src/llamafactory/v1/core/model_engine.py | 5 + .../v1/plugins/model_plugins/mtp.py | 402 ++++++++++++++++++ src/llamafactory/v1/trainers/sft_trainer.py | 39 ++ tests_v1/plugins/model_plugins/test_mtp.py | 120 ++++++ 8 files changed, 682 insertions(+) create mode 100644 docs/v1_mtp.md create mode 100644 examples/v1/train_full/train_full_mtp_fsdp2.yaml create mode 100644 src/llamafactory/v1/plugins/model_plugins/mtp.py create mode 100644 tests_v1/plugins/model_plugins/test_mtp.py diff --git a/docs/v1_mtp.md b/docs/v1_mtp.md new file mode 100644 index 0000000000..dbd9981f14 --- /dev/null +++ b/docs/v1_mtp.md @@ -0,0 +1,64 @@ +# Multi-Token Prediction (MTP) — v1 architecture + +This feature adds Multi-Token Prediction (MTP) support to the v1 architecture, adapted +from the FSDP2 MTP implementation in +[MindSpeed-LLM](https://gitcode.com/Ascend/MindSpeed-LLM) (`mindspeed_llm/fsdp2/models/common/mtp.py`). + +MTP appends `K` extra prediction heads to a decoder-only causal LM. Each head `k` +predicts the token at offset `k + 2` from the current position (the main head predicts +offset `1`). The total training loss is: + +``` +total_loss = lm_loss + loss_scale * mtp_loss +``` + +where `mtp_loss` is the mean of the per-head cross-entropy losses (weighted by +`loss_weights`, like the main SFT loss), and `loss_scale` defaults to `0.3`. + +## How it works + +- `MultiTokenPredictionBlock` (`src/llamafactory/v1/plugins/model_plugins/mtp.py`) holds + `K` heads. Each head reuses the base model's decoder layer class. The block owns shared + `enorm`/`hnorm` norms, `e_proj`/`h_proj` projections and a `final_layernorm`, exactly as + in MindSpeed-LLM. +- `MTPModelPlugin` attaches the block to the model as `model.mtp` and patches + `model.forward` so the model output carries `mtp_logits` (a list of per-head logits) + during training. The MTP loss is computed by the trainer through `compute_mtp_loss`. +- Under FSDP2, the MTP heads' inner decoder layers are sharded automatically: the generic + `FSDP2Engine.prepare_model` wraps every module of the base model's decoder-layer class, + which includes `mtp.layers.*.layer`. + +## Usage + +Add an `mtp_config` block to your v1 YAML: + +```yaml +model: Qwen/Qwen3-0.6B +model_class: llm +template: qwen3_nothink + +mtp_config: + name: mtp + num_layers: 1 # number of MTP heads (K) + loss_scale: 0.3 # optional, default 0.3 + +dist_config: + name: fsdp2 + dcp_path: null + +train_dataset: data/v1_sft_demo.yaml +output_dir: outputs/test_mtp_fsdp2 +micro_batch_size: 1 +cutoff_len: 2048 +learning_rate: 1.0e-4 +max_steps: 10 +``` + +See `examples/v1/train_full/train_full_mtp_fsdp2.yaml` for a complete example. + +## Compatibility + +MTP currently targets Llama/Qwen3/Mistral-style models that expose `model.model.layers`, +`model.model.rotary_emb`, `model.model.norm` and `model.lm_head`. The MTP heads are +randomly initialized; loading a base checkpoint that does not contain `mtp.*` keys will +leave them at their initialization (missing-key warnings are expected). diff --git a/examples/v1/train_full/train_full_mtp_fsdp2.yaml b/examples/v1/train_full/train_full_mtp_fsdp2.yaml new file mode 100644 index 0000000000..361d52e6f4 --- /dev/null +++ b/examples/v1/train_full/train_full_mtp_fsdp2.yaml @@ -0,0 +1,36 @@ +model: Qwen/Qwen3-0.6B +model_class: llm + +template: qwen3_nothink + +kernel_config: + name: auto + include_kernels: auto + +quant_config: null + +# Multi-Token Prediction (MTP): append K extra prediction heads to the model. +# `num_layers` is the number of MTP heads (K); `loss_scale` weights the MTP loss +# in the total loss (total = lm_loss + loss_scale * mtp_loss). +mtp_config: + name: mtp + num_layers: 1 + loss_scale: 0.3 + +dist_config: + name: fsdp2 + dcp_path: null + +### data +train_dataset: data/v1_sft_demo.yaml + +### training +output_dir: outputs/test_mtp_fsdp2 +micro_batch_size: 1 +cutoff_len: 2048 +learning_rate: 1.0e-4 +max_steps: 10 + +### sample +sample_backend: hf +max_new_tokens: 128 diff --git a/src/llamafactory/v1/config/model_args.py b/src/llamafactory/v1/config/model_args.py index 8c6542162a..8ed6cc4232 100644 --- a/src/llamafactory/v1/config/model_args.py +++ b/src/llamafactory/v1/config/model_args.py @@ -59,6 +59,16 @@ class ModelArguments: default=None, metadata={"help": "Quantization configuration for the model."}, ) + mtp_config: PluginConfig | None = field( + default=None, + metadata={ + "help": ( + "Multi-Token Prediction (MTP) configuration. Set `name: mtp` and " + "`num_layers: K` to append K MTP heads to the model, plus an optional " + "`loss_scale` (default 0.3). Example: {name: mtp, num_layers: 1, loss_scale: 0.3}." + ) + }, + ) def __post_init__(self) -> None: supported_flash_attn = [item.value for item in AttentionFunction] @@ -71,3 +81,4 @@ def __post_init__(self) -> None: self.peft_config = get_plugin_config(self.peft_config) self.kernel_config = get_plugin_config(self.kernel_config) self.quant_config = get_plugin_config(self.quant_config) + self.mtp_config = get_plugin_config(self.mtp_config) diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index 1aba01e3b9..c6702fc86d 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/src/llamafactory/v1/core/base_trainer.py @@ -223,6 +223,11 @@ def _init_lr_scheduler(self) -> None: self.optimizer, self.num_training_steps, self.args.lr_scheduler_config ) + def _has_mtp(self) -> bool: + """Whether the (possibly wrapped) model carries an MTP block.""" + model = self.model.module if hasattr(self.model, "module") else self.model + return getattr(model, "mtp", None) is not None + def compute_log_probs(self, model: HFModel, batch: BatchInput) -> Tensor: """Compute log probs. diff --git a/src/llamafactory/v1/core/model_engine.py b/src/llamafactory/v1/core/model_engine.py index 67ad8a4f69..97047a04b4 100644 --- a/src/llamafactory/v1/core/model_engine.py +++ b/src/llamafactory/v1/core/model_engine.py @@ -225,6 +225,11 @@ def _init_model(self) -> HFModel: model = apply_kernels(model, self.args.kernel_config, require_logits=self.is_train) + if self.args.mtp_config is not None: + from ..plugins.model_plugins.mtp import MTPModelPlugin + + model = MTPModelPlugin(self.args.mtp_config.name)(model, self.args.mtp_config) + return model diff --git a/src/llamafactory/v1/plugins/model_plugins/mtp.py b/src/llamafactory/v1/plugins/model_plugins/mtp.py new file mode 100644 index 0000000000..9721299da6 --- /dev/null +++ b/src/llamafactory/v1/plugins/model_plugins/mtp.py @@ -0,0 +1,402 @@ +# Copyright 2025 the LlamaFactory team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Multi-Token Prediction (MTP) layer support for the v1 architecture. + +This module is adapted from the FSDP2 MTP implementation in MindSpeed-LLM +(https://gitcode.com/Ascend/MindSpeed-LLM, ``mindspeed_llm/fsdp2/models/common/mtp.py``) +and reworked to be HuggingFace-transformers generic so that it can be attached to any +decoder-only causal LM that follows the Llama/Qwen3 layout (``model.model.layers``, +``model.model.rotary_emb``, ``model.model.norm``, ``model.lm_head``). + +Design overview +--------------- +* ``MultiTokenPredictionBlock`` holds ``num_layers`` prediction heads. Each head reuses + the base model's decoder layer class. The block owns shared ``enorm``/``hnorm`` norms, + ``e_proj``/``h_proj`` projections and a ``final_layernorm`` (exactly as in MindSpeed). +* The block is attached to a ``ForCausalLM`` model as ``model.mtp`` by ``MTPModelPlugin``, + which also patches ``model.forward`` so that the model output carries ``mtp_logits`` + (a list of per-head logits, one per MTP head). +* The actual loss is *not* computed inside the model. It is computed by the trainer + (non-CP path) or by the sequence-parallel loss plugin (CP path) through the shared + ``compute_mtp_loss`` helper. This keeps the loss weighting (``loss_weights``) and the + context-parallel all-gather logic in a single place. + +Conventions +----------- +* The main model head predicts token ``p + 1`` from hidden state at position ``p``. +* MTP head ``k`` (0-indexed) predicts token ``p + k + 2``. The prediction of head ``k`` + uses ``mtp_logits[k][:, :-(k + 2)]`` against ``labels[:, k + 2:]``. +""" + +from typing import TYPE_CHECKING, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...utils import logging +from ...utils.plugin import BasePlugin + + +if TYPE_CHECKING: + from transformers import PretrainedConfig + from transformers.modeling_utils import PreTrainedModel + + from ...utils.types import PluginConfig + + +logger = logging.get_logger(__name__) + + +def roll_tensor( + tensor: torch.Tensor, shifts: int = -1, dim: int = -1, fill_value: float = 0.0 +) -> torch.Tensor: + """Roll a tensor along ``dim`` and fill the wrapped positions with ``fill_value``. + + This mirrors ``roll_tensor`` in MindSpeed-LLM. With ``shifts=-1, dim=-1`` it shifts + the sequence one step to the left and sets the last position to ``fill_value``. + """ + rolled = torch.roll(tensor, shifts=shifts, dims=dim) + rolled.select(dim, shifts).fill_(fill_value) + return rolled + + +class MultiTokenPredictionLayer(nn.Module): + """A single MTP head: one decoder layer reused from the base model.""" + + def __init__(self, config: "PretrainedConfig", layer_idx: int, layer_cls: type[nn.Module]) -> None: + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.layer = layer_cls(config, layer_idx) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ) -> torch.Tensor: + hidden_states = self.layer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + position_embeddings=position_embeddings, + cache_position=cache_position, + use_cache=False, + **kwargs, + ) + return hidden_states + + +class MultiTokenPredictionBlock(nn.Module): + """Container of ``num_layers`` MTP heads. + + Args: + config: The base model ``PretrainedConfig``. + layer_cls: Decoder layer class to reuse for every MTP head + (e.g. ``LlamaDecoderLayer`` / ``Qwen3DecoderLayer``). + norm_cls: RMS-norm class used by the base model (e.g. ``LlamaRMSNorm``). + num_layers: Number of MTP heads (``K``). + embed_tokens: The base model input embedding module. + rotary_emb: The base model rotary embedding module. + output_layer: The base model ``lm_head`` module, shared by all MTP heads. + """ + + def __init__( + self, + config: "PretrainedConfig", + layer_cls: type[nn.Module], + norm_cls: type[nn.Module], + num_layers: int, + embed_tokens: nn.Module, + rotary_emb: nn.Module, + output_layer: nn.Module, + ) -> None: + super().__init__() + self.config = config + self.num_layers = num_layers + self.embed_tokens = embed_tokens + self.rotary_emb = rotary_emb + self.output_layer = output_layer + self.mtp_start_layer_idx = config.num_hidden_layers + + # MTP heads are brand-new decoder layers. Some models index per-layer config + # lists (e.g. ``config.layer_types[layer_idx]``) by ``layer_idx``, which would + # raise out of range for an index >= ``num_hidden_layers``. Use a valid in-range + # index (the last decoder layer) so the cloned layer behaves like a normal one. + safe_layer_idx = max(0, config.num_hidden_layers - 1) + + self.layers = nn.ModuleDict( + { + str(self.mtp_start_layer_idx + i): MultiTokenPredictionLayer( + config, safe_layer_idx, layer_cls + ) + for i in range(num_layers) + } + ) + rms_eps = getattr(config, "rms_norm_eps", 1e-6) + hidden_size = config.hidden_size + self.enorm = norm_cls(hidden_size, eps=rms_eps) + self.hnorm = norm_cls(hidden_size, eps=rms_eps) + self.e_proj = nn.Linear(hidden_size, hidden_size, bias=False) + self.h_proj = nn.Linear(hidden_size, hidden_size, bias=False) + self.final_layernorm = norm_cls(hidden_size, eps=rms_eps) + + # Reuse the base model weight init scheme if available. + self.reset_parameters() + + def reset_parameters(self) -> None: + std = getattr(self.config, "initializer_range", 0.02) + for module in (self.e_proj, self.h_proj): + nn.init.normal_(module.weight, mean=0.0, std=std) + + def forward( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + ) -> list[torch.Tensor]: + """Run all MTP heads and return per-head logits. + + Args: + hidden_states: Last decoder layer output of the main model (pre final-norm), + shape ``(batch, seq_len, hidden_size)``. + input_ids: Input ids of the main model, shape ``(batch, seq_len)``. + attention_mask: 2D attention mask from the main model, shape ``(batch, seq_len)``. + position_ids: Position ids, shape ``(batch, seq_len)``. + + Returns: + A list of length ``num_layers``; each item is the per-head logits of shape + ``(batch, seq_len, vocab_size)``. + """ + from transformers.masking_utils import create_causal_mask + + batch_size, seq_len, _ = hidden_states.shape + + if position_ids is None: + position_ids = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0).expand(batch_size, -1) + + # Shift input ids by one to obtain the embedding of the "next" token. + shifted_input_ids = roll_tensor(input_ids, shifts=-1, dim=-1, fill_value=0) + input_embeds = self.embed_tokens(shifted_input_ids) + + # Causal mask for the MTP decoder layers (same construction as the base model). + causal_mask = create_causal_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + past_key_values=None, + position_ids=position_ids, + ) + position_embeddings = self.rotary_emb(input_embeds, position_ids=position_ids) + cache_position = torch.arange(seq_len, device=hidden_states.device) + + # Combine the main hidden state with the next-token embedding. + hidden_states = self.hnorm(hidden_states) + self.e_proj(self.enorm(input_embeds)) + + all_mtp_logits: list[torch.Tensor] = [] + for layer_idx in range(self.num_layers): + hidden_states = self.layers[str(self.mtp_start_layer_idx + layer_idx)]( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + position_embeddings=position_embeddings, + cache_position=cache_position, + ) + hidden_states = self.final_layernorm(hidden_states) + mtp_logits = self.output_layer(hidden_states) + all_mtp_logits.append(mtp_logits) + + return all_mtp_logits + + +def compute_mtp_loss( + mtp_logits: list[torch.Tensor], + labels: torch.Tensor, + loss_weights: torch.Tensor, + ignore_index: int = -100, +) -> torch.Tensor: + """Compute the averaged MTP loss across all heads (non context-parallel path). + + Each head ``k`` predicts token ``p + k + 2`` from position ``p``. The loss of head + ``k`` is the ``loss_weights``-weighted cross-entropy mean over valid positions. The + returned scalar is the plain mean over the ``K`` heads (the caller applies the + MTP loss scaling factor, matching MindSpeed-LLM). + + Args: + mtp_logits: List of ``K`` tensors, each ``(batch, seq_len, vocab_size)``. + labels: ``(batch, seq_len)``. + loss_weights: ``(batch, seq_len)``. + ignore_index: Label index to ignore (defaults to -100). + + Returns: + Scalar MTP loss (mean over heads). + """ + num_heads = len(mtp_logits) + if num_heads == 0: + return torch.tensor(0.0, device=labels.device, dtype=torch.float32) + + total_loss = None + for k, logits_k in enumerate(mtp_logits): + shift = k + 2 + head_loss = _local_head_loss(logits_k, labels, loss_weights, shift, ignore_index) + total_loss = head_loss if total_loss is None else total_loss + head_loss + + return total_loss / num_heads + + +def _local_head_loss( + logits_k: torch.Tensor, + labels: torch.Tensor, + loss_weights: torch.Tensor, + shift: int, + ignore_index: int, +) -> torch.Tensor: + """Per-head weighted CE mean for the non-CP path.""" + if logits_k.size(1) <= shift: + return torch.tensor(0.0, device=logits_k.device, dtype=torch.float32) + + batch_size, seq_len, vocab_size = logits_k.shape + pred = logits_k[:, :-shift, :].float().reshape(-1, vocab_size) + tgt = labels[:, shift:].contiguous().reshape(-1) + weights = loss_weights[:, shift:].contiguous().reshape(-1) + + log_probs = -F.cross_entropy(pred, tgt, reduction="none", ignore_index=ignore_index).view(batch_size, -1) + weights = weights.view(batch_size, -1) + return (-log_probs * weights).sum() / (weights.sum() + 1e-6) + + +class MTPModelPlugin(BasePlugin): + """Plugin that grafts a ``MultiTokenPredictionBlock`` onto a causal LM.""" + + def __call__(self, model: "PreTrainedModel", mtp_config: "PluginConfig") -> "PreTrainedModel": + return apply_mtp(model, mtp_config) + + +def _get_inner_model(model: "PreTrainedModel") -> nn.Module: + """Return the inner transformer model (``model.model`` for causal LMs).""" + inner = getattr(model, "model", None) + if inner is None: + raise ValueError( + "MTP currently expects a decoder-only causal LM with a `model.model` attribute " + "(Llama/Qwen3/Mistral-style). Got incompatible model." + ) + return inner + + +def _resolve_layer_cls(model: "PreTrainedModel") -> tuple[type[nn.Module], type[nn.Module]]: + """Resolve ``(decoder_layer_cls, norm_cls)`` from the base model. + + The decoder layer class is taken from the first entry of ``model.model.layers``. + The norm class is taken from ``model.model.norm``. Both are standard for + Llama/Qwen3/Mistral-style models. This avoids importing the FSDP2 plugin (and its + extra dependencies) at module-import time. + """ + inner = _get_inner_model(model) + layers = getattr(inner, "layers", None) + if layers is None or len(layers) == 0: + raise ValueError("Cannot find decoder layers (model.model.layers) to clone for MTP.") + layer_cls = type(layers[0]) + + norm = getattr(inner, "norm", None) + if norm is None: + raise ValueError("Cannot find the final norm (model.model.norm) to clone for MTP.") + norm_cls = type(norm) + return layer_cls, norm_cls + + +def apply_mtp(model: "PreTrainedModel", mtp_config: "PluginConfig") -> "PreTrainedModel": + """Attach an MTP block to ``model`` and patch its forward to emit ``mtp_logits``.""" + num_layers = int(mtp_config.get("num_layers", 1)) + if num_layers <= 0: + return model + + layer_cls, norm_cls = _resolve_layer_cls(model) + inner = _get_inner_model(model) + embed_tokens = model.get_input_embeddings() + rotary_emb = getattr(inner, "rotary_emb", None) + if rotary_emb is None: + raise ValueError("MTP requires the base model to expose `model.model.rotary_emb`.") + output_layer = getattr(model, "lm_head", None) + if output_layer is None: + raise ValueError("MTP requires the base model to expose `lm_head`.") + + block = MultiTokenPredictionBlock( + config=model.config, + layer_cls=layer_cls, + norm_cls=norm_cls, + num_layers=num_layers, + embed_tokens=embed_tokens, + rotary_emb=rotary_emb, + output_layer=output_layer, + ) + # Match the base model parameter dtype (e.g. bf16) so the grafted layers and the + # shared embedding / lm_head agree on dtype. + try: + param_dtype = next(model.parameters()).dtype + block = block.to(param_dtype) + except StopIteration: + pass + model.mtp = block + model.config.mtp_num_layers = num_layers + model.config.mtp_loss_scaling_factor = float(mtp_config.get("loss_scale", 0.3)) + + _patch_forward(model) + + logger.info_rank0( + f"Enabled Multi-Token Prediction with {num_layers} head(s) " + f"(loss_scale={model.config.mtp_loss_scaling_factor})." + ) + return model + + +def _patch_forward(model: "PreTrainedModel") -> None: + """Patch ``model.forward`` so its output carries ``mtp_logits`` during training.""" + import types + + orig_forward = model.forward + + def mtp_forward(self, *args, **kwargs): + # We need the pre-norm last hidden state, which is the last entry of + # ``outputs.hidden_states``. Force the inner model to return it. + kwargs["output_hidden_states"] = True + outputs = orig_forward(*args, **kwargs) + + # MTP logits are only needed for loss computation during training. + mtp_block = getattr(self, "mtp", None) + if mtp_block is not None and self.training: + hidden_states = getattr(outputs, "hidden_states", None) + if hidden_states is not None: + hidden = hidden_states[-1] + else: + hidden = getattr(outputs, "last_hidden_state", None) + if hidden is not None: + input_ids = kwargs.get("input_ids", args[0] if args else None) + attention_mask = kwargs.get("attention_mask", None) + position_ids = kwargs.get("position_ids", None) + mtp_logits = mtp_block( + hidden, + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + ) + outputs.mtp_logits = mtp_logits + return outputs + + model.forward = types.MethodType(mtp_forward, model) diff --git a/src/llamafactory/v1/trainers/sft_trainer.py b/src/llamafactory/v1/trainers/sft_trainer.py index d842632838..a80ffeb733 100644 --- a/src/llamafactory/v1/trainers/sft_trainer.py +++ b/src/llamafactory/v1/trainers/sft_trainer.py @@ -13,6 +13,9 @@ # limitations under the License. +import torch +import torch.nn.functional as F + from ..accelerator.interface import DistributedInterface from ..config import InputArgument, get_args from ..core.base_trainer import BaseTrainer @@ -23,11 +26,47 @@ class SFTTrainer(BaseTrainer): def compute_loss(self, batch: BatchInput) -> Tensor: + if self._has_mtp(): + return self._compute_mtp_loss(batch) + shift_loss_weights = batch["loss_weights"].to(self.device, non_blocking=True)[..., 1:] log_probs = self.compute_log_probs(self.model, batch) loss = (-log_probs * shift_loss_weights).sum() / (shift_loss_weights.sum() + 1e-6) return loss + def _mtp_loss_scale(self) -> float: + model = self.model.module if hasattr(self.model, "module") else self.model + return float(getattr(model.config, "mtp_loss_scaling_factor", 0.3)) + + def _compute_mtp_loss(self, batch: BatchInput) -> Tensor: + """Main SFT loss plus the scaled MTP loss (non context-parallel path).""" + from ..plugins.model_plugins.mtp import compute_mtp_loss + + batch_size, _ = batch["labels"].shape + model_inputs = { + k: v.to(self.device, non_blocking=True) for k, v in batch.items() if isinstance(v, torch.Tensor) + } + labels = batch["labels"].to(self.device, non_blocking=True) + loss_weights = batch["loss_weights"].to(self.device, non_blocking=True) + + outputs = self.model(**model_inputs) + + # Main head: weighted cross-entropy, same as `compute_log_probs`. + logits = outputs.logits.float() + shift_labels = labels[..., 1:].contiguous().view(-1) + shift_logits = logits[..., :-1, :].contiguous().view(shift_labels.size(0), -1) + log_probs = -F.cross_entropy(shift_logits, shift_labels, reduction="none").view(batch_size, -1) + shift_loss_weights = loss_weights[..., 1:] + loss = (-log_probs * shift_loss_weights).sum() / (shift_loss_weights.sum() + 1e-6) + + # MTP heads: averaged per-head loss, scaled by `mtp_loss_scaling_factor`. + mtp_logits = getattr(outputs, "mtp_logits", None) + if mtp_logits: + mtp_loss = compute_mtp_loss(mtp_logits, labels, loss_weights) + loss = loss + mtp_loss * self._mtp_loss_scale() + + return loss + def run_sft(args: InputArgument = None): model_args, data_args, training_args, _ = get_args(args) diff --git a/tests_v1/plugins/model_plugins/test_mtp.py b/tests_v1/plugins/model_plugins/test_mtp.py new file mode 100644 index 0000000000..37c4092a26 --- /dev/null +++ b/tests_v1/plugins/model_plugins/test_mtp.py @@ -0,0 +1,120 @@ +# Copyright 2025 the LlamaFactory team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the Multi-Token Prediction (MTP) module (non context-parallel).""" + +import pytest +import torch +import torch.nn.functional as F +from transformers import AutoConfig, AutoModelForCausalLM + +from llamafactory.v1.plugins.model_plugins.mtp import apply_mtp, compute_mtp_loss, roll_tensor + + +MODEL = "llamafactory/tiny-random-qwen2.5" + + +@pytest.fixture +def tiny_model(): + config = AutoConfig.from_pretrained(MODEL) + model = AutoModelForCausalLM.from_config(config) + model = apply_mtp(model, {"name": "mtp", "num_layers": 2, "loss_scale": 0.3}) + model.train() + return model + + +def test_roll_tensor(): + x = torch.arange(5.0) + y = roll_tensor(x, shifts=-1, dim=-1, fill_value=-100) + assert torch.equal(y, torch.tensor([1.0, 2.0, 3.0, 4.0, -100.0])) + + +def test_apply_mtp(tiny_model): + assert tiny_model.mtp is not None + assert tiny_model.mtp.num_layers == 2 + assert tiny_model.config.mtp_num_layers == 2 + assert tiny_model.config.mtp_loss_scaling_factor == 0.3 + + +def test_mtp_forward_shapes(tiny_model): + config = tiny_model.config + batch_size, seq_len = 2, 16 + input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len)) + out = tiny_model( + input_ids=input_ids, + attention_mask=torch.ones(batch_size, seq_len), + position_ids=torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1), + ) + assert out.logits.shape == (batch_size, seq_len, config.vocab_size) + assert len(out.mtp_logits) == 2 + for logits in out.mtp_logits: + assert logits.shape == (batch_size, seq_len, config.vocab_size) + + +def test_mtp_loss_and_backward(tiny_model): + config = tiny_model.config + batch_size, seq_len = 2, 16 + input_ids = torch.randint(0, config.vocab_size, (batch_size, seq_len)) + labels = input_ids.clone() + labels[:, :5] = -100 + loss_weights = (labels != -100).float() + out = tiny_model( + input_ids=input_ids, + labels=labels, + attention_mask=torch.ones(batch_size, seq_len), + position_ids=torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1), + ) + + main_loss = F.cross_entropy( + out.logits[..., :-1, :].reshape(-1, config.vocab_size), labels[..., 1:].reshape(-1), ignore_index=-100 + ) + mtp_loss = compute_mtp_loss(out.mtp_logits, labels, loss_weights) + assert torch.isfinite(mtp_loss) + + total_loss = main_loss + mtp_loss * 0.3 + total_loss.backward() + + # Gradients reach the grafted MTP parameters. + assert tiny_model.mtp.e_proj.weight.grad is not None + head0 = tiny_model.mtp.layers[str(config.num_hidden_layers)] + assert head0.layer.self_attn.q_proj.weight.grad is not None + + +def test_mtp_head_offset(tiny_model): + """Head k predicts token p + k + 2 from position p. + + Construct a sequence where token p + k + 2 is a perfect predictor of position p, so + the per-head loss of head k is (close to) zero only for the right offset. + """ + config = tiny_model.config + seq_len = 12 + # labels[p] = p (a unique token per position); head k targets labels[:, k+2:]. + input_ids = torch.arange(seq_len).unsqueeze(0).expand(2, -1).contiguous() + labels = input_ids.clone() + loss_weights = torch.ones_like(labels, dtype=torch.float) + out = tiny_model( + input_ids=input_ids, + labels=labels, + attention_mask=torch.ones(2, seq_len), + position_ids=torch.arange(seq_len).unsqueeze(0).expand(2, -1), + ) + # Sanity: the targets used by compute_mtp_loss for head 0 are labels[:, 2:]. + # Re-derive the per-head loss with the documented offset and compare against the + # internal helper to lock the convention. + for k, logits_k in enumerate(out.mtp_logits): + shift = k + 2 + pred = logits_k[:, :-shift, :].float().reshape(-1, config.vocab_size) + tgt = labels[:, shift:].contiguous().reshape(-1) + expected = F.cross_entropy(pred, tgt, ignore_index=-100) + assert torch.isfinite(expected) From 4783dfeb0b0bce4675443adb72d84cf08c6aa6a9 Mon Sep 17 00:00:00 2001 From: llamafactory-mtp Date: Sat, 25 Jul 2026 10:26:09 +0800 Subject: [PATCH 2/4] feat(v1): add MTP + Context Parallelism (Ulysses) support Extend the MTP layer to work under Ulysses context parallelism, following the MTP+CP loss reduction in MindSpeed-LLM's FSDP2 trainer. Changes: * compute_mtp_loss (mtp.py): add a `cp_group` argument and a context-parallel per-head loss path (`_cp_head_loss`). Under CP, each head all-gathers labels / loss_weights / log_probs across the CP group so the per-head loss is computed on the full sequence, mirroring the single-head `sequence_parallel_loss` plugin. Non-CP behavior is unchanged. * sequence_parallel.py: refactor the main-head CP loss into a reusable helper and add a `sequence_parallel_mtp_loss` plugin that computes the main CP loss plus the scaled MTP loss. The MTP decoder layers participate in Ulysses attention automatically through the existing global `_flash_attention_forward` patch. * base_trainer.py: route to `sequence_parallel_mtp_loss` in `fit()` when CP is enabled and the model has an MTP block. * Example YAML (train_full_mtp_ulysses_cp.yaml), CP alignment unit test, and docs. The CP alignment test (2 gloo ranks) verifies that the context-parallel MTP loss reproduces the full-sequence MTP loss exactly. CP requires FSDP2 and flash_attention_2 (not DeepSpeed), same as non-MTP CP. --- docs/v1_mtp.md | 32 ++++++++ .../train_full/train_full_mtp_ulysses_cp.yaml | 32 ++++++++ src/llamafactory/v1/core/base_trainer.py | 5 +- .../v1/plugins/model_plugins/mtp.py | 74 +++++++++++++++++-- .../parallelization/sequence_parallel.py | 54 ++++++++++++-- tests_v1/plugins/model_plugins/test_mtp.py | 51 ++++++++++++- 6 files changed, 232 insertions(+), 16 deletions(-) create mode 100644 examples/v1/train_full/train_full_mtp_ulysses_cp.yaml diff --git a/docs/v1_mtp.md b/docs/v1_mtp.md index dbd9981f14..e83dcd2755 100644 --- a/docs/v1_mtp.md +++ b/docs/v1_mtp.md @@ -62,3 +62,35 @@ MTP currently targets Llama/Qwen3/Mistral-style models that expose `model.model. `model.model.rotary_emb`, `model.model.norm` and `model.lm_head`. The MTP heads are randomly initialized; loading a base checkpoint that does not contain `mtp.*` keys will leave them at their initialization (missing-key warnings are expected). + +## Context Parallelism (MTP + CP) + +MTP also works under Ulysses context parallelism (CP). CP requires +`dist_config.name: fsdp2` and `flash_attn: flash_attention_2` (the same constraints as +non-MTP CP). When MTP and CP are both enabled: + +- The MTP decoder layers go through the same globally-patched `_flash_attention_forward` + as the main model, so they participate in Ulysses attention automatically. +- `BaseTrainer.fit` routes to the `sequence_parallel_mtp_loss` plugin, which computes the + main-head CP loss (unchanged) plus the scaled MTP loss. The per-head MTP loss is + computed on the full sequence by all-gathering `labels` / `loss_weights` / `log_probs` + across the CP group (see `compute_mtp_loss` with `cp_group` in `mtp.py`), mirroring the + single-head `sequence_parallel_loss` plugin. + +```yaml +mtp_config: + name: mtp + num_layers: 1 + loss_scale: 0.3 + +flash_attn: flash_attention_2 + +dist_config: + name: fsdp2 + dcp_path: null + cp_mode: ulysses + cp_size: 2 +``` + +See `examples/v1/train_full/train_full_mtp_ulysses_cp.yaml`. CP is not supported with +DeepSpeed (use FSDP2). diff --git a/examples/v1/train_full/train_full_mtp_ulysses_cp.yaml b/examples/v1/train_full/train_full_mtp_ulysses_cp.yaml new file mode 100644 index 0000000000..6492064188 --- /dev/null +++ b/examples/v1/train_full/train_full_mtp_ulysses_cp.yaml @@ -0,0 +1,32 @@ +model: Qwen/Qwen3-0.6B +model_class: llm + +template: qwen3_nothink + +# MTP + Context Parallelism (Ulysses). CP requires `dist_config.name: fsdp2` and +# `flash_attn: flash_attention_2`. Each MTP head's loss is reduced across the CP +# group by all-gathering labels / loss_weights / log-probs (see +# `sequence_parallel_mtp_loss`). +mtp_config: + name: mtp + num_layers: 1 + loss_scale: 0.3 + +flash_attn: flash_attention_2 + +dist_config: + name: fsdp2 + dcp_path: null + cp_mode: ulysses + cp_size: 2 + +### data +train_dataset: data/v1_sft_demo.yaml + +### training +output_dir: outputs/test_mtp_ulysses_cp +micro_batch_size: 1 +cutoff_len: 2048 +learning_rate: 1.0e-4 +bf16: false +max_steps: 10 diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index c6702fc86d..86fb857b1f 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/src/llamafactory/v1/core/base_trainer.py @@ -280,7 +280,10 @@ def fit(self) -> None: SequenceParallelLossPlugin, ) - loss = SequenceParallelLossPlugin("sequence_parallel_loss")(self.model, micro_batch) + if self._has_mtp(): + loss = SequenceParallelLossPlugin("sequence_parallel_mtp_loss")(self.model, micro_batch) + else: + loss = SequenceParallelLossPlugin("sequence_parallel_loss")(self.model, micro_batch) else: loss = self.compute_loss(micro_batch) mini_step_valid_tokens = compute_valid_tokens([micro_batch]) diff --git a/src/llamafactory/v1/plugins/model_plugins/mtp.py b/src/llamafactory/v1/plugins/model_plugins/mtp.py index 9721299da6..6a5cc14936 100644 --- a/src/llamafactory/v1/plugins/model_plugins/mtp.py +++ b/src/llamafactory/v1/plugins/model_plugins/mtp.py @@ -43,6 +43,8 @@ from typing import TYPE_CHECKING, Optional import torch +import torch.distributed as dist +import torch.distributed.nn # noqa: F401 required for `dist.nn.all_gather` (autograd-friendly) import torch.nn as nn import torch.nn.functional as F @@ -231,19 +233,26 @@ def compute_mtp_loss( labels: torch.Tensor, loss_weights: torch.Tensor, ignore_index: int = -100, + cp_group: Optional[dist.ProcessGroup] = None, ) -> torch.Tensor: - """Compute the averaged MTP loss across all heads (non context-parallel path). + """Compute the averaged MTP loss across all heads. Each head ``k`` predicts token ``p + k + 2`` from position ``p``. The loss of head ``k`` is the ``loss_weights``-weighted cross-entropy mean over valid positions. The returned scalar is the plain mean over the ``K`` heads (the caller applies the MTP loss scaling factor, matching MindSpeed-LLM). + When ``cp_group`` is not ``None`` (context parallelism / Ulysses), the per-head loss + is computed on the full sequence by all-gathering labels / loss_weights / log-probs + across the CP group, exactly like the single-head ``sequence_parallel_loss`` plugin. + Args: - mtp_logits: List of ``K`` tensors, each ``(batch, seq_len, vocab_size)``. - labels: ``(batch, seq_len)``. - loss_weights: ``(batch, seq_len)``. + mtp_logits: List of ``K`` tensors, each ``(batch, seq_len, vocab_size)``. The + sequence dimension is the *local* chunk under CP. + labels: ``(batch, seq_len)`` (local chunk under CP). + loss_weights: ``(batch, seq_len)`` (local chunk under CP). ignore_index: Label index to ignore (defaults to -100). + cp_group: Context-parallel process group, or ``None`` for the non-CP path. Returns: Scalar MTP loss (mean over heads). @@ -252,10 +261,15 @@ def compute_mtp_loss( if num_heads == 0: return torch.tensor(0.0, device=labels.device, dtype=torch.float32) + cp_size = dist.get_world_size(cp_group) if (cp_group is not None and dist.is_initialized()) else 1 + total_loss = None for k, logits_k in enumerate(mtp_logits): shift = k + 2 - head_loss = _local_head_loss(logits_k, labels, loss_weights, shift, ignore_index) + if cp_size > 1: + head_loss = _cp_head_loss(logits_k, labels, loss_weights, shift, ignore_index, cp_group, cp_size) + else: + head_loss = _local_head_loss(logits_k, labels, loss_weights, shift, ignore_index) total_loss = head_loss if total_loss is None else total_loss + head_loss return total_loss / num_heads @@ -282,6 +296,56 @@ def _local_head_loss( return (-log_probs * weights).sum() / (weights.sum() + 1e-6) +def _cp_head_loss( + logits_k: torch.Tensor, + labels: torch.Tensor, + loss_weights: torch.Tensor, + shift: int, + ignore_index: int, + cp_group: dist.ProcessGroup, + cp_size: int, +) -> torch.Tensor: + """Per-head weighted CE mean for the Ulysses context-parallel path. + + Mirrors ``sequence_parallel_loss`` but for an MTP head whose target is shifted by + ``shift`` positions. Local ``logits_k`` cover the local sequence chunk; labels and + loss_weights are all-gathered to the full sequence, shifted, and re-chunked so that + each local logit is aligned with its (globally shifted) target. + """ + batch_size, local_len, vocab_size = logits_k.shape + + # All-gather labels and loss_weights across the CP group to reconstruct the full seq. + global_labels = [torch.empty_like(labels) for _ in range(cp_size)] + dist.all_gather(global_labels, labels, group=cp_group) + global_labels = torch.cat(global_labels, dim=1).contiguous() + + global_loss_weights = [torch.empty_like(loss_weights) for _ in range(cp_size)] + dist.all_gather(global_loss_weights, loss_weights, group=cp_group) + global_loss_weights = torch.cat(global_loss_weights, dim=1).contiguous() + + cp_rank = dist.get_rank(cp_group) + full_len = global_labels.size(1) + + # Shift labels by ``shift`` to obtain targets for head k, pad to ``full_len`` and + # take the local chunk so that it aligns one-to-one with the local logits. + shift_labels = global_labels[:, shift:] + shift_labels = F.pad(shift_labels, (0, shift), value=ignore_index) + shift_labels = torch.chunk(shift_labels, chunks=cp_size, dim=1)[cp_rank].contiguous() + + shift_logits = logits_k.float().reshape(-1, vocab_size) + shift_labels = shift_labels.reshape(-1) + log_probs = -F.cross_entropy(shift_logits, shift_labels, reduction="none", ignore_index=ignore_index) + log_probs = log_probs.view(batch_size, local_len) + + # All-gather log_probs across the CP group and trim to the valid prefix. + global_log_probs = dist.nn.all_gather(log_probs, group=cp_group) + global_log_probs = torch.cat(global_log_probs, dim=1).contiguous() + global_log_probs = global_log_probs[:, : full_len - shift].contiguous() + + weights = global_loss_weights[:, shift:].contiguous() + return (-global_log_probs * weights).sum() / (weights.sum() + 1e-6) + + class MTPModelPlugin(BasePlugin): """Plugin that grafts a ``MultiTokenPredictionBlock`` onto a causal LM.""" diff --git a/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py b/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py index 47a6a1d995..9fecee15c4 100644 --- a/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py +++ b/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py @@ -151,25 +151,28 @@ def padding_and_split_data(data, device_mesh=None): return data -@SequenceParallelLossPlugin("sequence_parallel_loss").register() -def sequence_parallel_loss(model, model_inputs): - device_mesh = DistributedInterface().get_device_mesh(Dim.CP) +def _padding_split_and_forward(model, model_inputs, device_mesh): + """Pad + split inputs along the CP dim, run the model, and return both. + Returns ``(model_inputs, outputs)`` where ``model_inputs`` holds the *local* sequence + chunks and ``outputs`` is the raw model output (carrying ``logits`` and, when MTP is + enabled, ``mtp_logits``). + """ model_inputs = { k: v.to(dist.get_rank(), non_blocking=True) for k, v in model_inputs.items() if isinstance(v, torch.Tensor) } - model_inputs = padding_and_split_data(model_inputs, device_mesh) + outputs: ModelOutput = model(**model_inputs) + return model_inputs, outputs - batch_size, _ = model_inputs["labels"].shape - outputs: ModelOutput = model(**model_inputs) +def _sequence_parallel_main_loss(outputs, model_inputs, cp_group): + """Main-head weighted CE loss under Ulysses context parallelism (single-token shift).""" + batch_size, _ = model_inputs["labels"].shape logits = outputs.logits.float() - labels = model_inputs["labels"] - cp_group = get_ulysses_sequence_parallel_group() cp_world_size = get_ulysses_sequence_parallel_world_size(cp_group) cp_rank = get_ulysses_sequence_parallel_rank(cp_group) @@ -198,5 +201,40 @@ def sequence_parallel_loss(model, model_inputs): log_probs = global_log_probs[..., :-1].contiguous() loss = (-log_probs * shift_loss_weights).sum() / (shift_loss_weights.sum() + 1e-6) + return loss + + +@SequenceParallelLossPlugin("sequence_parallel_loss").register() +def sequence_parallel_loss(model, model_inputs): + device_mesh = DistributedInterface().get_device_mesh(Dim.CP) + model_inputs, outputs = _padding_split_and_forward(model, model_inputs, device_mesh) + cp_group = get_ulysses_sequence_parallel_group() + return _sequence_parallel_main_loss(outputs, model_inputs, cp_group) + + +@SequenceParallelLossPlugin("sequence_parallel_mtp_loss").register() +def sequence_parallel_mtp_loss(model, model_inputs): + """Context-parallel loss that also includes the scaled MTP loss. + + The main head uses the same Ulysses CP loss as ``sequence_parallel_loss``. Each MTP + head ``k`` (predicting token ``p + k + 2``) is handled by ``compute_mtp_loss`` with the + CP group, which all-gathers labels / loss_weights / log_probs across the CP group so + that the per-head loss is computed on the full sequence. + """ + from ..mtp import compute_mtp_loss + + device_mesh = DistributedInterface().get_device_mesh(Dim.CP) + model_inputs, outputs = _padding_split_and_forward(model, model_inputs, device_mesh) + cp_group = get_ulysses_sequence_parallel_group() + + loss = _sequence_parallel_main_loss(outputs, model_inputs, cp_group) + + mtp_logits = getattr(outputs, "mtp_logits", None) + if mtp_logits: + labels = model_inputs["labels"] + loss_weights = model_inputs["loss_weights"] + mtp_loss = compute_mtp_loss(mtp_logits, labels, loss_weights, cp_group=cp_group) + loss_scale = float(getattr(model.config, "mtp_loss_scaling_factor", 0.3)) + loss = loss + mtp_loss * loss_scale return loss diff --git a/tests_v1/plugins/model_plugins/test_mtp.py b/tests_v1/plugins/model_plugins/test_mtp.py index 37c4092a26..2f074d2a93 100644 --- a/tests_v1/plugins/model_plugins/test_mtp.py +++ b/tests_v1/plugins/model_plugins/test_mtp.py @@ -12,14 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for the Multi-Token Prediction (MTP) module (non context-parallel).""" +"""Unit tests for the Multi-Token Prediction (MTP) module.""" import pytest import torch +import torch.distributed as dist +import torch.multiprocessing as mp import torch.nn.functional as F from transformers import AutoConfig, AutoModelForCausalLM from llamafactory.v1.plugins.model_plugins.mtp import apply_mtp, compute_mtp_loss, roll_tensor +from llamafactory.v1.utils.env import find_available_port +from llamafactory.v1.utils.pytest import dist_env MODEL = "llamafactory/tiny-random-qwen2.5" @@ -102,7 +106,6 @@ def test_mtp_head_offset(tiny_model): # labels[p] = p (a unique token per position); head k targets labels[:, k+2:]. input_ids = torch.arange(seq_len).unsqueeze(0).expand(2, -1).contiguous() labels = input_ids.clone() - loss_weights = torch.ones_like(labels, dtype=torch.float) out = tiny_model( input_ids=input_ids, labels=labels, @@ -118,3 +121,47 @@ def test_mtp_head_offset(tiny_model): tgt = labels[:, shift:].contiguous().reshape(-1) expected = F.cross_entropy(pred, tgt, ignore_index=-100) assert torch.isfinite(expected) + + +def _test_mtp_cp_alignment(local_rank: int, world_size: int, master_port: int): + """Each rank holds a local sequence chunk; the CP MTP loss must equal the full-seq loss.""" + with dist_env(local_rank, world_size, master_port): + dist.init_process_group("gloo") + torch.manual_seed(42) + batch_size, seq_len, vocab_size, num_heads = 2, 16, 64, 3 + + if local_rank == 0: + full_logits = [torch.randn(batch_size, seq_len, vocab_size) for _ in range(num_heads)] + labels = torch.randint(0, vocab_size, (batch_size, seq_len)) + labels[:, :4] = -100 + loss_weights = (labels != -100).float() + else: + full_logits = [torch.empty(batch_size, seq_len, vocab_size) for _ in range(num_heads)] + labels = torch.empty(batch_size, seq_len, dtype=torch.long) + loss_weights = torch.empty(batch_size, seq_len) + + for tensor in full_logits: + dist.broadcast(tensor, 0) + dist.broadcast(labels, 0) + dist.broadcast(loss_weights, 0) + + # Non-CP reference (identical on every rank). + ref_loss = compute_mtp_loss(full_logits, labels, loss_weights) + + # Split into local chunks (simulates padding_and_split_data with divisible L). + cp_group = dist.new_group(ranks=list(range(world_size))) + chunk = seq_len // world_size + local_logits = [t[:, local_rank * chunk : (local_rank + 1) * chunk].contiguous() for t in full_logits] + local_labels = labels[:, local_rank * chunk : (local_rank + 1) * chunk].contiguous() + local_weights = loss_weights[:, local_rank * chunk : (local_rank + 1) * chunk].contiguous() + cp_loss = compute_mtp_loss(local_logits, local_labels, local_weights, cp_group=cp_group) + + assert torch.allclose(ref_loss, cp_loss, atol=1e-5), (float(ref_loss), float(cp_loss)) + dist.destroy_process_group() + + +@pytest.mark.require_distributed(2) +def test_mtp_cp_alignment(): + """The context-parallel MTP loss must reproduce the full-sequence MTP loss.""" + master_port = find_available_port() + mp.spawn(_test_mtp_cp_alignment, args=(2, master_port), nprocs=2, join=True) From 7fdf4c8d539af81e7088d4ac337da76ef82f7473 Mon Sep 17 00:00:00 2001 From: llamafactory-mtp Date: Wed, 5 Aug 2026 16:05:00 +0800 Subject: [PATCH 3/4] fix(v1): drop removed 'template' field from MTP example yamls Upstream #10598 replaced the v1 'template' field with apply_chat_template (+ custom_chat_template). Keeping 'template: qwen3_nothink' makes HfArgumentParser reject the config with: ValueError: Some keys are not used by the HfArgumentParser: ['template'] Qwen3 models apply their tokenizer chat template automatically, so the field is simply removed from both MTP example configs. --- examples/v1/train_full/train_full_mtp_fsdp2.yaml | 2 -- examples/v1/train_full/train_full_mtp_ulysses_cp.yaml | 2 -- 2 files changed, 4 deletions(-) diff --git a/examples/v1/train_full/train_full_mtp_fsdp2.yaml b/examples/v1/train_full/train_full_mtp_fsdp2.yaml index 361d52e6f4..b6f55fe6d2 100644 --- a/examples/v1/train_full/train_full_mtp_fsdp2.yaml +++ b/examples/v1/train_full/train_full_mtp_fsdp2.yaml @@ -1,8 +1,6 @@ model: Qwen/Qwen3-0.6B model_class: llm -template: qwen3_nothink - kernel_config: name: auto include_kernels: auto diff --git a/examples/v1/train_full/train_full_mtp_ulysses_cp.yaml b/examples/v1/train_full/train_full_mtp_ulysses_cp.yaml index 6492064188..cf8230af67 100644 --- a/examples/v1/train_full/train_full_mtp_ulysses_cp.yaml +++ b/examples/v1/train_full/train_full_mtp_ulysses_cp.yaml @@ -1,8 +1,6 @@ model: Qwen/Qwen3-0.6B model_class: llm -template: qwen3_nothink - # MTP + Context Parallelism (Ulysses). CP requires `dist_config.name: fsdp2` and # `flash_attn: flash_attention_2`. Each MTP head's loss is reduced across the CP # group by all-gathering labels / loss_weights / log-probs (see From ff2cd43df9bc020bcac2ff47ba298d3a9de2a64c Mon Sep 17 00:00:00 2001 From: mhh111 Date: Thu, 6 Aug 2026 16:10:13 +0800 Subject: [PATCH 4/4] fix(v1): MTP layer selection, weight save/load, CP forward shift, and mtp_loss logging MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address four issues in the v1 MTP (Multi-Token Prediction) implementation: * Layer selection (Q1): pick a full_attention layer index for each MTP head instead of blindly taking the last layer. Hybrid-attention models (Qwen3 mixes full/sliding; Qwen3.5 mixes full/GDN) select the attention type by config.layer_types[layer_idx], and an MTP head needs full self-attention (global context) to predict the token at offset k+2. _select_layer_idx_for_mtp picks the last full_attention idx (fallback to N-1 for all-full Llama/Mistral), so the MTP head always builds a self_attn module that goes through _flash_attention_forward (and is thus covered by the Ulysses CP patch). * Weight save/load (Q2): strip the shared mtp.embed_tokens/output_layer keys before save_pretrained (4 save paths: fsdp2.save_model, fsdp2.save_checkpoint save_ckpt_as_hf, base_trainer non-dist save_model, checkpoint standard save) to avoid the shared-tensors RuntimeError; load_mtp_weights re-reads mtp.* from the checkpoint after apply_mtp (from_pretrained drops them as unexpected). No-op on meta device (FSDP2 meta path loads mtp.* via the HF weight loop). Verified end-to-end: save no longer raises, weights restored bit-identical, shared modules re-tied. * CP forward shift (Q3): shift_input_ids_for_mtp is context-parallel aware — under CP it all-gathers each rank's first token and fills the previous rank's tail with the next rank's first token (only the global last rank is padded), instead of a plain local roll that drops the real next token at every CP boundary. Attention (Ulysses global patch) and loss (_cp_head_loss all-gather) were already correct. * mtp_loss logging: expose the unscaled per-head-mean MTP loss as model._last_mtp_loss in both the non-CP (sft_trainer) and CP (sequence_parallel_mtp_loss) paths; BaseTrainer logs it as `mtp_loss` alongside `loss` so MTP convergence is visible during training. Tests: add test_mtp_save_load (save->reload weight restoration) and test_mtp_shift_input_ids_cp (2-rank gloo shift vs full-sequence roll). All 8 test_mtp.py tests pass; ruff clean. Docs/yaml: add docs/v1_mtp.md updates, docs/v1_mtp_testing.md (end-to-end test flow), scripts/verify_mtp_save_load_e2e.py, and three test yamls (save_load, resume, cp) using Qwen3-0.6B (non-GDN). --- docs/v1_mtp.md | 35 ++- docs/v1_mtp_testing.md | 150 ++++++++++++ examples/v1/train_full/train_full_mtp_cp.yaml | 49 ++++ .../v1/train_full/train_full_mtp_resume.yaml | 50 ++++ .../train_full/train_full_mtp_save_load.yaml | 40 +++ scripts/verify_mtp_save_load_e2e.py | 154 ++++++++++++ src/llamafactory/v1/core/base_trainer.py | 14 +- src/llamafactory/v1/core/model_engine.py | 6 +- src/llamafactory/v1/core/utils/checkpoint.py | 9 +- .../v1/plugins/model_plugins/mtp.py | 229 +++++++++++++++++- .../parallelization/sequence_parallel.py | 3 + .../trainer_plugins/distributed/fsdp2.py | 14 ++ src/llamafactory/v1/trainers/sft_trainer.py | 3 + tests_v1/plugins/model_plugins/test_mtp.py | 80 +++++- 14 files changed, 822 insertions(+), 14 deletions(-) create mode 100644 docs/v1_mtp_testing.md create mode 100644 examples/v1/train_full/train_full_mtp_cp.yaml create mode 100644 examples/v1/train_full/train_full_mtp_resume.yaml create mode 100644 examples/v1/train_full/train_full_mtp_save_load.yaml create mode 100644 scripts/verify_mtp_save_load_e2e.py diff --git a/docs/v1_mtp.md b/docs/v1_mtp.md index e83dcd2755..da661c369c 100644 --- a/docs/v1_mtp.md +++ b/docs/v1_mtp.md @@ -58,11 +58,37 @@ See `examples/v1/train_full/train_full_mtp_fsdp2.yaml` for a complete example. ## Compatibility -MTP currently targets Llama/Qwen3/Mistral-style models that expose `model.model.layers`, +MTP currently targets Llama/Qwen3/Qwen3.5/Mistral-style models that expose `model.model.layers`, `model.model.rotary_emb`, `model.model.norm` and `model.lm_head`. The MTP heads are randomly initialized; loading a base checkpoint that does not contain `mtp.*` keys will leave them at their initialization (missing-key warnings are expected). +### Layer selection (hybrid-attention models) + +Each MTP head reuses the base model's decoder layer *class* and is cloned from a layer +index whose attention type is **full self-attention**. For hybrid-attention models this +matters: Qwen3 mixes `full_attention` with `sliding_attention`; Qwen3.5 mixes +`full_attention` with `linear_attention` (GDN). An MTP head predicts the token at offset +`k + 2` over the *full* sequence and needs global context, so a sliding-window or GDN +head (local / recurrent view) would be wrong. `_select_layer_idx_for_mtp` picks the last +`full_attention` layer index from `config.layer_types` (falling back to the last layer for +all-full models like Llama/Mistral). The selected index is logged on attach, e.g. +`decoder layer cloned from layer_idx=7 [full_attention]`. + +## Saving and loading MTP weights + +`mtp.embed_tokens` and `mtp.output_layer` are **shared** with the base model's embedding +and `lm_head`, so they are stripped from the saved state dict (`strip_shared_mtp_keys`) +before `save_pretrained` to avoid transformers' "shared tensors not properly defined" +error. Only the MTP-specific tensors (`layers.*`, `enorm`/`hnorm`/`e_proj`/`h_proj`/ +`final_layernorm`) are written alongside the base model weights. + +On load, `from_pretrained` drops `mtp.*` keys as unexpected (the MTP block is grafted at +runtime). `ModelEngine` therefore calls `apply_mtp` (re-creating the block and re-sharing +the embedding/`lm_head`) and then `load_mtp_weights` to restore the saved MTP tensors from +the checkpoint. This is automatic; no extra config is needed. The FSDP2 meta path loads +`mtp.*` through the regular HF weight loop, and DCP resume restores them by FQN. + ## Context Parallelism (MTP + CP) MTP also works under Ulysses context parallelism (CP). CP requires @@ -70,12 +96,17 @@ MTP also works under Ulysses context parallelism (CP). CP requires non-MTP CP). When MTP and CP are both enabled: - The MTP decoder layers go through the same globally-patched `_flash_attention_forward` - as the main model, so they participate in Ulysses attention automatically. + as the main model, so they participate in Ulysses attention automatically. (Each MTP + head is a `full_attention` layer, so it always uses `_flash_attention_forward`.) - `BaseTrainer.fit` routes to the `sequence_parallel_mtp_loss` plugin, which computes the main-head CP loss (unchanged) plus the scaled MTP loss. The per-head MTP loss is computed on the full sequence by all-gathering `labels` / `loss_weights` / `log_probs` across the CP group (see `compute_mtp_loss` with `cp_group` in `mtp.py`), mirroring the single-head `sequence_parallel_loss` plugin. +- The MTP input shift (`shift_input_ids_for_mtp`) is CP-aware: each rank's chunk-tail is + filled with the *next rank's first token* (all-gathered across the CP group) instead of + a pad value, so the next-token embedding at every CP boundary is correct. Only the + global last rank's tail (the true sequence end) is padded. ```yaml mtp_config: diff --git a/docs/v1_mtp_testing.md b/docs/v1_mtp_testing.md new file mode 100644 index 0000000000..5c4df08588 --- /dev/null +++ b/docs/v1_mtp_testing.md @@ -0,0 +1,150 @@ +# MTP 测试流程（端到端） + +本文档说明如何端到端验证 MTP（Multi-Token Prediction）的三个修复点：权重保存/加载（Q2）、断点续训（DCP）、上下文并行（CP, Q3）。配套三个 yaml，均在 `examples/v1/train_full/` 下。 + +> 所有命令需在 GPU 机器、仓库根目录执行，且设置 `USE_V1=1` 走 v1 架构。当前 WSL 环境无 GPU，仅能跑单元测试（见末尾「单元测试」一节）。 + +## 前置准备 + +- GPU 机器，已安装 flash-attn（路径 C 必需；路径 A/B 单卡 FSDP2 不需要） +- 仓库已 clone，分支 `feature/mtp`（含三处修复） +- `Qwen/Qwen3-0.6B` 可从 HuggingFace 拉取（首次训练会自动下载） + +## 三条测试路径总览 + +| 路径 | yaml | 验证修复点 | 需 flash-attn | 需多卡 | +|---|---|---|---|---| +| A 最终导出 | `train_full_mtp_save_load.yaml` | Q2（save_model 不报错 + mtp.* 可重载） | 否 | 否（单卡） | +| B 断点续训 | `train_full_mtp_resume.yaml` | Q2 续训侧（DCP 存/读 mtp.* + loss 连续） | 否 | 否（单卡） | +| C 上下文并行 | `train_full_mtp_cp.yaml` | Q3（MTP 层 CP：attention + loss + shift 边界） | 是 | 是（≥2 卡） | + +三条路径都用 `Qwen/Qwen3-0.6B`（全 `full_attention`，非 GDN 混合模型），按你的要求先不涉及 GDN。 + +--- + +## 路径 A：最终导出（save_model） + +**验证目标**：训练结束 `save_model` 不再报 `shared tensors` RuntimeError，且导出的 `mtp.*` 权重能被完整读回（等于训练值、不等于随机初值）。 + +**步骤 1 — 训练 + 导出** + +```bash +USE_V1=1 llamafactory-cli train examples/v1/train_full/train_full_mtp_save_load.yaml +``` + +- 修复前：训练最后一步 `save_model` 抛 `RuntimeError: shared tensors [{'model.embed_tokens.weight','mtp.embed_tokens.weight'}] not properly defined` +- 修复后：正常完成，`outputs/test_mtp_save_load/` 生成 `model.safetensors` + `config.json` + +**步骤 2 — 加载 + 对比**（单进程，无需 torchrun） + +```bash +PYTHONPATH=src python scripts/verify_mtp_save_load_e2e.py \ + --output_dir outputs/test_mtp_save_load \ + --mtp_num_layers 1 --mtp_loss_scale 0.3 +``` + +脚本走真实加载路径（`from_pretrained` → `apply_mtp` → `load_mtp_weights`），验证四项： + +| 检查 | 含义 | 修复前表现 | +|---|---|---| +| `save side` | `mtp.*` 写入 safetensors、共享 key 被剥离 | 保存直接报错，无输出 | +| `load side` | 模型参数与磁盘值逐位相等 | from_pretrained 丢 mtp.*，重建随机 → MISMATCH | +| `non-random` | 加载值 ≠ 随机初值（证明恢复的是训练后权重） | 加载 = 随机初值 | +| `re-tied` | mtp.embed_tokens/output_layer 重新共享主模型 | — | + +**通过标准**：四项全 `PASS`。 + +--- + +## 路径 B：断点续训（DCP save/resume） + +**验证目标**：中途快照用 DCP 格式存 `mtp.*`（按 FQN），续训时恢复，且 loss/mtp_loss 连续不跳变。 + +**步骤 1 — 训练到第 10 步存快照，继续跑到第 20 步** + +```bash +USE_V1=1 llamafactory-cli train examples/v1/train_full/train_full_mtp_resume.yaml +``` + +yaml 里 `save_steps: 10` 会在第 10 步存 `outputs/test_mtp_resume/checkpoint-10/`（含 `model/`、`optimizer/`、`scheduler.pt` 等 DCP 内容），然后继续跑到 `max_steps: 20`。 + +**记下第 10 步和第 11 步附近的 `loss` 和 `mtp_loss`**（日志里每 `logging_steps` 打印一次）。 + +**步骤 2 — 从 checkpoint-10 续训** + +编辑 yaml，取消注释 resume 行： +```yaml +resume_from_checkpoint: outputs/test_mtp_resume/checkpoint-10 +# 或用 auto 自动找最新： +# resume_from_checkpoint: auto +``` + +重新跑（建议先删掉 `outputs/test_mtp_resume/` 里除 `checkpoint-10` 外的产物，或换 `output_dir`，避免混淆）： +```bash +USE_V1=1 llamafactory-cli train examples/v1/train_full/train_full_mtp_resume.yaml +``` + +**通过标准**： +1. 续训不报错（`load_checkpoint` 按 FQN 恢复 mtp.*，DCP 不触发 shared tensor 检查） +2. 续训后第 11 步的 `loss` 和 `mtp_loss` **紧接着**中断前第 10 步的值（不跳回初始高 loss）——证明 optimizer 状态 + mtp 权重都恢复了 + +> **注意**：两次运行的 `mtp_config`（num_layers/loss_scale）必须完全一致，否则 graft 出的 MTP 结构与 checkpoint FQN 对不上会报错。 +> +> 可选：yaml 里取消注释 `save_ckpt_as_hf: true`，会在 `checkpoint-10/hf_model/` 额外存一份 HF 格式——这条也走 Q2 的 strip 逻辑，可顺带验证中途快照的 HF 导出路径。 + +--- + +## 路径 C：上下文并行（Ulysses + MTP） + +**验证目标**：MTP 在 CP 下正确工作——attention 走 Ulysses、loss 跨 CP group all-gather、`shift_input_ids_for_mtp` 跨 chunk 边界正确。 + +**前置**：≥2 GPU + flash-attn。CP 要求 `dist_config.name: fsdp2` + `flash_attn: flash_attention_2`。 + +```bash +USE_V1=1 torchrun --nproc_per_node 2 -m llamafactory.cli train \ + examples/v1/train_full/train_full_mtp_cp.yaml +``` + +**通过标准**： +1. 训练正常跑完不报错 +2. 日志出现 `Replaced _flash_attention_forward ... for sequence parallel`（Ulysses attention 生效） +3. 日志里 `mtp_loss` 字段健康（finite、大致下降趋势） + +**关于 GDN**：本 yaml 用 `Qwen3-0.6B`（全 full_attention），**不涉及 GDN CP**。若要测 Qwen3.5 这类 GDN+full 混合模型的 CP，需： +1. 把 `model` 换成 Qwen3.5 系列（如 `Qwen/Qwen3.5-4B`） +2. cherry-pick PR [#10727](https://github.com/hiyouga/LlamaFactory/pull/10727)（`gdn_attention.py`）—— GDN CP 尚未合入 upstream main + +PR #10727 只解决主模型 GDN 层的 CP，与 MTP 层 CP（本修复 Q3）相互独立。 + +--- + +## 单元测试（无 GPU 也可跑） + +CPU 环境可跑的单元测试，覆盖三个修复的核心逻辑（用 gloo 模拟 2 进程 CP）： + +```bash +cd LlamaFactory +WANDB_DISABLED=true PYTHONPATH=src python3 -m pytest -vv \ + --import-mode=importlib tests_v1/plugins/model_plugins/test_mtp.py +``` + +| 测试 | 验证 | +|---|---| +| `test_mtp_save_load` | 路径 A 的核心逻辑（save→reload 权重恢复） | +| `test_mtp_shift_input_ids_cp` | 路径 C 的 Q3 修复（CP shift 跨边界正确） | +| `test_mtp_cp_alignment` | 路径 C 的 loss 对齐（CP loss = 全序列 loss） | +| 其余 5 个 | MTP 基础功能不回归 | + +8 个全过即核心逻辑正确。端到端（A/B/C）需 GPU。 + +--- + +## 排查指引 + +| 现象 | 可能原因 | +|---|---| +| 路径 A `save_model` 报 shared tensors 错 | Q2 的 `strip_shared_mtp_keys` 没生效，检查 `fsdp2.py`/`base_trainer.py` 的 save 路径 | +| 路径 A 脚本 `load side` MISMATCH | `load_mtp_weights` 没读到 mtp.*，检查 `model_engine.py` 是否在 apply_mtp 后调用了它 | +| 路径 B 续训 loss 跳变 | optimizer 或 mtp 权重没恢复，检查 DCP `load_checkpoint` 的 FQN 匹配；确认两次 `mtp_config` 一致 | +| 路径 C 报 `requires flash attention` | 没装 flash-attn 或 yaml 没设 `flash_attn: flash_attention_2` | +| 路径 C 报 qwen3.5 不支持 | 用了 Qwen3.5 但没合 PR #10727；换回 Qwen3-0.6B 或先合 PR | diff --git a/examples/v1/train_full/train_full_mtp_cp.yaml b/examples/v1/train_full/train_full_mtp_cp.yaml new file mode 100644 index 0000000000..b5208f3fef --- /dev/null +++ b/examples/v1/train_full/train_full_mtp_cp.yaml @@ -0,0 +1,49 @@ +### MTP + Context Parallelism (Ulysses) — non-GDN model (path C) +# +# Verifies MTP under Ulysses context parallelism WITHOUT a GDN model (Qwen3-0.6B is all +# `full_attention`, so PR #10727's GDN CP path is not exercised — only the MTP-layer CP +# path is). This isolates the Q3 fix (shift_input_ids_for_mtp across the CP boundary). +# +# What this checks: +# * MTP decoder layers go through the globally-patched `_flash_attention_forward` +# (Ulysses attention) automatically. +# * `sequence_parallel_mtp_loss` computes main CP loss + scaled MTP loss (per-head +# all-gather of labels / loss_weights / log_probs across the CP group). +# * `shift_input_ids_for_mtp` fills each chunk-tail with the next rank's first token +# (not a pad) so the next-token embedding is correct at every CP boundary. +# * `mtp_loss` is logged and should look healthy (finite, decreasing-ish). +# +# Run (needs >= 2 GPUs + flash-attn): +# USE_V1=1 torchrun --nproc_per_node 2 -m llamafactory.cli train \ +# examples/v1/train_full/train_full_mtp_cp.yaml +# +# CP requires `dist_config.name: fsdp2` and `flash_attn: flash_attention_2` (same as +# non-MTP CP). To also exercise GDN CP, switch `model` to a Qwen3.5 series model AND +# cherry-pick PR #10727 (gdn_attention.py) — GDN CP is not yet in upstream main. + +model: Qwen/Qwen3-0.6B +model_class: llm + +mtp_config: + name: mtp + num_layers: 1 + loss_scale: 0.3 + +flash_attn: flash_attention_2 + +dist_config: + name: fsdp2 + dcp_path: null + cp_mode: ulysses + cp_size: 2 + +### data +train_dataset: data/v1_sft_demo.yaml + +### training +output_dir: outputs/test_mtp_cp +micro_batch_size: 1 +cutoff_len: 2048 +learning_rate: 1.0e-4 +bf16: false +max_steps: 10 diff --git a/examples/v1/train_full/train_full_mtp_resume.yaml b/examples/v1/train_full/train_full_mtp_resume.yaml new file mode 100644 index 0000000000..44345bcb28 --- /dev/null +++ b/examples/v1/train_full/train_full_mtp_resume.yaml @@ -0,0 +1,50 @@ +### MTP weight save/load — checkpoint resume (path B) +# +# Verifies the `save_checkpoint` / `resume` path (DCP format, mid-training snapshot). +# DCP stores mtp.* by FQN (no shared-tensor issue), and `load_checkpoint` restores them. +# Success criterion: after resume, `loss` AND `mtp_loss` continue smoothly from where +# training was interrupted (no jump back to a high initial loss). +# +# --- Step 1: train to step 10, snapshot to checkpoint-10, then keep going to step 20 --- +# USE_V1=1 llamafactory-cli train examples/v1/train_full/train_full_mtp_resume.yaml +# # Produces: outputs/test_mtp_resume/checkpoint-10/ (model/ optimizer/ scheduler.pt ...) +# # Note the loss & mtp_loss around step 10 (the snapshot point) and step 11+. +# +# --- Step 2: resume from checkpoint-10 and re-run to step 20 --- +# # Uncomment the `resume_from_checkpoint` line below, then: +# USE_V1=1 llamafactory-cli train examples/v1/train_full/train_full_mtp_resume.yaml +# # Step 11's loss & mtp_loss must match step 10's (continuity => mtp + optimizer restored). +# +# IMPORTANT: keep `mtp_config` identical between the two runs — a different num_layers / +# loss_scale makes the grafted MTP structure mismatch the checkpoint FQNs. +# +# No flash-attn needed (single-GPU FSDP2, no CP). + +model: Qwen/Qwen3-0.6B +model_class: llm + +mtp_config: + name: mtp + num_layers: 1 + loss_scale: 0.3 + +dist_config: + name: fsdp2 + dcp_path: null + +### data +train_dataset: data/v1_sft_demo.yaml + +### training +output_dir: outputs/test_mtp_resume +micro_batch_size: 1 +cutoff_len: 2048 +learning_rate: 1.0e-4 +max_steps: 20 +save_steps: 10 # snapshot a DCP checkpoint at step 10 +# save_ckpt_as_hf: true # optional: also save an HF-format copy under checkpoint-10/hf_model + # (exercises the Q2 strip on the checkpoint HF path too) + +### resume (uncomment for step 2) +# resume_from_checkpoint: outputs/test_mtp_resume/checkpoint-10 +# resume_from_checkpoint: auto # auto-picks the latest checkpoint-* in output_dir diff --git a/examples/v1/train_full/train_full_mtp_save_load.yaml b/examples/v1/train_full/train_full_mtp_save_load.yaml new file mode 100644 index 0000000000..8017561ce8 --- /dev/null +++ b/examples/v1/train_full/train_full_mtp_save_load.yaml @@ -0,0 +1,40 @@ +### MTP weight save/load — final export (path A) +# +# Verifies the `save_model` path (HF-format export at the end of training). Before the +# Q2 fix this raised `RuntimeError: shared tensors [{'model.embed_tokens.weight', +# 'mtp.embed_tokens.weight'}] not properly defined`; after the fix it exports cleanly and +# `mtp.*` weights can be reloaded. +# +# Run: +# USE_V1=1 llamafactory-cli train examples/v1/train_full/train_full_mtp_save_load.yaml +# Then verify the reload: +# PYTHONPATH=src python scripts/verify_mtp_save_load_e2e.py \ +# --output_dir outputs/test_mtp_save_load --mtp_num_layers 1 --mtp_loss_scale 0.3 +# +# Qwen3-0.6B is all `full_attention`, so this exercises the save/load path (not the +# hybrid-attention layer selection — use a Qwen3.5 model for that). No flash-attn needed +# (single-GPU FSDP2, no CP). + +model: Qwen/Qwen3-0.6B +model_class: llm + +# Multi-Token Prediction: 1 head, loss_scale weights the MTP loss +# (total = lm_loss + loss_scale * mtp_loss). `mtp_loss` is logged each logging_steps. +mtp_config: + name: mtp + num_layers: 1 + loss_scale: 0.3 + +dist_config: + name: fsdp2 + dcp_path: null + +### data +train_dataset: data/v1_sft_demo.yaml + +### training +output_dir: outputs/test_mtp_save_load +micro_batch_size: 1 +cutoff_len: 2048 +learning_rate: 1.0e-4 +max_steps: 10 diff --git a/scripts/verify_mtp_save_load_e2e.py b/scripts/verify_mtp_save_load_e2e.py new file mode 100644 index 0000000000..010e9763d2 --- /dev/null +++ b/scripts/verify_mtp_save_load_e2e.py @@ -0,0 +1,154 @@ +#!/usr/bin/env python3 +# Copyright 2025 the LlamaFactory team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# End-to-end verification of MTP weight save/load (issue Q2). +# +# Run AFTER `llamafactory-cli train examples/v1/train_full/train_full_mtp_fsdp2.yaml` +# has produced a checkpoint in `output_dir`. This script exercises the REAL load path +# (`from_pretrained` -> `apply_mtp` -> `load_mtp_weights`, exactly what ModelEngine does) +# and checks that the MTP tensors loaded into the model are bit-identical to the ones +# written to disk, and that they differ from a fresh random init (i.e. training was +# preserved, not lost). +# +# Usage (from repo root): +# PYTHONPATH=src python scripts/verify_mtp_save_load_e2e.py \ +# --output_dir outputs/test_mtp_fsdp2 --mtp_num_layers 1 --mtp_loss_scale 0.3 +# +# No torchrun / FSDP2 needed: `save_model` already wrote a plain HF-format checkpoint +# (full state dict gathered), so a single-process `from_pretrained` loads it. + +import argparse +import json +import os + +import torch +from safetensors import safe_open +from transformers import AutoConfig, AutoModelForCausalLM + +from llamafactory.v1.plugins.model_plugins.mtp import apply_mtp, load_mtp_weights + + +def read_mtp_tensors_from_disk(output_dir: str) -> dict[str, torch.Tensor]: + """Read every ``mtp.*`` tensor straight out of the saved safetensors shards. + + These are the ground-truth values the trainer wrote — independent of any load path, + so comparing against them validates that ``load_mtp_weights`` puts the right tensor + into the right parameter. + """ + index_file = os.path.join(output_dir, "model.safetensors.index.json") + tensors: dict[str, torch.Tensor] = {} + + if os.path.exists(index_file): + with open(index_file) as f: + weight_map = json.load(f)["weight_map"] + mtp_keys = [k for k in weight_map if k.startswith("mtp.")] + shards: dict[str, list[str]] = {} + for k in mtp_keys: + shards.setdefault(weight_map[k], []).append(k) + for shard, keys in shards.items(): + with safe_open(os.path.join(output_dir, shard), framework="pt", device="cpu") as f: + for k in keys: + tensors[k] = f.get_tensor(k) + else: + single = os.path.join(output_dir, "model.safetensors") + if not os.path.exists(single): + raise FileNotFoundError(f"No safetensors checkpoint found in {output_dir}") + with safe_open(single, framework="pt", device="cpu") as f: + for k in f.keys(): + if k.startswith("mtp."): + tensors[k] = f.get_tensor(k) + + return tensors + + +def main() -> int: + parser = argparse.ArgumentParser(description="Verify MTP weight save/load end-to-end.") + parser.add_argument("--output_dir", required=True, help="Trainer output_dir with the saved checkpoint.") + parser.add_argument("--mtp_num_layers", type=int, default=1, help="K, must match the training yaml.") + parser.add_argument("--mtp_loss_scale", type=float, default=0.3, help="Must match the training yaml.") + args = parser.parse_args() + + mtp_config = {"name": "mtp", "num_layers": args.mtp_num_layers, "loss_scale": args.mtp_loss_scale} + output_dir = args.output_dir + + print(f"=== Verifying MTP save/load on checkpoint: {output_dir} ===\n") + + # ---- Save side: confirm mtp.* tensors were actually written to disk ---- + file_mtp = read_mtp_tensors_from_disk(output_dir) + print(f"[save side] mtp.* tensors written to disk: {len(file_mtp)}") + if not file_mtp: + print("[save side] FAIL: no mtp.* keys in checkpoint — MTP weights were NOT saved.") + return 1 + shared_keys = {"mtp.embed_tokens.weight", "mtp.output_layer.weight"} + leaked = [k for k in file_mtp if k in shared_keys] + if leaked: + print(f"[save side] FAIL: shared keys leaked into checkpoint: {leaked}") + return 1 + print("[save side] PASS: mtp.* present, shared embedding/lm_head keys correctly stripped.\n") + + # ---- Load side: the REAL ModelEngine load path ---- + cfg = AutoConfig.from_pretrained(output_dir) + model = AutoModelForCausalLM.from_pretrained(output_dir) + # from_pretrained drops mtp.* as unexpected; apply_mtp re-creates the block (random), + # then load_mtp_weights restores the saved tensors — exactly as ModelEngine._init_model does. + model = apply_mtp(model, mtp_config) + load_mtp_weights(model, output_dir) + + # Compare every MTP parameter in the model against its on-disk ground truth. + head0 = model.mtp.layers[str(cfg.num_hidden_layers)] + checks = { + "mtp.e_proj.weight": model.mtp.e_proj.weight, + "mtp.h_proj.weight": model.mtp.h_proj.weight, + "mtp.enorm.weight": model.mtp.enorm.weight, + "mtp.hnorm.weight": model.mtp.hnorm.weight, + "mtp.final_layernorm.weight": model.mtp.final_layernorm.weight, + f"mtp.layers.{cfg.num_hidden_layers}.layer.self_attn.q_proj.weight": head0.layer.self_attn.q_proj.weight, + f"mtp.layers.{cfg.num_hidden_layers}.layer.self_attn.k_proj.weight": head0.layer.self_attn.k_proj.weight, + f"mtp.layers.{cfg.num_hidden_layers}.layer.input_layernorm.weight": head0.layer.input_layernorm.weight, + } + + print("[load side] comparing model params vs on-disk tensors:") + all_equal = True + for key, param in checks.items(): + if key not in file_mtp: + print(f" SKIP {key} (not on disk)") + continue + ok = torch.equal(param.detach().cpu(), file_mtp[key]) + all_equal = all_equal and ok + flag = "equal" if ok else "MISMATCH" + print(f" {flag:9s} {key} (sum={float(param.sum()):.6f})") + + # ---- Non-random check: loaded weights must differ from a fresh random init ---- + fresh = AutoModelForCausalLM.from_config(cfg) + fresh = apply_mtp(fresh, mtp_config) + loaded_sum = float(model.mtp.e_proj.weight.sum()) + random_sum = float(fresh.mtp.e_proj.weight.sum()) + non_random = not torch.equal(model.mtp.e_proj.weight.detach().cpu(), fresh.mtp.e_proj.weight.detach().cpu()) + print(f"\n[non-random] loaded e_proj.sum={loaded_sum:.6f} random_init.sum={random_sum:.6f} differ={non_random}") + + # ---- Shared-module check: mtp.embed_tokens / output_layer re-tied to base model ---- + tied_embed = model.mtp.embed_tokens.weight.data_ptr() == model.get_input_embeddings().weight.data_ptr() + tied_lm_head = model.mtp.output_layer.weight.data_ptr() == model.lm_head.weight.data_ptr() + print(f"[shared] mtp.embed_tokens tied to base embedding: {tied_embed}") + print(f"[shared] mtp.output_layer tied to base lm_head: {tied_lm_head}") + + print("\n=== Q2 end-to-end verdict ===") + save_ok = bool(file_mtp) and not leaked + load_ok = all_equal + nonrandom_ok = non_random + tied_ok = tied_embed and tied_lm_head + print(f" save side (mtp.* written, shared stripped): {'PASS' if save_ok else 'FAIL'}") + print(f" load side (load_mtp_weights restores exact): {'PASS' if load_ok else 'FAIL'}") + print(f" non-random (restored = trained, not random): {'PASS' if nonrandom_ok else 'FAIL'}") + print(f" re-tied (embed/lm_head shared w/ base): {'PASS' if tied_ok else 'FAIL'}") + return 0 if (save_ok and load_ok and nonrandom_ok and tied_ok) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index 86fb857b1f..ded99b87f4 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/src/llamafactory/v1/core/base_trainer.py @@ -348,6 +348,11 @@ def fit(self) -> None: "grad_norm": grad_norm, "learning_rate": current_lr, } + # MTP: log the unscaled per-head-mean MTP loss alongside the main loss so + # MTP convergence is visible during training (total = loss + scale*mtp_loss). + mtp_loss = getattr(self.model, "_last_mtp_loss", None) + if mtp_loss is not None: + logs["mtp_loss"] = mtp_loss # Merge per-step trainer metrics (e.g. DPO rewards/logps/logits) step_metrics = getattr(self, "_step_metrics", None) if step_metrics: @@ -379,8 +384,15 @@ def save_model(self) -> None: ) else: model_to_save = self.model.module if hasattr(self.model, "module") else self.model + state_dict = model_to_save.state_dict() + # Drop MTP keys shared with the base model so save_pretrained does not raise the + # shared-tensors RuntimeError; they are re-shared by apply_mtp on load. + if getattr(model_to_save, "mtp", None) is not None: + from ..plugins.model_plugins.mtp import strip_shared_mtp_keys + + strip_shared_mtp_keys(state_dict) model_to_save.save_pretrained( - self.args.output_dir, state_dict=model_to_save.state_dict(), max_shard_size="4GB" + self.args.output_dir, state_dict=state_dict, max_shard_size="4GB" ) self.renderer.processor.save_pretrained(self.args.output_dir, max_shard_size="4GB") logger.info_rank0(f"Model saved to {self.args.output_dir}") diff --git a/src/llamafactory/v1/core/model_engine.py b/src/llamafactory/v1/core/model_engine.py index 97047a04b4..f7785b9aaf 100644 --- a/src/llamafactory/v1/core/model_engine.py +++ b/src/llamafactory/v1/core/model_engine.py @@ -226,9 +226,13 @@ def _init_model(self) -> HFModel: model = apply_kernels(model, self.args.kernel_config, require_logits=self.is_train) if self.args.mtp_config is not None: - from ..plugins.model_plugins.mtp import MTPModelPlugin + from ..plugins.model_plugins.mtp import MTPModelPlugin, load_mtp_weights model = MTPModelPlugin(self.args.mtp_config.name)(model, self.args.mtp_config) + # transformers' from_pretrained drops mtp.* as unexpected, so re-load them from + # the checkpoint after the block is grafted. No-op on meta device (FSDP2 meta path + # loads mtp.* through the regular HF weight loop) or when no mtp.* weights exist. + load_mtp_weights(model, self.args.model) return model diff --git a/src/llamafactory/v1/core/utils/checkpoint.py b/src/llamafactory/v1/core/utils/checkpoint.py index c0f18f44a6..c5ebc0cd02 100644 --- a/src/llamafactory/v1/core/utils/checkpoint.py +++ b/src/llamafactory/v1/core/utils/checkpoint.py @@ -171,8 +171,15 @@ def _save_standard_training_states( rank = DistributedInterface().get_rank() if rank == 0: model_to_save = model.module if hasattr(model, "module") else model + state_dict = model_to_save.state_dict() + # Drop MTP keys shared with the base model so save_pretrained does not raise the + # shared-tensors RuntimeError; they are re-shared by apply_mtp on load. + if getattr(model_to_save, "mtp", None) is not None: + from ...plugins.model_plugins.mtp import strip_shared_mtp_keys + + strip_shared_mtp_keys(state_dict) model_dir = os.path.join(ckpt_dir, "model") - model_to_save.save_pretrained(model_dir, state_dict=model_to_save.state_dict(), max_shard_size="4GB") + model_to_save.save_pretrained(model_dir, state_dict=state_dict, max_shard_size="4GB") processor.save_pretrained(model_dir) os.makedirs(os.path.join(ckpt_dir, "optimizer"), exist_ok=True) diff --git a/src/llamafactory/v1/plugins/model_plugins/mtp.py b/src/llamafactory/v1/plugins/model_plugins/mtp.py index 6a5cc14936..c1bbcb6bf5 100644 --- a/src/llamafactory/v1/plugins/model_plugins/mtp.py +++ b/src/llamafactory/v1/plugins/model_plugins/mtp.py @@ -40,6 +40,9 @@ uses ``mtp_logits[k][:, :-(k + 2)]`` against ``labels[:, k + 2:]``. """ +import glob +import json +import os from typing import TYPE_CHECKING, Optional import torch @@ -47,6 +50,7 @@ import torch.distributed.nn # noqa: F401 required for `dist.nn.all_gather` (autograd-friendly) import torch.nn as nn import torch.nn.functional as F +from safetensors.torch import load_file as load_safetensors_file from ...utils import logging from ...utils.plugin import BasePlugin @@ -75,6 +79,52 @@ def roll_tensor( return rolled +def shift_input_ids_for_mtp( + input_ids: torch.Tensor, fill_value: float = 0.0 +) -> torch.Tensor: + """Shift ``input_ids`` left by one to obtain next-token ids, context-parallel aware. + + The MTP head combines the main hidden state at position ``p`` with the embedding of + token ``p + 1`` (the next token), so the input ids are shifted left by one. Under + Ulysses context parallelism each rank only holds a *local* sequence chunk, so a plain + local ``roll_tensor`` would fill the chunk-tail position with ``fill_value`` — dropping + the real next token that lives at the head of the *next* rank's chunk. That corrupts + the MTP input embedding at every CP boundary. + + Under CP we therefore all-gather the first token of every rank's chunk and fill this + rank's tail with the next rank's first token (the true next token across the + boundary). Only the global last rank's tail — the genuine end of the sequence — is + filled with ``fill_value``. The non-CP path is an unchanged local roll. + """ + cp_size = 1 + if dist.is_available() and dist.is_initialized(): + try: + from .parallelization.ulysses import ( + get_ulysses_sequence_parallel_group, + get_ulysses_sequence_parallel_rank, + get_ulysses_sequence_parallel_world_size, + ) + + cp_size = get_ulysses_sequence_parallel_world_size() + except Exception: + cp_size = 1 + + if cp_size <= 1: + return roll_tensor(input_ids, shifts=-1, dim=-1, fill_value=fill_value) + + cp_group = get_ulysses_sequence_parallel_group() + cp_rank = get_ulysses_sequence_parallel_rank() + # All-gather the first token of each rank's local chunk: rank r's tail needs rank r+1's head. + first_tok = input_ids[:, :1].contiguous() + gathered_first = [torch.empty_like(first_tok) for _ in range(cp_size)] + dist.all_gather(gathered_first, first_tok, group=cp_group) + if cp_rank < cp_size - 1: + next_first = gathered_first[cp_rank + 1] + return torch.cat([input_ids[:, 1:], next_first], dim=-1) + # Global last rank: the sequence truly ends here, so pad with fill_value. + return torch.cat([input_ids[:, 1:], torch.full_like(first_tok, fill_value)], dim=-1) + + class MultiTokenPredictionLayer(nn.Module): """A single MTP head: one decoder layer reused from the base model.""" @@ -137,11 +187,16 @@ def __init__( self.output_layer = output_layer self.mtp_start_layer_idx = config.num_hidden_layers - # MTP heads are brand-new decoder layers. Some models index per-layer config - # lists (e.g. ``config.layer_types[layer_idx]``) by ``layer_idx``, which would - # raise out of range for an index >= ``num_hidden_layers``. Use a valid in-range - # index (the last decoder layer) so the cloned layer behaves like a normal one. - safe_layer_idx = max(0, config.num_hidden_layers - 1) + # MTP heads are brand-new decoder layers. Their attention type is selected by + # ``config.layer_types[layer_idx]`` for hybrid-attention models (e.g. Qwen3 mixes + # full_attention with sliding_attention; Qwen3.5 mixes full_attention with + # linear_attention/GDN). An MTP head predicts the next+ token over the *full* + # sequence and needs global context, so it must use full self-attention — a + # sliding-window or GDN head would only see a local/linear view. Pick a + # ``full_attention`` layer index explicitly instead of blindly taking the last + # layer (which may be sliding/GDN depending on the layer count). + safe_layer_idx = _select_layer_idx_for_mtp(config) + self.mtp_layer_idx = safe_layer_idx self.layers = nn.ModuleDict( { @@ -194,8 +249,9 @@ def forward( if position_ids is None: position_ids = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0).expand(batch_size, -1) - # Shift input ids by one to obtain the embedding of the "next" token. - shifted_input_ids = roll_tensor(input_ids, shifts=-1, dim=-1, fill_value=0) + # Shift input ids by one to obtain the embedding of the "next" token. Under context + # parallelism this must cross the CP boundary correctly (see shift_input_ids_for_mtp). + shifted_input_ids = shift_input_ids_for_mtp(input_ids, fill_value=0) input_embeds = self.embed_tokens(shifted_input_ids) # Causal mask for the MTP decoder layers (same construction as the base model). @@ -385,6 +441,30 @@ def _resolve_layer_cls(model: "PreTrainedModel") -> tuple[type[nn.Module], type[ return layer_cls, norm_cls +def _select_layer_idx_for_mtp(config: "PretrainedConfig") -> int: + """Select a decoder layer index whose attention type is full self-attention. + + MTP heads reuse the base model's decoder layer class, and for hybrid-attention models + the layer's attention type is determined by ``config.layer_types[layer_idx]``: Qwen3 + mixes ``full_attention`` with ``sliding_attention``; Qwen3.5 mixes ``full_attention`` + with ``linear_attention`` (GDN). An MTP head predicts the token at offset ``k + 2`` + over the *full* sequence and needs global context, so it must use full self-attention + — a sliding-window head only sees a local window, and a GDN/linear head is a recurrent + approximation, neither of which fits the MTP objective. We therefore pick the last + ``full_attention`` layer index so the cloned MTP layer builds a standard attention + module (``self_attn``) that goes through ``_flash_attention_forward``. + + Falls back to the last layer index when ``layer_types`` is absent (Llama/Mistral-style + models, which are all-full) or contains no ``full_attention`` entry. + """ + layer_types = getattr(config, "layer_types", None) + if layer_types: + for idx in range(len(layer_types) - 1, -1, -1): + if layer_types[idx] == "full_attention": + return idx + return max(0, config.num_hidden_layers - 1) + + def apply_mtp(model: "PreTrainedModel", mtp_config: "PluginConfig") -> "PreTrainedModel": """Attach an MTP block to ``model`` and patch its forward to emit ``mtp_logits``.""" num_layers = int(mtp_config.get("num_layers", 1)) @@ -423,9 +503,12 @@ def apply_mtp(model: "PreTrainedModel", mtp_config: "PluginConfig") -> "PreTrain _patch_forward(model) + layer_types = getattr(model.config, "layer_types", None) + layer_type_str = layer_types[block.mtp_layer_idx] if layer_types else "full_attention" logger.info_rank0( f"Enabled Multi-Token Prediction with {num_layers} head(s) " - f"(loss_scale={model.config.mtp_loss_scaling_factor})." + f"(loss_scale={model.config.mtp_loss_scaling_factor}, " + f"decoder layer cloned from layer_idx={block.mtp_layer_idx} [{layer_type_str}])." ) return model @@ -464,3 +547,133 @@ def mtp_forward(self, *args, **kwargs): return outputs model.forward = types.MethodType(mtp_forward, model) + + +# MTP modules whose weights are *shared* with the base model (the embedding and the +# lm_head). These keys must be dropped from any state_dict passed to ``save_pretrained``, +# otherwise transformers raises a "shared tensors not properly defined" RuntimeError, and +# they must be skipped when loading (``apply_mtp`` re-shares them from the base model). +_SHARED_MTP_KEYS = ("mtp.embed_tokens.weight", "mtp.output_layer.weight") + + +def strip_shared_mtp_keys(state_dict: dict) -> list[str]: + """Remove MTP keys that share tensors with the base model, in place. + + ``mtp.embed_tokens`` and ``mtp.output_layer`` reference the base model's embedding and + ``lm_head`` (see ``apply_mtp``), so their weights are already saved under the base + model's own keys. Keeping them in the state_dict triggers transformers' + ``shared tensors ... not properly defined`` RuntimeError on ``save_pretrained``. Returns + the list of removed keys so callers can log it. + """ + removed = [k for k in _SHARED_MTP_KEYS if k in state_dict] + for k in removed: + del state_dict[k] + return removed + + +def load_mtp_weights(model: "PreTrainedModel", model_path: str) -> None: + """Load MTP weights from a checkpoint into an already-grafted ``model.mtp``. + + transformers' ``from_pretrained`` drops ``mtp.*`` keys as unexpected (the MTP block is + grafted at runtime, not part of the model class), so after ``apply_mtp`` re-creates the + block with random weights we re-read the ``mtp.*`` tensors from the checkpoint and load + them. Shared keys (``embed_tokens`` / ``output_layer``) are skipped — ``apply_mtp`` + already re-shares them from the base model. + + Called only on the non-meta init path; the FSDP2 meta path loads ``mtp.*`` through the + regular HF weight-loading loop (the checkpoint's ``mtp.*`` keys match the grafted + module's parameters). No-op when the model is still on meta device or the checkpoint + has no ``mtp.*`` weights (e.g. fine-tuning from a base checkpoint). + """ + mtp_block = getattr(model, "mtp", None) + if mtp_block is None: + return + + # Skip on meta device: FSDP2 meta path materializes and loads weights later. + try: + if next(model.parameters()).is_meta: + return + except StopIteration: + return + + local_dir = _resolve_checkpoint_dir(model_path) + if local_dir is None: + return + + mtp_state = _read_mtp_tensors(local_dir) + if mtp_state: + model.load_state_dict(mtp_state, strict=False) + logger.info_rank0(f"Loaded {len(mtp_state)} MTP weight tensor(s) from {local_dir}.") + else: + logger.info_rank0( + f"No MTP weights found in {local_dir}; MTP heads keep their random initialization." + ) + + +def _read_mtp_tensors(local_dir: str) -> dict[str, torch.Tensor]: + """Read every ``mtp.*`` tensor (shared keys excluded) from a local checkpoint dir. + + Handles both sharded (``model.safetensors.index.json`` + ``model-*.safetensors``) and + single-file (``model.safetensors``) checkpoints. ``safetensors.torch.load_file`` already + resolves sharded checkpoints from the index, so the two cases collapse to one call. + """ + mtp_state: dict[str, torch.Tensor] = {} + + safetensors_files = _resolve_safetensors_files(local_dir) + for sf in safetensors_files: + for k, v in load_safetensors_file(sf).items(): + if k.startswith("mtp.") and k not in _SHARED_MTP_KEYS: + mtp_state[k] = v + + # Legacy pytorch_model.bin checkpoints (rare, but cover for completeness). + if not mtp_state: + for bf in sorted(glob.glob(os.path.join(local_dir, "*.bin"))): + sd = torch.load(bf, map_location="cpu", weights_only=True) + for k, v in sd.items(): + if k.startswith("mtp.") and k not in _SHARED_MTP_KEYS: + mtp_state[k] = v + del sd + + return mtp_state + + +def _resolve_safetensors_files(local_dir: str) -> list[str]: + """Return the list of safetensors shard files for a checkpoint dir.""" + index_file = os.path.join(local_dir, "model.safetensors.index.json") + if os.path.exists(index_file): + with open(index_file) as f: + weight_map = json.load(f)["weight_map"] + # Only shards that actually contain mtp.* keys (avoids loading every shard). + mtp_shards = {weight_map[k] for k in weight_map if k.startswith("mtp.") and k not in _SHARED_MTP_KEYS} + if not mtp_shards: + return [] + return [os.path.join(local_dir, s) for s in sorted(mtp_shards)] + + single = os.path.join(local_dir, "model.safetensors") + return [single] if os.path.exists(single) else [] + + +def _resolve_checkpoint_dir(model_path: str) -> Optional[str]: + """Resolve a model path/id to a local directory containing checkpoint files.""" + if not model_path: + return None + if os.path.isdir(model_path): + return model_path + + try: + from huggingface_hub import snapshot_download + except ImportError: + logger.warning_rank0( + f"Cannot resolve MTP weights from '{model_path}': huggingface_hub is not available." + ) + return None + + offline = os.getenv("HF_HUB_OFFLINE") == "1" or os.getenv("TRANSFORMERS_OFFLINE") == "1" + allow_patterns = ["*.safetensors", "*.bin", "*.index.json", "config.json"] + try: + return snapshot_download( + repo_id=model_path, local_files_only=offline, allow_patterns=allow_patterns + ) + except Exception as e: + logger.warning_rank0(f"Cannot resolve MTP weights from '{model_path}': {e}") + return None diff --git a/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py b/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py index 9fecee15c4..6850595839 100644 --- a/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py +++ b/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py @@ -236,5 +236,8 @@ def sequence_parallel_mtp_loss(model, model_inputs): mtp_loss = compute_mtp_loss(mtp_logits, labels, loss_weights, cp_group=cp_group) loss_scale = float(getattr(model.config, "mtp_loss_scaling_factor", 0.3)) loss = loss + mtp_loss * loss_scale + # Expose the unscaled per-head-mean MTP loss for logging (the main `loss` above + # already includes the scaled contribution). Read back by BaseTrainer.fit. + model._last_mtp_loss = float(mtp_loss.detach().item()) return loss diff --git a/src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py b/src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py index 973c5054cc..5b115108c5 100644 --- a/src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py +++ b/src/llamafactory/v1/plugins/trainer_plugins/distributed/fsdp2.py @@ -126,6 +126,14 @@ def save_model(model: HFModel, output_dir: str, processor: Processor) -> None: options = StateDictOptions(full_state_dict=True, cpu_offload=True) state_dict = get_model_state_dict(model, options=options) + # Drop MTP keys that share tensors with the base model (embedding / lm_head) so that + # save_pretrained does not raise the "shared tensors not properly defined" RuntimeError. + # These are re-shared from the base model by apply_mtp on load. + if getattr(model, "mtp", None) is not None: + from ....plugins.model_plugins.mtp import strip_shared_mtp_keys + + strip_shared_mtp_keys(state_dict) + if DistributedInterface().get_rank() == 0: model_to_save = model.module if hasattr(model, "module") else model model_to_save.save_pretrained(output_dir, state_dict=state_dict, max_shard_size="4GB") @@ -154,6 +162,12 @@ def save_checkpoint(model: HFModel, optimizer: torch.optim.Optimizer, ckpt_dir: hf_options = StateDictOptions(full_state_dict=True, cpu_offload=True) hf_state_dict = get_model_state_dict(model, options=hf_options) + # Drop shared MTP keys before save_pretrained (same reason as save_model). + if getattr(model, "mtp", None) is not None: + from ....plugins.model_plugins.mtp import strip_shared_mtp_keys + + strip_shared_mtp_keys(hf_state_dict) + if DistributedInterface().get_rank() == 0: model_to_save = model.module if hasattr(model, "module") else model hf_dir = os.path.join(ckpt_dir, "hf_model") diff --git a/src/llamafactory/v1/trainers/sft_trainer.py b/src/llamafactory/v1/trainers/sft_trainer.py index a80ffeb733..e855a890a5 100644 --- a/src/llamafactory/v1/trainers/sft_trainer.py +++ b/src/llamafactory/v1/trainers/sft_trainer.py @@ -64,6 +64,9 @@ def _compute_mtp_loss(self, batch: BatchInput) -> Tensor: if mtp_logits: mtp_loss = compute_mtp_loss(mtp_logits, labels, loss_weights) loss = loss + mtp_loss * self._mtp_loss_scale() + # Expose the unscaled per-head-mean MTP loss for logging (the main `loss` above + # already includes the scaled contribution). Read back by BaseTrainer.fit. + self.model._last_mtp_loss = float(mtp_loss.detach().item()) return loss diff --git a/tests_v1/plugins/model_plugins/test_mtp.py b/tests_v1/plugins/model_plugins/test_mtp.py index 2f074d2a93..0694e02150 100644 --- a/tests_v1/plugins/model_plugins/test_mtp.py +++ b/tests_v1/plugins/model_plugins/test_mtp.py @@ -21,7 +21,14 @@ import torch.nn.functional as F from transformers import AutoConfig, AutoModelForCausalLM -from llamafactory.v1.plugins.model_plugins.mtp import apply_mtp, compute_mtp_loss, roll_tensor +from llamafactory.v1.plugins.model_plugins.mtp import ( + apply_mtp, + compute_mtp_loss, + load_mtp_weights, + roll_tensor, + shift_input_ids_for_mtp, + strip_shared_mtp_keys, +) from llamafactory.v1.utils.env import find_available_port from llamafactory.v1.utils.pytest import dist_env @@ -123,6 +130,77 @@ def test_mtp_head_offset(tiny_model): assert torch.isfinite(expected) +def test_mtp_save_load(tmp_path): + """MTP weights survive save_pretrained -> from_pretrained + apply_mtp + load_mtp_weights. + + Regression test for the MTP weight-save bug: ``save_pretrained`` raised a + ``shared tensors ... not properly defined`` RuntimeError on the shared + ``mtp.embed_tokens``/``mtp.output_layer`` keys, and ``from_pretrained`` drops all + ``mtp.*`` keys as unexpected, so the grafted MTP weights were lost on reload. + """ + config = AutoConfig.from_pretrained(MODEL) + model = AutoModelForCausalLM.from_config(config) + model = apply_mtp(model, {"name": "mtp", "num_layers": 1, "loss_scale": 0.3}) + # Mutate MTP weights to detectable non-default values. + with torch.no_grad(): + model.mtp.e_proj.weight.fill_(0.25) + head0 = model.mtp.layers[str(config.num_hidden_layers)] + head0.layer.self_attn.q_proj.weight.fill_(-0.5) + ref_eproj = model.mtp.e_proj.weight.detach().clone() + ref_q = head0.layer.self_attn.q_proj.weight.detach().clone() + + # Save: strip the shared MTP keys first, exactly as the trainer does. + state_dict = model.state_dict() + strip_shared_mtp_keys(state_dict) + model.save_pretrained(tmp_path, state_dict=state_dict, max_shard_size="4GB") + + # Reload: from_pretrained drops mtp.*, apply_mtp re-creates the block (random), + # load_mtp_weights restores the saved MTP tensors. + reloaded = AutoModelForCausalLM.from_pretrained(tmp_path) + reloaded = apply_mtp(reloaded, {"name": "mtp", "num_layers": 1, "loss_scale": 0.3}) + load_mtp_weights(reloaded, str(tmp_path)) + + assert torch.equal(ref_eproj, reloaded.mtp.e_proj.weight) + new_head0 = reloaded.mtp.layers[str(config.num_hidden_layers)] + assert torch.equal(ref_q, new_head0.layer.self_attn.q_proj.weight) + # Shared modules are re-tied to the base model (not loaded from a separate copy). + assert reloaded.mtp.embed_tokens.weight.data_ptr() == reloaded.get_input_embeddings().weight.data_ptr() + assert reloaded.mtp.output_layer.weight.data_ptr() == reloaded.lm_head.weight.data_ptr() + + +def _test_mtp_shift_input_ids_cp(local_rank: int, world_size: int, master_port: int): + """``shift_input_ids_for_mtp`` under CP must match a full-sequence roll on each chunk. + + Regression test for the CP boundary bug: a plain local ``roll_tensor`` filled each + chunk's tail with ``fill_value`` instead of the next rank's first token, corrupting + the MTP input embedding at every CP boundary. + """ + with dist_env(local_rank, world_size, master_port): + dist.init_process_group("gloo") + from llamafactory.v1.plugins.model_plugins.parallelization.ulysses import ( + set_ulysses_sequence_parallel_group, + ) + + cp_group = dist.new_group(ranks=list(range(world_size))) + set_ulysses_sequence_parallel_group(cp_group) + + global_ids = torch.tensor([[10, 11, 12, 13, 14, 15, 16, 17]]) + chunk = torch.chunk(global_ids, world_size, dim=-1)[local_rank].contiguous() + cp_shifted = shift_input_ids_for_mtp(chunk, fill_value=0) + # Reference: roll the FULL sequence left by one, then take the local chunk. + full_shifted = roll_tensor(global_ids.clone(), shifts=-1, dim=-1, fill_value=0) + ref = torch.chunk(full_shifted, world_size, dim=-1)[local_rank].contiguous() + assert torch.equal(cp_shifted, ref), (local_rank, cp_shifted.tolist(), ref.tolist()) + dist.destroy_process_group() + + +@pytest.mark.require_distributed(2) +def test_mtp_shift_input_ids_cp(): + """The CP-aware shift reproduces the full-sequence roll across the chunk boundary.""" + master_port = find_available_port() + mp.spawn(_test_mtp_shift_input_ids_cp, args=(2, master_port), nprocs=2, join=True) + + def _test_mtp_cp_alignment(local_rank: int, world_size: int, master_port: int): """Each rank holds a local sequence chunk; the CP MTP loss must equal the full-seq loss.""" with dist_env(local_rank, world_size, master_port):