# coding=utf-8 # Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved. # # 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. """PyTorch Qwen3-VL-MOE model.""" from typing import Optional, Union import torch import torch.nn as nn from ...activations import ACT2FN from ...cache_utils import Cache from ...configuration_utils import PretrainedConfig from ...modeling_rope_utils import rope_config_validation from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import TransformersKwargs, logging from ..qwen3_moe.modeling_qwen3_moe import ( Qwen3MoeDecoderLayer, Qwen3MoePreTrainedModel, Qwen3MoeRMSNorm, load_balancing_loss_func, ) from ..qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig, Qwen3VLVisionConfig from ..qwen3_vl.modeling_qwen3_vl import ( Qwen3VLCausalLMOutputWithPast, Qwen3VLForConditionalGeneration, Qwen3VLModel, Qwen3VLTextAttention, Qwen3VLTextModel, Qwen3VLVisionModel, ) logger = logging.get_logger(__name__) class Qwen3VLMoeTextConfig(PretrainedConfig): r""" This is the configuration class to store the configuration of a [`Qwen3VLMoeTextModel`]. It is used to instantiate a Qwen3-VL-MOE model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of Qwen3-VL-30B-A3B-Instruct [Qwen/Qwen3-VL-30B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-30B-A3B-Instruct). Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: vocab_size (`int`, *optional*, defaults to 151936): Vocabulary size of the Qwen2MoE model. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling [`Qwen2MoeModel`] hidden_size (`int`, *optional*, defaults to 2048): Dimension of the hidden representations. intermediate_size (`int`, *optional*, defaults to 5632): Dimension of the MLP representations. num_hidden_layers (`int`, *optional*, defaults to 24): Number of hidden layers in the Transformer encoder. num_attention_heads (`int`, *optional*, defaults to 16): Number of attention heads for each attention layer in the Transformer encoder. num_key_value_heads (`int`, *optional*, defaults to 16): This is the number of key_value heads that should be used to implement Grouped Query Attention. If `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed by meanpooling all the original heads within that group. For more details checkout [this paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`. hidden_act (`str` or `function`, *optional*, defaults to `"silu"`): The non-linear activation function (function or string) in the decoder. max_position_embeddings (`int`, *optional*, defaults to 128000): The maximum sequence length that this model might ever be used with. initializer_range (`float`, *optional*, defaults to 0.02): The standard deviation of the truncated_normal_initializer for initializing all weight matrices. rms_norm_eps (`float`, *optional*, defaults to 1e-06): The epsilon used by the rms normalization layers. use_cache (`bool`, *optional*, defaults to `True`): Whether or not the model should return the last key/values attentions (not used by all models). Only relevant if `config.is_decoder=True`. tie_word_embeddings (`bool`, *optional*, defaults to `False`): Whether the model's input and output word embeddings should be tied. rope_theta (`float`, *optional*, defaults to 5000000.0): The base period of the RoPE embeddings. attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): Whether to use a bias in the query, key, value and output projection layers during self-attention. attention_dropout (`float`, *optional*, defaults to 0.0): The dropout ratio for the attention probabilities. decoder_sparse_step (`int`, *optional*, defaults to 1): The frequency of the MoE layer. moe_intermediate_size (`int`, *optional*, defaults to 1408): Intermediate size of the routed expert. num_experts_per_tok (`int`, *optional*, defaults to 4): Number of selected experts. num_experts (`int`, *optional*, defaults to 60): Number of routed experts. norm_topk_prob (`bool`, *optional*, defaults to `True`): Whether to normalize the topk probabilities. router_aux_loss_coef (`float`, *optional*, defaults to 0.001): The aux loss factor for the total loss. mlp_only_layers (`List[int]`, *optional*, defaults to `[]`): Indicate which layers use Qwen3VLMoeMLP rather than Qwen3VLMoeSparseMoeBlock The list contains layer index, from 0 to num_layers-1 if we have num_layers layers If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity. rope_scaling (`Dict`, *optional*): Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value accordingly. Expected contents: `rope_type` (`str`): The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope', 'llama3'], with 'default' being the original RoPE implementation. `factor` (`float`, *optional*): Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In most scaling types, a `factor` of x will enable the model to handle sequences of length x * original maximum pre-trained length. `original_max_position_embeddings` (`int`, *optional*): Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during pretraining. `attention_factor` (`float`, *optional*): Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention computation. If unspecified, it defaults to value recommended by the implementation, using the `factor` field to infer the suggested value. `beta_fast` (`float`, *optional*): Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear ramp function. If unspecified, it defaults to 32. `beta_slow` (`float`, *optional*): Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear ramp function. If unspecified, it defaults to 1. `short_factor` (`List[float]`, *optional*): Only used with 'longrope'. The scaling factor to be applied to short contexts (< `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden size divided by the number of attention heads divided by 2 `long_factor` (`List[float]`, *optional*): Only used with 'longrope'. The scaling factor to be applied to long contexts (< `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden size divided by the number of attention heads divided by 2 `low_freq_factor` (`float`, *optional*): Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE `high_freq_factor` (`float`, *optional*): Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE head_dim (`int`, *optional*): The dimension of the head. If not specified, will default to `hidden_size // num_attention_heads`. ```python >>> from transformers import Qwen3VLMoeForConditionalGeneration, Qwen3VLMoeConfig >>> # Initializing a Qwen3VLMoe style configuration >>> configuration = Qwen3VLMoeConfig() >>> # Initializing a model from the Qwen3-VL-30B-A3B style configuration >>> model = Qwen3VLMoeForConditionalGeneration(configuration) >>> # Accessing the model configuration >>> configuration = model.config ```""" model_type = "qwen3_vl_moe_text" base_config_key = "text_config" keys_to_ignore_at_inference = ["past_key_values"] # Default tensor parallel plan for base model `Qwen3VLMoe` base_model_tp_plan = { "layers.*.self_attn.q_proj": "colwise", "layers.*.self_attn.k_proj": "colwise", "layers.*.self_attn.v_proj": "colwise", "layers.*.self_attn.o_proj": "rowwise", "layers.*.mlp.gate_proj": "colwise", "layers.*.mlp.up_proj": "colwise", "layers.*.mlp.down_proj": "rowwise", } base_model_pp_plan = { "embed_tokens": (["input_ids"], ["inputs_embeds"]), "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), "norm": (["hidden_states"], ["hidden_states"]), } def __init__( self, vocab_size=151936, hidden_size=2048, intermediate_size=5632, num_hidden_layers=24, num_attention_heads=16, num_key_value_heads=16, hidden_act="silu", max_position_embeddings=128000, initializer_range=0.02, rms_norm_eps=1e-6, use_cache=True, tie_word_embeddings=False, rope_theta=5000000.0, attention_bias=False, attention_dropout=0.0, decoder_sparse_step=1, moe_intermediate_size=1408, num_experts_per_tok=4, num_experts=60, norm_topk_prob=True, router_aux_loss_coef=0.001, mlp_only_layers=None, rope_scaling=None, head_dim=None, **kwargs, ): self.vocab_size = vocab_size self.max_position_embeddings = max_position_embeddings self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.num_hidden_layers = num_hidden_layers self.num_attention_heads = num_attention_heads # for backward compatibility if num_key_value_heads is None: num_key_value_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.hidden_act = hidden_act self.initializer_range = initializer_range self.rms_norm_eps = rms_norm_eps self.use_cache = use_cache self.rope_theta = rope_theta self.attention_bias = attention_bias self.attention_dropout = attention_dropout self.rope_scaling = rope_scaling self.head_dim = head_dim or hidden_size // num_attention_heads rope_config_validation(self, ignore_keys={"mrope_section", "mrope_interleaved"}) # MoE arguments self.decoder_sparse_step = decoder_sparse_step self.moe_intermediate_size = moe_intermediate_size self.num_experts_per_tok = num_experts_per_tok self.num_experts = num_experts self.norm_topk_prob = norm_topk_prob self.router_aux_loss_coef = router_aux_loss_coef self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs) class Qwen3VLMoeVisionConfig(Qwen3VLVisionConfig): pass class Qwen3VLMoeConfig(Qwen3VLConfig): r""" This is the configuration class to store the configuration of a [`Qwen3VLMoeModel`]. It is used to instantiate a Qwen3-VL-MOE model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of Qwen3-VL-30B-A3B-Instruct [Qwen/Qwen3-VL-30B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-30B-A3B-Instruct). Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the documentation from [`PretrainedConfig`] for more information. Args: text_config (`Union[PreTrainedConfig, dict]`, *optional*, defaults to `Qwen3VLMoeTextConfig`): The config object or dictionary of the text backbone. vision_config (`Union[PreTrainedConfig, dict]`, *optional*, defaults to `Qwen3VLMoeVisionConfig`): The config object or dictionary of the vision backbone. image_token_id (`int`, *optional*, defaults to 151655): The image token index to encode the image prompt. video_token_id (`int`, *optional*, defaults to 151656): The video token index to encode the image prompt. vision_start_token_id (`int`, *optional*, defaults to 151652): The start token index to encode the image prompt. vision_end_token_id (`int`, *optional*, defaults to 151653): The end token index to encode the image prompt. tie_word_embeddings (`bool`, *optional*, defaults to `False`): Whether to tie the word embeddings. ```python >>> from transformers import Qwen3VLMoeForConditionalGeneration, Qwen3VLMoeConfig >>> # Initializing a Qwen3-VL-MOE style configuration >>> configuration = Qwen3VLMoeConfig() >>> # Initializing a model from the Qwen3-VL-30B-A3B style configuration >>> model = Qwen3VLMoeForConditionalGeneration(configuration) >>> # Accessing the model configuration >>> configuration = model.config ```""" model_type = "qwen3_vl_moe" sub_configs = {"vision_config": Qwen3VLMoeVisionConfig, "text_config": Qwen3VLMoeTextConfig} class Qwen3VLMoeTextRMSNorm(Qwen3MoeRMSNorm): pass class Qwen3VLMoeTextExperts(nn.Module): def __init__(self, config): super().__init__() self.num_experts = config.num_experts self.intermediate_size = config.moe_intermediate_size self.hidden_size = config.hidden_size self.expert_dim = self.intermediate_size self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, 2 * self.expert_dim)) self.down_proj = nn.Parameter(torch.empty((self.num_experts, self.expert_dim, self.hidden_size))) self.act_fn = ACT2FN[config.hidden_act] def forward( self, hidden_states: torch.Tensor, routing_weights: torch.Tensor, router_indices: torch.Tensor ) -> torch.Tensor: """ When training it is more efficient to just loop over the experts and compute the output for each expert as otherwise the memory would explode. For inference we can sacrifice some memory and compute the output for all experts at once. By repeating the inputs. Args: hidden_states (torch.Tensor): (batch_size * token_num, hidden_size) routing_weights (torch.Tensor): (batch_size * token_num, num_experts) router_indices (torch.Tensor): (batch_size * token_num, top_k) Returns: torch.Tensor """ batch_size = hidden_states.shape[0] hidden_states = hidden_states.reshape(-1, self.hidden_size) # (num_tokens, hidden_size) if self.training: next_states = torch.zeros_like(hidden_states, dtype=hidden_states.dtype, device=hidden_states.device) with torch.no_grad(): expert_mask = torch.nn.functional.one_hot(router_indices, num_classes=self.num_experts) expert_mask = expert_mask.permute(2, 1, 0) # we sum on the top_k and on the sequence length to get which experts # are hit this time around expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() for expert_idx in expert_hit[:]: with torch.no_grad(): _, token_idx = torch.where(expert_mask[expert_idx[0]]) current_state = hidden_states[token_idx] gate_up = current_state @ self.gate_up_proj[expert_idx] gate, up = gate_up.chunk(2, dim=-1) gated_output = up * self.act_fn(gate) out = gated_output @ self.down_proj[expert_idx] weighted_output = out[0] * routing_weights[token_idx, expert_idx, None] next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype)) next_states = next_states.view(batch_size, -1, self.hidden_size) else: hidden_states = hidden_states.repeat(self.num_experts, 1) hidden_states = hidden_states.view(self.num_experts, -1, self.hidden_size) gate_up = torch.bmm(hidden_states, self.gate_up_proj) gate, up = gate_up.chunk(2, dim=-1) # not supported for DTensors next_states = torch.bmm((up * self.act_fn(gate)), self.down_proj) next_states = next_states.reshape(self.num_experts, batch_size, -1, self.hidden_size) next_states = ( next_states * routing_weights.transpose(0, 1).view(self.num_experts, batch_size, -1)[..., None] ) next_states = next_states.sum(dim=0) return next_states class Qwen3VLMoeTextSparseMoeBlock(nn.Module): def __init__(self, config): super().__init__() self.hidden_size = config.hidden_size self.num_experts = config.num_experts self.top_k = config.num_experts_per_tok self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) self.experts = Qwen3VLMoeTextExperts(config) # since all the models use norm_topk_prob, we don't need to have a extra check for it # self.norm_topk_prob = config.norm_topk_prob def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: batch_size = hidden_states.shape[0] hidden_states = hidden_states.reshape(-1, self.hidden_size) router_logits = self.gate(hidden_states) routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float) routing_weights, router_indices = torch.topk(routing_weights, self.top_k, dim=-1) routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) routing_weights = routing_weights.to(hidden_states.dtype) router_weights = torch.zeros_like(router_logits).scatter_(1, router_indices, routing_weights) hidden_states = hidden_states.reshape(batch_size, -1, self.hidden_size) routed_out = self.experts(hidden_states, router_weights, router_indices) return routed_out, router_logits class Qwen3VLMoeTextAttention(Qwen3VLTextAttention): pass class Qwen3VLMoeTextDecoderLayer(Qwen3MoeDecoderLayer): pass class Qwen3VLMoePreTrainedModel(Qwen3MoePreTrainedModel): config: Qwen3VLMoeConfig _no_split_modules = ["Qwen3VLMoeTextDecoderLayer", "Qwen3VLMoeVisionBlock"] def _init_weights(self, module): """Initialize the weights.""" PreTrainedModel._init_weights(self, module) if hasattr(self.config, "initializer_range"): std = self.config.initializer_range else: std = getattr(self.config.get_text_config(), "initializer_range", 0.02) if isinstance(module, Qwen3VLMoeTextExperts): module.gate_up_proj.data.normal_(mean=0.0, std=std) module.down_proj.data.normal_(mean=0.0, std=std) class Qwen3VLMoeVisionModel(Qwen3VLVisionModel): pass class Qwen3VLMoeTextModel(Qwen3VLTextModel): pass class Qwen3VLMoeCausalLMOutputWithPast(Qwen3VLCausalLMOutputWithPast): aux_loss: Optional[torch.FloatTensor] = None class Qwen3VLMoeModel(Qwen3VLModel): pass class Qwen3VLMoeForConditionalGeneration(Qwen3VLForConditionalGeneration): def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, pixel_values: Optional[torch.Tensor] = None, pixel_values_videos: Optional[torch.FloatTensor] = None, image_grid_thw: Optional[torch.LongTensor] = None, video_grid_thw: Optional[torch.LongTensor] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs: Unpack[TransformersKwargs], ): 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]`. image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*): The temporal, height and width of feature shape of each image in LLM. video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*): The temporal, height and width of feature shape of each video in LLM. Example: ```python >>> from PIL import Image >>> import requests >>> from transformers import AutoProcessor, Qwen3VLMoeForConditionalGeneration >>> model = Qwen3VLMoeForConditionalGeneration.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct", dtype="auto", device_map="auto") >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-30B-A3B-Instruct") >>> messages = [ { "role": "user", "content": [ { "type": "image", "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg", }, {"type": "text", "text": "Describe this image in short."}, ], } ] >>> # Preparation for inference >>> inputs = processor.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt" ) >>> inputs = inputs.to(model.device) >>> # Generate >>> generated_ids = model.generate(**inputs, max_new_tokens=128) >>> generated_ids_trimmed = [ out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids) ] >>> processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "A woman in a plaid shirt sits on a sandy beach at sunset, smiling as she gives a high-five to a yellow Labrador Retriever wearing a harness. The ocean waves roll in the background." ```""" outputs = self.model( input_ids=input_ids, pixel_values=pixel_values, pixel_values_videos=pixel_values_videos, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, position_ids=position_ids, attention_mask=attention_mask, past_key_values=past_key_values, inputs_embeds=inputs_embeds, cache_position=cache_position, **kwargs, ) hidden_states = outputs[0] # 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.text_config.vocab_size) aux_loss = None if kwargs.get("output_router_logits", False): aux_loss = load_balancing_loss_func( outputs.router_logits, self.config.text_config.num_experts, self.config.text_config.num_experts_per_tok, attention_mask, ) if labels is not None: loss += self.config.text_config.router_aux_loss_coef * aux_loss.to( loss.device ) # make sure to reside in the same device return Qwen3VLMoeCausalLMOutputWithPast( loss=loss, aux_loss=aux_loss, logits=logits, past_key_values=outputs.past_key_values, rope_deltas=outputs.rope_deltas, ) __all__ = [ "Qwen3VLMoeConfig", "Qwen3VLMoeTextConfig", "Qwen3VLMoeVisionModel", "Qwen3VLMoeForConditionalGeneration", "Qwen3VLMoeModel", "Qwen3VLMoePreTrainedModel", "Qwen3VLMoeTextModel", ]