diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index f315cbb9..ba415eb6 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -418,6 +418,7 @@ class ArchitectureConfig(BaseModelConfig): # Vision shared fields (accessed as top-level config.X by tasks) mm_tokens_per_image: int | None = None image_token_id: int | None = None + video_token_id: int | None = None spatial_merge_size: int = 2 temporal_patch_size: int = 2 deepstack_visual_indexes: list[int] | None = None @@ -649,6 +650,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: "bloom", "qwen2", "qwen2_5_vl_text", + "qwen2_5_omni_text", "qwen2_moe", "qwen2_vl_text", ), diff --git a/src/mobius/_configs/_extractors.py b/src/mobius/_configs/_extractors.py index ac4c6e41..6c36a51a 100644 --- a/src/mobius/_configs/_extractors.py +++ b/src/mobius/_configs/_extractors.py @@ -156,7 +156,7 @@ def extract_vision_config(config, parent_config, model_type: str) -> dict: Either step can populate ``fields`` (which become kwargs for :class:`VisionConfig`), or a per-model hook can return a fully-formed dict to short-circuit. The dispatcher also lifts a fixed set of - "shared" vision fields (``image_token_id``, ``spatial_merge_size``, + "shared" vision fields (``image_token_id``, ``video_token_id``, ``spatial_merge_size``, ...) up to the top-level of the returned dict so callers can access them as ``config.image_token_id`` directly. """ @@ -182,6 +182,7 @@ def extract_vision_config(config, parent_config, model_type: str) -> dict: for shared in ( "mm_tokens_per_image", "image_token_id", + "video_token_id", "spatial_merge_size", "temporal_patch_size", "deepstack_visual_indexes", diff --git a/src/mobius/_configs/_sub_configs.py b/src/mobius/_configs/_sub_configs.py index f872d9a7..d4cfd6de 100644 --- a/src/mobius/_configs/_sub_configs.py +++ b/src/mobius/_configs/_sub_configs.py @@ -47,6 +47,7 @@ class VisionConfig: norm_eps: float = 1e-6 mm_tokens_per_image: int | None = None image_token_id: int | None = None + video_token_id: int | None = None # Pixtral / Mistral-3 vision fields model_type: str | None = None head_dim: int | None = None diff --git a/src/mobius/_configs/per_model/__init__.py b/src/mobius/_configs/per_model/__init__.py index fa21685a..44d70cae 100644 --- a/src/mobius/_configs/per_model/__init__.py +++ b/src/mobius/_configs/per_model/__init__.py @@ -31,5 +31,6 @@ _phi4mm_audio, _phi4mm_vision, _qwen3_asr_audio, + _qwen25_omni_vision, _sensevoice_audio, ) diff --git a/src/mobius/_configs/per_model/_qwen25_omni_vision.py b/src/mobius/_configs/per_model/_qwen25_omni_vision.py new file mode 100644 index 00000000..a86bf69f --- /dev/null +++ b/src/mobius/_configs/per_model/_qwen25_omni_vision.py @@ -0,0 +1,39 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Qwen2.5-Omni vision extractor (vision config lives under thinker_config).""" + +from __future__ import annotations + +from mobius._configs._extractors import register_vision_hook + + +@register_vision_hook("qwen2_5_omni_text") +def _qwen25_omni_vision(config, parent_config, model_type: str, fields: dict): + thinker = getattr(parent_config, "thinker_config", None) + if thinker is None: + return None + if isinstance(thinker, dict): + thinker = type("ThinkerConfig", (), thinker)() + vision = getattr(thinker, "vision_config", None) + if vision is None: + return None + if isinstance(vision, dict): + vision = type("VisionConfig", (), vision)() + + fields.update( + hidden_size=getattr(vision, "hidden_size", None), + intermediate_size=getattr(vision, "intermediate_size", None), + num_hidden_layers=getattr(vision, "depth", None), + num_attention_heads=getattr(vision, "num_heads", None), + patch_size=getattr(vision, "patch_size", None), + out_hidden_size=getattr(vision, "out_hidden_size", None), + in_channels=getattr(vision, "in_channels", 3), + spatial_merge_size=getattr(vision, "spatial_merge_size", 2), + temporal_patch_size=getattr(vision, "temporal_patch_size", 2), + fullatt_block_indexes=getattr(vision, "fullatt_block_indexes", None), + window_size=getattr(vision, "window_size", 112), + image_token_id=getattr(thinker, "image_token_id", None), + ) + fields["video_token_id"] = getattr(thinker, "video_token_id", None) + return None diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index d9fe9ef7..55b1251f 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -127,6 +127,7 @@ from mobius.models.qwen3_asr import Qwen3ASRForConditionalGeneration from mobius.models.qwen3_tts import Qwen3TTSForConditionalGeneration from mobius.models.qwen3_tts_tokenizer import Qwen3TTSTokenizerV2Model +from mobius.models.qwen25_omni import Qwen25OmniThinkerForConditionalGeneration from mobius.models.sam2 import Sam2VisionModel from mobius.models.segformer import SegformerForSemanticSegmentation from mobius.models.sensevoice_small import SenseVoiceSmallModel @@ -625,6 +626,11 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: task="speech-to-text", config_class=WhisperConfig, ), + # --- Omni --- + "qwen2_5_omni": ModelRegistration( + Qwen25OmniThinkerForConditionalGeneration, + task="qwen25-omni", + ), # --- Encoder-only --- "albert": ModelRegistration(BertModel, task="feature-extraction"), "bert": ModelRegistration(BertModel, task="feature-extraction"), diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index 4ef07186..24eb1d45 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -54,6 +54,8 @@ "PostNormDecoderLayer", "QuantizedEmbedding", "QuantizedLinear", + "Qwen25OmniAudioAttention", + "Qwen25OmniAudioEncoderLayer", "RMSNorm", "SelectiveScan", "SiLU", @@ -221,6 +223,10 @@ from mobius.components._qwen3_vl_vision import ( Qwen3VLVisionRotaryEmbedding as Qwen3VLVisionRotaryEmbedding, ) +from mobius.components._qwen25_omni_audio import ( + Qwen25OmniAudioAttention, + Qwen25OmniAudioEncoderLayer, +) from mobius.components._qwen25_vl_vision import ( Qwen2VLVisionBlock as Qwen2VLVisionBlock, ) diff --git a/src/mobius/components/_conv.py b/src/mobius/components/_conv.py index 4cbccd1f..6657841c 100644 --- a/src/mobius/components/_conv.py +++ b/src/mobius/components/_conv.py @@ -5,14 +5,9 @@ from __future__ import annotations -from typing import TYPE_CHECKING - import onnx_ir as ir from onnxscript import OpBuilder, nn -if TYPE_CHECKING: - pass - class Conv2d(nn.Module): """2D convolution with bias. @@ -53,6 +48,43 @@ def forward(self, op: OpBuilder, x: ir.Value): ) +class Conv1d(nn.Module): + """1D convolution with bias. + + Matches ``torch.nn.Conv1d`` with ``bias=True``. The default ``padding=0`` + follows PyTorch convention; callers should specify padding explicitly. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int = 3, + stride: int = 1, + padding: int = 0, + groups: int = 1, + ): + super().__init__() + self.weight = nn.Parameter((out_channels, in_channels // groups, kernel_size)) + self.bias = nn.Parameter((out_channels,)) + self._kernel_size = kernel_size + self._stride = stride + self._padding = padding + self._groups = groups + + def forward(self, op: OpBuilder, x: ir.Value): + p = self._padding + return op.Conv( + x, + self.weight, + self.bias, + kernel_shape=[self._kernel_size], + strides=[self._stride], + pads=[p, p], + group=self._groups, + ) + + class Conv2dNoBias(nn.Module): """2D convolution without bias.""" diff --git a/src/mobius/components/_qwen25_omni_audio.py b/src/mobius/components/_qwen25_omni_audio.py new file mode 100644 index 00000000..9d744b5d --- /dev/null +++ b/src/mobius/components/_qwen25_omni_audio.py @@ -0,0 +1,160 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Qwen2.5-Omni audio encoder components. + +Packed bidirectional transformer layers with LayerNorm. + +Reference: Transformers +https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py +""" + +from __future__ import annotations + +import onnx_ir as ir +from onnxscript import nn +from onnxscript._internal import builder + +from mobius._build_context import get_build_dtype +from mobius.components._common import LayerNorm, Linear + + +class Qwen25OmniAudioAttention(nn.Module): + """Bidirectional multi-head attention for Qwen2_5Omni audio encoder. + + Unlike WhisperAttention, all projections (Q, V, Out) have bias and K does not have bias. + No causal masking — the encoder uses full bidirectional attention. + """ + + def __init__(self, d_model: int, num_heads: int): + super().__init__() + self.q_proj = Linear(d_model, d_model, bias=True) + self.k_proj = Linear(d_model, d_model, bias=False) + self.v_proj = Linear(d_model, d_model, bias=True) + self.out_proj = Linear(d_model, d_model, bias=True) + self._num_heads = num_heads + self._head_dim = d_model // num_heads + + def forward( + self, + op: builder.OpBuilder, + hidden_states: ir.Value, + cu_seqlens: ir.Value, + ): + """Bidirectional self-attention. + + Args: + hidden_states: (batch, seq_len, d_model) + + Returns: + output: (batch, seq_len, d_model) + """ + seq_len = op.Shape(hidden_states, start=0, end=1) + packed_shape = op.Concat(seq_len, [self._num_heads, self._head_dim], axis=0) + q = op.Reshape(self.q_proj(op, hidden_states), packed_shape) + k = op.Reshape(self.k_proj(op, hidden_states), packed_shape) + v = op.Reshape(self.v_proj(op, hidden_states), packed_shape) + + # Build the block-diagonal mask represented by HF's cu_seqlens. + positions = op.Range(0, op.Squeeze(seq_len, [0]), 1) + segment_ids = op.Sub( + op.ReduceSum( + op.Cast( + op.GreaterOrEqual( + op.Unsqueeze(positions, [1]), + op.Unsqueeze(op.Cast(cu_seqlens, to=7), [0]), + ), + to=7, + ), + [1], + keepdims=False, + ), + 1, + ) + same_segment = op.Equal( + op.Unsqueeze(segment_ids, [1]), + op.Unsqueeze(segment_ids, [0]), + ) + attention_bias = op.Where( + same_segment, + op.CastLike(0.0, q), + op.CastLike(-1e9, q), + ) + attention_bias = op.Unsqueeze(attention_bias, [0, 1]) + + q = op.Unsqueeze(op.Transpose(q, perm=[1, 0, 2]), [0]) + k = op.Unsqueeze(op.Transpose(k, perm=[1, 0, 2]), [0]) + v = op.Unsqueeze(op.Transpose(v, perm=[1, 0, 2]), [0]) + attn_output = op.Attention( + q, + k, + v, + attention_bias, + q_num_heads=self._num_heads, + kv_num_heads=self._num_heads, + scale=float(self._head_dim**-0.5), + ) + attn_output = op.Transpose(op.Squeeze(attn_output, [0]), perm=[1, 0, 2]) + attn_output = op.Reshape(attn_output, op.Concat(seq_len, [-1], axis=0)) + return self.out_proj(op, attn_output) + + +class Qwen25OmniAudioEncoderLayer(nn.Module): + """Qwen25-Omni audio encoder layer. + + Pre-norm pattern: LayerNorm → self-attn → residual + → LayerNorm → FFN → residual. + Uses GELU activation in the FFN. + + Huggingface class: ``Qwen2_5OmniAudioEncoder`` + """ + + def __init__( + self, + d_model: int, + num_heads: int, + ffn_dim: int, + eps: float = 1e-5, + ): + super().__init__() + self.self_attn = Qwen25OmniAudioAttention(d_model, num_heads) + self.self_attn_layer_norm = LayerNorm(d_model, eps=eps) + self.fc1 = Linear(d_model, ffn_dim, bias=True) + self.fc2 = Linear(ffn_dim, d_model, bias=True) + self.final_layer_norm = LayerNorm(d_model, eps=eps) + + def forward( + self, + op: builder.OpBuilder, + hidden_states: ir.Value, + cu_seqlens: ir.Value, + ): + """Pre-norm encoder layer with bidirectional attention. + + Args: + hidden_states: (batch, seq_len, d_model) + + Returns: + hidden_states: (batch, seq_len, d_model) + """ + # Self-attention with pre-norm and residual + residual = hidden_states + hidden_states = self.self_attn_layer_norm(op, hidden_states) + hidden_states = self.self_attn(op, hidden_states, cu_seqlens) + hidden_states = op.Add(residual, hidden_states) + + # FFN with pre-norm, GELU, and residual + residual = hidden_states + hidden_states = self.final_layer_norm(op, hidden_states) + hidden_states = self.fc1(op, hidden_states) + hidden_states = op.Gelu(hidden_states) + hidden_states = self.fc2(op, hidden_states) + hidden_states = op.Add(residual, hidden_states) + if get_build_dtype() == ir.DataType.FLOAT16: + hidden_states = op.Clip( + hidden_states, + op.CastLike(-64504.0, hidden_states), + op.CastLike(64504.0, hidden_states), + ) + + return hidden_states diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 49c94dd1..1516759f 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -93,6 +93,7 @@ "Phi4MMMultiModalModel", "PhiCausalLMModel", "Qwen25VLCausalLMModel", + "Qwen25OmniThinkerForConditionalGeneration", "Qwen25VLDecoderModel", "Qwen25VLEmbeddingModel", "Qwen25VLTextModel", @@ -250,6 +251,7 @@ Qwen3TTSCodecEncoderModel, Qwen3TTSTokenizerV2Model, ) +from mobius.models.qwen25_omni import Qwen25OmniThinkerForConditionalGeneration from mobius.models.qwen35 import ( Qwen35CausalLMModel, Qwen35MoECausalLMModel, diff --git a/src/mobius/models/qwen25_omni.py b/src/mobius/models/qwen25_omni.py new file mode 100644 index 00000000..fa54e6d3 --- /dev/null +++ b/src/mobius/models/qwen25_omni.py @@ -0,0 +1,528 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Qwen2.5-Omni Thinker: audio + vision + text. + +Architecture (Thinker only): + - Audio encoder: Conv1d x2 → sinusoidal PE → 32 encoder layers → AvgPool → proj + - Vision encoder: Conv3d patch embed → 32 ViT blocks → patch merger + - Fusion: Audio/vision features replace placeholder token positions + - Text decoder: Qwen2 (no QK norm) + MRoPE + +Reference: https://huggingface.co/Qwen/Qwen2.5-Omni-7B +HuggingFace class: Qwen2_5OmniForConditionalGeneration +""" + +from __future__ import annotations + +import numpy as np +import onnx_ir as ir +import torch +from onnxscript import OpBuilder, nn + +from mobius._build_context import ep_capabilities +from mobius._configs import ArchitectureConfig +from mobius.components import ( + GatedMLP, + Qwen25OmniAudioEncoderLayer, + Qwen25VLPatchMerger, + Qwen25VLVisionAttention, + Qwen25VLVisionBlock, + Qwen25VLVisionModel, +) +from mobius.components._common import ( + Embedding, + LayerNorm, + Linear, + create_attention_bias, +) +from mobius.components._conv import Conv1d +from mobius.components._decoder import DecoderLayer +from mobius.components._rms_norm import RMSNorm +from mobius.components._rotary_embedding import initialize_rope + + +def _sinusoidal_position_embedding(max_positions: int, d_model: int) -> np.ndarray: + """Compute sinusoidal positional embeddings matching Qwen3-ASR. + + Uses log-timescale increments (different from Whisper which uses + alternating sin/cos layout). Layout: [sin_0..sin_n, cos_0..cos_n]. + """ + channels = d_model + log_timescale_increment = np.log(10000.0) / (channels // 2 - 1) + inv_timescales = np.exp( + -log_timescale_increment * np.arange(channels // 2, dtype=np.float32) + ) + scaled_time = ( + np.arange(max_positions, dtype=np.float32)[:, np.newaxis] + * inv_timescales[np.newaxis, :] + ) + # Layout: [sin, cos] matching HF SinusoidsPositionEmbedding + pe = np.concatenate([np.sin(scaled_time), np.cos(scaled_time)], axis=1).astype(np.float32) + return pe + + +class Qwen25OmniAudioEncoder(nn.Module): + """Qwen25-Omni audio encoder. + + Converts mel spectrogram to audio feature embeddings: + mel (batch, num_mel_bins, seq_len) + -> 2x Conv1d with GELU + -> sinusoidal position embeddings + -> N bidirectional encoder layers + -> AvgPool1d (2x downsample) + -> LayerNorm (ln_post) + -> Linear proj (d_model -> output_dim) + + Output: (batch, out_seq_len, output_dim) + """ + + def __init__(self, config: ArchitectureConfig): + super().__init__() + audio = config.audio + assert audio is not None + + d_model = audio.d_model or 1280 + self._d_model = d_model + num_mel_bin = audio.num_mel_bins or 128 + encoder_layers = audio.encoder_layers or 32 + encoder_heads = audio.encoder_attention_heads or 20 + encoder_ffn = audio.encoder_ffn_dim or 5120 + max_source_positions = audio.max_source_positions or 1500 + output_dim = audio.output_dim or 3584 + + # 2x Conv1d: mel -> d_model with GELU between them + self.conv1 = Conv1d(num_mel_bin, d_model, kernel_size=3, padding=1) + self.conv2 = Conv1d(d_model, d_model, kernel_size=3, stride=2, padding=1) + + # Sinusoidal positional embeddings (frozen) + pe_data = _sinusoidal_position_embedding(max_source_positions, d_model) + self.positional_embedding = nn.Parameter( + [max_source_positions, d_model], + name="positional_embedding.positional_embedding", + data=ir.tensor(pe_data), + ) + + # Encoder transformer layers + self.layers = nn.ModuleList( + [ + Qwen25OmniAudioEncoderLayer(d_model, encoder_heads, encoder_ffn) + for _ in range(encoder_layers) + ] + ) + + # Post-encoder normalization + self.ln_post = LayerNorm(d_model) + + # Output projection: d_model -> output_dim + self.proj = Linear(d_model, output_dim) + + def forward( + self, + op: OpBuilder, + input_features: ir.Value, + chunk_lengths: ir.Value, + pool_indices: ir.Value, + ): + """Encode pre-chunked mel spectrograms to packed audio features. + + Args: + input_features: (num_chunks, num_mel_bins, max_chunk_len) + chunk_lengths: Valid mel-frame count for each chunk. + pool_indices: Indices of the first token in each stride-2 pooling pair. + + Returns: + audio_features: (num_audio_tokens, output_dim) + """ + input_features = op.CastLike(input_features, self.conv1.weight) + + # Match HF's chunk padding mask before the stride-2 convolution. + chunk_seq_len = op.Shape(input_features, start=2, end=3) + chunk_positions = op.Range(0, op.Squeeze(chunk_seq_len, [0]), 1) + chunk_mask = op.Less( + op.Unsqueeze(chunk_positions, [0]), + op.Unsqueeze(chunk_lengths, [1]), + ) + chunk_mask = op.Unsqueeze(op.CastLike(chunk_mask, input_features), [1]) + + # (num_chunks, mel, time) -> (num_chunks, d_model, ceil(time / 2)) + hidden_states = op.Mul(op.Gelu(self.conv1(op, input_features)), chunk_mask) + hidden_states = op.Gelu(self.conv2(op, hidden_states)) + + # (num_chunks, d_model, time) -> (num_chunks, time, d_model) + hidden_states = op.Transpose(hidden_states, perm=[0, 2, 1]) + + # Add sinusoidal positional embeddings + seq_len = op.Shape(hidden_states, start=1, end=2) + pe_slice = op.Slice( + self.positional_embedding, + op.Constant(value_ints=[0]), + seq_len, + op.Constant(value_ints=[0]), + ) + hidden_states = op.Add(hidden_states, pe_slice) + + # Remove per-chunk padding and derive packed-attention boundaries. + after_conv_lengths = op.Add(op.Div(op.Sub(chunk_lengths, 1), 2), 1) + valid_mask = op.Less( + op.Unsqueeze(op.Range(0, op.Squeeze(seq_len, [0]), 1), [0]), + op.Unsqueeze(after_conv_lengths, [1]), + ) + valid_indices = op.Squeeze(op.NonZero(op.Reshape(valid_mask, [-1])), [0]) + hidden_states = op.Gather( + op.Reshape(hidden_states, [-1, self._d_model]), + valid_indices, + axis=0, + ) + cu_seqlens = op.Concat( + op.Constant(value_ints=[0]), + op.CumSum(after_conv_lengths, op.Constant(value_int=0)), + axis=0, + ) + + for layer in self.layers: + hidden_states = layer(op, hidden_states, cu_seqlens) + + # HF pools adjacent valid tokens using indices computed per original audio. + pooled_first = op.Gather(hidden_states, pool_indices, axis=0) + pooled_second = op.Gather(hidden_states, op.Add(pool_indices, 1), axis=0) + hidden_states = op.Mul( + op.Add(pooled_first, pooled_second), + op.CastLike(0.5, hidden_states), + ) + hidden_states = self.ln_post(op, hidden_states) + hidden_states = self.proj(op, hidden_states) + return hidden_states + + +class Qwen25OmniVisionAttention(Qwen25VLVisionAttention): + """Qwen2.5-Omni vision attention with separate Q/K/V checkpoint weights.""" + + def __init__(self, hidden_size: int, num_heads: int): + nn.Module.__init__(self) + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.q = Linear(hidden_size, hidden_size, bias=True) + self.k = Linear(hidden_size, hidden_size, bias=True) + self.v = Linear(hidden_size, hidden_size, bias=True) + self.proj = Linear(hidden_size, hidden_size, bias=True) + + def forward(self, op, hidden_states, cu_seqlens, cos, sin): + seq_len = op.Shape(hidden_states, start=0, end=1) + head_shape = op.Concat(seq_len, [self.num_heads, self.head_dim], axis=0) + q = self._apply_rotary(op, op.Reshape(self.q(op, hidden_states), head_shape), cos, sin) + k = self._apply_rotary(op, op.Reshape(self.k(op, hidden_states), head_shape), cos, sin) + v = op.Reshape(self.v(op, hidden_states), head_shape) + + if ep_capabilities().supports_packed_multi_head_attention: + output = self._emit_packed_mha(op, q, k, v, cu_seqlens, seq_len) + else: + output = self._emit_standard_attention(op, q, k, v, cu_seqlens, seq_len) + return self.proj(op, output) + + +class Qwen25OmniVisionBlock(Qwen25VLVisionBlock): + """Qwen2.5-Omni vision block with separate attention projections.""" + + def __init__(self, hidden_size: int, intermediate_size: int, num_heads: int): + super().__init__(hidden_size, intermediate_size, num_heads) + self.attn = Qwen25OmniVisionAttention(hidden_size, num_heads) + self.mlp = GatedMLP( + hidden_size, + intermediate_size, + activation="silu", + bias=True, + ) + + +class Qwen25OmniVisionModel(Qwen25VLVisionModel): + """Qwen2.5-Omni vision tower using the Omni checkpoint parameter layout.""" + + def __init__( + self, + depth: int, + hidden_size: int, + intermediate_size: int, + num_heads: int, + **kwargs, + ): + super().__init__( + depth, + hidden_size, + intermediate_size, + num_heads, + **kwargs, + ) + self.blocks = nn.ModuleList( + [ + Qwen25OmniVisionBlock(hidden_size, intermediate_size, num_heads) + for _ in range(depth) + ] + ) + self.merger = Qwen25VLPatchMerger( + out_hidden_size=kwargs.get("out_hidden_size") or hidden_size, + hidden_size=hidden_size, + spatial_merge_size=kwargs.get("spatial_merge_size", 2), + ) + + +class Qwen25OmniVisionEncoder(nn.Module): + """Qwen2.5-Omni vision encoder.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + vc = config.vision + assert vc is not None + + self.visual = Qwen25OmniVisionModel( + depth=vc.num_hidden_layers or 32, + hidden_size=vc.hidden_size or 1280, + intermediate_size=vc.intermediate_size or 3420, + num_heads=vc.num_attention_heads or 16, + patch_size=vc.patch_size or 14, + temporal_patch_size=vc.temporal_patch_size or 2, + in_channels=vc.in_channels or 3, + out_hidden_size=vc.out_hidden_size or 3584, + spatial_merge_size=vc.spatial_merge_size or 2, + fullatt_block_indexes=vc.fullatt_block_indexes or (7, 15, 23, 31), + window_size=vc.window_size or 112, + ) + + def forward(self, op: OpBuilder, pixel_values: ir.Value, image_grid_thw: ir.Value): + return self.visual(op, pixel_values, image_grid_thw) + + +class Qwen25OmniEmbeddingModel(nn.Module): + """Fuses text embedding with audio and image features. + + Replaces audio, image, and video placeholder tokens with encoder features. + + Inputs: + input_ids: (batch, seq_len) + audio_features: (num_audio_tokens, hidden_size) + image_features: (num_image_tokens, hidden_size) + video_features: (num_video_tokens, hidden_size) + + Output: + inputs_embeds: (batch, seq_len, hidden_size) + """ + + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + config.pad_token_id, + ) + + # Token IDs default to the values used by Qwen2.5-Omni-7B. + audio = config.audio + vision = config.vision + self._audio_token_id = (audio.audio_token_id if audio else None) or 151646 + self._image_token_id = (vision.image_token_id if vision else None) or 151655 + self._video_token_id = config.video_token_id or 151656 + + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + audio_features: ir.Value, + image_features: ir.Value, + video_features: ir.Value, + ): + inputs_embeds = self.embed_tokens(op, input_ids) + + # Fuse audio features at audio token positions. + inputs_embeds = self._replace_tokens( + op, inputs_embeds, input_ids, audio_features, self._audio_token_id + ) + + # Fuse image features at image token positions. + inputs_embeds = self._replace_tokens( + op, inputs_embeds, input_ids, image_features, self._image_token_id + ) + inputs_embeds = self._replace_tokens( + op, inputs_embeds, input_ids, video_features, self._video_token_id + ) + + return inputs_embeds + + def _replace_tokens(self, op, inputs_embeds, input_ids, features, token_id): + """Replace token positions with encoder features (masked_scatter equivalent).""" + mask = op.Equal(input_ids, op.Constant(value_int=token_id)) + mask_3d = op.Unsqueeze(mask, [-1]) + + # Pad with a zero row for safety (text-only case: no features). + feature_dim = op.Shape(features, start=1, end=2) + zero_shape = op.Concat(op.Constant(value_ints=[1]), feature_dim, axis=0) + zero_row = op.Expand(op.CastLike(0.0, features), zero_shape) + padded = op.Concat(zero_row, features, axis=0) + + # CumSum-based per-position gather index. Mask positions get the + # next feature row in order; non-mask positions get the zero row. + mask_int = op.Cast(mask, to=7) + flat = op.Reshape(mask_int, op.Constant(value_ints=[-1])) + indices = op.CumSum(flat, op.Constant(value_int=0)) + indices = op.Mul(indices, flat) + indices = op.Reshape(indices, op.Shape(input_ids)) + + gathered = op.Gather(padded, indices, axis=0) + return op.Where(mask_3d, gathered, inputs_embeds) + + +class Qwen25OmniDecoderModel(nn.Module): + """Qwen2.5-Omni text decoder: inputs_embeds → logits + KV cache. + + Standard Qwen2 decoder with MRoPE (3D position_ids). + No QK norm (unlike Qwen3-ASR which uses attn_qk_norm=True). + """ + + def __init__(self, config: ArchitectureConfig): + super().__init__() + self._dtype = config.dtype + self.layers = nn.ModuleList( + [DecoderLayer(config) for _ in range(config.num_hidden_layers)] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = initialize_rope(config) + self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False) + + def forward( + self, + op: OpBuilder, + inputs_embeds: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values=None, + ): + hidden_states = inputs_embeds + position_embeddings = ( + self.rotary_emb(op, position_ids) if self.rotary_emb is not None else None + ) + + attention_bias = create_attention_bias( + op, + input_ids=inputs_embeds, + attention_mask=attention_mask, + dtype=self._dtype, + ) + + present_key_values = [] + past_kvs = past_key_values or [None] * len(self.layers) + for layer, past_kv in zip(self.layers, past_kvs): + hidden_states, present_kv = layer( + op, + hidden_states=hidden_states, + attention_bias=attention_bias, + position_embeddings=position_embeddings, + past_key_value=past_kv, + ) + present_key_values.append(present_kv) + + hidden_states = self.norm(op, hidden_states) + logits = self.lm_head(op, hidden_states) + return logits, present_key_values + + +class Qwen25OmniThinkerForConditionalGeneration(nn.Module): + """Qwen2.5-Omni Thinker: composite audio + vision + text model. + + Builds four separate ONNX models: + + - ``decoder``: Qwen2.5 text decoder taking ``inputs_embeds`` + - ``vision_encoder``: Qwen2.5-VL ViT (pixel_values + grid_thw → image features) + - ``audio_tower``: 2x Conv1d + transformer audio tower (mel → audio features) + - ``embedding``: word embedding + multimodal feature fusion + + HuggingFace class: ``Qwen2_5OmniForConditionalGeneration`` (Thinker only — + the Talker / streaming code generation head is out of scope for now). + """ + + default_task: str = "qwen25-omni" + category: str = "Multimodal" + config_class: type = ArchitectureConfig + + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.config = config + self.decoder = Qwen25OmniDecoderModel(config) + self.embedding = Qwen25OmniEmbeddingModel(config) + self.vision_encoder: Qwen25OmniVisionEncoder | None = ( + Qwen25OmniVisionEncoder(config) if config.vision is not None else None + ) + self.audio_encoder: Qwen25OmniAudioEncoder | None = ( + Qwen25OmniAudioEncoder(config) if config.audio is not None else None + ) + + def forward(self, op: OpBuilder, **kwargs): + raise NotImplementedError( + "Qwen25OmniThinkerForConditionalGeneration is a multi-model split; the corresponding " + "Qwen25OmniTask builds each sub-module (decoder, embedding, vision_encoder, " + "audio_encoder) " + "separately." + ) + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Map HuggingFace weight names to ONNX module structure. + + HF Qwen2.5-Omni checkpoints prefix every Thinker key with ``thinker.``: + + - ``thinker.audio_tower.*`` → ``audio_encoder.*`` + - ``thinker.visual.*`` → ``vision_encoder.visual.*`` + - ``thinker.model.embed_tokens.*`` → ``embedding.embed_tokens.*`` + - ``thinker.model.layers.N.*`` and ``model.norm.*`` → ``decoder.*`` + - ``thinker.lm_head.*`` → ``decoder.lm_head.*`` + - ``thinker.model.rotary_emb.*`` → ``decoder.rotary_emb.*`` + + The Talker sub-tree (``talker.*``) and the audio-output codec head + are not consumed by this model and are silently dropped. + """ + cleaned: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + # Strip the thinker. prefix if present. + if key.startswith("thinker."): + key = key[len("thinker.") :] + + # Drop talker.* and any codec output keys — not part of Thinker. + if key.startswith(("talker.", "token2wav.", "code_predictor.")): + continue + + if key.startswith("audio_tower."): + if ".audio_bos_eos_token." not in key: + cleaned["audio_encoder." + key[len("audio_tower.") :]] = value + continue + + if key.startswith("visual."): + new_key = "vision_encoder." + key + new_key = new_key.replace(".merger.mlp.0.", ".merger.mlp_0.") + new_key = new_key.replace(".merger.mlp.2.", ".merger.mlp_2.") + cleaned[new_key] = value + continue + + if key.startswith("lm_head."): + cleaned["decoder." + key] = value + continue + + if key.startswith("model."): + inner = key[len("model.") :] + if inner.startswith("embed_tokens."): + cleaned["embedding." + inner] = value + continue + if inner.startswith(("layers.", "norm.", "rotary_emb.")): + cleaned["decoder." + inner] = value + continue + + cleaned[key] = value + + # Weight tying: ``embedding.embed_tokens.weight`` ↔ ``decoder.lm_head.weight``. + embed_key = "embedding.embed_tokens.weight" + lm_key = "decoder.lm_head.weight" + if getattr(self.config, "tie_word_embeddings", False): + if embed_key in cleaned and lm_key not in cleaned: + cleaned[lm_key] = cleaned[embed_key] + elif lm_key in cleaned and embed_key not in cleaned: + cleaned[embed_key] = cleaned[lm_key] + + return cleaned diff --git a/src/mobius/models/qwen25_omni_test.py b/src/mobius/models/qwen25_omni_test.py new file mode 100644 index 00000000..990bf74d --- /dev/null +++ b/src/mobius/models/qwen25_omni_test.py @@ -0,0 +1,106 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from mobius._configs import ArchitectureConfig +from mobius.models.qwen25_omni import Qwen25OmniThinkerForConditionalGeneration + + +def _hf_config(): + text = SimpleNamespace( + model_type="qwen2_5_omni_text", + vocab_size=256, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + hidden_act="silu", + rms_norm_eps=1e-6, + max_position_embeddings=128, + rope_parameters={ + "rope_type": "default", + "rope_theta": 1_000_000.0, + "mrope_section": [4, 2, 2], + }, + tie_word_embeddings=False, + ) + thinker = SimpleNamespace( + text_config=text, + audio_config=SimpleNamespace( + d_model=64, + encoder_layers=2, + encoder_attention_heads=4, + encoder_ffn_dim=128, + num_mel_bins=32, + max_source_positions=128, + n_window=8, + output_dim=64, + ), + vision_config=SimpleNamespace( + hidden_size=64, + intermediate_size=128, + depth=2, + num_heads=4, + patch_size=14, + temporal_patch_size=2, + in_channels=3, + out_hidden_size=64, + spatial_merge_size=2, + fullatt_block_indexes=[0], + window_size=112, + ), + audio_token_id=100, + image_token_id=101, + video_token_id=102, + ) + return text, SimpleNamespace(thinker_config=thinker, tie_word_embeddings=False) + + +def test_qwen25_omni_extracts_nested_thinker_config(): + text, parent = _hf_config() + config = ArchitectureConfig.from_transformers(text, parent_config=parent) + + assert config.attn_qkv_bias + assert config.audio is not None + assert config.audio.encoder_ffn_dim == 128 + assert config.audio.audio_token_id == 100 + assert config.vision is not None + assert config.vision.hidden_size == 64 + assert config.image_token_id == 101 + assert config.video_token_id == 102 + + +def test_qwen25_omni_preprocess_weights_routes_thinker_components(): + text, parent = _hf_config() + config = ArchitectureConfig.from_transformers(text, parent_config=parent) + model = Qwen25OmniThinkerForConditionalGeneration(config) + weight = torch.randn(1) + + processed = model.preprocess_weights( + { + "thinker.audio_tower.conv1.weight": weight, + "thinker.audio_tower.audio_bos_eos_token.weight": weight, + "thinker.visual.blocks.0.attn.q.weight": weight, + "thinker.visual.merger.mlp.0.weight": weight, + "thinker.model.embed_tokens.weight": weight, + "thinker.model.layers.0.self_attn.q_proj.bias": weight, + "thinker.lm_head.weight": weight, + "talker.model.layers.0.weight": weight, + "token2wav.dit.weight": weight, + } + ) + + assert set(processed) == { + "audio_encoder.conv1.weight", + "vision_encoder.visual.blocks.0.attn.q.weight", + "vision_encoder.visual.merger.mlp_0.weight", + "embedding.embed_tokens.weight", + "decoder.layers.0.self_attn.q_proj.bias", + "decoder.lm_head.weight", + } diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 2671c100..49e7ba55 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -31,6 +31,7 @@ "DFlashDraftTask", "Eagle3DraftTask", "Qwen35MtpTask", + "Qwen25OmniTask", "DenoisingTask", "FeatureExtractionTask", "FunASRSpeechLanguageTask", @@ -102,6 +103,7 @@ from mobius.tasks._multimodal import MultiModalTask from mobius.tasks._object_detection import ObjectDetectionTask from mobius.tasks._phi4mm_multimodal import Phi4MMMultiModalTask +from mobius.tasks._qwen25_omni import Qwen25OmniTask from mobius.tasks._qwen35_mtp import Qwen35MtpTask from mobius.tasks._qwen_image_vae import QwenImageVAETask from mobius.tasks._rnnt import RNNTTask @@ -144,6 +146,7 @@ "dflash-draft": DFlashDraftTask, "eagle3-draft": Eagle3DraftTask, "qwen35-mtp": Qwen35MtpTask, + "qwen25-omni": Qwen25OmniTask, "vae": VAETask, "qwen-image-vae": QwenImageVAETask, "vision-language": VisionLanguageTask, diff --git a/src/mobius/tasks/_qwen25_omni.py b/src/mobius/tasks/_qwen25_omni.py new file mode 100644 index 00000000..34e05eb3 --- /dev/null +++ b/src/mobius/tasks/_qwen25_omni.py @@ -0,0 +1,123 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Qwen2.5-Omni Thinker four-model split task.""" + +from __future__ import annotations + +from typing import ClassVar + +import onnx_ir as ir +from onnxscript import nn + +from mobius._configs import ArchitectureConfig +from mobius._model_package import ModelPackage +from mobius.tasks._base import ( + ComponentSpec, + _make_graph, + _make_model, + build_decoder_from_embeds, +) +from mobius.tasks._vision_language_3model import QwenVLTask + + +class Qwen25OmniTask(QwenVLTask): + """Build the Thinker's audio, vision, embedding, and decoder ONNX models.""" + + model_roles: ClassVar[dict[str, str]] = { + "audio_encoder": "encoder", + "vision_encoder": "encoder", + "embedding": "embedding", + "decoder": "decoder", + } + components = ComponentSpec( + audio_encoder="audio_encoder", + vision_encoder="vision_encoder", + embedding="embedding", + decoder="decoder", + ) + + def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage: + self._validate_components(module) + models = { + "audio_encoder": self._build_audio(module.audio_encoder, config), + "vision_encoder": self._build_vision(module.vision_encoder, config), + "embedding": self._build_embedding(module.embedding, config), + "decoder": build_decoder_from_embeds(module.decoder, config, mrope=True), + } + return ModelPackage(models, config=config) + + def _build_audio(self, audio_encoder: nn.Module, config: ArchitectureConfig) -> ir.Model: + """Build packed audio chunks into packed LLM audio tokens.""" + num_chunks = ir.SymbolicDim("num_audio_chunks") + chunk_len = ir.SymbolicDim("audio_chunk_len") + num_audio_tokens = ir.SymbolicDim("num_audio_tokens") + n_mels = (config.audio.num_mel_bins if config.audio else None) or 128 + + graph, builder = _make_graph(name="audio_encoder") + input_features = builder.input( + "input_features", + dtype=ir.DataType.FLOAT, + shape=[num_chunks, n_mels, chunk_len], + ) + chunk_lengths = builder.input( + "chunk_lengths", + dtype=ir.DataType.INT64, + shape=[num_chunks], + ) + pool_indices = builder.input( + "pool_indices", + dtype=ir.DataType.INT64, + shape=[num_audio_tokens], + ) + audio_features = audio_encoder( + builder.op, + input_features, + chunk_lengths, + pool_indices, + ) + builder.add_output(audio_features, "audio_features") + return _make_model(graph) + + def _build_embedding( + self, + embedding: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build text embedding and three-modality feature replacement.""" + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + num_audio_tokens = ir.SymbolicDim("num_audio_tokens") + num_image_tokens = ir.SymbolicDim("num_image_tokens") + num_video_tokens = ir.SymbolicDim("num_video_tokens") + + graph, builder = _make_graph(name="embedding") + input_ids = builder.input( + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], + ) + audio_features = builder.input( + "audio_features", + dtype=config.dtype, + shape=[num_audio_tokens, config.hidden_size], + ) + image_features = builder.input( + "image_features", + dtype=config.dtype, + shape=[num_image_tokens, config.hidden_size], + ) + video_features = builder.input( + "video_features", + dtype=config.dtype, + shape=[num_video_tokens, config.hidden_size], + ) + inputs_embeds = embedding( + builder.op, + input_ids, + audio_features, + image_features, + video_features, + ) + builder.add_output(inputs_embeds, "inputs_embeds") + return _make_model(graph) diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index 8053003a..1548fbd8 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -2862,6 +2862,72 @@ def test_3model_pipeline_runs_with_ort(self): assert logits.shape[1] == seq_len +class TestBuildGraphQwen25Omni: + """Verify the Qwen2.5-Omni Thinker four-model split.""" + + def _omni_config(self): + return _base_config( + model_type="qwen2_5_omni_text", + attn_qkv_bias=True, + hidden_act="silu", + mrope_section=[4, 2, 2], + audio=AudioConfig( + d_model=64, + encoder_layers=2, + encoder_attention_heads=4, + encoder_ffn_dim=128, + num_mel_bins=32, + max_source_positions=128, + output_dim=64, + audio_token_id=100, + n_window=8, + ), + vision=VisionConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + patch_size=14, + temporal_patch_size=2, + in_channels=3, + out_hidden_size=64, + spatial_merge_size=2, + fullatt_block_indexes=[0], + window_size=112, + image_token_id=101, + video_token_id=102, + ), + image_token_id=101, + video_token_id=102, + ) + + def test_package_builds_four_models(self): + from mobius.models import Qwen25OmniThinkerForConditionalGeneration + from mobius.tasks import Qwen25OmniTask + + config = self._omni_config() + module = Qwen25OmniThinkerForConditionalGeneration(config) + package = build_from_module(module, config, task=Qwen25OmniTask()) + + assert set(package) == { + "audio_encoder", + "vision_encoder", + "embedding", + "decoder", + } + assert {value.name for value in package["audio_encoder"].graph.inputs} == { + "input_features", + "chunk_lengths", + "pool_indices", + } + assert {value.name for value in package["embedding"].graph.inputs} == { + "input_ids", + "audio_features", + "image_features", + "video_features", + } + + class TestBuildGraphFunASR: """Verify Fun-ASR-Nano 3-model split with FunASRSpeechLanguageTask.""" @@ -4938,6 +5004,7 @@ def test_jamba_preprocess_weights_moe_renames(self): "mms", "qwen3_asr", "qwen3_forced_aligner", + "qwen2_5_omni", "qwen3_tts", "qwen3_tts_tokenizer_12hz", "whisper",