diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 33dcd0551..b27d90488 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1805,6 +1805,7 @@ class TestQwen2_5OmniPipeline: assert s.owns_tokenizer is True assert s.engine_output_type == "latent" assert s.requires_multimodal_data is True + assert s.hf_config_name == "thinker_config" def test_talker(self): p = resolve_pipeline_config("qwen2_5_omni") @@ -1813,6 +1814,7 @@ class TestQwen2_5OmniPipeline: s = p.get_stage(1) assert isinstance(s, StagePipelineConfig) assert s.input_sources == (0,) + assert s.hf_config_name == "talker_config" assert s.sampling_constraints["stop_token_ids"] == [8294] # thinker2talker was removed: qwen2_5_omni has no async_chunk support, # so sync_process_input_func always wins and custom_process_input_func @@ -1829,6 +1831,7 @@ class TestQwen2_5OmniPipeline: assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.final_output_type == "audio" assert s.engine_output_type == "audio" + assert s.hf_config_name == "thinker_config" class TestQwen3TTSPipeline: diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index ca491275e..abd67bf14 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -12,7 +12,7 @@ from unittest.mock import Mock import pytest from omegaconf import OmegaConf from pydantic import ValidationError -from transformers import PretrainedConfig +from transformers import PretrainedConfig, Qwen3OmniMoeConfig from vllm.engine.arg_utils import EngineArgs from vllm_omni.config.model import OmniModelConfig @@ -244,6 +244,28 @@ def test_stage_specific_text_config_override(): assert omni_config.get_head_size() == talker_head_dim +def test_non_override_ar_stage_inputs_embeds_size_matches_hidden_size(): + """A stage without an embedding-size override keeps the default input width.""" + vllm_config = EngineArgs().create_model_config() + vllm_config.hf_config = Qwen3OmniMoeConfig(enable_audio_output=True) + + omni_config = OmniModelConfig.from_vllm_model_config( + vllm_config, + model_stage="talker", + hf_config_name="talker_config", + ) + + assert ( + getattr( + omni_config.hf_config.talker_config, + "embedding_size", + None, + ) + is None + ) + assert omni_config.get_inputs_embeds_size() == omni_config.hf_text_config.hidden_size + + def test_stage_configs_path_field(): """OmniEngineArgs with stage_configs_path should construct without error.""" args = OmniEngineArgs(stage_configs_path="/some/path.yaml") diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index d18cabcf5..f830fc37e 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -178,6 +178,14 @@ class OmniModelConfig(ModelConfig): return override return super().embedding_size + def get_inputs_embeds_size(self) -> int: + if self.hf_config_name is not None: + stage_config = getattr(self.hf_config, self.hf_config_name, None) + override = getattr(stage_config, "embedding_size", None) + if override is not None: + return override + return super().get_inputs_embeds_size() + def get_model_arch_config(self): # For multi-stage omni models, use a stage-aware convertor so that # only the correct stage's quantization config is surfaced. diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py index 4ce5ef207..de79e1aae 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py @@ -29,6 +29,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig( final_output_type="text", owns_tokenizer=True, requires_multimodal_data=True, + hf_config_name="thinker_config", engine_output_type="latent", sampling_constraints={"detokenize": True}, custom_process_next_stage_input_func=f"{_PROC}.thinker2talker_full_payload", @@ -38,6 +39,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig( model_stage="talker", execution_type=StageExecutionType.LLM_AR, input_sources=(0,), + hf_config_name="talker_config", engine_output_type="latent", sync_process_input_func=f"{_PROC}.thinker2talker_token_only", custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload", @@ -53,6 +55,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig( input_sources=(1,), final_output=True, final_output_type="audio", + hf_config_name="thinker_config", engine_output_type="audio", sync_process_input_func=f"{_PROC}.talker2code2wav_token_only", sampling_constraints={"detokenize": True}, @@ -75,6 +78,7 @@ QWEN2_5_OMNI_THINKER_ONLY_PIPELINE = PipelineConfig( final_output_type="text", owns_tokenizer=True, requires_multimodal_data=True, + hf_config_name="thinker_config", engine_output_type="latent", sampling_constraints={"detokenize": True}, custom_process_next_stage_input_func=f"{_PROC}.thinker2talker_full_payload", diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index b2437ca02..cb6e57579 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -303,7 +303,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32) # each model stage has their own hidden size self.hidden_size = self.model_config.hf_text_config.hidden_size - self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False) + # The width of embeddings received by this stage. It can differ from + # the stage's internal hidden size (e.g., Qwen2.5-Omni Talker receives + # 3584-wide Thinker states and projects them to a 896-wide hidden state). + self.inputs_embeds_size = self.model_config.get_inputs_embeds_size() + self.inputs_embeds = self._make_buffer( + self.max_num_tokens, + self.inputs_embeds_size, + dtype=self.dtype, + numpy=False, + ) # Initialize KV cache manager (preserve vllm_config fallback behavior) self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config) self._async_chunk = getattr(self.model_config, "async_chunk", False)