diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 7078c9125..ee63aea29 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -1089,6 +1089,7 @@ def _per_layer_value(attribute: str) -> int | None: "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 5f7e14027..31231899a 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. """ diff --git a/src/mobius/_configs/per_model/__init__.py b/src/mobius/_configs/per_model/__init__.py index f4342d356..fda72ffbb 100644 --- a/src/mobius/_configs/per_model/__init__.py +++ b/src/mobius/_configs/per_model/__init__.py @@ -41,6 +41,7 @@ _phi4mm_vision, _phi_vision, _qwen3_asr_audio, + _qwen25_omni_vision, _sensenova_u1_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 000000000..8a9d22f89 --- /dev/null +++ b/src/mobius/_configs/per_model/_qwen25_omni_vision.py @@ -0,0 +1,43 @@ +# 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", getattr(thinker, "image_token_index", None) + ), + ) + fields["video_token_id"] = getattr( + thinker, "video_token_id", getattr(thinker, "video_token_index", None) + ) + return None diff --git a/src/mobius/_configs/per_model/_qwen3_asr_audio.py b/src/mobius/_configs/per_model/_qwen3_asr_audio.py index 73dcd7772..fabdbf7fd 100644 --- a/src/mobius/_configs/per_model/_qwen3_asr_audio.py +++ b/src/mobius/_configs/per_model/_qwen3_asr_audio.py @@ -41,7 +41,9 @@ def _qwen3_asr_audio(config, parent_config, model_type: str, fields: dict): n_window_infer=getattr(ac, "n_window_infer", None), ) # Special tokens from thinker config - fields["audio_token_id"] = getattr(tc, "audio_token_id", None) + fields["audio_token_id"] = getattr( + tc, "audio_token_id", getattr(tc, "audio_token_index", None) + ) fields["audio_start_token_id"] = getattr(tc, "audio_start_token_id", None) fields["audio_end_token_id"] = getattr(tc, "audio_end_token_id", None) fields["classify_num"] = getattr(tc, "classify_num", None) diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 869518b0f..1cd51caef 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -225,6 +225,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.sensenova_u1 import SenseNovaU1Model @@ -962,6 +963,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", + ), "moonshine": ModelRegistration( MoonshineForConditionalGeneration, task="speech-to-text", diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index d9cee1d06..4543d77b0 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -111,6 +111,8 @@ "PixtralProjector", "QuantizedEmbedding", "QuantizedLinear", + "Qwen25OmniAudioAttention", + "Qwen25OmniAudioEncoderLayer", "RadioVisionModel", "RMSNorm", "RMSNormBias", @@ -416,6 +418,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 e8cabed97..c2d72c909 100644 --- a/src/mobius/components/_conv.py +++ b/src/mobius/components/_conv.py @@ -6,14 +6,10 @@ from __future__ import annotations from collections.abc import Sequence -from typing import TYPE_CHECKING import onnx_ir as ir from onnxscript import OpBuilder, nn -if TYPE_CHECKING: - pass - def _pair(value: int | Sequence[int], name: str) -> tuple[int, int]: if isinstance(value, int): @@ -103,6 +99,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 000000000..ce20c17a4 --- /dev/null +++ b/src/mobius/components/_qwen25_omni_audio.py @@ -0,0 +1,157 @@ +# 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 OpBuilder, nn + +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: OpBuilder, + hidden_states: ir.Value, + cu_seqlens: ir.Value, + ): + """Bidirectional self-attention. + + Args: + hidden_states: (total_tokens, d_model) + cu_seqlens: Cumulative lengths of the independent audio chunks. + + Returns: + output: (total_tokens, d_model) + """ + seq_len = op.Shape(hidden_states, start=0, end=1) + # Rank-3 Attention splits the projected hidden dimension into heads. + q = op.Unsqueeze(self.q_proj(op, hidden_states), [0]) + k = op.Unsqueeze(self.k_proj(op, hidden_states), [0]) + v = op.Unsqueeze(self.v_proj(op, hidden_states), [0]) + + # 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]) + + 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.Squeeze(attn_output, [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: OpBuilder, + hidden_states: ir.Value, + cu_seqlens: ir.Value, + ): + """Pre-norm encoder layer with bidirectional attention. + + Args: + hidden_states: (total_tokens, d_model) + cu_seqlens: Cumulative lengths of the independent audio chunks. + + Returns: + hidden_states: (total_tokens, 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/integrations/transformers/_builder.py b/src/mobius/integrations/transformers/_builder.py index 3dbe632f6..7e4fa1183 100644 --- a/src/mobius/integrations/transformers/_builder.py +++ b/src/mobius/integrations/transformers/_builder.py @@ -221,7 +221,8 @@ def _select_primary_config(hf_config): parent_config = hf_config model_type = hf_config.model_type - if hasattr(hf_config, "talker_config"): + # Omni exports Thinker; its sibling Talker has a separate decoder config. + if hasattr(hf_config, "talker_config") and model_type != "qwen2_5_omni": hf_config = hf_config.talker_config elif hasattr(hf_config, "thinker_config"): thinker = hf_config.thinker_config diff --git a/src/mobius/integrations/transformers/_builder_test.py b/src/mobius/integrations/transformers/_builder_test.py index 4f02da31a..f8176761e 100644 --- a/src/mobius/integrations/transformers/_builder_test.py +++ b/src/mobius/integrations/transformers/_builder_test.py @@ -391,6 +391,17 @@ def build_module(_module, config, *args, **kwargs): assert built_configs[0].image_token_id == 248056 +def test_qwen3_tts_primary_config_remains_talker() -> None: + talker = SimpleNamespace(model_type="qwen3_tts_talker") + hf_config = SimpleNamespace(model_type="qwen3_tts", talker_config=talker) + + primary, parent, model_type = transformers_builder._select_primary_config(hf_config) + + assert primary is talker + assert parent is hf_config + assert model_type == "qwen3_tts" + + def test_vibevoice_architecture_dispatch_keeps_native_asr_separate() -> None: """Native ASR has its own model_type and cannot fall through to VibeVoice TTS.""" from mobius.models import ( diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 11b9ef922..e4acc012b 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -157,6 +157,7 @@ "Phi4SigLIPModel", "PhiCausalLMModel", "Qwen25VLCausalLMModel", + "Qwen25OmniThinkerForConditionalGeneration", "Qwen25VLDecoderModel", "Qwen25VLEmbeddingModel", "Qwen25VLTextModel", @@ -412,6 +413,7 @@ Qwen4ExpCausalLMModel, Qwen4ExpForConditionalGeneration, ) +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 000000000..f5a00ceb9 --- /dev/null +++ b/src/mobius/models/qwen25_omni.py @@ -0,0 +1,526 @@ +# 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 ( + DecoderLayer, + Embedding, + GatedMLP, + LayerNorm, + Linear, + Qwen25OmniAudioEncoderLayer, + Qwen25VLPatchMerger, + Qwen25VLVisionAttention, + Qwen25VLVisionBlock, + Qwen25VLVisionModel, + RMSNorm, + create_attention_bias, + initialize_rope, +) +from mobius.components._conv import Conv1d + + +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, eps=1e-5) + + # 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_encoder``: 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 000000000..6c865d41e --- /dev/null +++ b/src/mobius/models/qwen25_omni_test.py @@ -0,0 +1,138 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +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 + + +@pytest.mark.parametrize("raw_json", [False, True]) +def test_qwen25_omni_full_checkpoint_builds_thinker(monkeypatch, raw_json): + from transformers import Qwen2_5OmniConfig + + from mobius.integrations.transformers import _builder + from mobius.integrations.transformers._config_resolver import _dict_to_pretrained_config + + _, parent = _hf_config() + hf_config = Qwen2_5OmniConfig( + thinker_config={ + name: vars(value) if isinstance(value, SimpleNamespace) else value + for name, value in vars(parent.thinker_config).items() + } + ) + if raw_json: + hf_config = _dict_to_pretrained_config(hf_config.to_dict()) + assert hf_config.talker_config is not None + monkeypatch.setattr( + _builder, "_load_transformers_config", lambda *args, **kwargs: (hf_config, raw_json) + ) + + package = _builder.build_transformers_model("test/qwen25-omni", load_weights=False) + + assert set(package) == {"audio_encoder", "vision_encoder", "embedding", "decoder"} + assert package.config.vocab_size == 256 + assert package.config.hidden_size == 64 + assert package.config.audio.audio_token_id == 100 + assert package.config.image_token_id == 101 + assert package.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 1677cd01e..36abc9ce4 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -35,6 +35,7 @@ "DraftTargetCausalLMTask", "Eagle3DraftTask", "Qwen35MtpTask", + "Qwen25OmniTask", "Qwen4ExpCausalLMTask", "Qwen4ExpVisionLanguageTask", "DenoisingTask", @@ -194,6 +195,7 @@ Qwen4ExpCausalLMTask, Qwen4ExpVisionLanguageTask, ) +from mobius.tasks._qwen25_omni import Qwen25OmniTask from mobius.tasks._qwen35_mtp import Qwen35MtpTask from mobius.tasks._qwen_image import QwenImageDenoisingTask from mobius.tasks._qwen_image_text_encoder import QwenImageTextEncoderTask @@ -275,6 +277,7 @@ "dflash-draft": DFlashDraftTask, "eagle3-draft": Eagle3DraftTask, "qwen35-mtp": Qwen35MtpTask, + "qwen25-omni": Qwen25OmniTask, "qwen4-exp-text-generation": Qwen4ExpCausalLMTask, "qwen4-exp-vision-language": Qwen4ExpVisionLanguageTask, "vae": VAETask, diff --git a/src/mobius/tasks/_qwen25_omni.py b/src/mobius/tasks/_qwen25_omni.py new file mode 100644 index 000000000..5caf41c0e --- /dev/null +++ b/src/mobius/tasks/_qwen25_omni.py @@ -0,0 +1,126 @@ +# 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) + for name in ("audio_encoder", "vision_encoder"): + if getattr(module, name) is None: + raise ValueError(f"Qwen25OmniTask requires a non-None {name}.") + 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=config.dtype, + 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/_test_configs.py b/tests/_test_configs.py index 51401b492..885b4caa4 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -3625,6 +3625,44 @@ def vl_overrides(model_type: str) -> dict: }, True, ), + # --- Qwen2.5-Omni Thinker (audio + vision + embedding + decoder) --- + ( + "qwen2_5_omni", + { + "model_type": "qwen2_5_omni_text", + "attn_qkv_bias": True, + "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, + }, + True, + ), # --- Qwen3-ASR (speech-language, 3-model split) --- ( "qwen3_asr", diff --git a/tests/build_graph/speech_test.py b/tests/build_graph/speech_test.py index 58a44ac81..d1fd1372b 100644 --- a/tests/build_graph/speech_test.py +++ b/tests/build_graph/speech_test.py @@ -43,7 +43,64 @@ ) _SPEECH_MODEL_PARAMS = _make_params(SPEECH_CONFIGS) + + +class TestBuildGraphQwen25Omni: + """Verify the Qwen2.5-Omni Thinker four-model split.""" + + def _omni_config(self): + overrides = next( + overrides + for model_type, overrides, _ in SPEECH_CONFIGS + if model_type == "qwen2_5_omni" + ) + return _base_config(**overrides) + + @pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] + ) + def test_package_builds_four_models(self, dtype): + from mobius.models import Qwen25OmniThinkerForConditionalGeneration + from mobius.tasks import Qwen25OmniTask + + config = self._omni_config() + config.dtype = dtype + 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 package["audio_encoder"].graph.inputs[0].dtype == dtype + assert {value.name for value in package["embedding"].graph.inputs} == { + "input_ids", + "audio_features", + "image_features", + "video_features", + } + + @pytest.mark.parametrize("missing", ["audio", "vision"]) + def test_missing_encoder_rejected(self, missing): + from mobius.models import Qwen25OmniThinkerForConditionalGeneration + from mobius.tasks import Qwen25OmniTask + + config = self._omni_config() + setattr(config, missing, None) + module = Qwen25OmniThinkerForConditionalGeneration(config) + with pytest.raises(ValueError, match=f"non-None {missing}_encoder"): + Qwen25OmniTask().build(module, config) + + _SPEECH_TASK_KEYS: dict[str, set[str]] = { + "qwen25-omni": {"audio_encoder", "vision_encoder", "embedding", "decoder"}, "speech-to-text": {"encoder", "decoder"}, "speech-language": {"audio_encoder", "embedding", "decoder"}, "codec": {"decoder", "encoder"}, diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index c7cf8f73a..b05ad73c4 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -46,6 +46,211 @@ logger = logging.getLogger(__name__) +def test_qwen2_5_omni_thinker_synthetic_parity(): + """Exercise all four Thinker components with native HF weights, offline.""" + from _test_configs import SPEECH_CONFIGS + from transformers import ( + Qwen2_5OmniThinkerConfig, + Qwen2_5OmniThinkerForConditionalGeneration, + ) + from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import ( + chunk_and_pad_features, + get_pool_indices, + ) + + from mobius._builder import build_from_module + from mobius._testing.ort_inference import OnnxModelSession + + overrides = next(o for mt, o, _ in SPEECH_CONFIGS if mt == "qwen2_5_omni") + config = _base_config(**overrides) + audio, vision = config.audio, config.vision + hf_config = Qwen2_5OmniThinkerConfig( + audio_config={ + name: getattr(audio, name) + for name in ( + "d_model", + "encoder_layers", + "encoder_attention_heads", + "encoder_ffn_dim", + "num_mel_bins", + "max_source_positions", + "n_window", + "output_dim", + ) + }, + vision_config={ + "depth": vision.num_hidden_layers, + "num_heads": vision.num_attention_heads, + **{ + name: getattr(vision, name) + for name in ( + "hidden_size", + "intermediate_size", + "patch_size", + "temporal_patch_size", + "in_channels", + "out_hidden_size", + "spatial_merge_size", + "fullatt_block_indexes", + "window_size", + ) + }, + }, + text_config={ + **{ + name: getattr(config, name) + for name in ( + "vocab_size", + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "max_position_embeddings", + "rms_norm_eps", + "hidden_act", + ) + }, + "rope_parameters": { + "rope_type": "default", + "rope_theta": config.rope_theta, + "mrope_section": config.mrope_section, + }, + }, + audio_token_id=audio.audio_token_id, + image_token_id=config.image_token_id, + video_token_id=config.video_token_id, + ) + torch.manual_seed(42) + reference = Qwen2_5OmniThinkerForConditionalGeneration._from_config( + hf_config, attn_implementation="eager" + ).eval() + module = registry.get("qwen2_5_omni")(config) + package = build_from_module(module, config, task="qwen25-omni") + weights = module.preprocess_weights(reference.state_dict()) + for model in package.values(): + missing = { + name + for name, value in model.graph.initializers.items() + if value.const_value is None and name not in weights + } + assert not missing, f"Unmapped HF parameters: {missing}" + package.apply_weights(weights) + + # Unequal, odd-length audio clips include a tail chunk and independent clips. + session = OnnxModelSession(package["audio_encoder"]) + try: + for lengths in ([7], [35, 13]): + feature_lens = torch.tensor(lengths) + features = torch.randn(audio.num_mel_bins, sum(lengths)) + padded, chunk_lengths = chunk_and_pad_features( + features, feature_lens, audio.n_window + ) + actual = session.run( + { + "input_features": padded.numpy(), + "chunk_lengths": chunk_lengths.numpy(), + "pool_indices": get_pool_indices(feature_lens).numpy(), + } + )["audio_features"] + with torch.no_grad(): + expected = reference.audio_tower( + features, feature_lens=feature_lens + ).last_hidden_state.numpy() + np.testing.assert_allclose(actual, expected, atol=1e-5, rtol=1e-4) + finally: + session.close() + + # A non-square image and multi-frame video exercise window/full attention. + grid = torch.tensor([[1, 4, 10], [2, 4, 4]]) + pixels = torch.randn( + int(grid.prod(-1).sum()), + vision.in_channels * vision.temporal_patch_size * vision.patch_size**2, + ) + with torch.no_grad(): + expected = reference.visual(pixels, grid_thw=grid).pooler_output.numpy() + session = OnnxModelSession(package["vision_encoder"]) + try: + actual = session.run({"pixel_values": pixels.numpy(), "image_grid_thw": grid.numpy()})[ + "image_features" + ] + np.testing.assert_allclose(actual, expected, atol=1e-5, rtol=1e-4) + finally: + session.close() + + # Features must advance across batch rows, with separate image/video streams. + ids = torch.tensor([[5, 100, 101, 102], [101, 102, 100, 6]]) + session = OnnxModelSession(package["embedding"]) + try: + media = { + name: torch.randn(2, config.hidden_size) for name in ("audio", "image", "video") + } + feeds = { + "input_ids": ids.numpy(), + **{f"{name}_features": value.numpy() for name, value in media.items()}, + } + actual_embeds = session.run(feeds)["inputs_embeds"] + with torch.no_grad(): + expected_embeds = reference.model.embed_tokens(ids) + for name, token_id in (("audio", 100), ("image", 101), ("video", 102)): + expected_embeds[ids == token_id] = media[name] + np.testing.assert_allclose(actual_embeds, expected_embeds.numpy(), atol=0, rtol=0) + text_ids = np.array([[7], [8]], dtype=np.int64) + text_embeds = session.run( + { + "input_ids": text_ids, + **{ + f"{name}_features": np.empty((0, config.hidden_size), np.float32) + for name in media + }, + } + )["inputs_embeds"] + with torch.no_grad(): + expected_text = reference.model.embed_tokens(torch.from_numpy(text_ids)).numpy() + np.testing.assert_allclose(text_embeds, expected_text, atol=0, rtol=0) + finally: + session.close() + + # Compare full prefill logits and a cached decode with three distinct MRoPE axes. + session = OnnxModelSession(package["decoder"]) + cache = None + feeds = {} + for layer in range(config.num_hidden_layers): + for kind in ("key", "value"): + feeds[f"past_key_values.{layer}.{kind}"] = np.zeros( + (2, config.num_key_value_heads, 0, config.head_dim), np.float32 + ) + try: + for embeds, offset in ((actual_embeds, 0), (text_embeds, 4)): + length = embeds.shape[1] + positions = np.broadcast_to( + np.arange(offset, offset + length), (3, 2, length) + ).copy() + positions[1] += 1 + positions[2] += 2 + mask = np.ones((2, offset + length), np.int64) + feeds.update(inputs_embeds=embeds, attention_mask=mask, position_ids=positions) + outputs = session.run(feeds) + with torch.no_grad(): + result = reference.model( + inputs_embeds=torch.from_numpy(embeds), + attention_mask=torch.from_numpy(mask), + position_ids=torch.from_numpy(positions), + past_key_values=cache, + use_cache=True, + ) + expected = reference.lm_head(result.last_hidden_state).numpy() + np.testing.assert_allclose(outputs["logits"], expected, atol=1e-5, rtol=1e-4) + cache = result.past_key_values + for layer in range(config.num_hidden_layers): + for kind in ("key", "value"): + feeds[f"past_key_values.{layer}.{kind}"] = outputs[ + f"present.{layer}.{kind}" + ] + finally: + session.close() + + def test_vibevoice_synthetic_pipeline_parity(): """Run the dedicated eight-stage continuous-token parity harness.""" from mobius.models.vibevoice_test import run_vibevoice_synthetic_stage_parity