[Bugfix][Model] Validate DefaultModelLoader / LoadConfig and fail with clear errors (#45196)

Signed-off-by: Ting Sun <[email protected]>

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)

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)

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.
 

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

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: