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 new file mode 100644 index 000000000..e2970dcb2 --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tests for code predictor dtype alignment (fix for #2385). + +Verifies that the code predictor handles dtype mismatches between input +tensors and model parameters without raising RuntimeError. This can happen +when model weights are loaded in float16/bfloat16 but upstream modules +produce float32 hidden states. +""" + +from __future__ import annotations + +import importlib.util +import os +import sys +import types +from unittest.mock import MagicMock, patch + +import torch + +# Direct file import to avoid vllm_omni.__init__ patch dependencies. +_BASE = os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + os.pardir, + os.pardir, + "vllm_omni", + "model_executor", + "models", + "qwen3_tts", +) + + +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) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def _build_mock_modules() -> dict[str, object]: + """Build the dict of modules to inject into sys.modules.""" + platforms_mock = MagicMock() + platforms_mock.current_omni_platform.supports_torch_inductor.return_value = False + + logger_mock = MagicMock() + logger_mock.init_logger = lambda name: MagicMock() + + vllm_config_mod = MagicMock() + vllm_config_mod.set_current_vllm_config = lambda cfg: MagicMock(__enter__=MagicMock(), __exit__=MagicMock()) + + weight_utils_mock = MagicMock() + weight_utils_mock.default_weight_loader = lambda p, w: None + + pkg = types.ModuleType("vllm_omni.model_executor.models.qwen3_tts") + pkg.__path__ = [os.path.abspath(_BASE)] + + return { + "vllm_omni": MagicMock(), + "vllm_omni.platforms": platforms_mock, + "vllm.logger": logger_mock, + "vllm.config": MagicMock(), + "vllm.config.vllm": vllm_config_mod, + "vllm.model_executor.model_loader.weight_utils": weight_utils_mock, + "vllm_omni.model_executor": types.ModuleType("vllm_omni.model_executor"), + "vllm_omni.model_executor.models": types.ModuleType("vllm_omni.model_executor.models"), + "vllm_omni.model_executor.models.qwen3_tts": pkg, + } + + +def _load_target_classes(): + """Load config and code predictor modules with mocked dependencies. + + Uses patch.dict to ensure sys.modules is always restored, even on failure. + """ + mocks = _build_mock_modules() + with patch.dict(sys.modules, mocks): + config_mod = _load_module( + "vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts", + "configuration_qwen3_tts.py", + ) + sys.modules["vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts"] = config_mod + + cp_mod = _load_module( + "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm", + "qwen3_tts_code_predictor_vllm.py", + ) + + return config_mod, cp_mod + + +_config_mod, _cp_mod = _load_target_classes() + +Qwen3TTSTalkerCodePredictorConfig = _config_mod.Qwen3TTSTalkerCodePredictorConfig +Qwen3TTSTalkerConfig = _config_mod.Qwen3TTSTalkerConfig +CodePredictorWrapper = _cp_mod.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM +CodePredictorModel = _cp_mod.Qwen3TTSTalkerCodePredictorModelVLLM + + +def _make_tiny_config() -> tuple: + """Create minimal configs for a tiny code predictor model.""" + cp_config = Qwen3TTSTalkerCodePredictorConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + num_code_groups=4, + rms_norm_eps=1e-6, + ) + talker_config = Qwen3TTSTalkerConfig( + hidden_size=32, + num_code_groups=4, + ) + return cp_config, talker_config + + +def _make_vllm_config(max_num_seqs: int = 4) -> MagicMock: + """Create a mock VllmConfig with scheduler_config.""" + vllm_config = MagicMock() + vllm_config.scheduler_config.max_num_seqs = max_num_seqs + return vllm_config + + +class TestCodePredictorDtypeAlignment: + """Test that code predictor buffers match model parameter dtype.""" + + def test_ensure_buffers_uses_given_dtype(self) -> None: + """_ensure_buffers should create proj_buf with the given dtype.""" + cp_config, talker_config = _make_tiny_config() + vllm_config = _make_vllm_config() + + predictor = CodePredictorWrapper( + vllm_config=vllm_config, + config=cp_config, + talker_config=talker_config, + ) + + # Create buffer in float16 + predictor._ensure_buffers(torch.device("cpu"), torch.float16) + assert predictor._proj_buf is not None + assert predictor._proj_buf.dtype == torch.float16 + + # Re-create buffer in float32 (different dtype triggers re-allocation) + predictor._ensure_buffers(torch.device("cpu"), torch.float32) + assert predictor._proj_buf.dtype == torch.float32 + + def test_warmup_aligns_buffer_to_model_params(self) -> None: + """_warmup_buckets should align proj_buf dtype to model parameters.""" + cp_config, talker_config = _make_tiny_config() + vllm_config = _make_vllm_config(max_num_seqs=2) + + predictor = CodePredictorWrapper( + vllm_config=vllm_config, + config=cp_config, + talker_config=talker_config, + ) + + # Cast model to float16 (simulating vLLM loading weights in half precision) + predictor = predictor.to(torch.float16) + + # Pre-create proj_buf with WRONG dtype (float32) — simulating the bug + predictor._ensure_buffers(torch.device("cpu"), torch.float32) + assert predictor._proj_buf.dtype == torch.float32 + + # Simulate _setup_compile having cached model dtype and compiled forward + predictor._model_dtype = torch.float16 + predictor._compiled_model_fwd = predictor.model.forward + + # _warmup_buckets should fix the dtype mismatch + predictor._warmup_buckets() + + assert predictor._proj_buf.dtype == torch.float16 + + def test_setup_compile_caches_model_dtype(self) -> None: + """_setup_compile should cache model parameter dtype.""" + cp_config, talker_config = _make_tiny_config() + vllm_config = _make_vllm_config(max_num_seqs=2) + + predictor = CodePredictorWrapper( + vllm_config=vllm_config, + config=cp_config, + talker_config=talker_config, + ) + predictor = predictor.to(torch.float16) + + assert predictor._model_dtype is None + predictor._setup_compile() + assert predictor._model_dtype == torch.float16 + + def test_forward_with_mismatched_input_dtype(self) -> None: + """forward() should not crash when inputs are float32 but model is float16.""" + cp_config, talker_config = _make_tiny_config() + vllm_config = _make_vllm_config(max_num_seqs=2) + + predictor = CodePredictorWrapper( + vllm_config=vllm_config, + config=cp_config, + talker_config=talker_config, + ) + + # Model in float16 + predictor = predictor.to(torch.float16) + + bsz = 1 + num_groups = cp_config.num_code_groups + hidden = talker_config.hidden_size + + # Inputs in float32 (simulating the dtype mismatch from #2385) + layer0_code = torch.zeros(bsz, dtype=torch.long) + layer0_embed = torch.randn(bsz, hidden, dtype=torch.float32) + last_talker_hidden = torch.randn(bsz, hidden, dtype=torch.float32) + + # This should NOT raise RuntimeError about dtype mismatch + result = predictor( + layer0_code=layer0_code, + layer0_embed=layer0_embed, + last_talker_hidden=last_talker_hidden, + do_sample=False, + ) + + assert result.shape == (bsz, num_groups) + assert result.dtype == torch.long + + +class TestCodePredictorModelDtype: + """Test the inner model forward with different dtypes.""" + + def test_model_forward_float16(self) -> None: + """Inner model forward should work in float16.""" + cp_config, _ = _make_tiny_config() + model = CodePredictorModel(cp_config, talker_hidden_size=32).to(torch.float16) + + bsz, seq_len = 1, 4 + inputs = torch.randn(bsz, seq_len, 32, dtype=torch.float16) + pos_ids = torch.arange(seq_len).unsqueeze(0).expand(bsz, -1) + + output = model(inputs, pos_ids) + assert output.dtype == torch.float16 + assert output.shape == (bsz, seq_len, 32) + + def test_model_forward_float32(self) -> None: + """Inner model forward should work in float32.""" + cp_config, _ = _make_tiny_config() + model = CodePredictorModel(cp_config, talker_hidden_size=32).to(torch.float32) + + bsz, seq_len = 1, 4 + inputs = torch.randn(bsz, seq_len, 32, dtype=torch.float32) + pos_ids = torch.arange(seq_len).unsqueeze(0).expand(bsz, -1) + + output = model(inputs, pos_ids) + assert output.dtype == torch.float32 + assert output.shape == (bsz, seq_len, 32) diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index 11c0369e8..1e84eaeba 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py @@ -348,6 +348,7 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): # Pre-allocated buffers (lazily initialized on first forward). self._proj_buf: torch.Tensor | None = None + self._model_dtype: torch.dtype | None = None # torch.compile + warmup state (lazily initialized in _setup_compile). self._compiled_model_fwd = None @@ -404,6 +405,10 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): """Lazily set up torch.compile with manual CUDA graph capture.""" if self._compiled_model_fwd is not None: return + # Cache model parameter dtype so forward() doesn't need to query it + # on every call. Also ensures warmup buffers match model precision + # even when upstream modules produce a different dtype (#2385). + self._model_dtype = next(self.model.parameters()).dtype self._lm_heads_list = list(self.lm_head) self._codec_embeds_list = list(self.model.codec_embedding) if not current_omni_platform.supports_torch_inductor(): @@ -443,6 +448,9 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): max_seq = self._num_groups + 1 device = next(self.model.parameters()).device + # Ensure proj_buf matches model parameter dtype to avoid dtype + # mismatch during warmup compilation (see #2385). + self._ensure_buffers(device, self._model_dtype) proj_buf = self._proj_buf for bsz in self._bucket_sizes: # position_ids: [batch, seq_len] for HF-style RoPE @@ -499,13 +507,15 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): bsz = int(layer0_code.shape[0]) num_groups = self._num_groups device = layer0_code.device - dtype = layer0_embed.dtype all_codes = torch.empty(bsz, num_groups, dtype=torch.long, device=device) all_codes[:, 0] = layer0_code.reshape(bsz) - self._ensure_buffers(device, dtype) + # _setup_compile caches _model_dtype on first call; use it for buffers + # so they always match model weight precision (#2385). self._setup_compile() + dtype = self._model_dtype + self._ensure_buffers(device, dtype) proj_buf = self._proj_buf max_seq = self._num_groups + 1 @@ -525,8 +535,8 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): padded_bsz = self._padded_bsz(bsz) proj_buf[:padded_bsz].zero_() - proj_buf[:bsz, 0, :] = projection(last_talker_hidden.reshape(bsz, 1, -1)).reshape(bsz, -1) - proj_buf[:bsz, 1, :] = projection(layer0_embed.reshape(bsz, 1, -1)).reshape(bsz, -1) + proj_buf[:bsz, 0, :] = projection(last_talker_hidden.reshape(bsz, 1, -1).to(dtype)).reshape(bsz, -1) + proj_buf[:bsz, 1, :] = projection(layer0_embed.reshape(bsz, 1, -1).to(dtype)).reshape(bsz, -1) full_pos_ids = self._bucket_pos_ids.get(padded_bsz) if full_pos_ids is None: full_pos_ids = torch.arange(max_seq, device=device, dtype=torch.long).unsqueeze(0).expand(padded_bsz, -1)