diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 10077d24a..8dec9c5a9 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -10,6 +10,7 @@ from types import SimpleNamespace from unittest.mock import Mock import pytest +from omegaconf import OmegaConf from pydantic import ValidationError from transformers import PretrainedConfig from vllm.engine.arg_utils import EngineArgs @@ -17,6 +18,7 @@ from vllm.engine.arg_utils import EngineArgs from vllm_omni.config.model import OmniModelConfig from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.stage_init_utils import build_engine_args_dict pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -314,3 +316,15 @@ def test_strip_single_engine_args_model_does_not_trigger_warning(mocker): assert "compilation_config" in warned_args assert "tensor_parallel_size" not in warned_args assert "model" not in warned_args + + +# For https://github.com/vllm-project/vllm-omni/issues/3293 +def test_tensor_parallel_size_none_is_handled(): + """Ensure the tensor parallel size of None isn't forwarded.""" + engine_args = OmegaConf.create({"stage_id": 0, "engine_args": {"tensor_parallel_size": None}}) + args = build_engine_args_dict( + engine_args, + model="snu-aidas/Dynin-Omni", + ) + assert isinstance(args, dict) + assert "tensor_parallel_size" not in args diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index eb25f1577..277edb5eb 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -397,6 +397,13 @@ def build_engine_args_dict( ) -> dict[str, Any]: """Build the normalized engine args dict for one stage.""" engine_args = stage_config.engine_args + # HACK (Alex) Tensor parallel size should not be passed as None; + # remove it if this is the case so that we fall back to default + # creation from vLLM's engine args. + # NOTE: This will be fixed more generically in ongoing work for engine arg filtering. + if "tensor_parallel_size" in engine_args and engine_args["tensor_parallel_size"] is None: + del engine_args["tensor_parallel_size"] + stage_type = getattr(stage_config, "stage_type", "llm") stage_id = stage_config.stage_id