From 9d4b87f4f0bb489c66c466b0b8ae6bbb212f0a2f Mon Sep 17 00:00:00 2001 From: Ting SUN Date: Thu, 18 Jun 2026 05:46:33 +0800 Subject: [PATCH] [Bugfix][Model] Validate DefaultModelLoader / LoadConfig and fail with clear errors (#45196) Signed-off-by: Ting Sun --- .../model_loader/test_registry.py | 55 +++++++++++++++++++ tests/test_config.py | 11 ++++ vllm/config/load.py | 9 ++- vllm/engine/arg_utils.py | 5 +- .../model_loader/default_loader.py | 35 +++++++++++- 5 files changed, 108 insertions(+), 7 deletions(-) diff --git a/tests/model_executor/model_loader/test_registry.py b/tests/model_executor/model_loader/test_registry.py index 020988ccac1..95b797bb514 100644 --- a/tests/model_executor/model_loader/test_registry.py +++ b/tests/model_executor/model_loader/test_registry.py @@ -8,6 +8,7 @@ from vllm.config import ModelConfig from vllm.config.load import LoadConfig from vllm.model_executor.model_loader import get_model_loader, register_model_loader from vllm.model_executor.model_loader.base_loader import BaseModelLoader +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader @register_model_loader("custom_load_format") @@ -33,3 +34,57 @@ def test_invalid_model_loader(): @register_model_loader("invalid_load_format") class InValidModelLoader: pass + + +def test_default_loader_rejects_zero_num_threads(): + # num_threads=0 used to fail late in ThreadPoolExecutor ("max_workers must be > 0"). + with pytest.raises(ValueError, match="num_threads"): + DefaultModelLoader( + LoadConfig( + model_loader_extra_config={ + "enable_multithread_load": True, + "num_threads": 0, + } + ) + ) + + +def test_default_loader_rejects_multithread_with_non_lazy_strategy(): + # The multi-thread loader ignores safetensors_load_strategy; reject the + # combination instead of silently dropping the requested strategy. + with pytest.raises(ValueError, match="does not support"): + DefaultModelLoader( + LoadConfig( + safetensors_load_strategy="torchao", + model_loader_extra_config={"enable_multithread_load": True}, + ) + ) + + +def test_default_loader_explicit_safetensors_does_not_misread_pt(tmp_path): + # Explicit safetensors must not fall back to a .pt and open it as safetensors. + (tmp_path / "model.pt").write_bytes(b"\x00\x00\x00\x00") + loader = DefaultModelLoader(LoadConfig(load_format="safetensors")) + with pytest.raises(RuntimeError, match="Cannot find any model weights"): + loader._prepare_weights( + str(tmp_path), + None, + None, + fall_back_to_pt=True, + allow_patterns_overrides=None, + ) + + +def test_default_loader_hf_still_falls_back_to_pt(tmp_path): + # Control: load_format="hf" still picks up .pt weights via fallback. + (tmp_path / "model.pt").write_bytes(b"\x00\x00\x00\x00") + loader = DefaultModelLoader(LoadConfig(load_format="hf")) + _, files, use_safetensors = loader._prepare_weights( + str(tmp_path), + None, + None, + fall_back_to_pt=True, + allow_patterns_overrides=None, + ) + assert use_safetensors is False + assert any(f.endswith("model.pt") for f in files) diff --git a/tests/test_config.py b/tests/test_config.py index e6c122de83d..d992ac29696 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1557,3 +1557,14 @@ def test_ir_op_priority_ctx(): # context restored even after exception assert ir.ops.rms_norm.get_priority() == ["vllm_c", "native"] assert ir.ops.fused_add_rms_norm.get_priority() == ["native"] + + +def test_load_config_rejects_invalid_safetensors_load_strategy(): + with pytest.raises(pydantic.ValidationError): + LoadConfig(safetensors_load_strategy="not_a_real_strategy") + + +@pytest.mark.parametrize("bad_load_format", [None, 123]) +def test_load_config_rejects_non_string_load_format(bad_load_format): + with pytest.raises(pydantic.ValidationError): + LoadConfig(load_format=bad_load_format) diff --git a/vllm/config/load.py b/vllm/config/load.py index ed591a2299f..f1066c2b9ad 100644 --- a/vllm/config/load.py +++ b/vllm/config/load.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal, TypeAlias from pydantic import Field, field_validator @@ -11,12 +11,11 @@ from vllm.utils.hashing import safe_hash DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS = 8 DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE = 16 * 1024 * 1024 +SafetensorsLoadStrategy: TypeAlias = Literal["lazy", "eager", "prefetch", "torchao"] if TYPE_CHECKING: - from vllm.model_executor.model_loader import LoadFormats from vllm.model_executor.model_loader.tensorizer import TensorizerConfig else: - LoadFormats = Any TensorizerConfig = Any logger = init_logger(__name__) @@ -26,7 +25,7 @@ logger = init_logger(__name__) class LoadConfig: """Configuration for loading the model weights.""" - load_format: str | LoadFormats = "auto" + load_format: str = "auto" """ The format of the model weights to load. @@ -59,7 +58,7 @@ class LoadConfig: download_dir: str | None = None """Directory to download and load the weights, default to the default cache directory of Hugging Face.""" - safetensors_load_strategy: str | None = None + safetensors_load_strategy: SafetensorsLoadStrategy | None = None """ Specifies the loading strategy for safetensors weights. diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 85a59bebf8a..9172a8728a0 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -73,6 +73,7 @@ from vllm.config.cache import ( ) from vllm.config.device import Device from vllm.config.kernel import IrOpPriorityConfig, LinearBackend, MoEBackend +from vllm.config.load import SafetensorsLoadStrategy from vllm.config.lora import MaxLoRARanks from vllm.config.mamba import MambaBackendEnum from vllm.config.model import ( @@ -427,7 +428,9 @@ class EngineArgs: allowed_local_media_path: str = ModelConfig.allowed_local_media_path allowed_media_domains: list[str] | None = ModelConfig.allowed_media_domains download_dir: str | None = LoadConfig.download_dir - safetensors_load_strategy: str | None = LoadConfig.safetensors_load_strategy + safetensors_load_strategy: SafetensorsLoadStrategy | None = ( + LoadConfig.safetensors_load_strategy + ) safetensors_prefetch_num_threads: int = LoadConfig.safetensors_prefetch_num_threads safetensors_prefetch_block_size: int = LoadConfig.safetensors_prefetch_block_size load_format: str | LoadFormats = LoadConfig.load_format diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index 43d5d4a4496..3ea76f4d9b3 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -76,6 +76,11 @@ class DefaultModelLoader(BaseModelLoader): self.local_expert_ids: set[int] | None = None extra_config = load_config.model_loader_extra_config + if not isinstance(extra_config, dict): + raise ValueError( + f"model_loader_extra_config must be a dict for load format " + f"{load_config.load_format}, got {type(extra_config).__name__}" + ) allowed_keys = { "enable_multithread_load", "num_threads", @@ -90,10 +95,36 @@ class DefaultModelLoader(BaseModelLoader): f"{unexpected_keys}" ) + enable_multithread_load = extra_config.get("enable_multithread_load", False) + if not isinstance(enable_multithread_load, bool): + raise ValueError( + f"enable_multithread_load must be a bool, got " + f"{type(enable_multithread_load).__name__}" + ) + num_threads = extra_config.get("num_threads") + if num_threads is not None and not ( + isinstance(num_threads, int) and num_threads > 0 + ): + raise ValueError( + f"num_threads must be a positive integer, got {num_threads!r}" + ) + self.enable_weights_track: bool | None = extra_config.get( "enable_weights_track", None ) + # The multi-thread loader ignores safetensors_load_strategy, so reject + # the combination instead of silently dropping the requested strategy. + if extra_config.get("enable_multithread_load") and ( + load_config.safetensors_load_strategy not in (None, "lazy") + ): + raise ValueError( + "enable_multithread_load does not support " + "safetensors_load_strategy=" + f"{load_config.safetensors_load_strategy!r}; the multi-thread " + "loader only implements the default lazy strategy." + ) + def _prepare_weights( self, model_name_or_path: str, @@ -152,7 +183,9 @@ class DefaultModelLoader(BaseModelLoader): else: raise ValueError(f"Unknown load_format: {load_format}") - if fall_back_to_pt: + # Don't fall back to .pt for explicit safetensors formats; otherwise a + # .pt file is matched and later opened as safetensors. + if fall_back_to_pt and not use_safetensors: allow_patterns += ["*.pt"] if allow_patterns_overrides is not None: