diff --git a/src/transformers/generation/candidate_generator.py b/src/transformers/generation/candidate_generator.py index c31d27e85805..c643fd7ea9d4 100644 --- a/src/transformers/generation/candidate_generator.py +++ b/src/transformers/generation/candidate_generator.py @@ -1512,7 +1512,7 @@ def get_candidates( # We need the last invalidated tokens, as well as the new one in a single passcfor efficiency mtp_input_ids = input_ids[:, -num_last_main_model_tokens - self.num_mtp_layers :] - mtp_position_ids = model_kwargs["position_ids"][:, -num_last_main_model_tokens - self.num_mtp_layers :] + mtp_position_ids = model_kwargs["position_ids"][..., -num_last_main_model_tokens - self.num_mtp_layers :] mtp_attention_mask = model_kwargs["attention_mask"][:, -num_last_main_model_tokens - self.num_mtp_layers :] last_hidden_states = self.full_seq_last_hidden_states[ :, -num_last_main_model_tokens - self.num_mtp_layers :, : @@ -1520,7 +1520,7 @@ def get_candidates( # If we have a single layer, or validated enough, we can simply pass the new tokens else: mtp_input_ids = input_ids[:, -num_last_main_model_tokens:] - mtp_position_ids = model_kwargs["position_ids"][:, -num_last_main_model_tokens:] + mtp_position_ids = model_kwargs["position_ids"][..., -num_last_main_model_tokens:] mtp_attention_mask = model_kwargs["attention_mask"][:, -num_last_main_model_tokens:] candidate_ids, candidate_logits, _ = self.mtp_model( diff --git a/src/transformers/modeling_layers.py b/src/transformers/modeling_layers.py index bbecda8f3f32..c09e782714fa 100644 --- a/src/transformers/modeling_layers.py +++ b/src/transformers/modeling_layers.py @@ -502,18 +502,27 @@ def forward( for i, mtp_layer in enumerate(self.layers): # We need to recompute those every layer since they change inputs_embeds = self.embed_tokens(input_ids).to(last_hidden_states.device) + + # Support both 2D position_ids and MRoPE-style nD ids (e.g. Qwen3.5 uses [4, B, L]) + if position_ids.ndim == 2: + text_position_ids = position_ids + rope_position_ids = position_ids + else: + text_position_ids = position_ids[0] # first component is text positions + rope_position_ids = position_ids[1:] if position_ids.shape[0] > 1 else position_ids + position_embeddings = ( - self.rotary_emb(inputs_embeds, position_ids=position_ids) if self.rotary_emb is not None else None + self.rotary_emb(inputs_embeds, position_ids=rope_position_ids) if self.rotary_emb is not None else None ) # In full generality, we may need to recompute masks for every layer due to the position offset of each layer - masks = self.create_masks_for_mtp_layer(i, inputs_embeds, mtp_cache, position_ids) + masks = self.create_masks_for_mtp_layer(i, inputs_embeds, mtp_cache, text_position_ids) last_hidden_states = mtp_layer( inputs_embeds, last_hidden_states, position_embeddings=position_embeddings, - position_ids=position_ids, + position_ids=text_position_ids, past_key_values=mtp_cache, **masks, **kwargs, @@ -549,7 +558,8 @@ def forward( # Roll by 1 and append for next layer input_ids = torch.cat([input_ids[:, 1:], next_mtp_token], dim=-1) attention_mask = torch.cat([attention_mask[:, 1:], attention_mask.new_ones(batch_size, 1)], dim=-1) # type: ignore - position_ids = torch.cat([position_ids[:, 1:], position_ids[:, -1:] + 1], dim=-1) + next_position_ids = position_ids[..., -1:] + 1 + position_ids = torch.cat([position_ids[..., 1:], next_position_ids], dim=-1) # Need to cat ful_ids as well for the processors if full_input_ids is not None: @@ -564,12 +574,15 @@ def from_pretrained(cls, main_model: PreTrainedModel, device_map=None, **kwargs) pretrained_model_name_or_path = main_model.config.name_or_path num_hidden_layers = main_model.config.get_text_config().num_hidden_layers # Heuristic: the main model should have the mtp layer patterns under `_keys_to_ignore_on_load_unexpected` to avoid - # loading them by default, so use it to later load the correct keys from the checkpoints + # loading them by default, so use it to later load the correct keys from the checkpoints. + # Filter out non-MTP patterns (e.g. vision tower weights for multimodal models like Qwen3.5). mtp_patterns = main_model._keys_to_ignore_on_load_unexpected.copy() # type: ignore - # Due to different released checkpoints, only keep the ones with layer number >= num_hidden_layers - otherwise - # mtp layers in a smaller checkpoints could be wrongly added as a 2nd mtp layer of a bigger checkpoint final_mtp_patterns = [] for pattern in mtp_patterns: + has_mtp_keyword = "mtp" in pattern.lower() + has_layer_pattern = re.search(r"layers?\.\d+", pattern, re.IGNORECASE) is not None + if not has_mtp_keyword and not has_layer_pattern: + continue match_object = re.search(r"\.(\d+)", pattern) if match_object is not None and int(match_object.group(1)) < num_hidden_layers: continue diff --git a/src/transformers/models/qwen3_5/configuration_qwen3_5.py b/src/transformers/models/qwen3_5/configuration_qwen3_5.py index c33663287cea..a89330333426 100644 --- a/src/transformers/models/qwen3_5/configuration_qwen3_5.py +++ b/src/transformers/models/qwen3_5/configuration_qwen3_5.py @@ -38,6 +38,9 @@ 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. + mtp_num_hidden_layers (`int`, *optional*, defaults to 0): + Number of Multi-Token Prediction (MTP) layers. When set to 0, MTP is disabled. + This field is exposed as `num_mtp_layers` via `attribute_map`. ```python >>> from transformers import Qwen3_5TextModel, Qwen3_5TextConfig @@ -105,7 +108,10 @@ class Qwen3_5TextConfig(PreTrainedConfig): bos_token_id: int | None = None eos_token_id: int | list[int] | None = None base_config_key = "text_config" + + attribute_map = {"num_mtp_layers": "mtp_num_hidden_layers"} ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} + mtp_num_hidden_layers: int = 0 def __post_init__(self, **kwargs): kwargs.setdefault("partial_rotary_factor", 0.25) # assign default for BC @@ -120,6 +126,13 @@ def __post_init__(self, **kwargs): super().__post_init__(**kwargs) + # MTP layers always use full attention + @property + def mtp_layer_types(self): + if self.num_mtp_layers is not None: + return ["full_attention"] * self.num_mtp_layers + return None + @auto_docstring(checkpoint="Qwen/Qwen3.5-27B") @strict @@ -152,6 +165,10 @@ class Qwen3_5VisionConfig(PreTrainedConfig): @strict class Qwen3_5Config(PreTrainedConfig): r""" + mtp_num_hidden_layers (`int`, *optional*, defaults to 0): + Number of Multi-Token Prediction (MTP) layers. When set to 0, MTP is disabled. + This field is exposed as `num_mtp_layers` via `attribute_map`. + Example: ```python @@ -180,6 +197,10 @@ class Qwen3_5Config(PreTrainedConfig): vision_end_token_id: int = 248054 tie_word_embeddings: bool = False + attribute_map = {"num_mtp_layers": "mtp_num_hidden_layers"} + + mtp_num_hidden_layers: int = 0 + def __post_init__(self, **kwargs): if isinstance(self.vision_config, dict): # old ckpt with incorrect model type -> override manually @@ -196,5 +217,12 @@ def __post_init__(self, **kwargs): super().__post_init__(**kwargs) + # MTP layers always use full attention + @property + def mtp_layer_types(self): + if self.num_mtp_layers is not None: + return ["full_attention"] * self.num_mtp_layers + return None + __all__ = ["Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5VisionConfig"] diff --git a/src/transformers/models/qwen3_5/modular_qwen3_5.py b/src/transformers/models/qwen3_5/modular_qwen3_5.py index 655ce31ce1b2..867f6ed84721 100644 --- a/src/transformers/models/qwen3_5/modular_qwen3_5.py +++ b/src/transformers/models/qwen3_5/modular_qwen3_5.py @@ -81,6 +81,9 @@ 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. + mtp_num_hidden_layers (`int`, *optional*, defaults to 0): + Number of Multi-Token Prediction (MTP) layers. When set to 0, MTP is disabled. + This field is exposed as `num_mtp_layers` via `attribute_map`. ```python >>> from transformers import Qwen3_5TextModel, Qwen3_5TextConfig @@ -99,6 +102,8 @@ class Qwen3_5TextConfig(Qwen3NextConfig): model_type = "qwen3_5_text" base_config_key = "text_config" + attribute_map = {"num_mtp_layers": "mtp_num_hidden_layers"} + base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", @@ -123,6 +128,7 @@ class Qwen3_5TextConfig(Qwen3NextConfig): intermediate_size: int = 12288 num_hidden_layers: int = 32 num_key_value_heads: int = 4 + mtp_num_hidden_layers: int = 0 decoder_sparse_step = AttributeError() norm_topk_prob = AttributeError() @@ -138,6 +144,13 @@ def __post_init__(self, **kwargs): super().__post_init__(**kwargs) del self.mlp_only_layers + # MTP layers always use full attention + @property + def mtp_layer_types(self): + if self.num_mtp_layers is not None: + return ["full_attention"] * self.num_mtp_layers + return None + @auto_docstring(checkpoint="Qwen/Qwen3.5-27B") @strict @@ -156,6 +169,10 @@ class Qwen3_5VisionConfig(Qwen3VLVisionConfig): @strict class Qwen3_5Config(Qwen3VLConfig): r""" + mtp_num_hidden_layers (`int`, *optional*, defaults to 0): + Number of Multi-Token Prediction (MTP) layers. When set to 0, MTP is disabled. + This field is exposed as `num_mtp_layers` via `attribute_map`. + Example: ```python @@ -176,6 +193,17 @@ class Qwen3_5Config(Qwen3VLConfig): vision_start_token_id: int = 248053 vision_end_token_id: int = 248054 + attribute_map = {"num_mtp_layers": "mtp_num_hidden_layers"} + + mtp_num_hidden_layers: int = 0 + + # MTP layers always use full attention + @property + def mtp_layer_types(self): + if self.num_mtp_layers is not None: + return ["full_attention"] * self.num_mtp_layers + return None + class Qwen3_5VisionRotaryEmbedding(Qwen3VLVisionRotaryEmbedding): pass diff --git a/tests/generation/test_utils.py b/tests/generation/test_utils.py index fd5dac37c00c..25066886f10a 100644 --- a/tests/generation/test_utils.py +++ b/tests/generation/test_utils.py @@ -2488,6 +2488,11 @@ def test_generate_with_mtp(self): ): self.skipTest("No MTP keys registered") + # Assisted generation (use_mtp) is blocked for stateful models, whose recurrent state + # cannot be rolled back during candidate verification + if getattr(model_class, "_is_stateful", False): + self.skipTest("MTP generation is not supported for stateful models") + config.get_text_config().num_mtp_layers = 1 model = model_class(config).to(torch_device).eval() mtp_model = MtpModel(model, num_mtp_layers=1)