From 348dff5af4420cbbfd79bbc1dcd35a1e31722853 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:33:01 -0700 Subject: [PATCH 1/7] Add Qwen2.5 Omni Talker export Preserve independent Thinker/Talker configs and export the codec embedding plus Talker decoder graphs. Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- src/mobius/_builder.py | 16 +++- src/mobius/_configs/_base.py | 13 +++ src/mobius/models/__init__.py | 8 +- src/mobius/models/qwen25_omni.py | 112 +++++++++++++++++++++++--- src/mobius/models/qwen25_omni_test.py | 38 ++++++++- src/mobius/tasks/_qwen25_omni.py | 82 ++++++++++++++++++- tests/build_graph_test.py | 51 +++++++++++- 7 files changed, 300 insertions(+), 20 deletions(-) diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py index a080f36aa..91df10adf 100644 --- a/src/mobius/_builder.py +++ b/src/mobius/_builder.py @@ -476,7 +476,11 @@ def build( model_type = hf_config.model_type parent_config = hf_config - if hasattr(hf_config, "talker_config"): + if model_type == "qwen2_5_omni" and hasattr(hf_config, "thinker_config"): + thinker = hf_config.thinker_config + if hasattr(thinker, "text_config"): + hf_config = thinker.text_config + elif hasattr(hf_config, "talker_config"): hf_config = hf_config.talker_config elif hasattr(hf_config, "thinker_config"): thinker = hf_config.thinker_config @@ -567,7 +571,15 @@ def build( if dtype is not None: dtype = resolve_dtype(dtype) - config = dataclasses.replace(config, dtype=dtype) + talker_config = getattr(config, "talker", None) + if talker_config is not None: + config = dataclasses.replace( + config, + dtype=dtype, + talker=dataclasses.replace(talker_config, dtype=dtype), + ) + else: + config = dataclasses.replace(config, dtype=dtype) if output_layer_indices is not None: # Opt-in: emit additional `hidden_states.{k}` ONNX outputs for each diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index ba415eb6e..9e63a9a00 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -316,6 +316,7 @@ class BaseModelConfig: vocab_size: int = DEFAULT_INT hidden_size: int = DEFAULT_INT + embedding_size: int | None = None intermediate_size: int = DEFAULT_INT num_hidden_layers: int = DEFAULT_INT num_attention_heads: int = DEFAULT_INT @@ -436,6 +437,9 @@ class ArchitectureConfig(BaseModelConfig): mrope_section: list[int] | None = None mrope_interleaved: bool = False + # Qwen2.5-Omni uses independent Thinker and Talker decoder dimensions. + talker: ArchitectureConfig | None = None + # Standalone vision config image_size: int = 224 patch_size: int = 16 @@ -589,6 +593,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: num_hidden_layers=_as_int(num_hidden_layers), vocab_size=getattr(config, "vocab_size", None) or 0, hidden_size=_as_int(hidden_size), + embedding_size=getattr(config, "embedding_size", None), intermediate_size=( getattr(config, "intermediate_size", None) or getattr(config, "n_inner", None) @@ -651,6 +656,7 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: "qwen2", "qwen2_5_vl_text", "qwen2_5_omni_text", + "qwen2_5_omni_talker", "qwen2_moe", "qwen2_vl_text", ), @@ -982,6 +988,13 @@ def from_transformers(cls, config, parent_config=None) -> ArchitectureConfig: num_semantic_quantizers=getattr(ec, "num_semantic_quantizers", 1), ) + if model_type == "qwen2_5_omni_text" and parent_config is not None: + talker_config = getattr(parent_config, "talker_config", None) + if talker_config is not None: + if isinstance(talker_config, dict): + talker_config = type("TalkerConfig", (), talker_config)() + options["talker"] = ArchitectureConfig.from_transformers(talker_config) + # Model dtype resolved = _resolve_dtype(config) if resolved is not None: diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 1516759fb..4a8bc2268 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -94,6 +94,8 @@ "PhiCausalLMModel", "Qwen25VLCausalLMModel", "Qwen25OmniThinkerForConditionalGeneration", + "Qwen25OmniTalkerForConditionalGeneration", + "Qwen25OmniTalkerModel", "Qwen25VLDecoderModel", "Qwen25VLEmbeddingModel", "Qwen25VLTextModel", @@ -251,7 +253,11 @@ Qwen3TTSCodecEncoderModel, Qwen3TTSTokenizerV2Model, ) -from mobius.models.qwen25_omni import Qwen25OmniThinkerForConditionalGeneration +from mobius.models.qwen25_omni import ( + Qwen25OmniTalkerForConditionalGeneration, + Qwen25OmniTalkerModel, + Qwen25OmniThinkerForConditionalGeneration, +) from mobius.models.qwen35 import ( Qwen35CausalLMModel, Qwen35MoECausalLMModel, diff --git a/src/mobius/models/qwen25_omni.py b/src/mobius/models/qwen25_omni.py index fa54e6d3f..9f2a85b48 100644 --- a/src/mobius/models/qwen25_omni.py +++ b/src/mobius/models/qwen25_omni.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Qwen2.5-Omni Thinker: audio + vision + text. +"""Qwen2.5-Omni: Thinker and Talker models. Architecture (Thinker only): - Audio encoder: Conv1d x2 → sinusoidal PE → 32 encoder layers → AvgPool → proj @@ -424,18 +424,100 @@ def forward( return logits, present_key_values +class Qwen25OmniTalkerModel(nn.Module): + """Talker backbone with codec embedding, MRoPE decoder layers, and final norm.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + embedding_size = config.embedding_size or config.hidden_size + self._dtype = config.dtype + self.embed_tokens = Embedding(config.vocab_size, embedding_size) + 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) + + 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) + + return self.norm(op, hidden_states), present_key_values + + +class Qwen25OmniTalkerForConditionalGeneration(nn.Module): + """Generate codec-token logits from Thinker-width input embeddings.""" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + embedding_size = config.embedding_size or config.hidden_size + self.thinker_to_talker_proj = Linear( + embedding_size, config.hidden_size, bias=True + ) + self.model = Qwen25OmniTalkerModel(config) + self.codec_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, + ): + # The host combines Thinker reply states, text embeddings, and codec + # embeddings in the shared embedding space before this projection. + hidden_states = self.thinker_to_talker_proj(op, inputs_embeds) + hidden_states, present_key_values = self.model( + op, + inputs_embeds=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + return self.codec_head(op, hidden_states), present_key_values + + class Qwen25OmniThinkerForConditionalGeneration(nn.Module): - """Qwen2.5-Omni Thinker: composite audio + vision + text model. + """Qwen2.5-Omni composite Thinker and optional Talker model. - Builds four separate ONNX models: + Builds four Thinker models and, when configured, two Talker 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 + - ``talker_embedding``: codec token embedding in the Thinker-width space + - ``talker``: projection + speech-token decoder + codec logits - HuggingFace class: ``Qwen2_5OmniForConditionalGeneration`` (Thinker only — - the Talker / streaming code generation head is out of scope for now). + HuggingFace class: ``Qwen2_5OmniForConditionalGeneration``. """ default_task: str = "qwen25-omni" @@ -453,12 +535,17 @@ def __init__(self, config: ArchitectureConfig): self.audio_encoder: Qwen25OmniAudioEncoder | None = ( Qwen25OmniAudioEncoder(config) if config.audio is not None else None ) + self.talker: Qwen25OmniTalkerForConditionalGeneration | None = ( + Qwen25OmniTalkerForConditionalGeneration(config.talker) + if config.talker 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) " + "audio_encoder, talker, and talker_embedding) " "separately." ) @@ -476,17 +563,22 @@ def preprocess_weights( - ``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. + Talker keys already align with the nested ``talker.*`` module and are + retained when audio output is enabled. """ cleaned: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): + if key.startswith("talker."): + if self.talker is not None: + cleaned[key] = value + continue + # 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.")): + # Token2wav is exported independently from the Talker. + if key.startswith(("token2wav.", "code_predictor.")): continue if key.startswith("audio_tower."): diff --git a/src/mobius/models/qwen25_omni_test.py b/src/mobius/models/qwen25_omni_test.py index 990bf74d4..a07c70e96 100644 --- a/src/mobius/models/qwen25_omni_test.py +++ b/src/mobius/models/qwen25_omni_test.py @@ -59,7 +59,28 @@ def _hf_config(): image_token_id=101, video_token_id=102, ) - return text, SimpleNamespace(thinker_config=thinker, tie_word_embeddings=False) + talker = SimpleNamespace( + model_type="qwen2_5_omni_talker", + vocab_size=8448, + embedding_size=64, + hidden_size=32, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + hidden_act="silu", + rms_norm_eps=1e-6, + max_position_embeddings=128, + rope_scaling={"rope_type": "default", "mrope_section": [2, 1, 1]}, + rope_theta=1_000_000.0, + ) + return text, SimpleNamespace( + model_type="qwen2_5_omni", + thinker_config=thinker, + talker_config=talker, + tie_word_embeddings=False, + ) def test_qwen25_omni_extracts_nested_thinker_config(): @@ -74,6 +95,12 @@ def test_qwen25_omni_extracts_nested_thinker_config(): assert config.vision.hidden_size == 64 assert config.image_token_id == 101 assert config.video_token_id == 102 + assert config.hidden_size == 64 + assert config.talker is not None + assert config.talker.embedding_size == 64 + assert config.talker.hidden_size == 32 + assert config.talker.vocab_size == 8448 + assert config.talker.attn_qkv_bias def test_qwen25_omni_preprocess_weights_routes_thinker_components(): @@ -91,7 +118,10 @@ def test_qwen25_omni_preprocess_weights_routes_thinker_components(): "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, + "talker.thinker_to_talker_proj.weight": weight, + "talker.model.embed_tokens.weight": weight, + "talker.model.layers.0.self_attn.q_proj.weight": weight, + "talker.codec_head.weight": weight, "token2wav.dit.weight": weight, } ) @@ -103,4 +133,8 @@ def test_qwen25_omni_preprocess_weights_routes_thinker_components(): "embedding.embed_tokens.weight", "decoder.layers.0.self_attn.q_proj.bias", "decoder.lm_head.weight", + "talker.thinker_to_talker_proj.weight", + "talker.model.embed_tokens.weight", + "talker.model.layers.0.self_attn.q_proj.weight", + "talker.codec_head.weight", } diff --git a/src/mobius/tasks/_qwen25_omni.py b/src/mobius/tasks/_qwen25_omni.py index 34e05eb3c..9b3de839b 100644 --- a/src/mobius/tasks/_qwen25_omni.py +++ b/src/mobius/tasks/_qwen25_omni.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Qwen2.5-Omni Thinker four-model split task.""" +"""Qwen2.5-Omni Thinker and Talker split task.""" from __future__ import annotations @@ -18,17 +18,23 @@ _make_model, build_decoder_from_embeds, ) +from mobius.tasks._cache_utils import ( + _make_kv_cache_inputs, + _register_kv_cache_outputs, +) from mobius.tasks._vision_language_3model import QwenVLTask class Qwen25OmniTask(QwenVLTask): - """Build the Thinker's audio, vision, embedding, and decoder ONNX models.""" + """Build the Thinker models and optional Talker ONNX models.""" model_roles: ClassVar[dict[str, str]] = { "audio_encoder": "encoder", "vision_encoder": "encoder", "embedding": "embedding", "decoder": "decoder", + "talker_embedding": "embedding", + "talker": "decoder", } components = ComponentSpec( audio_encoder="audio_encoder", @@ -39,14 +45,84 @@ class Qwen25OmniTask(QwenVLTask): def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage: self._validate_components(module) - models = { + models: dict[str, ir.Model] = { "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), } + if module.talker is not None and config.talker is not None: + models["talker_embedding"] = self._build_talker_embedding( + module.talker.model.embed_tokens, config.talker + ) + models["talker"] = self._build_talker(module.talker, config.talker) return ModelPackage(models, config=config) + def _build_talker_embedding( + self, + embedding: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build codec-token embeddings in the shared Thinker-width space.""" + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + graph, builder = _make_graph(name="talker_embedding") + input_ids = builder.input( + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], + ) + codec_embeds = embedding(builder.op, input_ids) + builder.add_output(codec_embeds, "codec_embeds") + return _make_model(graph) + + def _build_talker( + self, + talker: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build speech-token logits from preconstructed shared-space embeddings.""" + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + past_seq_len = ir.SymbolicDim("past_sequence_len") + embedding_size = config.embedding_size or config.hidden_size + + graph, builder = _make_graph(name="talker") + inputs_embeds = builder.input( + "inputs_embeds", + dtype=config.dtype, + shape=[batch, seq_len, embedding_size], + ) + attention_mask = builder.input( + "attention_mask", + dtype=ir.DataType.INT64, + shape=[batch, "past_seq_len + seq_len"], + ) + position_ids = builder.input( + "position_ids", + dtype=ir.DataType.INT64, + shape=[3, batch, seq_len], + ) + past_key_values = _make_kv_cache_inputs( + builder, + config.num_hidden_layers, + config.num_key_value_heads, + config.head_dim, + config.dtype, + batch, + past_seq_len, + ) + logits, present_key_values = talker( + builder.op, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + builder.add_output(logits, "logits") + _register_kv_cache_outputs(builder, present_key_values) + return _make_model(graph) + 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") diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index 1548fbd85..13b014e4e 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -2863,9 +2863,22 @@ def test_3model_pipeline_runs_with_ort(self): class TestBuildGraphQwen25Omni: - """Verify the Qwen2.5-Omni Thinker four-model split.""" + """Verify the Qwen2.5-Omni Thinker and Talker model split.""" def _omni_config(self): + talker_config = _base_config( + model_type="qwen2_5_omni_talker", + vocab_size=128, + embedding_size=64, + hidden_size=32, + intermediate_size=64, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + attn_qkv_bias=True, + hidden_act="silu", + mrope_section=[2, 1, 1], + ) return _base_config( model_type="qwen2_5_omni_text", attn_qkv_bias=True, @@ -2899,9 +2912,10 @@ def _omni_config(self): ), image_token_id=101, video_token_id=102, + talker=talker_config, ) - def test_package_builds_four_models(self): + def test_package_builds_talker_models(self): from mobius.models import Qwen25OmniThinkerForConditionalGeneration from mobius.tasks import Qwen25OmniTask @@ -2914,6 +2928,8 @@ def test_package_builds_four_models(self): "vision_encoder", "embedding", "decoder", + "talker_embedding", + "talker", } assert {value.name for value in package["audio_encoder"].graph.inputs} == { "input_features", @@ -2926,6 +2942,37 @@ def test_package_builds_four_models(self): "image_features", "video_features", } + assert {value.name for value in package["talker_embedding"].graph.inputs} == { + "input_ids" + } + assert {value.name for value in package["talker"].graph.inputs} >= { + "inputs_embeds", + "attention_mask", + "position_ids", + "past_key_values.0.key", + "past_key_values.0.value", + } + assert {value.name for value in package["talker"].graph.outputs} >= { + "logits", + "present.0.key", + "present.0.value", + } + + def test_package_omits_talker_models_when_disabled(self): + from mobius.models import Qwen25OmniThinkerForConditionalGeneration + from mobius.tasks import Qwen25OmniTask + + config = self._omni_config() + config.talker = None + module = Qwen25OmniThinkerForConditionalGeneration(config) + package = build_from_module(module, config, task=Qwen25OmniTask()) + + assert set(package) == { + "audio_encoder", + "vision_encoder", + "embedding", + "decoder", + } class TestBuildGraphFunASR: From 6d4f7f968dbc5ac742ddc29a3876a37fed6ae868 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:45:33 -0700 Subject: [PATCH 2/7] Fix Qwen2.5 Omni Talker formatting Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- src/mobius/models/qwen25_omni.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/mobius/models/qwen25_omni.py b/src/mobius/models/qwen25_omni.py index 9f2a85b48..accca4e34 100644 --- a/src/mobius/models/qwen25_omni.py +++ b/src/mobius/models/qwen25_omni.py @@ -478,9 +478,7 @@ class Qwen25OmniTalkerForConditionalGeneration(nn.Module): def __init__(self, config: ArchitectureConfig): super().__init__() embedding_size = config.embedding_size or config.hidden_size - self.thinker_to_talker_proj = Linear( - embedding_size, config.hidden_size, bias=True - ) + self.thinker_to_talker_proj = Linear(embedding_size, config.hidden_size, bias=True) self.model = Qwen25OmniTalkerModel(config) self.codec_head = Linear(config.hidden_size, config.vocab_size, bias=False) From d95ab1523df46e3aa9b4514bbb6553f611a46c4a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:59:59 -0700 Subject: [PATCH 3/7] Add Qwen2.5 Omni test coverage metadata Register the specialized tiny graph config and public checkpoint used by L1 and L2 coverage checks. Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- src/mobius/_registry.py | 1 + tests/_test_configs.py | 53 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+) diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 55b1251f9..9eb95be22 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -979,6 +979,7 @@ def _create_default_registry() -> ModelRegistry: # --- Speech --- "whisper": "openai/whisper-tiny", + "qwen2_5_omni": "Qwen/Qwen2.5-Omni-7B", "qwen3_asr": "Qwen/Qwen3-ASR-0.6B", "fun_asr": "justinchuby/Fun-ASR-Nano-2512", "sensevoice_small": "mlx-community/SenseVoiceSmall", diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 9fb16b002..d905c0c40 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -2350,6 +2350,59 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: }, True, ), + # --- Qwen2.5-Omni (audio + vision + text Thinker, speech Talker) --- + ( + "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, + "talker": ArchitectureConfig( + model_type="qwen2_5_omni_talker", + vocab_size=128, + embedding_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + attn_qkv_bias=True, + hidden_act="silu", + mrope_section=[2, 1, 1], + rope_type="default", + ), + }, + True, + ), # --- Qwen3-ForcedAligner (speech-language, same class as ASR) --- ( "qwen3_forced_aligner", From cb46e5013b36751dc9a6a6c5972d8afe7bd8a209 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:26:33 -0700 Subject: [PATCH 4/7] Fix Qwen2.5 Omni Thinker parity fixture Disable the optional Talker graphs in the Thinker-only synthetic parity test so its strict weight mapping checks only the four components represented by the Hugging Face reference model. Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- tests/synthetic_parity_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index b05ad73c4..66840231a 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -63,6 +63,7 @@ def test_qwen2_5_omni_thinker_synthetic_parity(): overrides = next(o for mt, o, _ in SPEECH_CONFIGS if mt == "qwen2_5_omni") config = _base_config(**overrides) + config.talker = None audio, vision = config.audio, config.vision hf_config = Qwen2_5OmniThinkerConfig( audio_config={ From 9a99d55b7a326600d484d4c7cd77758e0ade1746 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:37:12 -0700 Subject: [PATCH 5/7] Update Qwen2.5 Omni L1 package contract Expect the optional Talker embedding and decoder in the representative six-model speech package across supported graph dtypes. Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- tests/build_graph/speech_test.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/build_graph/speech_test.py b/tests/build_graph/speech_test.py index d1fd1372b..77f209eac 100644 --- a/tests/build_graph/speech_test.py +++ b/tests/build_graph/speech_test.py @@ -46,7 +46,7 @@ class TestBuildGraphQwen25Omni: - """Verify the Qwen2.5-Omni Thinker four-model split.""" + """Verify the Qwen2.5-Omni Thinker and Talker six-model split.""" def _omni_config(self): overrides = next( @@ -59,7 +59,7 @@ def _omni_config(self): @pytest.mark.parametrize( "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] ) - def test_package_builds_four_models(self, dtype): + def test_package_builds_six_models(self, dtype): from mobius.models import Qwen25OmniThinkerForConditionalGeneration from mobius.tasks import Qwen25OmniTask @@ -73,6 +73,8 @@ def test_package_builds_four_models(self, dtype): "vision_encoder", "embedding", "decoder", + "talker_embedding", + "talker", } assert {value.name for value in package["audio_encoder"].graph.inputs} == { "input_features", @@ -100,7 +102,14 @@ def test_missing_encoder_rejected(self, missing): _SPEECH_TASK_KEYS: dict[str, set[str]] = { - "qwen25-omni": {"audio_encoder", "vision_encoder", "embedding", "decoder"}, + "qwen25-omni": { + "audio_encoder", + "vision_encoder", + "embedding", + "decoder", + "talker_embedding", + "talker", + }, "speech-to-text": {"encoder", "decoder"}, "speech-language": {"audio_encoder", "embedding", "decoder"}, "codec": {"decoder", "encoder"}, From 0fc9b69b8ed18b0d18127cdd4cfc862b4561a2fd Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 18 Sep 2026 09:49:43 -0700 Subject: [PATCH 6/7] Document Qwen2.5 Omni golden coverage gap Exclude the six-stage multimodal Talker pipeline from the generic L4/L5 golden runner while retaining its dedicated L1 graph and L3 Thinker parity coverage. Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- tests/model_coverage_test.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index 7167e670e..255082ba8 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -294,6 +294,9 @@ def _all_registered_with_test_id() -> dict[str, str]: "molmo": "VL model — requires image inputs", "phi4_multimodal": "VL model (14B) — needs GPU for golden", "phi4mm": "VL model (14B) — needs GPU for golden", + "qwen2_5_omni": "Six-stage audio, vision, Thinker, and Talker pipeline — L1 graph " + "and L3 Thinker parity are covered; the generic L4/L5 runner cannot orchestrate " + "multimodal inputs and codec-token generation", "qwen2_5_vl": "VL model — requires image inputs", "qwen2_vl": "VL model — requires image inputs", "qwen3_5": "VL model — hybrid VL, requires image inputs", From 6214c69a4d6b75c5c8e734524931d0a72944ee56 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 18 Sep 2026 11:06:38 -0700 Subject: [PATCH 7/7] Address Qwen2.5 Omni Talker review feedback Expose normalized Thinker hidden states for Talker orchestration, align the directly traced codec embedding weight, preserve non-interleaved MRoPE sections, and stamp Talker cache metadata for fused attention exports. Signed-off-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> --- src/mobius/_configs/_base.py | 10 ++-- src/mobius/models/qwen25_omni.py | 7 ++- src/mobius/models/qwen25_omni_test.py | 24 ++++++++- src/mobius/tasks/_qwen25_omni.py | 71 +++++++++++++++++++++++++-- 4 files changed, 101 insertions(+), 11 deletions(-) diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 68b2ef06d..5ebff24d8 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -338,11 +338,11 @@ def _extract_mrope_fields(config) -> dict: ) if mrope_interleaved: result["mrope_interleaved"] = True - section = rope_scaling.get("mrope_section", None) or rope_parameters.get( - "mrope_section", None - ) - if section is not None: - result["mrope_section"] = section + section = rope_scaling.get("mrope_section", None) or rope_parameters.get( + "mrope_section", None + ) + if section is not None: + result["mrope_section"] = section return result diff --git a/src/mobius/models/qwen25_omni.py b/src/mobius/models/qwen25_omni.py index b8cc13a41..bf0530038 100644 --- a/src/mobius/models/qwen25_omni.py +++ b/src/mobius/models/qwen25_omni.py @@ -419,7 +419,7 @@ def forward( hidden_states = self.norm(op, hidden_states) logits = self.lm_head(op, hidden_states) - return logits, present_key_values + return logits, hidden_states, present_key_values class Qwen25OmniTalkerModel(nn.Module): @@ -564,6 +564,11 @@ def preprocess_weights( """ cleaned: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): + if key == "talker.model.embed_tokens.weight": + # The embedding component traces this nested module directly, + # so its initializer has no outer Talker scopes. + cleaned["embed_tokens.weight"] = value + continue if key.startswith("talker."): if self.talker is not None: cleaned[key] = value diff --git a/src/mobius/models/qwen25_omni_test.py b/src/mobius/models/qwen25_omni_test.py index e3ecac25f..852579fdd 100644 --- a/src/mobius/models/qwen25_omni_test.py +++ b/src/mobius/models/qwen25_omni_test.py @@ -104,6 +104,7 @@ def test_qwen25_omni_extracts_nested_thinker_config(): assert config.talker.hidden_size == 32 assert config.talker.vocab_size == 8448 assert config.talker.attn_qkv_bias + assert config.talker.mrope_section == [2, 1, 1] @pytest.mark.parametrize("raw_json", [False, True]) @@ -175,7 +176,7 @@ def test_qwen25_omni_preprocess_weights_routes_thinker_components(): "decoder.layers.0.self_attn.q_proj.bias", "decoder.lm_head.weight", "talker.thinker_to_talker_proj.weight", - "talker.model.embed_tokens.weight", + "embed_tokens.weight", "talker.model.layers.0.self_attn.q_proj.weight", "talker.codec_head.weight", } @@ -196,6 +197,19 @@ def test_qwen25_omni_package_builds_talker_models(): "talker", } assert {value.name for value in package["talker_embedding"].graph.inputs} == {"input_ids"} + assert set(package["talker_embedding"].graph.initializers) == {"embed_tokens.weight"} + decoder_outputs = {value.name: value for value in package["decoder"].graph.outputs} + assert set(decoder_outputs) >= { + "logits", + "hidden_states", + "present.0.key", + "present.0.value", + } + assert [str(dim) for dim in decoder_outputs["hidden_states"].shape] == [ + "batch", + "sequence_len", + "64", + ] assert {value.name for value in package["talker"].graph.inputs} >= { "inputs_embeds", "attention_mask", @@ -208,6 +222,14 @@ def test_qwen25_omni_package_builds_talker_models(): "present.0.key", "present.0.value", } + talker_outputs = {value.name: value for value in package["talker"].graph.outputs} + for name in ("present.0.key", "present.0.value"): + assert [str(dim) for dim in talker_outputs[name].shape] == [ + "batch", + "2", + "past_sequence_len + sequence_len", + "8", + ] def test_qwen25_omni_package_omits_talker_models_when_disabled(): diff --git a/src/mobius/tasks/_qwen25_omni.py b/src/mobius/tasks/_qwen25_omni.py index 48968d6e3..6c2275917 100644 --- a/src/mobius/tasks/_qwen25_omni.py +++ b/src/mobius/tasks/_qwen25_omni.py @@ -16,7 +16,6 @@ ComponentSpec, _make_graph, _make_model, - build_decoder_from_embeds, ) from mobius.tasks._cache_utils import ( _make_kv_cache_inputs, @@ -52,7 +51,7 @@ def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage: "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), + "decoder": self._build_decoder(module.decoder, config), } if module.talker is not None and config.talker is not None: models["talker_embedding"] = self._build_talker_embedding( @@ -61,6 +60,61 @@ def build(self, module: nn.Module, config: ArchitectureConfig) -> ModelPackage: models["talker"] = self._build_talker(module.talker, config.talker) return ModelPackage(models, config=config) + def _build_decoder( + self, + decoder: nn.Module, + config: ArchitectureConfig, + ) -> ir.Model: + """Build Thinker logits, normalized hidden states, and KV cache.""" + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + past_seq_len = ir.SymbolicDim("past_sequence_len") + graph, builder = _make_graph(name="decoder") + inputs_embeds = builder.input( + "inputs_embeds", + dtype=config.dtype, + shape=[batch, seq_len, config.hidden_size], + ) + attention_mask = builder.input( + "attention_mask", + dtype=ir.DataType.INT64, + shape=[batch, "past_sequence_len + sequence_len"], + ) + position_ids = builder.input( + "position_ids", + dtype=ir.DataType.INT64, + shape=[3, batch, seq_len], + ) + past_key_values = _make_kv_cache_inputs( + builder, + config.num_hidden_layers, + config.num_key_value_heads, + config.head_dim, + config.dtype, + batch, + past_seq_len, + ) + logits, hidden_states, present_key_values = decoder( + builder.op, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + ) + builder.add_output(logits, "logits") + builder.add_output(hidden_states, "hidden_states") + _register_kv_cache_outputs( + builder, + present_key_values, + batch=batch, + num_kv_heads=config.num_key_value_heads, + key_head_dim=config.head_dim, + value_head_dim=config.head_dim, + total_seq_len="past_sequence_len + sequence_len", + dtype=config.dtype, + ) + return _make_model(graph) + def _build_talker_embedding( self, embedding: nn.Module, @@ -99,7 +153,7 @@ def _build_talker( attention_mask = builder.input( "attention_mask", dtype=ir.DataType.INT64, - shape=[batch, "past_seq_len + seq_len"], + shape=[batch, "past_sequence_len + sequence_len"], ) position_ids = builder.input( "position_ids", @@ -123,7 +177,16 @@ def _build_talker( past_key_values=past_key_values, ) builder.add_output(logits, "logits") - _register_kv_cache_outputs(builder, present_key_values) + _register_kv_cache_outputs( + builder, + present_key_values, + batch=batch, + num_kv_heads=config.num_key_value_heads, + key_head_dim=config.head_dim, + value_head_dim=config.head_dim, + total_seq_len="past_sequence_len + sequence_len", + dtype=config.dtype, + ) return _make_model(graph) def _build_audio(self, audio_encoder: nn.Module, config: ArchitectureConfig) -> ir.Model: