From a8f1394699f6955282e7248138f2b2cefadee518 Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Sun, 22 Mar 2026 04:27:28 -0600 Subject: [PATCH] [Core] Simplify OmniModelConfig Initialization (#1768) Signed-off-by: Alex Brooks --- tests/engine/test_arg_utils.py | 134 +++++++++++++++++++++ vllm_omni/config/model.py | 212 ++++++++++++++++++++------------- vllm_omni/engine/arg_utils.py | 99 ++------------- 3 files changed, 275 insertions(+), 170 deletions(-) create mode 100644 tests/engine/test_arg_utils.py diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py new file mode 100644 index 000000000..7ba1cebec --- /dev/null +++ b/tests/engine/test_arg_utils.py @@ -0,0 +1,134 @@ +""" +Tests for Omni config utils. For stability, these tests should largely be +invariant to the specific attributes of vLLM config except in cases where we +explicitly patch values that differ from vLLM. +""" + +import inspect +from unittest.mock import Mock + +import pytest +from pydantic import ValidationError +from transformers import PretrainedConfig +from vllm.engine.arg_utils import EngineArgs + +from vllm_omni.config.model import OmniModelConfig +from vllm_omni.engine.arg_utils import OmniEngineArgs + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_sync_config_is_omni(): + """Ensure create_model_config gives the right type.""" + cfg = OmniEngineArgs().create_model_config() + assert isinstance(cfg, OmniModelConfig) + + +def test_multimodal_kwarg_overrides(): + """Ensure that overrides in the multimodal config are preserved.""" + # Get a different value than the default for a multimodal field + sig = inspect.signature(OmniEngineArgs) + default_mm_cache = sig.parameters["mm_processor_cache_gb"].default + override_val = default_mm_cache + 1 + + # NOTE: This needs to be a model that resolves to supports_multimodal=True + # in vLLM, otherwise we won't have an MM config + cfg = OmniEngineArgs( + model="Qwen/Qwen2-VL-2B-Instruct", + mm_processor_cache_gb=override_val, + ).create_model_config() + + # Ensure that the override was applied correctly + assert cfg.multimodal_config is not None + assert cfg.multimodal_config.mm_processor_cache_gb == override_val + + +def test_from_vllm_config_validates_invalid_omni_kwargs(): + """Ensure omni-specific field validation catches invalid keys.""" + model_config = EngineArgs().create_model_config() + with pytest.raises(ValueError, match="Unexpected omni kwarg"): + OmniModelConfig.from_vllm_model_config(model_config, foo="bar") + + +def test_from_vllm_config_validates_bad_omni_kwarg_types(): + """Ensure omni-specific field validation catches type errors.""" + model_config = EngineArgs().create_model_config() + with pytest.raises(ValidationError): + OmniModelConfig.from_vllm_model_config(model_config, stage_id="not_an_int") + + +def test_default_all_values_are_initialized(): + """Ensure omni-specific field initializes all fields""" + model_config = EngineArgs().create_model_config() + cfg = OmniModelConfig.from_vllm_model_config(model_config) + + # Test a primitive + assert cfg.model_stage == "thinker" + # Test a field initialized with a default factory + assert cfg.stage_connector_config == { + "name": "SharedMemoryConnector", + "extra": {}, + } + + # Ensure that hf_config is initialized on model_config in the vLLM by ModelConfig's + # __post_init__, and that the hf_config is copied over to the OmniModelConfig; + # we explicitly set this since the field sets init=False + assert isinstance(model_config.hf_config, PretrainedConfig) + assert cfg.hf_config is model_config.hf_config + + # Ensure that we can convert it to a string; this will convert + # all attributes, so should raise if we have attributes that are + # not initialized correctly, e.g., due to default factories + str(cfg) + + +def test_qwen3_tts_codec_frame_rate_patching(): + """Ensure the patch for qwen3 tts is applied correctly when creating the omni config.""" + # Create a vLLM ModelConfig + vllm_config = EngineArgs().create_model_config() + + # Create a mock talking config with a dummy value for position_id_per_seconds + mock_talker_config = Mock() + mock_talker_config.position_id_per_seconds = 12.3 + vllm_config.hf_config.talker_config = mock_talker_config + + # Ensure creating the config for a Qwen3TTSTalkerForConditionalGenerationARVLLM + # model calls the patch func to apply position_id_per_seconds from the talker + # config to the config's codec_frame_rate_hz + omni_config = OmniModelConfig.from_vllm_model_config( + vllm_config, + model_arch="Qwen3TTSTalkerForConditionalGenerationARVLLM", + ) + + # Verify codec_frame_rate_hz was patched + assert omni_config.codec_frame_rate_hz == 12.3 + + +def test_stage_specific_text_config_override(): + """Ensure dependent attributes are updated when using stage-specific config.""" + vllm_config = EngineArgs().create_model_config() + vllm_config.disable_sliding_window = True + + # Switch the created hf text config with a mock whose + # values we want to pull through the text config helper + stage_text_config = vllm_config.hf_text_config + vllm_config.hf_text_config = Mock() + stage_text_config.sliding_window = 4096 + stage_text_config.attention_chunk_size = 2048 + + # Move the stage config's text config getter & thinker config + mock_stage_config = Mock() + mock_stage_config.get_text_config.return_value = stage_text_config + vllm_config.hf_config.thinker_config = mock_stage_config + + # Ensure that create from a vLLM config correctly pulls the + # expected values off of the thinker config & swaps the text config + omni_config = OmniModelConfig.from_vllm_model_config( + vllm_config, + hf_config_name="thinker_config", + ) + + assert omni_config.hf_text_config is stage_text_config + assert omni_config.attention_chunk_size == 2048 + assert omni_config.max_model_len == 4096 + assert omni_config.hf_text_config.sliding_window is None diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index f837b8d45..db558f64a 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -1,13 +1,11 @@ -from dataclasses import field +from dataclasses import MISSING, field from typing import Any -from pydantic import ConfigDict +from pydantic import ConfigDict, TypeAdapter from vllm.config import ModelConfig -from vllm.config.multimodal import MMCacheType, MMEncoderTPMode from vllm.config.utils import config from vllm.logger import init_logger from vllm.transformers_utils.config import get_hf_text_config -from vllm.v1.attention.backends.registry import AttentionBackendEnum import vllm_omni.model_executor.models as me_models @@ -18,31 +16,38 @@ logger = init_logger(__name__) class OmniModelConfig(ModelConfig): """Configuration for Omni models, extending the base ModelConfig. - This configuration class extends the base vLLM ModelConfig with - omni-specific fields for multi-stage pipeline processing. + This configuration class extends the base vLLM ModelConfig with + omni-specific fields for multi-stage pipeline processing. - Attributes: - stage_id: Identifier for the stage in a multi-stage pipeline (default: 0) - async_chunk: If set to True, perform async chunk - model_stage: Stage type identifier, e.g., "thinker" or "talker" - (default: "thinker") - model_arch: Model architecture name - (default: "Qwen2_5OmniForConditionalGeneration") - worker_type: Model Type, e.g., "ar" or "generation" - engine_output_type: Optional output type specification for the engine. - Used to route outputs to appropriate processors (e.g., "image", - "audio", "latents"). If None, output type is inferred. - stage_connector_config: Stage connector configuration dictionary. - Contains "name" (connector name), "extra" (extra connector config). - task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base). - If not specified, will be inferred from model path. + Attributes: + hf_config: The model's HF Transformers config (default: None) + hf_text_config: The sub text_config of the model's hf_config (default: None) + stage_id: Identifier for the stage in a multi-stage pipeline (default: 0) + async_chunk: If set to True, perform async chunk + model_stage: Stage type identifier, e.g., "thinker" or "talker" + (default: "thinker") + model_arch: Model architecture name + (default: "Qwen2_5OmniForConditionalGeneration") + worker_type: Model Type, e.g., "ar" or "generation" + engine_output_type: Optional output type specification for the engine. + Used to route outputs to appropriate processors (e.g., "image", + "audio", "latents"). If None, output type is inferred. + stage_connector_config: Stage connector configuration dictionary. + Contains "name" (connector name), "extra" (extra connector config). + task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base). + If not specified, will be inferred from model path. - Example: - >>> config = OmniModelConfig( - ... stage_id=0, - ... model_stage="thinker", - ... model_arch="Qwen2_5OmniForConditionalGeneration" - ... ) + + The correct way to initialize this class is via vLLM config, as most + of the logic for handling values is in the ModelConfig's __post_init__. + + Example: + >>> config = OmniModelConfig.from_vllm_model_config( + ... vllm_config, + ... stage_id=0, + ... model_stage="thinker", + ... model_arch="Qwen2_5OmniForConditionalGeneration" + ... ) """ stage_id: int = 0 @@ -100,59 +105,27 @@ class OmniModelConfig(ModelConfig): ) return get_hf_text_config(self.hf_config) - def __post_init__( - self, - # Multimodal config init vars - language_model_only: bool, - limit_mm_per_prompt: dict[str, int | dict[str, int]] | None, - enable_mm_embeds: bool | None, - media_io_kwargs: dict[str, dict[str, Any]] | None, - mm_processor_kwargs: dict[str, Any] | None, - mm_processor_cache_gb: float | None, - mm_processor_cache_type: MMCacheType | None, - mm_shm_cache_max_object_size_mb: int | None, - mm_encoder_only: bool | None, - mm_encoder_tp_mode: MMEncoderTPMode | None, - mm_encoder_attn_backend: AttentionBackendEnum | str | None, - interleave_mm_strings: bool | None, - skip_mm_profiling: bool | None, - video_pruning_rate: float | None, - ) -> None: - # Call parent's __post_init__ to handle all standard ModelConfig initialization - super().__post_init__( - language_model_only=language_model_only, - limit_mm_per_prompt=limit_mm_per_prompt, - enable_mm_embeds=enable_mm_embeds, - media_io_kwargs=media_io_kwargs, - mm_processor_kwargs=mm_processor_kwargs, - mm_processor_cache_gb=mm_processor_cache_gb, - mm_processor_cache_type=mm_processor_cache_type, - mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb, - mm_encoder_only=mm_encoder_only, - mm_encoder_tp_mode=mm_encoder_tp_mode, - mm_encoder_attn_backend=mm_encoder_attn_backend, - interleave_mm_strings=interleave_mm_strings, - skip_mm_profiling=skip_mm_profiling, - video_pruning_rate=video_pruning_rate, - ) + def _patch_qwen3_tts(self): + """Patches the value of `position_id_per_seconds` in Qwen3's + TTS's talker_config into the this class's codec_frame_rate_hz. + """ + talker_cfg = getattr(self.hf_config, "talker_config", None) + if isinstance(talker_cfg, dict): + pos_per_sec = talker_cfg.get("position_id_per_seconds") + else: + pos_per_sec = getattr(talker_cfg, "position_id_per_seconds", None) + if pos_per_sec is not None: + try: + fps = float(pos_per_sec) + except Exception: + fps = None + if fps is not None and fps > 0: + self.codec_frame_rate_hz = fps - # Qwen3-TTS: infer codec frame rate from the model config for online serving. - if self.codec_frame_rate_hz is None and self.model_arch == "Qwen3TTSTalkerForConditionalGenerationARVLLM": - talker_cfg = getattr(self.hf_config, "talker_config", None) - if isinstance(talker_cfg, dict): - pos_per_sec = talker_cfg.get("position_id_per_seconds") - else: - pos_per_sec = getattr(talker_cfg, "position_id_per_seconds", None) - if pos_per_sec is not None: - try: - fps = float(pos_per_sec) - except Exception: - fps = None - if fps is not None and fps > 0: - self.codec_frame_rate_hz = fps - - # Override hf_text_config with omni-specific logic for multi-stage models - # (e.g., thinker_config, talker_config) + def _maybe_override_text_config(self): + """Override hf_text_config with omni-specific logic for multi-stage + models (e.g., thinker_config, talker_config). + """ new_hf_text_config = self.draw_hf_text_config() if new_hf_text_config is not self.hf_text_config: self.hf_text_config = new_hf_text_config @@ -161,5 +134,84 @@ class OmniModelConfig(ModelConfig): # Recalculate max_model_len since it depends on hf_text_config self.max_model_len = self.get_and_verify_max_len(self.original_max_model_len) # Reset sliding_window if needed - if self.disable_sliding_window: + if self.disable_sliding_window and self.hf_text_config is not None: self.hf_text_config.sliding_window = None + + @classmethod + def from_vllm_model_config(cls, model_config: ModelConfig, **omni_kwargs): + """Create OmniModelConfig from an existing vLLM ModelConfig + and additional Omni specific kwargs. + + NOTE: The validation and __post_init__ for ModelConfig is expensive; + to avoid calling it a second time, we explicitly retrieve defaults + from dataclass attributes for values not passed to omni_kwargs, + and use that to initialize a __new__ instance. This is significantly + faster than creating the OmniModelConfig directly from the ModelConfig, + and saves us from having to pass all kwargs to the OmniModelConfig. + """ + # Add missing defaults to the omni kwargs and ensure values are valid + cls.add_defaults_to_omni_kwargs(omni_kwargs) + cls._validate_omni_fields(**omni_kwargs) + + # Allocate the new omni config and copy the model config & omni fields + omni_cfg = object.__new__(cls) + omni_cfg.__dict__.update(model_config.__dict__) + omni_cfg.__dict__.update(omni_kwargs) + + # Apply any model specific patches or necessary overrides + if ( + omni_cfg.codec_frame_rate_hz is None + and omni_cfg.model_arch == "Qwen3TTSTalkerForConditionalGenerationARVLLM" + ): + omni_cfg._patch_qwen3_tts() + + omni_cfg._maybe_override_text_config() + + if omni_cfg.hf_config is not None: + omni_cfg.hf_config.architectures = omni_cfg.architectures + + return omni_cfg + + @classmethod + def _validate_omni_fields(cls, **omni_kwargs): + """Validate omni-specific fields; we use TypeAdapters here to quickly + validate only omni kwargs to avoid rerunning validation on the + ModelConfig. + + NOTE: This assumes add_defaults_to_omni_kwargs has already been called, + so that all omni fields are present in the provided omni_kwargs. + """ + omni_fields = set(cls.__dataclass_fields__) - set(ModelConfig.__dataclass_fields__) + + for key, value in omni_kwargs.items(): + if key not in omni_fields: + raise ValueError(f"Unexpected omni kwarg: {key}") + + field_type = cls.__dataclass_fields__[key].type + if field_type is not Any: + TypeAdapter(field_type).validate_python(value) + + # We should not have any uninitialized keys + uninitialized_fields = omni_fields - omni_kwargs.keys() + if len(uninitialized_fields): + logger.error(f"The following OmniModelConfig keys were not initialized: {uninitialized_fields}") + + @classmethod + def add_defaults_to_omni_kwargs(cls, omni_kwargs): + """Because we init the OmniModelConfig with __new__ to sidestep expensive + validation, we need to be careful to ensure fields with default factories + are initialized, otherwise we will get an AttributeError when we use it. + + To work around this issue, we explicitly add defaults to the omni_kwargs + dict provided to ensure all fields are defined correctly. + + NOTE: omni_kwargs are mutated in place. + """ + omni_fields = set(cls.__dataclass_fields__) - set(ModelConfig.__dataclass_fields__) + + for field_name in omni_fields - set(omni_kwargs.keys()): + field_def = cls.__dataclass_fields__[field_name] + if field_def.default_factory is not MISSING: + omni_kwargs[field_name] = field_def.default_factory() + elif field_def.default is not MISSING: + omni_kwargs[field_name] = field_def.default diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index fd1845c2c..4d517532d 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -3,10 +3,8 @@ import dataclasses from dataclasses import dataclass, field from typing import Any -import vllm.envs as envs from vllm.engine.arg_utils import EngineArgs from vllm.logger import init_logger -from vllm.transformers_utils.gguf_utils import is_gguf from vllm_omni.config import OmniModelConfig from vllm_omni.plugins import load_omni_general_plugins @@ -68,6 +66,9 @@ class OmniEngineArgs(EngineArgs): engine_output_type: Optional output type specification for the engine. Used to route outputs to appropriate processors (e.g., "image", "audio", "latents"). If None, output type is inferred. + hf_config_name: Optional key for HF config subkey to be extracted + for this stage, e.g., talker_config; If None, the default + HF config will be used. custom_process_next_stage_input_func: Optional path to a custom function for processing inputs from previous stages If None, default processing is used. @@ -113,38 +114,9 @@ class OmniEngineArgs(EngineArgs): Returns: OmniModelConfig instance with all configuration fields set """ - # GGUF files need a specific model loader path in vLLM. - if is_gguf(self.model): - self.quantization = self.load_format = "gguf" - - if not envs.VLLM_ENABLE_V1_MULTIPROCESSING: - logger.warning( - "The global random seed is set to %d. Since " - "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may " - "affect the random state of the Python process that " - "launched vLLM.", - self.seed, - ) - # register omni models to avoid model not found error self._ensure_omni_models_registered() - # Keep compatibility when async args are constructed from partial payloads. - language_model_only = getattr(self, "language_model_only", False) - limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {}) - enable_mm_embeds = getattr(self, "enable_mm_embeds", False) - interleave_mm_strings = getattr(self, "interleave_mm_strings", False) - media_io_kwargs = getattr(self, "media_io_kwargs", {}) - skip_mm_profiling = getattr(self, "skip_mm_profiling", False) - mm_processor_kwargs = getattr(self, "mm_processor_kwargs", None) - mm_processor_cache_gb = getattr(self, "mm_processor_cache_gb", 4) - mm_processor_cache_type = getattr(self, "mm_processor_cache_type", None) - mm_shm_cache_max_object_size_mb = getattr(self, "mm_shm_cache_max_object_size_mb", 128) - mm_encoder_only = getattr(self, "mm_encoder_only", False) - mm_encoder_tp_mode = getattr(self, "mm_encoder_tp_mode", "weights") - mm_encoder_attn_backend = getattr(self, "mm_encoder_attn_backend", None) - video_pruning_rate = getattr(self, "video_pruning_rate", 0.0) - # Build stage_connector_config from stage_connector_spec stage_connector_config = { "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"), @@ -152,63 +124,12 @@ class OmniEngineArgs(EngineArgs): } stage_connector_config["extra"]["stage_id"] = self.stage_id - # Create OmniModelConfig directly from engine args - # Note: We pass the actual init parameters matching vLLM's EngineArgs.create_model_config() - omni_config = OmniModelConfig( - # Base ModelConfig fields (matching vLLM's EngineArgs.create_model_config) - model=self.model, - model_weights=self.model_weights, - hf_config_path=self.hf_config_path, - runner=self.runner, - convert=self.convert, - tokenizer=self.tokenizer, - tokenizer_mode=self.tokenizer_mode, - trust_remote_code=self.trust_remote_code, - allowed_local_media_path=self.allowed_local_media_path, - allowed_media_domains=self.allowed_media_domains, - dtype=self.dtype, - seed=self.seed, - revision=self.revision, - code_revision=self.code_revision, - hf_token=self.hf_token, - hf_overrides=self.hf_overrides, - tokenizer_revision=self.tokenizer_revision, - max_model_len=self.max_model_len, - quantization=self.quantization, - allow_deprecated_quantization=self.allow_deprecated_quantization, - enforce_eager=self.enforce_eager, - enable_return_routed_experts=self.enable_return_routed_experts, - max_logprobs=self.max_logprobs, - logprobs_mode=self.logprobs_mode, - disable_sliding_window=self.disable_sliding_window, - disable_cascade_attn=self.disable_cascade_attn, - skip_tokenizer_init=self.skip_tokenizer_init, - enable_prompt_embeds=self.enable_prompt_embeds, - served_model_name=self.served_model_name, - language_model_only=language_model_only, - limit_mm_per_prompt=limit_mm_per_prompt, - enable_mm_embeds=enable_mm_embeds, - interleave_mm_strings=interleave_mm_strings, - media_io_kwargs=media_io_kwargs, - skip_mm_profiling=skip_mm_profiling, - config_format=self.config_format, - mm_processor_kwargs=mm_processor_kwargs, - mm_processor_cache_gb=mm_processor_cache_gb, - mm_processor_cache_type=mm_processor_cache_type, - mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb, - mm_encoder_only=mm_encoder_only, - mm_encoder_tp_mode=mm_encoder_tp_mode, - mm_encoder_attn_backend=mm_encoder_attn_backend, - pooler_config=self.pooler_config, - generation_config=self.generation_config, - override_generation_config=self.override_generation_config, - enable_sleep_mode=self.enable_sleep_mode, - model_impl=self.model_impl, - override_attention_dtype=self.override_attention_dtype, - logits_processors=self.logits_processors, - video_pruning_rate=video_pruning_rate, - io_processor_plugin=self.io_processor_plugin, - # Omni-specific fields + # Build the vLLM config first, then use it to create the Omni config. + model_config = super().create_model_config() + + omni_config = OmniModelConfig.from_vllm_model_config( + model_config=model_config, + # All kwargs below are Omni specific stage_id=self.stage_id, async_chunk=self.async_chunk, model_stage=self.model_stage, @@ -221,6 +142,4 @@ class OmniEngineArgs(EngineArgs): omni_kv_config=self.omni_kv_config, task_type=self.task_type, ) - omni_config.hf_config.architectures = omni_config.architectures - return omni_config