From 809ef6e4796c9c4fda3eadd90190365cd940e41b Mon Sep 17 00:00:00 2001 From: Curnane Date: Fri, 5 Jun 2026 10:28:04 +0800 Subject: [PATCH 1/4] feat(qwen3.5+trainer): MTP training loss plumbing aligned with #46229 Wire up Multi-Token Prediction (MTP) training for Qwen3.5 by reusing the inference-side MtpLayer from #46229 instead of the inline MTP definition from #45638, and add the matching Trainer hook that combines the auxiliary loss. Model side (Qwen3.5): * CausalLMOutputWithPast gains an mtp_loss field carrying the unweighted MTP loss (averaged over MTP layers and non-ignored tokens). * Qwen3_5TextConfig exposes num_nextn_predict_layers (canonical MTP layer count, aligned with #46229) and output_mtp_loss (training-only switch). * Qwen3_5ForCausalLM dynamically instantiates MtpLayer copies that mirror the main decoder layer architecture (full / linear attention mix, RMSNorm variant) and recomputes the per-layer loss with properly sliced inputs, position embeddings, and labels. Trainer side: * TrainingArguments gains mtp_loss_coef (default 0.0, BC-safe). * Trainer.compute_loss adds the unweighted output.mtp_loss * mtp_loss_coef to the main loss at the very end, leaving the model forward signature clean and the combination policy centralized in the trainer. Supersedes #45638's inline-MTP approach in favor of a single shared MtpLayer implementation that the inference path already uses. Refs: #46229, #45638 --- src/transformers/modeling_outputs.py | 5 + .../models/qwen3_5/configuration_qwen3_5.py | 23 +- .../models/qwen3_5/modeling_qwen3_5.py | 178 ++++++++++++ .../models/qwen3_5/modular_qwen3_5.py | 266 +++++++++++++++++- src/transformers/trainer.py | 21 ++ src/transformers/training_args.py | 22 ++ 6 files changed, 513 insertions(+), 2 deletions(-) diff --git a/src/transformers/modeling_outputs.py b/src/transformers/modeling_outputs.py index 4f3ac25fa58f..47985104043a 100755 --- a/src/transformers/modeling_outputs.py +++ b/src/transformers/modeling_outputs.py @@ -632,6 +632,10 @@ class CausalLMOutputWithPast(ModelOutput): Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + mtp_loss (`torch.FloatTensor` of shape `(1,)`, *optional*): + Raw, unweighted multi-token prediction loss. Returned by models that support MTP training + (e.g. Qwen3.5) when `output_mtp_loss=True` is passed or `config.output_mtp_loss=True` is + set. Callers (e.g. `Trainer`) decide how to combine this with the main `loss`. """ loss: torch.FloatTensor | None = None @@ -639,6 +643,7 @@ class CausalLMOutputWithPast(ModelOutput): past_key_values: Cache | None = None hidden_states: tuple[torch.FloatTensor, ...] | None = None attentions: tuple[torch.FloatTensor, ...] | None = None + mtp_loss: torch.FloatTensor | None = None @dataclass diff --git a/src/transformers/models/qwen3_5/configuration_qwen3_5.py b/src/transformers/models/qwen3_5/configuration_qwen3_5.py index ae9eb8f86c6d..13a9b1d70f6a 100644 --- a/src/transformers/models/qwen3_5/configuration_qwen3_5.py +++ b/src/transformers/models/qwen3_5/configuration_qwen3_5.py @@ -18,10 +18,11 @@ # See the License for the specific language governing permissions and # limitations under the License. from huggingface_hub.dataclasses import strict +from ...utils import auto_docstring from ...configuration_utils import PreTrainedConfig from ...modeling_rope_utils import RopeParameters -from ...utils import auto_docstring +from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available @auto_docstring(checkpoint="Qwen/Qwen3.5-27B") @@ -38,6 +39,17 @@ class Qwen3_5TextConfig(PreTrainedConfig): Number of key heads used in linear attention layers. linear_num_value_heads (`int`, *optional*, defaults to 32): Number of value heads used in linear attention layers. + num_nextn_predict_layers (`int`, *optional*, defaults to 0): + Number of Multi-Token Prediction (MTP) layers. `0` (default) disables MTP and + allocates no extra parameters, keeping behavior identical to a non-MTP model. + When `> 0`, that many `MtpLayer` instances are appended to the language model + head and trained jointly with the main model. Name aligned with the inference + side in #46229. + output_mtp_loss (`bool`, *optional*, defaults to False): + Training-only switch. When `True` and `num_nextn_predict_layers > 0`, the model + forward returns an additional unweighted `mtp_loss` tensor alongside `loss`. + Callers (e.g. `Trainer`) combine it with `loss` using their own weighting + coefficient. ```python >>> from transformers import Qwen3_5TextModel, Qwen3_5TextConfig @@ -101,6 +113,15 @@ class Qwen3_5TextConfig(PreTrainedConfig): base_config_key = "text_config" ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} + # === Multi-Token Prediction (MTP) training config === + # Aligned with #46229: num_nextn_predict_layers is the canonical name across HF MTP modules. + # 0 (default) means MTP is disabled and no extra weights are allocated, preserving BC. + num_nextn_predict_layers: int = 0 + # Training-only switch: when True, the model forward returns a separate raw `mtp_loss` + # tensor in the output (unweighted). Callers (e.g. `Trainer`) are expected to combine it + # with the main loss using their own coefficient (e.g. `TrainingArguments.mtp_loss_coef`). + output_mtp_loss: bool = False + def __post_init__(self, **kwargs): kwargs.setdefault("partial_rotary_factor", 0.25) # assign default for BC if self.layer_types is None: diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 7bacea8a8f41..14e87ed18847 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -1618,6 +1618,7 @@ class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel, GenerationMixin): _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} config: Qwen3_5TextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] + _no_split_modules = ["Qwen3_5DecoderLayer", "Qwen3_5VisionBlock", "MtpLayer"] def __init__(self, config): super().__init__(config) @@ -1625,9 +1626,166 @@ def __init__(self, config): self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + # === Multi-Token Prediction (MTP) training head === + # Aligned with #46229: when `num_nextn_predict_layers > 0`, lazily construct + # MTP layers that share the main model's embed/lm_head/rotary embeddings. + # MTP is intentionally a separate, lightweight module list (not a PreTrainedModel) + # to keep the training forward simple; inference-side `MtpLayerStack` (#46229) + # builds on the same `MtpLayer` primitive. + self._mtp_layers: nn.ModuleList | None = None + if getattr(config, "num_nextn_predict_layers", 0) > 0: + self._init_mtp_layers() + # Initialize weights and apply final processing self.post_init() + def _init_mtp_layers(self) -> None: + import dataclasses + + from ...generation.mtp import MtpLayer + + text_config = self.config.get_text_config() + num = int(getattr(text_config, "num_nextn_predict_layers", 0)) + if num <= 0: + self._mtp_layers = None + return + + # MTP layers are addressed at layer_idx = num_hidden_layers + k, but the + # model's strict-config validation requires `len(layer_types) == num_hidden_layers`. + # Build a transient copy of the text config that *does* know about the MTP + # layers, so the inner decoder-layer constructor can index layer_types. + mtp_layer_types = list(text_config.layer_types or []) + ["full_attention"] * num + mtp_text_config = dataclasses.replace( + text_config, + num_hidden_layers=text_config.num_hidden_layers + num, + layer_types=mtp_layer_types, + ) + # Carry over the attention implementation choice; the inner decoder layer reads + # `config._attn_implementation` and we don't want to lose the user's selection. + mtp_text_config._attn_implementation = getattr(text_config, "_attn_implementation", None) + + # Discover the decoder layer class & first norm class from the main model so + # the MTP layer mirrors the exact same architecture (full attention / + # linear attention mix, RMSNorm variant, etc.). + reference_layer = self.model.layers[0] + layer_cls = type(reference_layer) + norm_cls = next( + (type(m) for n, m in reference_layer.named_modules() if n.endswith("input_layernorm")), + None, + ) + if norm_cls is None: + raise ValueError( + "Qwen3.5 MTP training requires the main decoder layer to expose an `input_layernorm` module." + ) + + self._mtp_layers = nn.ModuleList( + [MtpLayer(mtp_text_config, layer_cls, norm_cls, text_config.num_hidden_layers + k) for k in range(num)] + ) + + def _compute_mtp_loss( + self, + input_ids: torch.Tensor, + main_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + position_ids: torch.Tensor | None, + labels: torch.Tensor, + ) -> torch.Tensor: + """Compute raw (unweighted) MTP loss over all configured MTP layers. + + For each MTP layer k (0-indexed), the model predicts `labels[t+k+1]` from + `main_hidden_states[:, t, :]` (the main hidden state at time t). We feed + `input_ids[:, t+1+k]` (shifted by k+1) as the next-token embedding. + + The returned loss is the *raw* mean of per-MTP-layer losses; callers + (e.g. `Trainer`) decide the weighting via their own coefficient. + """ + device = main_hidden_states.device + if self._mtp_layers is None or len(self._mtp_layers) == 0: + return torch.tensor(0.0, device=device, requires_grad=True) + if labels is None: + return torch.tensor(0.0, device=device, requires_grad=True) + + # We do not have access to the cached position_embeddings from the main forward, + # so we recompute them here. Cost is one RoPE pass per MTP layer; minor vs. the + # attention itself, and avoids leaking internals of the main model. + embed = self.model.embed_tokens + rotary = self.model.rotary_emb + + # `text_position_ids` is what the main forward passes to layers (shape (B, S)). + # The main model expands 1D/2D position_ids into the (4, B, S) mrope form first + # and then slices [0] to get the text positions. We mirror that here so the + # rotary embedding call below gets the right shape. + if position_ids is None: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(4, input_ids.shape[0], -1) + ) + elif position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(4, *position_ids.shape) + if position_ids.ndim == 3 and position_ids.shape[0] == 4: + text_position_ids = position_ids[0] + else: + text_position_ids = position_ids + + total_loss = torch.tensor(0.0, device=device) + current_hidden = main_hidden_states + + for i, mtp_layer in enumerate(self._mtp_layers): + shift = i + 1 + # Shifted input: token at position t+shift becomes the "next token" prediction + # for the main hidden at position t. + shifted_input_ids = torch.roll(input_ids, -shift, dims=1) + input_embeds = embed(shifted_input_ids) + + # The MTP layer attends to the first `seq_len - shift` positions; later + # positions are invalid labels and would leak info beyond the end of the + # main sequence. We slice the inputs *and* recompute position embeddings + # for the sliced length so RoPE dimensions line up exactly. + seq_len = input_embeds.shape[1] + mtp_input_embeds = input_embeds[:, : seq_len - shift] + mtp_position_ids = text_position_ids[:, : seq_len - shift] + mtp_position_embeddings = rotary(mtp_input_embeds, position_ids=mtp_position_ids) + + if attention_mask is not None: + if attention_mask.ndim == 4: + mtp_attention_mask = attention_mask[..., : seq_len - shift, : seq_len - shift] + elif attention_mask.ndim == 2: + mtp_attention_mask = attention_mask[:, : seq_len - shift] + else: + mtp_attention_mask = attention_mask + else: + mtp_attention_mask = None + + mtp_hidden = mtp_layer( + inputs_embeds=mtp_input_embeds, + previous_hidden_state=current_hidden[:, : seq_len - shift], + position_embeddings=mtp_position_embeddings, + attention_mask=mtp_attention_mask, + position_ids=mtp_position_ids, + past_key_values=None, + ) + + # Project to vocab & compute next-token loss vs. labels shifted by `shift`. + # The mtp_hidden was computed over the first `seq_len - shift` positions, + # so labels must be sliced to match (after rolling). Tail positions are + # masked to -100 so they never contribute even if length happened to match. + mtp_logits = self.lm_head(mtp_hidden) + shifted_labels = torch.roll(labels, -shift, dims=1)[:, : seq_len - shift].clone() + shifted_labels[:, -shift:] = -100 # ignore the rolled tail + + layer_loss = self.loss_function( + logits=mtp_logits, + labels=shifted_labels, + vocab_size=self.config.vocab_size, + ) + total_loss = total_loss + layer_loss + # The next MTP layer should be conditioned on this layer's hidden states, + # not the main hidden states. Concatenate along time so positions line up. + current_hidden = mtp_hidden + + return total_loss / len(self._mtp_layers) + @can_return_tuple @auto_docstring def forward( @@ -1640,6 +1798,7 @@ def forward( labels: torch.LongTensor | None = None, use_cache: bool | None = None, logits_to_keep: int | torch.Tensor = 0, + output_mtp_loss: bool | None = None, **kwargs: Unpack[TransformersKwargs], ) -> CausalLMOutputWithPast: r""" @@ -1647,6 +1806,11 @@ def forward( Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + output_mtp_loss (`bool`, *optional*): + If True, the model will compute and return a separate raw `mtp_loss` tensor in the output + (unweighted). When `None` (default), falls back to `config.output_mtp_loss`. + Combining the main `loss` with `mtp_loss` is the caller's responsibility (typically the + `Trainer` via a coefficient like `mtp_loss_coef`). Example: @@ -1683,12 +1847,26 @@ def forward( if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + mtp_loss = None + # Determine whether to compute MTP loss: explicit kwarg wins, else config default. + if output_mtp_loss is None: + output_mtp_loss = bool(getattr(self.config, "output_mtp_loss", False)) + if output_mtp_loss and labels is not None and self._mtp_layers is not None: + mtp_loss = self._compute_mtp_loss( + input_ids=input_ids, + main_hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + labels=labels, + ) + return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, + mtp_loss=mtp_loss, ) diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py index 91457c58f3ce..6bd13e91d682 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -28,7 +28,12 @@ GenericForTokenClassification, GradientCheckpointingLayer, ) -from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, SequenceClassifierOutputWithPast +from ...modeling_outputs import ( + BaseModelOutputWithPast, + BaseModelOutputWithPooling, + CausalLMOutputWithPast, + SequenceClassifierOutputWithPast, +) from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging @@ -74,6 +79,17 @@ class Qwen3_5TextConfig(Qwen3NextConfig): Number of key heads used in linear attention layers. linear_num_value_heads (`int`, *optional*, defaults to 32): Number of value heads used in linear attention layers. + num_nextn_predict_layers (`int`, *optional*, defaults to 0): + Number of Multi-Token Prediction (MTP) layers. `0` (default) disables MTP and + allocates no extra parameters, keeping behavior identical to a non-MTP model. + When `> 0`, that many `MtpLayer` instances are appended to the language model + head and trained jointly with the main model. Name aligned with the inference + side in #46229. + output_mtp_loss (`bool`, *optional*, defaults to False): + Training-only switch. When `True` and `num_nextn_predict_layers > 0`, the model + forward returns an additional unweighted `mtp_loss` tensor alongside `loss`. + Callers (e.g. `Trainer`) combine it with `loss` using their own weighting + coefficient. ```python >>> from transformers import Qwen3_5TextModel, Qwen3_5TextConfig @@ -121,6 +137,15 @@ class Qwen3_5TextConfig(Qwen3NextConfig): output_router_logits = AttributeError() router_aux_loss_coef = AttributeError() + # === Multi-Token Prediction (MTP) training config === + # Aligned with #46229: num_nextn_predict_layers is the canonical name across HF MTP modules. + # 0 (default) means MTP is disabled and no extra weights are allocated, preserving BC. + num_nextn_predict_layers: int = 0 + # Training-only switch: when True, the model forward returns a separate raw `mtp_loss` + # tensor in the output (unweighted). Callers (e.g. `Trainer`) are expected to combine it + # with the main loss using their own coefficient (e.g. `TrainingArguments.mtp_loss_coef`). + output_mtp_loss: bool = False + def __post_init__(self, **kwargs): super().__post_init__(**kwargs) del self.mlp_only_layers @@ -667,10 +692,249 @@ def forward( class Qwen3_5ForCausalLM(Qwen3ForCausalLM): config: Qwen3_5TextConfig _keys_to_ignore_on_load_unexpected = [r"^mtp.*", r"^model.visual.*"] + _no_split_modules = ["Qwen3_5DecoderLayer", "Qwen3_5VisionBlock", "MtpLayer"] def __init__(self, config): super().__init__(config) self.model = Qwen3_5TextModel(config) + # === Multi-Token Prediction (MTP) training head === + # Aligned with #46229: when `num_nextn_predict_layers > 0`, lazily construct + # MTP layers that share the main model's embed/lm_head/rotary embeddings. + # MTP is intentionally a separate, lightweight module list (not a PreTrainedModel) + # to keep the training forward simple; inference-side `MtpLayerStack` (#46229) + # builds on the same `MtpLayer` primitive. + self._mtp_layers: nn.ModuleList | None = None + if getattr(config, "num_nextn_predict_layers", 0) > 0: + self._init_mtp_layers() + + def _init_mtp_layers(self) -> None: + import dataclasses + + from ...generation.mtp import MtpLayer + + text_config = self.config.get_text_config() + num = int(getattr(text_config, "num_nextn_predict_layers", 0)) + if num <= 0: + self._mtp_layers = None + return + + # MTP layers are addressed at layer_idx = num_hidden_layers + k, but the + # model's strict-config validation requires `len(layer_types) == num_hidden_layers`. + # Build a transient copy of the text config that *does* know about the MTP + # layers, so the inner decoder-layer constructor can index layer_types. + mtp_layer_types = list(text_config.layer_types or []) + ["full_attention"] * num + mtp_text_config = dataclasses.replace( + text_config, + num_hidden_layers=text_config.num_hidden_layers + num, + layer_types=mtp_layer_types, + ) + # Carry over the attention implementation choice; the inner decoder layer reads + # `config._attn_implementation` and we don't want to lose the user's selection. + mtp_text_config._attn_implementation = getattr(text_config, "_attn_implementation", None) + + # Discover the decoder layer class & first norm class from the main model so + # the MTP layer mirrors the exact same architecture (full attention / + # linear attention mix, RMSNorm variant, etc.). + reference_layer = self.model.layers[0] + layer_cls = type(reference_layer) + norm_cls = next( + (type(m) for n, m in reference_layer.named_modules() if n.endswith("input_layernorm")), + None, + ) + if norm_cls is None: + raise ValueError( + "Qwen3.5 MTP training requires the main decoder layer to expose an `input_layernorm` module." + ) + + self._mtp_layers = nn.ModuleList( + [MtpLayer(mtp_text_config, layer_cls, norm_cls, text_config.num_hidden_layers + k) for k in range(num)] + ) + + def _compute_mtp_loss( + self, + input_ids: torch.Tensor, + main_hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None, + position_ids: torch.Tensor | None, + labels: torch.Tensor, + ) -> torch.Tensor: + """Compute raw (unweighted) MTP loss over all configured MTP layers. + + For each MTP layer k (0-indexed), the model predicts `labels[t+k+1]` from + `main_hidden_states[:, t, :]` (the main hidden state at time t). We feed + `input_ids[:, t+1+k]` (shifted by k+1) as the next-token embedding. + + The returned loss is the *raw* mean of per-MTP-layer losses; callers + (e.g. `Trainer`) decide the weighting via their own coefficient. + """ + device = main_hidden_states.device + if self._mtp_layers is None or len(self._mtp_layers) == 0: + return torch.tensor(0.0, device=device, requires_grad=True) + if labels is None: + return torch.tensor(0.0, device=device, requires_grad=True) + + # We do not have access to the cached position_embeddings from the main forward, + # so we recompute them here. Cost is one RoPE pass per MTP layer; minor vs. the + # attention itself, and avoids leaking internals of the main model. + embed = self.model.embed_tokens + rotary = self.model.rotary_emb + + # `text_position_ids` is what the main forward passes to layers (shape (B, S)). + # The main model expands 1D/2D position_ids into the (4, B, S) mrope form first + # and then slices [0] to get the text positions. We mirror that here so the + # rotary embedding call below gets the right shape. + if position_ids is None: + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, 1, -1) + .expand(4, input_ids.shape[0], -1) + ) + elif position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(4, *position_ids.shape) + if position_ids.ndim == 3 and position_ids.shape[0] == 4: + text_position_ids = position_ids[0] + else: + text_position_ids = position_ids + + total_loss = torch.tensor(0.0, device=device) + current_hidden = main_hidden_states + + for i, mtp_layer in enumerate(self._mtp_layers): + shift = i + 1 + # Shifted input: token at position t+shift becomes the "next token" prediction + # for the main hidden at position t. + shifted_input_ids = torch.roll(input_ids, -shift, dims=1) + input_embeds = embed(shifted_input_ids) + + # The MTP layer attends to the first `seq_len - shift` positions; later + # positions are invalid labels and would leak info beyond the end of the + # main sequence. We slice the inputs *and* recompute position embeddings + # for the sliced length so RoPE dimensions line up exactly. + seq_len = input_embeds.shape[1] + mtp_input_embeds = input_embeds[:, : seq_len - shift] + mtp_position_ids = text_position_ids[:, : seq_len - shift] + mtp_position_embeddings = rotary(mtp_input_embeds, position_ids=mtp_position_ids) + + if attention_mask is not None: + if attention_mask.ndim == 4: + mtp_attention_mask = attention_mask[..., : seq_len - shift, : seq_len - shift] + elif attention_mask.ndim == 2: + mtp_attention_mask = attention_mask[:, : seq_len - shift] + else: + mtp_attention_mask = attention_mask + else: + mtp_attention_mask = None + + mtp_hidden = mtp_layer( + inputs_embeds=mtp_input_embeds, + previous_hidden_state=current_hidden[:, : seq_len - shift], + position_embeddings=mtp_position_embeddings, + attention_mask=mtp_attention_mask, + position_ids=mtp_position_ids, + past_key_values=None, + ) + + # Project to vocab & compute next-token loss vs. labels shifted by `shift`. + # The mtp_hidden was computed over the first `seq_len - shift` positions, + # so labels must be sliced to match (after rolling). Tail positions are + # masked to -100 so they never contribute even if length happened to match. + mtp_logits = self.lm_head(mtp_hidden) + shifted_labels = torch.roll(labels, -shift, dims=1)[:, : seq_len - shift].clone() + shifted_labels[:, -shift:] = -100 # ignore the rolled tail + + layer_loss = self.loss_function( + logits=mtp_logits, + labels=shifted_labels, + vocab_size=self.config.vocab_size, + ) + total_loss = total_loss + layer_loss + # The next MTP layer should be conditioned on this layer's hidden states, + # not the main hidden states. Concatenate along time so positions line up. + current_hidden = mtp_hidden + + return total_loss / len(self._mtp_layers) + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + output_mtp_loss: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> CausalLMOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + output_mtp_loss (`bool`, *optional*): + If True, the model will compute and return a separate raw `mtp_loss` tensor in the output + (unweighted). When `None` (default), falls back to `config.output_mtp_loss`. + Combining the main `loss` with `mtp_loss` is the caller's responsibility (typically the + `Trainer` via a coefficient like `mtp_loss_coef`). + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen3_5ForCausalLM + + >>> model = Qwen3_5ForCausalLM.from_pretrained("Qwen/Qwen3_5-8B") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3_5-8B") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + mtp_loss = None + # Determine whether to compute MTP loss: explicit kwarg wins, else config default. + if output_mtp_loss is None: + output_mtp_loss = bool(getattr(self.config, "output_mtp_loss", False)) + if output_mtp_loss and labels is not None and self._mtp_layers is not None: + mtp_loss = self._compute_mtp_loss( + input_ids=input_ids, + main_hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + labels=labels, + ) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + mtp_loss=mtp_loss, + ) class Qwen3_5ForTokenClassification(GenericForTokenClassification, Qwen3_5PreTrainedModel): diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py index ad66f1a23b0c..0a7483af5d4c 100755 --- a/src/transformers/trainer.py +++ b/src/transformers/trainer.py @@ -2030,6 +2030,27 @@ def compute_loss( loss_scale //= pc.tp_size loss *= loss_scale if self.args.n_gpu <= 1 else self.args.n_gpu + # Multi-Token Prediction (MTP) auxiliary loss. + # Strategy mirrors #46229: the model returns the *unweighted* MTP loss (averaged + # over layers and non-ignored tokens) as a separate field on its output, and the + # Trainer is responsible for combining it with the main loss. This keeps the + # model forward signature clean and lets users tune the coefficient from the + # training command line without subclassing the model. + # + # We only do anything if both: (1) the user actually set a non-zero coefficient, + # and (2) the model returned a non-None `mtp_loss` tensor. Models that don't + # support MTP just pass through unchanged (BC-safe). + mtp_coef = getattr(self.args, "mtp_loss_coef", 0.0) + if mtp_coef and mtp_coef != 0.0: + mtp_loss = None + if isinstance(outputs, dict): + mtp_loss = outputs.get("mtp_loss") + else: + # Dataclass / ModelOutput / namedtuple — try attribute access first + mtp_loss = getattr(outputs, "mtp_loss", None) + if mtp_loss is not None: + loss = loss + mtp_coef * mtp_loss + return (loss, outputs) if return_outputs else loss def compute_loss_context_manager(self) -> contextlib.ExitStack: diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py index 328e3eca4aa9..d70bcfa2c215 100644 --- a/src/transformers/training_args.py +++ b/src/transformers/training_args.py @@ -273,6 +273,15 @@ class TrainingArguments: Label smoothing factor to prevent overconfidence. Replaces hard 0/1 targets with soft targets: 0 becomes `ε/num_labels` and 1 becomes `1 - ε + ε/num_labels`, where ε = `label_smoothing_factor`. Zero means no smoothing. Typical range: 0.0 to 0.1. + mtp_loss_coef (`float`, *optional*, defaults to 0.0): + Coefficient applied to the Multi-Token Prediction (MTP) auxiliary loss returned + by the model as `output.mtp_loss`. The total loss becomes + `loss + mtp_loss_coef * mtp_loss`. The model is expected to return the + *unweighted* MTP loss (averaged over MTP layers and non-ignored tokens); the + Trainer is responsible for weighting. Set to `0.0` (default) to disable, which + keeps behavior identical for models that do not return an MTP loss. Typical + range: 0.0 to 1.0; common values in the DeepSeek / Qwen3.5 MTP literature are + 0.1 to 0.3. > Mixed Precision Training @@ -846,6 +855,19 @@ class TrainingArguments: label_smoothing_factor: float = field( default=0.0, metadata={"help": "Label smoothing factor to prevent overconfidence. Zero means no smoothing."} ) + mtp_loss_coef: float = field( + default=0.0, + metadata={ + "help": ( + "Coefficient applied to the Multi-Token Prediction auxiliary loss returned by the model " + "as `output.mtp_loss`. The total loss becomes `loss + mtp_loss_coef * mtp_loss`. The model " + "is expected to return the *unweighted* MTP loss (avg over MTP layers and non-ignored " + "tokens); the Trainer combines it. Set to 0.0 (default) to disable, preserving BC for " + "models that do not return an MTP loss. Typical range: 0.0 to 1.0; common values in the " + "DeepSeek / Qwen3.5 MTP literature are 0.1 to 0.3." + ) + }, + ) # --- Mixed Precision --- bf16: bool = field( From 487745a295a02b9bd49936c84efdc05af7b97cde Mon Sep 17 00:00:00 2001 From: Curnane Date: Fri, 5 Jun 2026 18:34:55 +0800 Subject: [PATCH 2/4] Fix check_code_quality and check_repository_consistency for qwen3_5 MTP - Remove unused imports (is_causal_conv1d_available, is_flash_linear_attention_available) from configuration_qwen3_5.py - Move _init_mtp_layers and _compute_mtp_loss after forward() in both modular_qwen3_5.py and modeling_qwen3_5.py to match modular generator output - Remove extra blank lines inside __init__ to match generated code style --- .../models/qwen3_5/configuration_qwen3_5.py | 1 - .../models/qwen3_5/modeling_qwen3_5.py | 167 +++++++++--------- .../models/qwen3_5/modular_qwen3_5.py | 165 +++++++++-------- 3 files changed, 164 insertions(+), 169 deletions(-) diff --git a/src/transformers/models/qwen3_5/configuration_qwen3_5.py b/src/transformers/models/qwen3_5/configuration_qwen3_5.py index 13a9b1d70f6a..7cdd48332adf 100644 --- a/src/transformers/models/qwen3_5/configuration_qwen3_5.py +++ b/src/transformers/models/qwen3_5/configuration_qwen3_5.py @@ -22,7 +22,6 @@ from ...configuration_utils import PreTrainedConfig from ...modeling_rope_utils import RopeParameters -from ...utils.import_utils import is_causal_conv1d_available, is_flash_linear_attention_available @auto_docstring(checkpoint="Qwen/Qwen3.5-27B") diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 14e87ed18847..1994fac2624f 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -1625,7 +1625,6 @@ def __init__(self, config): self.model = Qwen3_5TextModel(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) - # === Multi-Token Prediction (MTP) training head === # Aligned with #46229: when `num_nextn_predict_layers > 0`, lazily construct # MTP layers that share the main model's embed/lm_head/rotary embeddings. @@ -1638,6 +1637,88 @@ def __init__(self, config): # Initialize weights and apply final processing self.post_init() + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + output_mtp_loss: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> CausalLMOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + output_mtp_loss (`bool`, *optional*): + If True, the model will compute and return a separate raw `mtp_loss` tensor in the output + (unweighted). When `None` (default), falls back to `config.output_mtp_loss`. + Combining the main `loss` with `mtp_loss` is the caller's responsibility (typically the + `Trainer` via a coefficient like `mtp_loss_coef`). + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen3_5ForCausalLM + + >>> model = Qwen3_5ForCausalLM.from_pretrained("Qwen/Qwen3_5-8B") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3_5-8B") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + mtp_loss = None + # Determine whether to compute MTP loss: explicit kwarg wins, else config default. + if output_mtp_loss is None: + output_mtp_loss = bool(getattr(self.config, "output_mtp_loss", False)) + if output_mtp_loss and labels is not None and self._mtp_layers is not None: + mtp_loss = self._compute_mtp_loss( + input_ids=input_ids, + main_hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + labels=labels, + ) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + mtp_loss=mtp_loss, + ) def _init_mtp_layers(self) -> None: import dataclasses @@ -1786,90 +1867,6 @@ def _compute_mtp_loss( return total_loss / len(self._mtp_layers) - @can_return_tuple - @auto_docstring - def forward( - self, - input_ids: torch.LongTensor | None = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - inputs_embeds: torch.FloatTensor | None = None, - labels: torch.LongTensor | None = None, - use_cache: bool | None = None, - logits_to_keep: int | torch.Tensor = 0, - output_mtp_loss: bool | None = None, - **kwargs: Unpack[TransformersKwargs], - ) -> CausalLMOutputWithPast: - r""" - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - output_mtp_loss (`bool`, *optional*): - If True, the model will compute and return a separate raw `mtp_loss` tensor in the output - (unweighted). When `None` (default), falls back to `config.output_mtp_loss`. - Combining the main `loss` with `mtp_loss` is the caller's responsibility (typically the - `Trainer` via a coefficient like `mtp_loss_coef`). - - Example: - - ```python - >>> from transformers import AutoTokenizer, Qwen3_5ForCausalLM - - >>> model = Qwen3_5ForCausalLM.from_pretrained("Qwen/Qwen3_5-8B") - >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3_5-8B") - - >>> prompt = "Hey, are you conscious? Can you talk to me?" - >>> inputs = tokenizer(prompt, return_tensors="pt") - - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." - ```""" - outputs: BaseModelOutputWithPast = self.model( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - **kwargs, - ) - - hidden_states = outputs.last_hidden_state - # Only compute necessary logits, and do not upcast them to float if we are not computing the loss - slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep - logits = self.lm_head(hidden_states[:, slice_indices, :]) - - loss = None - if labels is not None: - loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) - - mtp_loss = None - # Determine whether to compute MTP loss: explicit kwarg wins, else config default. - if output_mtp_loss is None: - output_mtp_loss = bool(getattr(self.config, "output_mtp_loss", False)) - if output_mtp_loss and labels is not None and self._mtp_layers is not None: - mtp_loss = self._compute_mtp_loss( - input_ids=input_ids, - main_hidden_states=hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - labels=labels, - ) - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - mtp_loss=mtp_loss, - ) - - class Qwen3_5ForTokenClassification(GenericForTokenClassification, Qwen3_5PreTrainedModel): config: Qwen3_5Config diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py index 6bd13e91d682..5414aa18a84c 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -707,6 +707,88 @@ def __init__(self, config): if getattr(config, "num_nextn_predict_layers", 0) > 0: self._init_mtp_layers() + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + use_cache: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + output_mtp_loss: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> CausalLMOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + output_mtp_loss (`bool`, *optional*): + If True, the model will compute and return a separate raw `mtp_loss` tensor in the output + (unweighted). When `None` (default), falls back to `config.output_mtp_loss`. + Combining the main `loss` with `mtp_loss` is the caller's responsibility (typically the + `Trainer` via a coefficient like `mtp_loss_coef`). + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen3_5ForCausalLM + + >>> model = Qwen3_5ForCausalLM.from_pretrained("Qwen/Qwen3_5-8B") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3_5-8B") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + mtp_loss = None + # Determine whether to compute MTP loss: explicit kwarg wins, else config default. + if output_mtp_loss is None: + output_mtp_loss = bool(getattr(self.config, "output_mtp_loss", False)) + if output_mtp_loss and labels is not None and self._mtp_layers is not None: + mtp_loss = self._compute_mtp_loss( + input_ids=input_ids, + main_hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + labels=labels, + ) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + mtp_loss=mtp_loss, + ) + def _init_mtp_layers(self) -> None: import dataclasses @@ -854,89 +936,6 @@ def _compute_mtp_loss( return total_loss / len(self._mtp_layers) - @can_return_tuple - @auto_docstring - def forward( - self, - input_ids: torch.LongTensor | None = None, - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values: Cache | None = None, - inputs_embeds: torch.FloatTensor | None = None, - labels: torch.LongTensor | None = None, - use_cache: bool | None = None, - logits_to_keep: int | torch.Tensor = 0, - output_mtp_loss: bool | None = None, - **kwargs: Unpack[TransformersKwargs], - ) -> CausalLMOutputWithPast: - r""" - labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): - Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., - config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored - (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. - output_mtp_loss (`bool`, *optional*): - If True, the model will compute and return a separate raw `mtp_loss` tensor in the output - (unweighted). When `None` (default), falls back to `config.output_mtp_loss`. - Combining the main `loss` with `mtp_loss` is the caller's responsibility (typically the - `Trainer` via a coefficient like `mtp_loss_coef`). - - Example: - - ```python - >>> from transformers import AutoTokenizer, Qwen3_5ForCausalLM - - >>> model = Qwen3_5ForCausalLM.from_pretrained("Qwen/Qwen3_5-8B") - >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3_5-8B") - - >>> prompt = "Hey, are you conscious? Can you talk to me?" - >>> inputs = tokenizer(prompt, return_tensors="pt") - - >>> # Generate - >>> generate_ids = model.generate(inputs.input_ids, max_length=30) - >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] - "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." - ```""" - outputs: BaseModelOutputWithPast = self.model( - input_ids=input_ids, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - inputs_embeds=inputs_embeds, - use_cache=use_cache, - **kwargs, - ) - - hidden_states = outputs.last_hidden_state - slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep - logits = self.lm_head(hidden_states[:, slice_indices, :]) - - loss = None - if labels is not None: - loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) - - mtp_loss = None - # Determine whether to compute MTP loss: explicit kwarg wins, else config default. - if output_mtp_loss is None: - output_mtp_loss = bool(getattr(self.config, "output_mtp_loss", False)) - if output_mtp_loss and labels is not None and self._mtp_layers is not None: - mtp_loss = self._compute_mtp_loss( - input_ids=input_ids, - main_hidden_states=hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - labels=labels, - ) - - return CausalLMOutputWithPast( - loss=loss, - logits=logits, - past_key_values=outputs.past_key_values, - hidden_states=outputs.hidden_states, - attentions=outputs.attentions, - mtp_loss=mtp_loss, - ) - - class Qwen3_5ForTokenClassification(GenericForTokenClassification, Qwen3_5PreTrainedModel): config: Qwen3_5Config From 9529206ca2925a87a8cb27d1dd2ce3c498871da6 Mon Sep 17 00:00:00 2001 From: Curnane Date: Fri, 5 Jun 2026 18:35:00 +0800 Subject: [PATCH 3/4] Update test_auto_docstring for CausalLMOutputWithPast.mtp_loss field --- tests/utils/test_auto_docstring.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/utils/test_auto_docstring.py b/tests/utils/test_auto_docstring.py index a38d2fc3f62d..79d6f83cb06e 100644 --- a/tests/utils/test_auto_docstring.py +++ b/tests/utils/test_auto_docstring.py @@ -519,6 +519,9 @@ def test_dummy_model_complete_docstring(self): Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. + - **mtp_loss** (`torch.FloatTensor` of shape `(1,)`, *optional*) -- Raw, unweighted multi-token prediction loss. Returned by models that support MTP training + (e.g. Qwen3.5) when `output_mtp_loss=True` is passed or `config.output_mtp_loss=True` is + set. Callers (e.g. `Trainer`) decide how to combine this with the main `loss`. Example: From 976337169eae34729fa2b413ae3b8dbe6feb65ea Mon Sep 17 00:00:00 2001 From: Curnane Date: Fri, 5 Jun 2026 19:15:04 +0800 Subject: [PATCH 4/4] Apply ruff formatting to qwen3_5 files for CI compliance --- src/transformers/models/qwen3_5/configuration_qwen3_5.py | 2 +- src/transformers/models/qwen3_5/modeling_qwen3_5.py | 2 ++ src/transformers/models/qwen3_5/modular_qwen3_5.py | 1 + 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/src/transformers/models/qwen3_5/configuration_qwen3_5.py b/src/transformers/models/qwen3_5/configuration_qwen3_5.py index 7cdd48332adf..3d28aa23d5e7 100644 --- a/src/transformers/models/qwen3_5/configuration_qwen3_5.py +++ b/src/transformers/models/qwen3_5/configuration_qwen3_5.py @@ -18,10 +18,10 @@ # See the License for the specific language governing permissions and # limitations under the License. from huggingface_hub.dataclasses import strict -from ...utils import auto_docstring from ...configuration_utils import PreTrainedConfig from ...modeling_rope_utils import RopeParameters +from ...utils import auto_docstring @auto_docstring(checkpoint="Qwen/Qwen3.5-27B") diff --git a/src/transformers/models/qwen3_5/modeling_qwen3_5.py b/src/transformers/models/qwen3_5/modeling_qwen3_5.py index 1994fac2624f..cad881e2600e 100644 --- a/src/transformers/models/qwen3_5/modeling_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modeling_qwen3_5.py @@ -1637,6 +1637,7 @@ def __init__(self, config): # Initialize weights and apply final processing self.post_init() + @can_return_tuple @auto_docstring def forward( @@ -1867,6 +1868,7 @@ def _compute_mtp_loss( return total_loss / len(self._mtp_layers) + class Qwen3_5ForTokenClassification(GenericForTokenClassification, Qwen3_5PreTrainedModel): config: Qwen3_5Config diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py index 5414aa18a84c..fada2ed6bf41 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -936,6 +936,7 @@ def _compute_mtp_loss( return total_loss / len(self._mtp_layers) + class Qwen3_5ForTokenClassification(GenericForTokenClassification, Qwen3_5PreTrainedModel): config: Qwen3_5Config