A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix][Model] Validate DefaultModelLoader / LoadConfig and fail with clear errors (#45196)
Signed-off-by: Ting Sun <suntcrick@gmail.com>
Ting SUN committed
Jun 18, 2026 at 05:46 UTC
9d4b87f4f0bb489c66c466b0b8ae6bbb212f0a2f
5 files changed
+108
-7
tests/model_executor/model_loader/test_registry.py
+55
@@ -8,6 +8,7 @@ from vllm.config import ModelConfig
8
from vllm.config.load import LoadConfig
9
from vllm.model_executor.model_loader import get_model_loader, register_model_loader
10
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
11
+from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
12
13
14
@register_model_loader("custom_load_format")
@@ -33,3 +34,57 @@ def test_invalid_model_loader():
34
@register_model_loader("invalid_load_format")
35
class InValidModelLoader:
36
pass
37
+
38
+
39
+def test_default_loader_rejects_zero_num_threads():
40
+ # num_threads=0 used to fail late in ThreadPoolExecutor ("max_workers must be > 0").
41
+ with pytest.raises(ValueError, match="num_threads"):
42
+ DefaultModelLoader(
43
+ LoadConfig(
44
+ model_loader_extra_config={
45
+ "enable_multithread_load": True,
46
+ "num_threads": 0,
47
+ }
48
+ )
49
+ )
50
+
51
+
52
+def test_default_loader_rejects_multithread_with_non_lazy_strategy():
53
+ # The multi-thread loader ignores safetensors_load_strategy; reject the
54
+ # combination instead of silently dropping the requested strategy.
55
+ with pytest.raises(ValueError, match="does not support"):
56
+ DefaultModelLoader(
57
+ LoadConfig(
58
+ safetensors_load_strategy="torchao",
59
+ model_loader_extra_config={"enable_multithread_load": True},
60
+ )
61
+ )
62
+
63
+
64
+def test_default_loader_explicit_safetensors_does_not_misread_pt(tmp_path):
65
+ # Explicit safetensors must not fall back to a .pt and open it as safetensors.
66
+ (tmp_path / "model.pt").write_bytes(b"\x00\x00\x00\x00")
67
+ loader = DefaultModelLoader(LoadConfig(load_format="safetensors"))
68
+ with pytest.raises(RuntimeError, match="Cannot find any model weights"):
69
+ loader._prepare_weights(
70
+ str(tmp_path),
71
+ None,
72
+ None,
73
+ fall_back_to_pt=True,
74
+ allow_patterns_overrides=None,
75
+ )
76
+
77
+
78
+def test_default_loader_hf_still_falls_back_to_pt(tmp_path):
79
+ # Control: load_format="hf" still picks up .pt weights via fallback.
80
+ (tmp_path / "model.pt").write_bytes(b"\x00\x00\x00\x00")
81
+ loader = DefaultModelLoader(LoadConfig(load_format="hf"))
82
+ _, files, use_safetensors = loader._prepare_weights(
83
+ str(tmp_path),
84
+ None,
85
+ None,
86
+ fall_back_to_pt=True,
87
+ allow_patterns_overrides=None,
88
+ )
89
+ assert use_safetensors is False
90
+ assert any(f.endswith("model.pt") for f in files)
tests/test_config.py
+11
@@ -1557,3 +1557,14 @@ def test_ir_op_priority_ctx():
1557
# context restored even after exception
1558
assert ir.ops.rms_norm.get_priority() == ["vllm_c", "native"]
1559
assert ir.ops.fused_add_rms_norm.get_priority() == ["native"]
1560
+
1561
+
1562
+def test_load_config_rejects_invalid_safetensors_load_strategy():
1563
+ with pytest.raises(pydantic.ValidationError):
1564
+ LoadConfig(safetensors_load_strategy="not_a_real_strategy")
1565
+
1566
+
1567
+@pytest.mark.parametrize("bad_load_format", [None, 123])
1568
+def test_load_config_rejects_non_string_load_format(bad_load_format):
1569
+ with pytest.raises(pydantic.ValidationError):
1570
+ LoadConfig(load_format=bad_load_format)
vllm/config/load.py
+4
-5
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import TYPE_CHECKING, Any
4
+from typing import TYPE_CHECKING, Any, Literal, TypeAlias
5
6
from pydantic import Field, field_validator
7
@@ -11,12 +11,11 @@ from vllm.utils.hashing import safe_hash
11
12
DEFAULT_SAFETENSORS_PREFETCH_NUM_THREADS = 8
13
DEFAULT_SAFETENSORS_PREFETCH_BLOCK_SIZE = 16 * 1024 * 1024
14
+SafetensorsLoadStrategy: TypeAlias = Literal["lazy", "eager", "prefetch", "torchao"]
15
16
if TYPE_CHECKING:
16
- from vllm.model_executor.model_loader import LoadFormats
17
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
18
else:
19
- LoadFormats = Any
19
TensorizerConfig = Any
20
21
logger = init_logger(__name__)
@@ -26,7 +25,7 @@ logger = init_logger(__name__)
25
class LoadConfig:
26
"""Configuration for loading the model weights."""
27
29
- load_format: str | LoadFormats = "auto"
28
+ load_format: str = "auto"
29
"""
30
The format of the model weights to load.
31
@@ -59,7 +58,7 @@ class LoadConfig:
58
download_dir: str | None = None
59
"""Directory to download and load the weights, default to the default
60
cache directory of Hugging Face."""
62
- safetensors_load_strategy: str | None = None
61
+ safetensors_load_strategy: SafetensorsLoadStrategy | None = None
62
"""
63
Specifies the loading strategy for safetensors weights.
64
vllm/engine/arg_utils.py
+4
-1
@@ -73,6 +73,7 @@ from vllm.config.cache import (
73
)
74
from vllm.config.device import Device
75
from vllm.config.kernel import IrOpPriorityConfig, LinearBackend, MoEBackend
76
+from vllm.config.load import SafetensorsLoadStrategy
77
from vllm.config.lora import MaxLoRARanks
78
from vllm.config.mamba import MambaBackendEnum
79
from vllm.config.model import (
@@ -427,7 +428,9 @@ class EngineArgs:
428
allowed_local_media_path: str = ModelConfig.allowed_local_media_path
429
allowed_media_domains: list[str] | None = ModelConfig.allowed_media_domains
430
download_dir: str | None = LoadConfig.download_dir
430
- safetensors_load_strategy: str | None = LoadConfig.safetensors_load_strategy
431
+ safetensors_load_strategy: SafetensorsLoadStrategy | None = (
432
+ LoadConfig.safetensors_load_strategy
433
+ )
434
safetensors_prefetch_num_threads: int = LoadConfig.safetensors_prefetch_num_threads
435
safetensors_prefetch_block_size: int = LoadConfig.safetensors_prefetch_block_size
436
load_format: str | LoadFormats = LoadConfig.load_format
vllm/model_executor/model_loader/default_loader.py
+34
-1
@@ -76,6 +76,11 @@ class DefaultModelLoader(BaseModelLoader):
76
self.local_expert_ids: set[int] | None = None
77
78
extra_config = load_config.model_loader_extra_config
79
+ if not isinstance(extra_config, dict):
80
+ raise ValueError(
81
+ f"model_loader_extra_config must be a dict for load format "
82
+ f"{load_config.load_format}, got {type(extra_config).__name__}"
83
+ )
84
allowed_keys = {
85
"enable_multithread_load",
86
"num_threads",
@@ -90,10 +95,36 @@ class DefaultModelLoader(BaseModelLoader):
95
f"{unexpected_keys}"
96
)
97
98
+ enable_multithread_load = extra_config.get("enable_multithread_load", False)
99
+ if not isinstance(enable_multithread_load, bool):
100
+ raise ValueError(
101
+ f"enable_multithread_load must be a bool, got "
102
+ f"{type(enable_multithread_load).__name__}"
103
+ )
104
+ num_threads = extra_config.get("num_threads")
105
+ if num_threads is not None and not (
106
+ isinstance(num_threads, int) and num_threads > 0
107
+ ):
108
+ raise ValueError(
109
+ f"num_threads must be a positive integer, got {num_threads!r}"
110
+ )
111
+
112
self.enable_weights_track: bool | None = extra_config.get(
113
"enable_weights_track", None
114
)
115
116
+ # The multi-thread loader ignores safetensors_load_strategy, so reject
117
+ # the combination instead of silently dropping the requested strategy.
118
+ if extra_config.get("enable_multithread_load") and (
119
+ load_config.safetensors_load_strategy not in (None, "lazy")
120
+ ):
121
+ raise ValueError(
122
+ "enable_multithread_load does not support "
123
+ "safetensors_load_strategy="
124
+ f"{load_config.safetensors_load_strategy!r}; the multi-thread "
125
+ "loader only implements the default lazy strategy."
126
+ )
127
+
128
def _prepare_weights(
129
self,
130
model_name_or_path: str,
@@ -152,7 +183,9 @@ class DefaultModelLoader(BaseModelLoader):
183
else:
184
raise ValueError(f"Unknown load_format: {load_format}")
185
155
- if fall_back_to_pt:
186
+ # Don't fall back to .pt for explicit safetensors formats; otherwise a
187
+ # .pt file is matched and later opened as safetensors.
188
+ if fall_back_to_pt and not use_safetensors:
189
allow_patterns += ["*.pt"]
190
191
if allow_patterns_overrides is not None: