From 1413eaab7c82e4ab96ac375561e776eb3284f8ef Mon Sep 17 00:00:00 2001 From: Deep Shah Date: Mon, 7 Sep 2026 19:25:27 -0700 Subject: [PATCH] [Bugfix][Qwen3-Omni] Read code-predictor RoPE theta from rope_parameters (#7228) Signed-off-by: Deep Shah Co-authored-by: amy-why-3459 --- .../qwen3_tts/test_code_predictor_dtype.py | 83 ++++++++++++++++++- .../models/common/qwen3_code_predictor.py | 9 +- 2 files changed, 89 insertions(+), 3 deletions(-) diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index 619f2169a..438a5ec1d 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Tests for code predictor dtype alignment (fix for #2385). @@ -20,6 +20,9 @@ import pytest import torch import torch.nn.functional as F from pytest_mock import MockerFixture +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( + Qwen3OmniMoeTalkerCodePredictorConfig, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -41,6 +44,7 @@ _COMMON = os.path.join(_MODELS, "common") def _load_module(name: str, filename: str): path = os.path.abspath(os.path.join(_BASE, filename)) spec = importlib.util.spec_from_file_location(name, path) + assert spec is not None and spec.loader is not None mod = importlib.util.module_from_spec(spec) sys.modules[name] = mod # register before exec (needed for dataclasses etc.) spec.loader.exec_module(mod) @@ -83,7 +87,7 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]: vllm_parallel_mock = mocker.MagicMock() vllm_parallel_mock.VocabParallelEmbedding = torch.nn.Embedding custom_op_mock = types.ModuleType("vllm_omni.diffusion.layers.custom_op") - custom_op_mock.CustomOp = NativeCustomOp + setattr(custom_op_mock, "CustomOp", NativeCustomOp) return { "vllm_omni": mocker.MagicMock(), @@ -119,6 +123,7 @@ def _load_target_classes(mocker: MockerFixture): common_spec = importlib.util.spec_from_file_location( "vllm_omni.model_executor.models.common.qwen3_code_predictor", common_cp_path ) + assert common_spec is not None and common_spec.loader is not None common_cp_mod = importlib.util.module_from_spec(common_spec) sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] = common_cp_mod common_spec.loader.exec_module(common_cp_mod) @@ -209,6 +214,80 @@ def test_npu_custom_ops_use_fused_norm_and_cached_rope(mocker: MockerFixture, lo torch.testing.assert_close(sin, rotary.sin_cached[position_ids].to(torch.float16)) +@pytest.mark.parametrize( + ("rope_kwargs", "expected_theta"), + [ + pytest.param( + {"rope_theta": 1_000_000.0}, + 1_000_000.0, + id="serialized_checkpoint", + ), + pytest.param( + { + "rope_theta": 10_000.0, + "rope_parameters": {"rope_type": "default", "rope_theta": 1_000_000.0}, + }, + 1_000_000.0, + id="nested_precedence", + ), + pytest.param({}, 10_000.0, id="default"), + ], +) +def test_code_predictor_rotary_uses_qwen3_omni_rope_parameters( + loaded_target_classes, + rope_kwargs, + expected_theta, +) -> None: + """Follow Transformers 5 deserialization and precedence for Qwen3-Omni.""" + _ = loaded_target_classes + config = Qwen3OmniMoeTalkerCodePredictorConfig( + hidden_size=16, + num_attention_heads=1, + head_dim=16, + num_code_groups=16, + **rope_kwargs, + ) + assert not hasattr(config, "rope_theta") + assert config.rope_parameters["rope_theta"] == expected_theta + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + + rotary = common_mod._RotaryEmbedding(config) + + expected = 1.0 / (expected_theta ** (torch.arange(0, 16, 2, dtype=torch.float32) / 16)) + torch.testing.assert_close(rotary.inv_freq, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("theta", [10_000.0, 1_000_000.0]) +@pytest.mark.parametrize( + "legacy_top_level_only", + [False, True], + ids=["normalized", "legacy_top_level"], +) +def test_code_predictor_rotary_preserves_qwen3_tts_rope_theta( + loaded_target_classes, + theta, + legacy_top_level_only, +) -> None: + """Keep the shared predictor compatible with current and legacy Qwen3-TTS.""" + config_class = loaded_target_classes[0] + config = config_class( + hidden_size=16, + num_attention_heads=1, + head_dim=16, + num_code_groups=16, + rope_theta=theta, + ) + assert config.rope_parameters["rope_theta"] == theta + if legacy_top_level_only: + delattr(config, "rope_parameters") + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + + rotary = common_mod._RotaryEmbedding(config) + + expected = 1.0 / (theta ** (torch.arange(0, 16, 2, dtype=torch.float32) / 16)) + torch.testing.assert_close(rotary.inv_freq, expected, rtol=0, atol=0) + + class TestCodePredictorDtypeAlignment: """Test that code predictor buffers match model parameter dtype.""" diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 0308cda7b..60b4d4525 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Qwen3 Code Predictor -- optimized re-prefill, no KV cache. Shared by Qwen3-Omni and Qwen3-TTS talker models. @@ -97,7 +100,11 @@ class _RotaryEmbedding(CustomOp): "head_dim", config.hidden_size // config.num_attention_heads, ) - rope_theta = getattr(config, "rope_theta", 10000.0) + rope_parameters = getattr(config, "rope_parameters", None) or {} + rope_theta = rope_parameters.get( + "rope_theta", + getattr(config, "rope_theta", 10000.0), + ) inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) self.register_buffer("inv_freq", inv_freq, persistent=False)