A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 23:45 UTC
[Core] Support model loader plugins (#21067)
Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>
22quinn committed
Jul 24, 2025 at 01:49 UTC
610852a42327836eb804c1a893fcb9bb6d52c29c
9 files changed
+159
-86
tests/fastsafetensors_loader/test_fastsafetensors_loader.py
+1
-3
@@ -2,7 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from vllm import SamplingParams
5
-from vllm.config import LoadFormat
5
6
test_model = "openai-community/gpt2"
7
@@ -17,7 +16,6 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
16
17
18
def test_model_loader_download_files(vllm_runner):
20
- with vllm_runner(test_model,
21
- load_format=LoadFormat.FASTSAFETENSORS) as llm:
19
+ with vllm_runner(test_model, load_format="fastsafetensors") as llm:
20
deserialized_outputs = llm.generate(prompts, sampling_params)
21
assert deserialized_outputs
tests/model_executor/model_loader/__init__.py
tests/model_executor/model_loader/test_registry.py
new
+37
@@ -0,0 +1,37 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+from torch import nn
6
+
7
+from vllm.config import LoadConfig, ModelConfig
8
+from vllm.model_executor.model_loader import (get_model_loader,
9
+ register_model_loader)
10
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
11
+
12
+
13
+@register_model_loader("custom_load_format")
14
+class CustomModelLoader(BaseModelLoader):
15
+
16
+ def __init__(self, load_config: LoadConfig) -> None:
17
+ super().__init__(load_config)
18
+
19
+ def download_model(self, model_config: ModelConfig) -> None:
20
+ pass
21
+
22
+ def load_weights(self, model: nn.Module,
23
+ model_config: ModelConfig) -> None:
24
+ pass
25
+
26
+
27
+def test_register_model_loader():
28
+ load_config = LoadConfig(load_format="custom_load_format")
29
+ assert isinstance(get_model_loader(load_config), CustomModelLoader)
30
+
31
+
32
+def test_invalid_model_loader():
33
+ with pytest.raises(ValueError):
34
+
35
+ @register_model_loader("invalid_load_format")
36
+ class InValidModelLoader:
37
+ pass
tests/runai_model_streamer_test/test_runai_model_streamer_loader.py
+4
-3
@@ -2,9 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from vllm import SamplingParams
5
-from vllm.config import LoadConfig, LoadFormat
5
+from vllm.config import LoadConfig
6
from vllm.model_executor.model_loader import get_model_loader
7
8
+load_format = "runai_streamer"
9
test_model = "openai-community/gpt2"
10
11
prompts = [
@@ -18,7 +19,7 @@ sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
19
20
21
def get_runai_model_loader():
21
- load_config = LoadConfig(load_format=LoadFormat.RUNAI_STREAMER)
22
+ load_config = LoadConfig(load_format=load_format)
23
return get_model_loader(load_config)
24
25
@@ -28,6 +29,6 @@ def test_get_model_loader_with_runai_flag():
29
30
31
def test_runai_model_loader_download_files(vllm_runner):
31
- with vllm_runner(test_model, load_format=LoadFormat.RUNAI_STREAMER) as llm:
32
+ with vllm_runner(test_model, load_format=load_format) as llm:
33
deserialized_outputs = llm.generate(prompts, sampling_params)
34
assert deserialized_outputs
vllm/config.py
+6
-24
@@ -65,7 +65,7 @@ if TYPE_CHECKING:
65
from vllm.model_executor.layers.quantization import QuantizationMethods
66
from vllm.model_executor.layers.quantization.base_config import (
67
QuantizationConfig)
68
- from vllm.model_executor.model_loader import BaseModelLoader
68
+ from vllm.model_executor.model_loader import LoadFormats
69
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
70
71
ConfigType = type[DataclassInstance]
@@ -78,6 +78,7 @@ else:
78
QuantizationConfig = Any
79
QuantizationMethods = Any
80
BaseModelLoader = Any
81
+ LoadFormats = Any
82
TensorizerConfig = Any
83
ConfigType = type
84
HfOverrides = Union[dict[str, Any], Callable[[type], type]]
@@ -1773,29 +1774,12 @@ class CacheConfig:
1774
logger.warning("Possibly too large swap space. %s", msg)
1775
1776
1776
-class LoadFormat(str, enum.Enum):
1777
- AUTO = "auto"
1778
- PT = "pt"
1779
- SAFETENSORS = "safetensors"
1780
- NPCACHE = "npcache"
1781
- DUMMY = "dummy"
1782
- TENSORIZER = "tensorizer"
1783
- SHARDED_STATE = "sharded_state"
1784
- GGUF = "gguf"
1785
- BITSANDBYTES = "bitsandbytes"
1786
- MISTRAL = "mistral"
1787
- RUNAI_STREAMER = "runai_streamer"
1788
- RUNAI_STREAMER_SHARDED = "runai_streamer_sharded"
1789
- FASTSAFETENSORS = "fastsafetensors"
1790
-
1791
-
1777
@config
1778
@dataclass
1779
class LoadConfig:
1780
"""Configuration for loading the model weights."""
1781
1797
- load_format: Union[str, LoadFormat,
1798
- "BaseModelLoader"] = LoadFormat.AUTO.value
1782
+ load_format: Union[str, LoadFormats] = "auto"
1783
"""The format of the model weights to load:\n
1784
- "auto" will try to load the weights in the safetensors format and fall
1785
back to the pytorch bin format if safetensors format is not available.\n
@@ -1816,7 +1800,8 @@ class LoadConfig:
1800
- "gguf" will load weights from GGUF format files (details specified in
1801
https://github.com/ggml-org/ggml/blob/master/docs/gguf.md).\n
1802
- "mistral" will load weights from consolidated safetensors files used by
1819
- Mistral models."""
1803
+ Mistral models.
1804
+ - Other custom values can be supported via plugins."""
1805
download_dir: Optional[str] = None
1806
"""Directory to download and load the weights, default to the default
1807
cache directory of Hugging Face."""
@@ -1864,10 +1849,7 @@ class LoadConfig:
1849
return hash_str
1850
1851
def __post_init__(self):
1867
- if isinstance(self.load_format, str):
1868
- load_format = self.load_format.lower()
1869
- self.load_format = LoadFormat(load_format)
1870
-
1852
+ self.load_format = self.load_format.lower()
1853
if self.ignore_patterns is not None and len(self.ignore_patterns) > 0:
1854
logger.info(
1855
"Ignoring the following patterns when downloading weights: %s",
vllm/engine/arg_utils.py
+13
-15
@@ -26,13 +26,12 @@ from vllm.config import (BlockSize, CacheConfig, CacheDType, CompilationConfig,
26
DetailedTraceModules, Device, DeviceConfig,
27
DistributedExecutorBackend, GuidedDecodingBackend,
28
GuidedDecodingBackendV1, HfOverrides, KVEventsConfig,
29
- KVTransferConfig, LoadConfig, LoadFormat,
30
- LogprobsMode, LoRAConfig, ModelConfig, ModelDType,
31
- ModelImpl, MultiModalConfig, ObservabilityConfig,
32
- ParallelConfig, PoolerConfig, PrefixCachingHashAlgo,
33
- SchedulerConfig, SchedulerPolicy, SpeculativeConfig,
34
- TaskOption, TokenizerMode, VllmConfig, get_attr_docs,
35
- get_field)
29
+ KVTransferConfig, LoadConfig, LogprobsMode,
30
+ LoRAConfig, ModelConfig, ModelDType, ModelImpl,
31
+ MultiModalConfig, ObservabilityConfig, ParallelConfig,
32
+ PoolerConfig, PrefixCachingHashAlgo, SchedulerConfig,
33
+ SchedulerPolicy, SpeculativeConfig, TaskOption,
34
+ TokenizerMode, VllmConfig, get_attr_docs, get_field)
35
from vllm.logger import init_logger
36
from vllm.platforms import CpuArchEnum, current_platform
37
from vllm.plugins import load_general_plugins
@@ -47,10 +46,12 @@ from vllm.utils import (STR_DUAL_CHUNK_FLASH_ATTN_VAL, FlexibleArgumentParser,
46
if TYPE_CHECKING:
47
from vllm.executor.executor_base import ExecutorBase
48
from vllm.model_executor.layers.quantization import QuantizationMethods
49
+ from vllm.model_executor.model_loader import LoadFormats
50
from vllm.usage.usage_lib import UsageContext
51
else:
52
ExecutorBase = Any
53
QuantizationMethods = Any
54
+ LoadFormats = Any
55
UsageContext = Any
56
57
logger = init_logger(__name__)
@@ -276,7 +277,7 @@ class EngineArgs:
277
trust_remote_code: bool = ModelConfig.trust_remote_code
278
allowed_local_media_path: str = ModelConfig.allowed_local_media_path
279
download_dir: Optional[str] = LoadConfig.download_dir
279
- load_format: str = LoadConfig.load_format
280
+ load_format: Union[str, LoadFormats] = LoadConfig.load_format
281
config_format: str = ModelConfig.config_format
282
dtype: ModelDType = ModelConfig.dtype
283
kv_cache_dtype: CacheDType = CacheConfig.cache_dtype
@@ -547,9 +548,7 @@ class EngineArgs:
548
title="LoadConfig",
549
description=LoadConfig.__doc__,
550
)
550
- load_group.add_argument("--load-format",
551
- choices=[f.value for f in LoadFormat],
552
- **load_kwargs["load_format"])
551
+ load_group.add_argument("--load-format", **load_kwargs["load_format"])
552
load_group.add_argument("--download-dir",
553
**load_kwargs["download_dir"])
554
load_group.add_argument("--model-loader-extra-config",
@@ -864,10 +863,9 @@ class EngineArgs:
863
864
# NOTE: This is to allow model loading from S3 in CI
865
if (not isinstance(self, AsyncEngineArgs) and envs.VLLM_CI_USE_S3
867
- and self.model in MODELS_ON_S3
868
- and self.load_format == LoadFormat.AUTO): # noqa: E501
866
+ and self.model in MODELS_ON_S3 and self.load_format == "auto"):
867
self.model = f"{MODEL_WEIGHTS_S3_BUCKET}/{self.model}"
870
- self.load_format = LoadFormat.RUNAI_STREAMER
868
+ self.load_format = "runai_streamer"
869
870
return ModelConfig(
871
model=self.model,
@@ -1299,7 +1297,7 @@ class EngineArgs:
1297
#############################################################
1298
# Unsupported Feature Flags on V1.
1299
1302
- if self.load_format == LoadFormat.SHARDED_STATE.value:
1300
+ if self.load_format == "sharded_state":
1301
_raise_or_fallback(
1302
feature_name=f"--load_format {self.load_format}",
1303
recommend_to_remove=False)
vllm/model_executor/model_loader/__init__.py
+87
-27
@@ -1,11 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from typing import Optional
4
+from typing import Literal, Optional
5
6
from torch import nn
7
8
-from vllm.config import LoadConfig, LoadFormat, ModelConfig, VllmConfig
8
+from vllm.config import LoadConfig, ModelConfig, VllmConfig
9
+from vllm.logger import init_logger
10
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
11
from vllm.model_executor.model_loader.bitsandbytes_loader import (
12
BitsAndBytesModelLoader)
@@ -20,34 +21,92 @@ from vllm.model_executor.model_loader.tensorizer_loader import TensorizerLoader
21
from vllm.model_executor.model_loader.utils import (
22
get_architecture_class_name, get_model_architecture, get_model_cls)
23
24
+logger = init_logger(__name__)
25
+
26
+# Reminder: Please update docstring in `LoadConfig`
27
+# if a new load format is added here
28
+LoadFormats = Literal[
29
+ "auto",
30
+ "bitsandbytes",
31
+ "dummy",
32
+ "fastsafetensors",
33
+ "gguf",
34
+ "mistral",
35
+ "npcache",
36
+ "pt",
37
+ "runai_streamer",
38
+ "runai_streamer_sharded",
39
+ "safetensors",
40
+ "sharded_state",
41
+ "tensorizer",
42
+]
43
+_LOAD_FORMAT_TO_MODEL_LOADER: dict[str, type[BaseModelLoader]] = {
44
+ "auto": DefaultModelLoader,
45
+ "bitsandbytes": BitsAndBytesModelLoader,
46
+ "dummy": DummyModelLoader,
47
+ "fastsafetensors": DefaultModelLoader,
48
+ "gguf": GGUFModelLoader,
49
+ "mistral": DefaultModelLoader,
50
+ "npcache": DefaultModelLoader,
51
+ "pt": DefaultModelLoader,
52
+ "runai_streamer": RunaiModelStreamerLoader,
53
+ "runai_streamer_sharded": ShardedStateLoader,
54
+ "safetensors": DefaultModelLoader,
55
+ "sharded_state": ShardedStateLoader,
56
+ "tensorizer": TensorizerLoader,
57
+}
58
+
59
+
60
+def register_model_loader(load_format: str):
61
+ """Register a customized vllm model loader.
62
+
63
+ When a load format is not supported by vllm, you can register a customized
64
+ model loader to support it.
65
+
66
+ Args:
67
+ load_format (str): The model loader format name.
68
+
69
+ Examples:
70
+ >>> from vllm.config import LoadConfig
71
+ >>> from vllm.model_executor.model_loader import get_model_loader, register_model_loader
72
+ >>> from vllm.model_executor.model_loader.base_loader import BaseModelLoader
73
+ >>>
74
+ >>> @register_model_loader("my_loader")
75
+ ... class MyModelLoader(BaseModelLoader):
76
+ ... def download_model(self):
77
+ ... pass
78
+ ...
79
+ ... def load_weights(self):
80
+ ... pass
81
+ >>>
82
+ >>> load_config = LoadConfig(load_format="my_loader")
83
+ >>> type(get_model_loader(load_config))
84
+ <class 'MyModelLoader'>
85
+ """ # noqa: E501
86
+
87
+ def _wrapper(model_loader_cls):
88
+ if load_format in _LOAD_FORMAT_TO_MODEL_LOADER:
89
+ logger.warning(
90
+ "Load format `%s` is already registered, and will be "
91
+ "overwritten by the new loader class `%s`.", load_format,
92
+ model_loader_cls)
93
+ if not issubclass(model_loader_cls, BaseModelLoader):
94
+ raise ValueError("The model loader must be a subclass of "
95
+ "`BaseModelLoader`.")
96
+ _LOAD_FORMAT_TO_MODEL_LOADER[load_format] = model_loader_cls
97
+ logger.info("Registered model loader `%s` with load format `%s`",
98
+ model_loader_cls, load_format)
99
+ return model_loader_cls
100
+
101
+ return _wrapper
102
+
103
104
def get_model_loader(load_config: LoadConfig) -> BaseModelLoader:
105
"""Get a model loader based on the load format."""
26
- if isinstance(load_config.load_format, type):
27
- return load_config.load_format(load_config)
28
-
29
- if load_config.load_format == LoadFormat.DUMMY:
30
- return DummyModelLoader(load_config)
31
-
32
- if load_config.load_format == LoadFormat.TENSORIZER:
33
- return TensorizerLoader(load_config)
34
-
35
- if load_config.load_format == LoadFormat.SHARDED_STATE:
36
- return ShardedStateLoader(load_config)
37
-
38
- if load_config.load_format == LoadFormat.BITSANDBYTES:
39
- return BitsAndBytesModelLoader(load_config)
40
-
41
- if load_config.load_format == LoadFormat.GGUF:
42
- return GGUFModelLoader(load_config)
43
-
44
- if load_config.load_format == LoadFormat.RUNAI_STREAMER:
45
- return RunaiModelStreamerLoader(load_config)
46
-
47
- if load_config.load_format == LoadFormat.RUNAI_STREAMER_SHARDED:
48
- return ShardedStateLoader(load_config, runai_model_streamer=True)
49
-
50
- return DefaultModelLoader(load_config)
106
+ load_format = load_config.load_format
107
+ if load_format not in _LOAD_FORMAT_TO_MODEL_LOADER:
108
+ raise ValueError(f"Load format `{load_format}` is not supported")
109
+ return _LOAD_FORMAT_TO_MODEL_LOADER[load_format](load_config)
110
111
112
def get_model(*,
@@ -66,6 +125,7 @@ __all__ = [
125
"get_architecture_class_name",
126
"get_model_architecture",
127
"get_model_cls",
128
+ "register_model_loader",
129
"BaseModelLoader",
130
"BitsAndBytesModelLoader",
131
"GGUFModelLoader",
vllm/model_executor/model_loader/default_loader.py
+9
-9
@@ -13,7 +13,7 @@ from torch import nn
13
from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
14
15
from vllm import envs
16
-from vllm.config import LoadConfig, LoadFormat, ModelConfig
16
+from vllm.config import LoadConfig, ModelConfig
17
from vllm.logger import init_logger
18
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
19
from vllm.model_executor.model_loader.weight_utils import (
@@ -104,19 +104,19 @@ class DefaultModelLoader(BaseModelLoader):
104
use_safetensors = False
105
index_file = SAFE_WEIGHTS_INDEX_NAME
106
# Some quantized models use .pt files for storing the weights.
107
- if load_format == LoadFormat.AUTO:
107
+ if load_format == "auto":
108
allow_patterns = ["*.safetensors", "*.bin"]
109
- elif (load_format == LoadFormat.SAFETENSORS
110
- or load_format == LoadFormat.FASTSAFETENSORS):
109
+ elif (load_format == "safetensors"
110
+ or load_format == "fastsafetensors"):
111
use_safetensors = True
112
allow_patterns = ["*.safetensors"]
113
- elif load_format == LoadFormat.MISTRAL:
113
+ elif load_format == "mistral":
114
use_safetensors = True
115
allow_patterns = ["consolidated*.safetensors"]
116
index_file = "consolidated.safetensors.index.json"
117
- elif load_format == LoadFormat.PT:
117
+ elif load_format == "pt":
118
allow_patterns = ["*.pt"]
119
- elif load_format == LoadFormat.NPCACHE:
119
+ elif load_format == "npcache":
120
allow_patterns = ["*.bin"]
121
else:
122
raise ValueError(f"Unknown load_format: {load_format}")
@@ -178,7 +178,7 @@ class DefaultModelLoader(BaseModelLoader):
178
hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
179
source.model_or_path, source.revision, source.fall_back_to_pt,
180
source.allow_patterns_overrides)
181
- if self.load_config.load_format == LoadFormat.NPCACHE:
181
+ if self.load_config.load_format == "npcache":
182
# Currently np_cache only support *.bin checkpoints
183
assert use_safetensors is False
184
weights_iterator = np_cache_weights_iterator(
@@ -189,7 +189,7 @@ class DefaultModelLoader(BaseModelLoader):
189
self.load_config.use_tqdm_on_load,
190
)
191
elif use_safetensors:
192
- if self.load_config.load_format == LoadFormat.FASTSAFETENSORS:
192
+ if self.load_config.load_format == "fastsafetensors":
193
weights_iterator = fastsafetensors_weights_iterator(
194
hf_weights_files,
195
self.load_config.use_tqdm_on_load,
vllm/model_executor/model_loader/sharded_state_loader.py
+2
-5
@@ -32,12 +32,9 @@ class ShardedStateLoader(BaseModelLoader):
32
33
DEFAULT_PATTERN = "model-rank-{rank}-part-{part}.safetensors"
34
35
- def __init__(self,
36
- load_config: LoadConfig,
37
- runai_model_streamer: bool = False):
35
+ def __init__(self, load_config: LoadConfig):
36
super().__init__(load_config)
37
40
- self.runai_model_streamer = runai_model_streamer
38
extra_config = ({} if load_config.model_loader_extra_config is None
39
else load_config.model_loader_extra_config.copy())
40
self.pattern = extra_config.pop("pattern", self.DEFAULT_PATTERN)
@@ -152,7 +149,7 @@ class ShardedStateLoader(BaseModelLoader):
149
150
def iterate_over_files(
151
self, paths) -> Generator[tuple[str, torch.Tensor], None, None]:
155
- if self.runai_model_streamer:
152
+ if self.load_config.load_format == "runai_streamer_sharded":
153
yield from runai_safetensors_weights_iterator(paths, True)
154
else:
155
from safetensors.torch import safe_open