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 15, 11:00 UTC
[Misc] Split model loader (#17712)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
May 7, 2025 at 12:42 UTC
822de7fb9426dced9cba62f4d470176831cb60f0
21 files changed
+1681
-1566
tests/runai_model_streamer_test/test_runai_model_streamer_loader.py
+2
-3
@@ -2,8 +2,7 @@
2
3
from vllm import SamplingParams
4
from vllm.config import LoadConfig, LoadFormat
5
-from vllm.model_executor.model_loader.loader import (RunaiModelStreamerLoader,
6
- get_model_loader)
5
+from vllm.model_executor.model_loader import get_model_loader
6
7
test_model = "openai-community/gpt2"
8
@@ -24,7 +23,7 @@ def get_runai_model_loader():
23
24
def test_get_model_loader_with_runai_flag():
25
model_loader = get_runai_model_loader()
27
- assert isinstance(model_loader, RunaiModelStreamerLoader)
26
+ assert model_loader.__class__.__name__ == "RunaiModelStreamerLoader"
27
28
29
def test_runai_model_loader_download_files(vllm_runner):
tests/test_sharded_state_loader.py
+1
-1
@@ -10,7 +10,7 @@ import torch
10
from huggingface_hub import snapshot_download
11
12
from vllm import LLM, SamplingParams
13
-from vllm.model_executor.model_loader.loader import ShardedStateLoader
13
+from vllm.model_executor.model_loader import ShardedStateLoader
14
15
prompts = [
16
"Hello, my name is",
tests/utils.py
+1
-1
@@ -29,7 +29,7 @@ from vllm.distributed import (ensure_model_parallel_initialized,
29
init_distributed_environment)
30
from vllm.engine.arg_utils import AsyncEngineArgs
31
from vllm.entrypoints.openai.cli_args import make_arg_parser
32
-from vllm.model_executor.model_loader.loader import get_model_loader
32
+from vllm.model_executor.model_loader import get_model_loader
33
from vllm.platforms import current_platform
34
from vllm.transformers_utils.tokenizer import get_tokenizer
35
from vllm.utils import (FlexibleArgumentParser, GB_bytes,
vllm/config.py
+1
-1
@@ -54,7 +54,7 @@ if TYPE_CHECKING:
54
from vllm.executor.executor_base import ExecutorBase
55
from vllm.model_executor.layers.quantization.base_config import (
56
QuantizationConfig)
57
- from vllm.model_executor.model_loader.loader import BaseModelLoader
57
+ from vllm.model_executor.model_loader import BaseModelLoader
58
59
ConfigType = type[DataclassInstance]
60
else:
vllm/model_executor/model_loader/__init__.py
+53
-5
@@ -2,19 +2,67 @@
2
3
from torch import nn
4
5
-from vllm.config import VllmConfig
6
-from vllm.model_executor.model_loader.loader import (BaseModelLoader,
7
- get_model_loader)
5
+from vllm.config import LoadConfig, LoadFormat, VllmConfig
6
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
7
+from vllm.model_executor.model_loader.bitsandbytes_loader import (
8
+ BitsAndBytesModelLoader)
9
+from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
10
+from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader
11
+from vllm.model_executor.model_loader.gguf_loader import GGUFModelLoader
12
+from vllm.model_executor.model_loader.runai_streamer_loader import (
13
+ RunaiModelStreamerLoader)
14
+from vllm.model_executor.model_loader.sharded_state_loader import (
15
+ ShardedStateLoader)
16
+from vllm.model_executor.model_loader.tensorizer_loader import TensorizerLoader
17
from vllm.model_executor.model_loader.utils import (
18
get_architecture_class_name, get_model_architecture)
19
20
21
+def get_model_loader(load_config: LoadConfig) -> BaseModelLoader:
22
+ """Get a model loader based on the load format."""
23
+ if isinstance(load_config.load_format, type):
24
+ return load_config.load_format(load_config)
25
+
26
+ if load_config.load_format == LoadFormat.DUMMY:
27
+ return DummyModelLoader(load_config)
28
+
29
+ if load_config.load_format == LoadFormat.TENSORIZER:
30
+ return TensorizerLoader(load_config)
31
+
32
+ if load_config.load_format == LoadFormat.SHARDED_STATE:
33
+ return ShardedStateLoader(load_config)
34
+
35
+ if load_config.load_format == LoadFormat.BITSANDBYTES:
36
+ return BitsAndBytesModelLoader(load_config)
37
+
38
+ if load_config.load_format == LoadFormat.GGUF:
39
+ return GGUFModelLoader(load_config)
40
+
41
+ if load_config.load_format == LoadFormat.RUNAI_STREAMER:
42
+ return RunaiModelStreamerLoader(load_config)
43
+
44
+ if load_config.load_format == LoadFormat.RUNAI_STREAMER_SHARDED:
45
+ return ShardedStateLoader(load_config, runai_model_streamer=True)
46
+
47
+ return DefaultModelLoader(load_config)
48
+
49
+
50
def get_model(*, vllm_config: VllmConfig) -> nn.Module:
51
loader = get_model_loader(vllm_config.load_config)
52
return loader.load_model(vllm_config=vllm_config)
53
54
55
__all__ = [
18
- "get_model", "get_model_loader", "BaseModelLoader",
19
- "get_architecture_class_name", "get_model_architecture"
56
+ "get_model",
57
+ "get_model_loader",
58
+ "get_architecture_class_name",
59
+ "get_model_architecture",
60
+ "BaseModelLoader",
61
+ "BitsAndBytesModelLoader",
62
+ "GGUFModelLoader",
63
+ "DefaultModelLoader",
64
+ "DummyModelLoader",
65
+ "RunaiModelStreamerLoader",
66
+ "ShardedStateLoader",
67
+ "TensorizerLoader",
68
]
vllm/model_executor/model_loader/base_loader.py
new
+23
@@ -0,0 +1,23 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from abc import ABC, abstractmethod
3
+
4
+import torch.nn as nn
5
+
6
+from vllm.config import LoadConfig, ModelConfig, VllmConfig
7
+
8
+
9
+class BaseModelLoader(ABC):
10
+ """Base class for model loaders."""
11
+
12
+ def __init__(self, load_config: LoadConfig):
13
+ self.load_config = load_config
14
+
15
+ @abstractmethod
16
+ def download_model(self, model_config: ModelConfig) -> None:
17
+ """Download a model so that it can be immediately loaded."""
18
+ raise NotImplementedError
19
+
20
+ @abstractmethod
21
+ def load_model(self, *, vllm_config: VllmConfig) -> nn.Module:
22
+ """Load a model with the given configurations."""
23
+ raise NotImplementedError
vllm/model_executor/model_loader/bitsandbytes_loader.py
new
+568
@@ -0,0 +1,568 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# ruff: noqa: SIM117
3
+import copy
4
+import fnmatch
5
+import glob
6
+import itertools
7
+import math
8
+import os
9
+from typing import Any, Callable, Dict, Generator, List, Optional, Tuple
10
+
11
+import numpy as np
12
+import torch
13
+from huggingface_hub import HfApi
14
+from torch import nn
15
+from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
16
+
17
+from vllm.config import LoadConfig, ModelConfig, VllmConfig
18
+from vllm.distributed import (get_tensor_model_parallel_rank,
19
+ get_tensor_model_parallel_world_size)
20
+# yapf: enable
21
+from vllm.logger import init_logger
22
+# yapf conflicts with isort for this block
23
+# yapf: disable
24
+from vllm.model_executor.layers.linear import (LinearBase,
25
+ MergedColumnParallelLinear,
26
+ QKVParallelLinear,
27
+ ReplicatedLinear,
28
+ RowParallelLinear)
29
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
30
+from vllm.model_executor.model_loader.utils import (ParamMapping,
31
+ initialize_model,
32
+ set_default_torch_dtype)
33
+from vllm.model_executor.model_loader.weight_utils import (
34
+ download_safetensors_index_file_from_hf, download_weights_from_hf,
35
+ filter_duplicate_safetensors_files, filter_files_not_needed_for_inference,
36
+ pt_weights_iterator, safetensors_weights_iterator)
37
+from vllm.model_executor.utils import set_weight_attrs
38
+from vllm.platforms import current_platform
39
+
40
+logger = init_logger(__name__)
41
+
42
+
43
+class BitsAndBytesModelLoader(BaseModelLoader):
44
+ """Model loader to load model weights with BitAndBytes quantization."""
45
+
46
+ possible_config_file_names = ["adapter_config.json"]
47
+
48
+ def __init__(self, load_config: LoadConfig):
49
+ super().__init__(load_config)
50
+
51
+ # Save the module names without sharding.
52
+ self.unsharded_weights_modules: List[str] = []
53
+ # Save the module names that are sharded by column.
54
+ self.column_sharded_weights_modules: List[str] = []
55
+ # Store all module names (from transformers) that support
56
+ # BNB quantization.
57
+ self.target_modules: List[str] = []
58
+ # mapping weight names from transformers to vllm.
59
+ self.weight_mapper: Callable = lambda name: name
60
+
61
+ def _get_weight_files(
62
+ self,
63
+ model_name_or_path: str,
64
+ allowed_patterns: List[str],
65
+ revision: Optional[str] = None,
66
+ ) -> Tuple[str, List[str], str]:
67
+ """Retrieve weight files. Download the files if necessary.
68
+
69
+ Return the weight files and the file pattern."""
70
+ is_local = os.path.isdir(model_name_or_path)
71
+
72
+ if is_local:
73
+ for pattern in allowed_patterns:
74
+ weight_files = glob.glob(
75
+ os.path.join(model_name_or_path, pattern))
76
+ if weight_files:
77
+ return model_name_or_path, weight_files, pattern
78
+ else:
79
+ hf_api = HfApi()
80
+ repo_files = hf_api.list_repo_files(repo_id=model_name_or_path)
81
+ for pattern in allowed_patterns:
82
+ matching_files = fnmatch.filter(repo_files, pattern)
83
+ if matching_files:
84
+ hf_folder = download_weights_from_hf(
85
+ model_name_or_path,
86
+ self.load_config.download_dir,
87
+ [pattern],
88
+ revision,
89
+ ignore_patterns=self.load_config.ignore_patterns,
90
+ )
91
+ return hf_folder, glob.glob(
92
+ os.path.join(hf_folder, pattern)), pattern
93
+
94
+ raise RuntimeError(
95
+ f"No model weights found in: `{model_name_or_path}`")
96
+
97
+ def _prepare_weights(self, model_name_or_path: str,
98
+ revision: Optional[str]) -> Tuple[List[str], bool]:
99
+ """Prepare weight files for the model."""
100
+
101
+ allowed_patterns = ["*.safetensors", "*.bin", "*.pt"]
102
+
103
+ hf_folder, hf_weights_files, matched_pattern = self._get_weight_files(
104
+ model_name_or_path, allowed_patterns, revision)
105
+
106
+ use_safetensors = matched_pattern == "*.safetensors"
107
+ is_local = os.path.isdir(model_name_or_path)
108
+ index_file = SAFE_WEIGHTS_INDEX_NAME
109
+ if use_safetensors:
110
+ # For models like Mistral-7B-Instruct-v0.3
111
+ # there are both sharded safetensors files and a consolidated
112
+ # safetensors file. Using both breaks.
113
+ # Here, we download the `model.safetensors.index.json` and filter
114
+ # any files not found in the index.
115
+ if not is_local:
116
+ download_safetensors_index_file_from_hf(
117
+ model_name_or_path,
118
+ index_file,
119
+ self.load_config.download_dir,
120
+ revision,
121
+ )
122
+ hf_weights_files = filter_duplicate_safetensors_files(
123
+ hf_weights_files, hf_folder, index_file)
124
+ else:
125
+ hf_weights_files = filter_files_not_needed_for_inference(
126
+ hf_weights_files)
127
+
128
+ if len(hf_weights_files) == 0:
129
+ raise RuntimeError(
130
+ f"Cannot find any model weights with `{model_name_or_path}`")
131
+
132
+ return hf_weights_files, use_safetensors
133
+
134
+ def _hf_weight_iter(self, hf_weights_files, use_safetensors: bool):
135
+ if use_safetensors:
136
+ iterator = safetensors_weights_iterator(
137
+ hf_weights_files,
138
+ self.load_config.use_tqdm_on_load,
139
+ )
140
+ else:
141
+ iterator = pt_weights_iterator(
142
+ hf_weights_files,
143
+ self.load_config.use_tqdm_on_load,
144
+ self.load_config.pt_load_map_location,
145
+ )
146
+ for org_name, param in iterator:
147
+ # mapping weight names from transformers to vllm while preserving
148
+ # original names.
149
+ mapped_name = self.weight_mapper(org_name)
150
+ yield org_name, mapped_name, param
151
+
152
+ def _get_quantized_weights_iterator(
153
+ self,
154
+ model_name_or_path: str,
155
+ revision: Optional[str],
156
+ pre_quant: bool,
157
+ load_8bit: bool,
158
+ ) -> Tuple[Generator[Tuple[str, torch.Tensor], None, None], Dict[str,
159
+ Any]]:
160
+ """Get an iterator to the model weights with bitsandbytes quantization,
161
+ as well as the quantization state dictionary."""
162
+
163
+ # only load the bitsandbytes module when needed
164
+ try:
165
+ import bitsandbytes
166
+
167
+ if bitsandbytes.__version__ < "0.45.3":
168
+ raise ImportError("bitsandbytes version is wrong. Please "
169
+ "install bitsandbytes>=0.45.3.")
170
+ except ImportError as err:
171
+ raise ImportError("Please install bitsandbytes>=0.45.3 via "
172
+ "`pip install bitsandbytes>=0.45.3` to use "
173
+ "bitsandbytes quantizer.") from err
174
+
175
+ hf_weights_files, use_safetensors = self._prepare_weights(
176
+ model_name_or_path, revision)
177
+
178
+ quant_state_dict: Dict[str, Any] = {}
179
+
180
+ if pre_quant:
181
+ if load_8bit:
182
+ return self._quantized_8bit_generator(
183
+ hf_weights_files, use_safetensors,
184
+ quant_state_dict), quant_state_dict
185
+ else:
186
+ return self._quantized_4bit_generator(
187
+ hf_weights_files, use_safetensors,
188
+ quant_state_dict), quant_state_dict
189
+
190
+ return self._unquantized_generator(hf_weights_files, use_safetensors,
191
+ quant_state_dict), quant_state_dict
192
+
193
+ def _is_8bit_weight_name(self, weight_name: str):
194
+ quantized_suffix = {".scb", ".weight_format"}
195
+ return any(weight_name.lower().endswith(suffix)
196
+ for suffix in quantized_suffix)
197
+
198
+ def _is_4bit_weight_name(self, weight_name: str):
199
+ quantized_suffix = {
200
+ "absmax",
201
+ "quant_map",
202
+ "nested_absmax",
203
+ "nested_quant_map",
204
+ "bitsandbytes",
205
+ }
206
+ suffix = weight_name.split(".")[-1]
207
+ return any(q_suffix in suffix for q_suffix in quantized_suffix)
208
+
209
+ def _quantized_8bit_generator(self, hf_weights_files, use_safetensors,
210
+ quant_state_dict) -> Generator:
211
+ for (
212
+ org_weight_name,
213
+ mapped_weight_name,
214
+ weight_tensor,
215
+ ) in self._hf_weight_iter(hf_weights_files, use_safetensors):
216
+ if not mapped_weight_name.lower().endswith(".scb"):
217
+ continue
218
+
219
+ weight_key = mapped_weight_name.lower().replace(".scb", ".weight")
220
+ quant_state_dict[weight_key] = weight_tensor
221
+
222
+ for (
223
+ org_weight_name,
224
+ mapped_weight_name,
225
+ weight_tensor,
226
+ ) in self._hf_weight_iter(hf_weights_files, use_safetensors):
227
+ if self._is_8bit_weight_name(mapped_weight_name):
228
+ continue
229
+
230
+ if mapped_weight_name in quant_state_dict:
231
+ set_weight_attrs(weight_tensor, {"load_in_8bit": True})
232
+ yield org_weight_name, weight_tensor
233
+ else:
234
+ yield org_weight_name, weight_tensor
235
+
236
+ def _quantized_4bit_generator(self, hf_weights_files, use_safetensors,
237
+ quant_state_dict) -> Generator:
238
+ from bitsandbytes.functional import QuantState
239
+
240
+ # First iterate over all quant state weights
241
+ weight_iterator = self._hf_weight_iter(hf_weights_files,
242
+ use_safetensors)
243
+ temp_state_dict = {}
244
+ for (
245
+ org_weight_name,
246
+ mapped_weight_name,
247
+ weight_tensor,
248
+ ) in weight_iterator:
249
+ if not self._is_4bit_weight_name(mapped_weight_name):
250
+ continue
251
+ # bitsandbytes library requires
252
+ # weight.quant_state.bitsandbytes__* in CPU
253
+ if "quant_state.bitsandbytes" in mapped_weight_name:
254
+ temp_state_dict[mapped_weight_name] = weight_tensor.cpu().data
255
+ else:
256
+ temp_state_dict[mapped_weight_name] = weight_tensor
257
+
258
+ # Closure to parse quant_state for each prequant weight
259
+ def _parse_quant_state(param_name: str,
260
+ temp_state_dict: Dict) -> QuantState:
261
+ quant_state = {}
262
+ for k in temp_state_dict:
263
+ if param_name + "." in k:
264
+ quant_state[k] = temp_state_dict[k]
265
+
266
+ return QuantState.from_dict(quant_state,
267
+ device=current_platform.device_type)
268
+
269
+ # Second iterate over all prequant and normal weights
270
+ # pre quantized weights would have a quant_state
271
+ for (
272
+ org_weight_name,
273
+ mapped_weight_name,
274
+ weight_tensor,
275
+ ) in self._hf_weight_iter(hf_weights_files, use_safetensors):
276
+ if self._is_4bit_weight_name(mapped_weight_name):
277
+ continue
278
+
279
+ if (f"{mapped_weight_name}.quant_state.bitsandbytes__nf4"
280
+ in temp_state_dict) or (
281
+ f"{mapped_weight_name}.quant_state.bitsandbytes__fp4"
282
+ in temp_state_dict):
283
+ quant_state = _parse_quant_state(mapped_weight_name,
284
+ temp_state_dict)
285
+ quant_state_dict[mapped_weight_name] = quant_state
286
+ yield org_weight_name, weight_tensor
287
+ else:
288
+ yield org_weight_name, weight_tensor
289
+
290
+ def _unquantized_generator(self, hf_weights_files, use_safetensors,
291
+ quant_state_dict) -> Generator:
292
+ from bitsandbytes.functional import quantize_4bit
293
+
294
+ tp_size = get_tensor_model_parallel_world_size()
295
+ tp_rank = get_tensor_model_parallel_rank()
296
+
297
+ for (
298
+ org_weight_name,
299
+ mapped_weight_name,
300
+ weight_tensor,
301
+ ) in self._hf_weight_iter(hf_weights_files, use_safetensors):
302
+ if any(target_module in mapped_weight_name
303
+ for target_module in self.target_modules
304
+ ) and mapped_weight_name.endswith(".weight"):
305
+ # Without sharding
306
+ if any(
307
+ mapped_weight_name.startswith(module)
308
+ for module in self.unsharded_weights_modules):
309
+ weight_sub_tensor = weight_tensor
310
+ # Shard by column
311
+ elif any(
312
+ mapped_weight_name.startswith(module)
313
+ for module in self.column_sharded_weights_modules):
314
+ total_size = weight_tensor.size(-1)
315
+ start_index = total_size // tp_size * tp_rank
316
+ end_index = total_size // tp_size * (tp_rank + 1)
317
+ weight_sub_tensor = weight_tensor[...,
318
+ start_index:end_index]
319
+ # Weights have fused on disk. In this case, we assume that the
320
+ # weight and module use same name.
321
+ elif any(
322
+ mapped_weight_name.startswith(module)
323
+ for module in self.maybe_fused_weights_modules):
324
+ # special case for fused weights
325
+ # get the size of each shard weight tensor
326
+ total_shard_sizes = next(
327
+ (sizes for module, sizes in
328
+ self.maybe_fused_weights_modules.items()
329
+ if mapped_weight_name.startswith(module)))
330
+ total_size = weight_tensor.size(0)
331
+ assert total_size == sum(total_shard_sizes)
332
+ # get the start/end index of each shard weight tensor
333
+ total_start_index = list(
334
+ itertools.accumulate([0] + total_shard_sizes))[:-1]
335
+ shard_weights_index = [(
336
+ idx + size // tp_size * tp_rank,
337
+ idx + size // tp_size * (tp_rank + 1),
338
+ ) for idx, size in zip(total_start_index,
339
+ total_shard_sizes)]
340
+ # slice and reorder the weight tensor
341
+ weight_tensor = [
342
+ weight_tensor[start_index:end_index, ...]
343
+ for start_index, end_index in shard_weights_index
344
+ ]
345
+ weight_sub_tensor = torch.cat(weight_tensor, dim=0)
346
+ # Shard by row
347
+ else:
348
+ total_size = weight_tensor.size(0)
349
+ start_index = total_size // tp_size * tp_rank
350
+ end_index = total_size // tp_size * (tp_rank + 1)
351
+ weight_sub_tensor = weight_tensor[start_index:end_index,
352
+ ...]
353
+
354
+ # bitsandbytes requires data in GPU
355
+ if weight_sub_tensor.is_cuda:
356
+ loaded_weight = weight_sub_tensor
357
+ else:
358
+ loaded_weight = weight_sub_tensor.cuda()
359
+
360
+ # remove the following after the issue is fixed:
361
+ # https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1342
362
+ if loaded_weight.is_contiguous() is False:
363
+ loaded_weight = loaded_weight.contiguous()
364
+
365
+ with set_default_torch_dtype(torch.float32):
366
+ processed_weight, quant_state = quantize_4bit(
367
+ loaded_weight,
368
+ compress_statistics=True,
369
+ quant_type="nf4",
370
+ )
371
+
372
+ quant_state_dict[mapped_weight_name] = quant_state
373
+ else:
374
+ processed_weight = weight_tensor
375
+ yield org_weight_name, processed_weight
376
+
377
+ def _get_bnb_target_modules(self, model: nn.Module) -> None:
378
+
379
+ for name, module in model.named_modules():
380
+ if isinstance(module, (LinearBase, )):
381
+ if modules_info := self.modules_mapping.get_sub_modules(name):
382
+ # Map vllm's names to transformers's names.
383
+ rep_name, sub_modules = modules_info
384
+ for sub_name in sub_modules:
385
+ self.target_modules.append(
386
+ name.replace(rep_name, sub_name))
387
+ # Add original module name even if the module has stacked map,
388
+ # in case model has a mixture of disk-merged and disk-splitted
389
+ # weights with same last name.
390
+ self.target_modules.append(name)
391
+
392
+ assert (self.target_modules
393
+ ), "vllm currently does not support BNB quantization for"
394
+ f" {type(model).__name__}"
395
+
396
+ def _load_weights(self, model_config: ModelConfig,
397
+ model: nn.Module) -> None:
398
+ if not hasattr(model, "load_weights"):
399
+ raise AttributeError(
400
+ "The required method 'load_weights' is not defined in class"
401
+ f" {type(model).__name__}.")
402
+
403
+ if not hasattr(model, "packed_modules_mapping"):
404
+ raise AttributeError(
405
+ f"Model {type(model).__name__} does not support BitsAndBytes "
406
+ "quantization yet. No 'packed_modules_mapping' found.")
407
+
408
+ self.modules_mapping = ParamMapping(
409
+ copy.deepcopy(model.packed_modules_mapping))
410
+
411
+ # For some models like Molmo, we need to use hf_to_vllm_mapper
412
+ # to ensure correct loading of weights.
413
+ if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None):
414
+ self.weight_mapper = lambda name: hf_to_vllm_mapper._map_name(name)
415
+
416
+ # Modules whose weights might have fused on disk
417
+ # we need their output_sizes to make shard in flight correctly with TP
418
+ self.maybe_fused_weights_modules: Dict[str, List[int]] = {}
419
+ self._get_bnb_target_modules(model)
420
+ for name, module in model.named_modules():
421
+ # Some modules like `ReplicatedLinear` should not have their weights
422
+ # sharded. The reason for implementing it this way is to avoid new
423
+ # static variable in the model implementation.
424
+ if isinstance(module, (ReplicatedLinear, )):
425
+ self.unsharded_weights_modules.append(name)
426
+ # `QKVParallelLinear` and `MergedColumnParallelLinear` might have
427
+ # fused weights on disk. We need to use the output sizes of these
428
+ # modules to shard the weights correctly.
429
+ elif isinstance(module,
430
+ (QKVParallelLinear, MergedColumnParallelLinear)):
431
+ self.maybe_fused_weights_modules[name] = module.output_sizes
432
+ # In TP, these weights are partitioned along the column
433
+ # dimension (dim=-1)
434
+ elif isinstance(module, (RowParallelLinear, )):
435
+ self.column_sharded_weights_modules.append(name)
436
+
437
+ self.model_type = type(model).__name__
438
+
439
+ logger.info("Loading weights with BitsAndBytes quantization. "
440
+ "May take a while ...")
441
+
442
+ quant_config = getattr(model_config.hf_config, "quantization_config",
443
+ None)
444
+
445
+ pre_quant = False
446
+ if quant_config is not None:
447
+ quant_method = quant_config.get("quant_method")
448
+ if quant_method == "bitsandbytes":
449
+ pre_quant = True
450
+ else:
451
+ raise ValueError(
452
+ f"BitsAndBytes loader does not support {quant_method} "
453
+ "quantization")
454
+
455
+ # The quant_states in pre_quantized models cannot work with a split
456
+ # weight tensor. So TP does not work with pre_quantized bnb models.
457
+ if pre_quant and get_tensor_model_parallel_world_size() > 1:
458
+ raise ValueError(
459
+ "Prequant BitsAndBytes models with tensor parallelism is not "
460
+ "supported. Please try with pipeline parallelism.")
461
+
462
+ load_8bit = False
463
+ if pre_quant:
464
+ load_8bit = quant_config.get("load_in_8bit", False)
465
+
466
+ qweight_iterator, quant_state_dict = (
467
+ self._get_quantized_weights_iterator(model_config.model,
468
+ model_config.revision,
469
+ pre_quant, load_8bit))
470
+
471
+ weights_to_load = {name for name, _ in model.named_parameters()}
472
+ loaded_weights = model.load_weights(qweight_iterator)
473
+ # Some models may have weights loading tracker unimplemented.
474
+ if loaded_weights is not None:
475
+ weights_not_loaded = weights_to_load - loaded_weights
476
+ if weights_not_loaded:
477
+ raise ValueError("Following weights were not initialized from "
478
+ f"checkpoint: {weights_not_loaded}")
479
+
480
+ torch.cuda.empty_cache()
481
+
482
+ param_dict = dict(model.named_parameters())
483
+ stacked_quant_state_dict: Dict[str, Dict[int, Any]] = {}
484
+ # TODO: Change this lazy import to normal import
485
+ # after the checks are updated to run on a new version
486
+ from vllm.model_executor.models.utils import is_pp_missing_parameter
487
+
488
+ for quant_param_name in quant_state_dict:
489
+ if is_pp_missing_parameter(quant_param_name, model):
490
+ continue
491
+
492
+ non_stacked_param_name = quant_param_name
493
+
494
+ shard_index = 0
495
+ for shard_name, (
496
+ weight_name,
497
+ index,
498
+ ) in self.modules_mapping.inverse_packed_mapping.items():
499
+ # Some models, such as MiniCPM V2.5/2.6, contain both
500
+ # module names 'kv_proj' and 'qkv_proj'. To prevent 'kv_proj'
501
+ # from being incorrectly identified as being present in
502
+ # 'vpm.encoder.layers.0.self_attn.qkv_proj.weight
503
+ shard_pos = quant_param_name.find(shard_name)
504
+ can_correct_rename = (shard_pos
505
+ > 0) and (quant_param_name[shard_pos - 1]
506
+ == ".")
507
+ # If the quant_param_name is packed, it won't occur in the
508
+ # param_dict before renaming.
509
+ new_quant_param_name = quant_param_name.replace(
510
+ shard_name, weight_name)
511
+ need_rename = (quant_param_name not in param_dict) \
512
+ and (new_quant_param_name in param_dict)
513
+ if can_correct_rename and need_rename:
514
+ shard_index = index
515
+ quant_param_name = new_quant_param_name
516
+ break
517
+
518
+ # Models like Clip/Siglip may skip some layers in initialization,
519
+ # causing unused quant_param_name in state_dict.
520
+ if quant_param_name not in param_dict:
521
+ continue
522
+
523
+ if quant_param_name not in stacked_quant_state_dict:
524
+ stacked_quant_state_dict[quant_param_name] = {}
525
+
526
+ stacked_quant_state_dict[quant_param_name][shard_index] = (
527
+ quant_state_dict[non_stacked_param_name])
528
+
529
+ # save quant_states and offsets as the attributes of the parameters
530
+ for param_name, param in param_dict.items():
531
+ if param_name in stacked_quant_state_dict:
532
+ quant_states = stacked_quant_state_dict[param_name]
533
+ set_weight_attrs(param, {"bnb_quant_state": quant_states})
534
+
535
+ pack_ratio = getattr(param, "pack_factor", -1)
536
+ if pack_ratio == -1:
537
+ raise ValueError(
538
+ f"pack_factor not set for parameter {param_name}.")
539
+
540
+ num_elements = [0] * len(quant_states)
541
+ for seq, quant_state in quant_states.items():
542
+ num_elements[seq] = (math.prod(quant_state.shape) //
543
+ pack_ratio)
544
+
545
+ offsets = np.concatenate(([0], np.cumsum(num_elements)))
546
+ # Make torch infer_schema happy
547
+ offsets = torch.tensor(offsets).cpu()
548
+ set_weight_attrs(param, {"bnb_shard_offsets": offsets})
549
+
550
+ if load_8bit:
551
+ set_weight_attrs(
552
+ param, {"matmul_state": [None] * len(quant_states)})
553
+
554
+ def download_model(self, model_config: ModelConfig) -> None:
555
+ self._prepare_weights(model_config.model, model_config.revision)
556
+
557
+ def load_model(self, vllm_config: VllmConfig) -> nn.Module:
558
+ device_config = vllm_config.device_config
559
+ model_config = vllm_config.model_config
560
+
561
+ with set_default_torch_dtype(model_config.dtype):
562
+ with torch.device(device_config.device):
563
+
564
+ model = initialize_model(vllm_config=vllm_config)
565
+
566
+ self._load_weights(model_config, model)
567
+
568
+ return model.eval()
vllm/model_executor/model_loader/default_loader.py
new
+293
@@ -0,0 +1,293 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import dataclasses
3
+import glob
4
+import os
5
+import time
6
+from typing import Generator, Iterable, List, Optional, Tuple, cast
7
+
8
+import huggingface_hub
9
+import torch
10
+from torch import nn
11
+from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
12
+
13
+from vllm.config import LoadConfig, LoadFormat, ModelConfig, VllmConfig
14
+from vllm.envs import VLLM_USE_MODELSCOPE
15
+from vllm.logger import init_logger
16
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
17
+from vllm.model_executor.model_loader.utils import (
18
+ initialize_model, process_weights_after_loading, set_default_torch_dtype)
19
+from vllm.model_executor.model_loader.weight_utils import (
20
+ download_safetensors_index_file_from_hf, download_weights_from_hf,
21
+ fastsafetensors_weights_iterator, filter_duplicate_safetensors_files,
22
+ filter_files_not_needed_for_inference, get_lock, np_cache_weights_iterator,
23
+ pt_weights_iterator, safetensors_weights_iterator)
24
+from vllm.platforms import current_platform
25
+
26
+logger = init_logger(__name__)
27
+
28
+
29
+class DefaultModelLoader(BaseModelLoader):
30
+ """Model loader that can load different file types from disk."""
31
+
32
+ @dataclasses.dataclass
33
+ class Source:
34
+ """A source for weights."""
35
+
36
+ model_or_path: str
37
+ """The model ID or path."""
38
+
39
+ revision: Optional[str]
40
+ """The optional model revision."""
41
+
42
+ prefix: str = ""
43
+ """A prefix to prepend to all weights."""
44
+
45
+ fall_back_to_pt: bool = True
46
+ """Whether .pt weights can be used."""
47
+
48
+ allow_patterns_overrides: Optional[list[str]] = None
49
+ """If defined, weights will load exclusively using these patterns."""
50
+
51
+ counter_before_loading_weights: float = 0.0
52
+ counter_after_loading_weights: float = 0.0
53
+
54
+ def __init__(self, load_config: LoadConfig):
55
+ super().__init__(load_config)
56
+ if load_config.model_loader_extra_config:
57
+ raise ValueError(f"Model loader extra config is not supported for "
58
+ f"load format {load_config.load_format}")
59
+
60
+ def _maybe_download_from_modelscope(
61
+ self, model: str, revision: Optional[str]) -> Optional[str]:
62
+ """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True.
63
+
64
+ Returns the path to the downloaded model, or None if the model is not
65
+ downloaded from ModelScope."""
66
+ if VLLM_USE_MODELSCOPE:
67
+ # download model from ModelScope hub,
68
+ # lazy import so that modelscope is not required for normal use.
69
+ # pylint: disable=C.
70
+ from modelscope.hub.snapshot_download import snapshot_download
71
+
72
+ if not os.path.exists(model):
73
+ # Use file lock to prevent multiple processes from
74
+ # downloading the same model weights at the same time.
75
+ with get_lock(model, self.load_config.download_dir):
76
+ model_path = snapshot_download(
77
+ model_id=model,
78
+ cache_dir=self.load_config.download_dir,
79
+ local_files_only=huggingface_hub.constants.
80
+ HF_HUB_OFFLINE,
81
+ revision=revision,
82
+ ignore_file_pattern=self.load_config.ignore_patterns,
83
+ )
84
+ else:
85
+ model_path = model
86
+ return model_path
87
+ return None
88
+
89
+ def _prepare_weights(
90
+ self,
91
+ model_name_or_path: str,
92
+ revision: Optional[str],
93
+ fall_back_to_pt: bool,
94
+ allow_patterns_overrides: Optional[list[str]],
95
+ ) -> Tuple[str, List[str], bool]:
96
+ """Prepare weights for the model.
97
+
98
+ If the model is not local, it will be downloaded."""
99
+ model_name_or_path = (self._maybe_download_from_modelscope(
100
+ model_name_or_path, revision) or model_name_or_path)
101
+
102
+ is_local = os.path.isdir(model_name_or_path)
103
+ load_format = self.load_config.load_format
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:
108
+ allow_patterns = ["*.safetensors", "*.bin"]
109
+ elif (load_format == LoadFormat.SAFETENSORS
110
+ or load_format == LoadFormat.FASTSAFETENSORS):
111
+ use_safetensors = True
112
+ allow_patterns = ["*.safetensors"]
113
+ elif load_format == LoadFormat.MISTRAL:
114
+ use_safetensors = True
115
+ allow_patterns = ["consolidated*.safetensors"]
116
+ index_file = "consolidated.safetensors.index.json"
117
+ elif load_format == LoadFormat.PT:
118
+ allow_patterns = ["*.pt"]
119
+ elif load_format == LoadFormat.NPCACHE:
120
+ allow_patterns = ["*.bin"]
121
+ else:
122
+ raise ValueError(f"Unknown load_format: {load_format}")
123
+
124
+ if fall_back_to_pt:
125
+ allow_patterns += ["*.pt"]
126
+
127
+ if allow_patterns_overrides is not None:
128
+ allow_patterns = allow_patterns_overrides
129
+
130
+ if not is_local:
131
+ hf_folder = download_weights_from_hf(
132
+ model_name_or_path,
133
+ self.load_config.download_dir,
134
+ allow_patterns,
135
+ revision,
136
+ ignore_patterns=self.load_config.ignore_patterns,
137
+ )
138
+ else:
139
+ hf_folder = model_name_or_path
140
+
141
+ hf_weights_files: List[str] = []
142
+ for pattern in allow_patterns:
143
+ hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
144
+ if len(hf_weights_files) > 0:
145
+ if pattern == "*.safetensors":
146
+ use_safetensors = True
147
+ break
148
+
149
+ if use_safetensors:
150
+ # For models like Mistral-7B-Instruct-v0.3
151
+ # there are both sharded safetensors files and a consolidated
152
+ # safetensors file. Using both breaks.
153
+ # Here, we download the `model.safetensors.index.json` and filter
154
+ # any files not found in the index.
155
+ if not is_local:
156
+ download_safetensors_index_file_from_hf(
157
+ model_name_or_path,
158
+ index_file,
159
+ self.load_config.download_dir,
160
+ revision,
161
+ )
162
+ hf_weights_files = filter_duplicate_safetensors_files(
163
+ hf_weights_files, hf_folder, index_file)
164
+ else:
165
+ hf_weights_files = filter_files_not_needed_for_inference(
166
+ hf_weights_files)
167
+
168
+ if len(hf_weights_files) == 0:
169
+ raise RuntimeError(
170
+ f"Cannot find any model weights with `{model_name_or_path}`")
171
+
172
+ return hf_folder, hf_weights_files, use_safetensors
173
+
174
+ def _get_weights_iterator(
175
+ self, source: "Source"
176
+ ) -> Generator[Tuple[str, torch.Tensor], None, None]:
177
+ """Get an iterator for the model weights based on the load format."""
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:
182
+ # Currently np_cache only support *.bin checkpoints
183
+ assert use_safetensors is False
184
+ weights_iterator = np_cache_weights_iterator(
185
+ source.model_or_path,
186
+ self.load_config.download_dir,
187
+ hf_folder,
188
+ hf_weights_files,
189
+ self.load_config.use_tqdm_on_load,
190
+ )
191
+ elif use_safetensors:
192
+ if self.load_config.load_format == LoadFormat.FASTSAFETENSORS:
193
+ weights_iterator = fastsafetensors_weights_iterator(
194
+ hf_weights_files,
195
+ self.load_config.use_tqdm_on_load,
196
+ )
197
+ else:
198
+ weights_iterator = safetensors_weights_iterator(
199
+ hf_weights_files,
200
+ self.load_config.use_tqdm_on_load,
201
+ )
202
+ else:
203
+ weights_iterator = pt_weights_iterator(
204
+ hf_weights_files,
205
+ self.load_config.use_tqdm_on_load,
206
+ self.load_config.pt_load_map_location,
207
+ )
208
+
209
+ if current_platform.is_tpu():
210
+ # In PyTorch XLA, we should call `xm.mark_step` frequently so that
211
+ # not too many ops are accumulated in the XLA program.
212
+ import torch_xla.core.xla_model as xm
213
+
214
+ def _xla_weights_iterator(iterator: Generator):
215
+ for weights in iterator:
216
+ yield weights
217
+ xm.mark_step()
218
+
219
+ weights_iterator = _xla_weights_iterator(weights_iterator)
220
+
221
+ elif current_platform.is_hpu():
222
+ import habana_frameworks.torch.core as htcore
223
+
224
+ def _hpu_weights_iterator(iterator: Generator):
225
+ for weights in iterator:
226
+ yield weights
227
+ htcore.mark_step()
228
+
229
+ weights_iterator = _hpu_weights_iterator(weights_iterator)
230
+
231
+ if self.counter_before_loading_weights == 0.0:
232
+ self.counter_before_loading_weights = time.perf_counter()
233
+ # Apply the prefix.
234
+ return ((source.prefix + name, tensor)
235
+ for (name, tensor) in weights_iterator)
236
+
237
+ def get_all_weights(
238
+ self,
239
+ model_config: ModelConfig,
240
+ model: nn.Module,
241
+ ) -> Generator[Tuple[str, torch.Tensor], None, None]:
242
+ primary_weights = DefaultModelLoader.Source(
243
+ model_config.model,
244
+ model_config.revision,
245
+ prefix="",
246
+ fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load",
247
+ True),
248
+ allow_patterns_overrides=getattr(model, "allow_patterns_overrides",
249
+ None),
250
+ )
251
+ yield from self._get_weights_iterator(primary_weights)
252
+
253
+ secondary_weights = cast(
254
+ Iterable[DefaultModelLoader.Source],
255
+ getattr(model, "secondary_weights", ()),
256
+ )
257
+ for source in secondary_weights:
258
+ yield from self._get_weights_iterator(source)
259
+
260
+ def download_model(self, model_config: ModelConfig) -> None:
261
+ self._prepare_weights(model_config.model,
262
+ model_config.revision,
263
+ fall_back_to_pt=True,
264
+ allow_patterns_overrides=None)
265
+
266
+ def load_model(self, vllm_config: VllmConfig) -> nn.Module:
267
+ device_config = vllm_config.device_config
268
+ model_config = vllm_config.model_config
269
+ target_device = torch.device(device_config.device)
270
+ with set_default_torch_dtype(model_config.dtype):
271
+ with target_device:
272
+ model = initialize_model(vllm_config=vllm_config)
273
+
274
+ weights_to_load = {name for name, _ in model.named_parameters()}
275
+ loaded_weights = model.load_weights(
276
+ self.get_all_weights(model_config, model))
277
+ self.counter_after_loading_weights = time.perf_counter()
278
+ logger.info(
279
+ "Loading weights took %.2f seconds",
280
+ self.counter_after_loading_weights -
281
+ self.counter_before_loading_weights)
282
+ # We only enable strict check for non-quantized models
283
+ # that have loaded weights tracking currently.
284
+ if model_config.quantization is None and loaded_weights is not None:
285
+ weights_not_loaded = weights_to_load - loaded_weights
286
+ if weights_not_loaded:
287
+ raise ValueError(
288
+ "Following weights were not initialized from "
289
+ f"checkpoint: {weights_not_loaded}")
290
+
291
+ process_weights_after_loading(model, model_config, target_device)
292
+
293
+ return model.eval()
vllm/model_executor/model_loader/dummy_loader.py
new
+37
@@ -0,0 +1,37 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import torch
3
+import torch.nn as nn
4
+
5
+from vllm.config import LoadConfig, ModelConfig, VllmConfig
6
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
7
+from vllm.model_executor.model_loader.utils import (
8
+ initialize_model, process_weights_after_loading, set_default_torch_dtype)
9
+from vllm.model_executor.model_loader.weight_utils import (
10
+ initialize_dummy_weights)
11
+
12
+
13
+class DummyModelLoader(BaseModelLoader):
14
+ """Model loader that will set model weights to random values."""
15
+
16
+ def __init__(self, load_config: LoadConfig):
17
+ super().__init__(load_config)
18
+ if load_config.model_loader_extra_config:
19
+ raise ValueError(f"Model loader extra config is not supported for "
20
+ f"load format {load_config.load_format}")
21
+
22
+ def download_model(self, model_config: ModelConfig) -> None:
23
+ pass # Nothing to download
24
+
25
+ def load_model(self, vllm_config: VllmConfig) -> nn.Module:
26
+ device_config = vllm_config.device_config
27
+ model_config = vllm_config.model_config
28
+ target_device = torch.device(device_config.device)
29
+ with set_default_torch_dtype(model_config.dtype):
30
+ with target_device:
31
+ model = initialize_model(vllm_config=vllm_config)
32
+ # NOTE(woosuk): For accurate performance evaluation, we assign
33
+ # random values to the weights.
34
+ initialize_dummy_weights(model)
35
+
36
+ process_weights_after_loading(model, model_config, target_device)
37
+ return model.eval()
vllm/model_executor/model_loader/gguf_loader.py
new
+113
@@ -0,0 +1,113 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import os
3
+from typing import Dict, Generator, Tuple
4
+
5
+import gguf
6
+import torch
7
+import torch.nn as nn
8
+from transformers import AutoModelForCausalLM
9
+
10
+from vllm.config import LoadConfig, ModelConfig, VllmConfig
11
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
12
+from vllm.model_executor.model_loader.utils import (
13
+ initialize_model, process_weights_after_loading, set_default_torch_dtype)
14
+from vllm.model_executor.model_loader.weight_utils import (
15
+ get_gguf_extra_tensor_names, gguf_quant_weights_iterator)
16
+
17
+
18
+class GGUFModelLoader(BaseModelLoader):
19
+ """
20
+ Model loader that can load GGUF files. This is useful for loading models
21
+ that are quantized with GGUF and saved in the GGUF format. This loader
22
+ supports loading both full models and sharded models.
23
+ """
24
+
25
+ def __init__(self, load_config: LoadConfig):
26
+ super().__init__(load_config)
27
+ if load_config.model_loader_extra_config:
28
+ raise ValueError(f"Model loader extra config is not supported for "
29
+ f"load format {load_config.load_format}")
30
+
31
+ def _prepare_weights(self, model_name_or_path: str):
32
+ if os.path.isfile(model_name_or_path):
33
+ return model_name_or_path
34
+ else:
35
+ raise ValueError(f"{model_name_or_path} is not a file.")
36
+
37
+ def _get_gguf_weights_map(self, model_config: ModelConfig):
38
+ """
39
+ GGUF uses this naming convention for their tensors from HF checkpoint:
40
+ `blk.N.BB.weight` and `blk.N.BB.bias`
41
+ where N signifies the block number of a layer, and BB signifies the
42
+ attention/mlp layer components.
43
+ See "Standardized tensor names" in
44
+ https://github.com/ggerganov/ggml/blob/master/docs/gguf.md for details.
45
+ """
46
+ config = model_config.hf_config
47
+ model_type = config.model_type
48
+ gguf_to_hf_name_map = {}
49
+ # hack: ggufs have a different name than transformers
50
+ if model_type == "cohere":
51
+ model_type = "command-r"
52
+ if model_type in ("deepseek_v3", "deepseek_v2"):
53
+ model_type = "deepseek2"
54
+ # GGUF layer map assumes that we will have a merged expert weights
55
+ # so we need to map them manually
56
+ for idx in range(config.num_hidden_layers):
57
+ gguf_to_hf_name_map[f"blk.{idx}.exp_probs_b.bias"] = \
58
+ f"model.layers.{idx}.mlp.gate.e_score_correction_bias"
59
+ gguf_to_hf_name_map[f"blk.{idx}.ffn_down_exps.weight"] = \
60
+ f"model.layers.{idx}.mlp.experts.0.down_proj.weight"
61
+ gguf_to_hf_name_map[f"blk.{idx}.ffn_gate_exps.weight"] = \
62
+ f"model.layers.{idx}.mlp.experts.0.gate_proj.weight"
63
+ gguf_to_hf_name_map[f"blk.{idx}.ffn_up_exps.weight"] = \
64
+ f"model.layers.{idx}.mlp.experts.0.up_proj.weight"
65
+
66
+ arch = None
67
+ for key, value in gguf.MODEL_ARCH_NAMES.items():
68
+ if value == model_type:
69
+ arch = key
70
+ break
71
+ if arch is None:
72
+ raise RuntimeError(f"Unknown gguf model_type: {model_type}")
73
+ num_layers = config.num_hidden_layers
74
+ name_map = gguf.get_tensor_name_map(arch, num_layers)
75
+ with torch.device("meta"):
76
+ dummy_model = AutoModelForCausalLM.from_config(
77
+ config, trust_remote_code=model_config.trust_remote_code)
78
+ state_dict = dummy_model.state_dict()
79
+
80
+ for hf_name in state_dict:
81
+ name, suffix = hf_name.rsplit(".", 1)
82
+ gguf_name = name_map.get_name(name)
83
+ gguf_to_hf_name_map[f"{gguf_name}.{suffix}"] = hf_name
84
+ return gguf_to_hf_name_map
85
+
86
+ def _get_weights_iterator(
87
+ self, model_name_or_path: str, gguf_to_hf_name_map: Dict[str, str]
88
+ ) -> Generator[Tuple[str, torch.Tensor], None, None]:
89
+ return gguf_quant_weights_iterator(model_name_or_path,
90
+ gguf_to_hf_name_map)
91
+
92
+ def download_model(self, model_config: ModelConfig) -> None:
93
+ self._prepare_weights(model_config.model)
94
+
95
+ def load_model(self, vllm_config: VllmConfig) -> nn.Module:
96
+ device_config = vllm_config.device_config
97
+ model_config = vllm_config.model_config
98
+ local_model_path = self._prepare_weights(model_config.model)
99
+ gguf_weights_map = self._get_gguf_weights_map(model_config)
100
+ # we can only know if tie word embeddings after mapping weights
101
+ if "lm_head.weight" in get_gguf_extra_tensor_names(
102
+ local_model_path, gguf_weights_map):
103
+ model_config.hf_config.update({"tie_word_embeddings": True})
104
+
105
+ target_device = torch.device(device_config.device)
106
+ with set_default_torch_dtype(model_config.dtype):
107
+ with target_device:
108
+ model = initialize_model(vllm_config=vllm_config)
109
+ model.load_weights(
110
+ self._get_weights_iterator(local_model_path, gguf_weights_map))
111
+
112
+ process_weights_after_loading(model, model_config, target_device)
113
+ return model
vllm/model_executor/model_loader/loader.py
deleted
-1544
@@ -1,1544 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-
3
-# ruff: noqa: SIM117
4
-import collections
5
-import copy
6
-import dataclasses
7
-import fnmatch
8
-import glob
9
-import inspect
10
-import itertools
11
-import math
12
-import os
13
-import time
14
-import warnings
15
-from abc import ABC, abstractmethod
16
-from contextlib import contextmanager
17
-from typing import (Any, Callable, Dict, Generator, Iterable, List, Optional,
18
- Tuple, cast)
19
-
20
-import gguf
21
-import huggingface_hub
22
-import numpy as np
23
-import torch
24
-from huggingface_hub import HfApi
25
-from torch import nn
26
-from transformers import AutoModelForCausalLM
27
-from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
28
-
29
-from vllm.attention import Attention
30
-from vllm.config import (LoadConfig, LoadFormat, ModelConfig, ParallelConfig,
31
- VllmConfig, set_current_vllm_config)
32
-from vllm.distributed import (get_tensor_model_parallel_rank,
33
- get_tensor_model_parallel_world_size)
34
-from vllm.envs import VLLM_USE_MODELSCOPE
35
-from vllm.logger import init_logger
36
-# yapf conflicts with isort for this block
37
-# yapf: disable
38
-from vllm.model_executor.layers.linear import (LinearBase,
39
- MergedColumnParallelLinear,
40
- QKVCrossParallelLinear,
41
- QKVParallelLinear,
42
- ReplicatedLinear,
43
- RowParallelLinear)
44
-# yapf: enable
45
-from vllm.model_executor.layers.quantization.base_config import (
46
- QuantizeMethodBase)
47
-from vllm.model_executor.model_loader.tensorizer import (
48
- TensorizerConfig, is_vllm_tensorized, load_with_tensorizer,
49
- serialize_vllm_model, tensorizer_weights_iterator)
50
-from vllm.model_executor.model_loader.utils import (ParamMapping,
51
- configure_quant_config,
52
- get_model_architecture,
53
- set_default_torch_dtype)
54
-from vllm.model_executor.model_loader.weight_utils import (
55
- download_safetensors_index_file_from_hf, download_weights_from_hf,
56
- fastsafetensors_weights_iterator, filter_duplicate_safetensors_files,
57
- filter_files_not_needed_for_inference, get_gguf_extra_tensor_names,
58
- get_lock, gguf_quant_weights_iterator, initialize_dummy_weights,
59
- np_cache_weights_iterator, pt_weights_iterator,
60
- runai_safetensors_weights_iterator, safetensors_weights_iterator)
61
-from vllm.model_executor.utils import set_weight_attrs
62
-from vllm.platforms import current_platform
63
-from vllm.transformers_utils.s3_utils import glob as s3_glob
64
-from vllm.transformers_utils.utils import is_s3
65
-from vllm.utils import is_pin_memory_available
66
-
67
-
68
-@contextmanager
69
-def device_loading_context(module: torch.nn.Module,
70
- target_device: torch.device):
71
- if target_device.type == "cpu":
72
- # If target is CPU, no need to move anything
73
- yield module
74
- return
75
-
76
- original_device_states: Dict[str, torch.device] = {}
77
-
78
- # Store original device states and move parameters to GPU if they're on CPU
79
- for name, p in module.named_parameters():
80
- if p.device.type == "cpu":
81
- original_device_states[name] = p.device
82
- p.data = p.data.to(target_device)
83
- # Parameters already on target device are not touched
84
-
85
- try:
86
- yield module
87
-
88
- finally:
89
- # Restore parameters to their original devices, ignoring new parameters
90
- pin_memory = is_pin_memory_available()
91
- for name, p in module.named_parameters():
92
- if name in original_device_states:
93
- original_device: torch.device = original_device_states[name]
94
- if original_device.type == "cpu":
95
- # `torch.empty_like` does not support `pin_memory` argument
96
- cpu_data = torch.empty_strided(
97
- size=p.data.size(),
98
- stride=p.data.stride(),
99
- dtype=p.data.dtype,
100
- layout=p.data.layout,
101
- device="cpu",
102
- pin_memory=pin_memory,
103
- )
104
- cpu_data.copy_(p.data)
105
- p.data = cpu_data
106
- else:
107
- p.data = p.data.to(original_device)
108
- # New parameters or parameters already on target device are untouched
109
-
110
-
111
-logger = init_logger(__name__)
112
-
113
-
114
-def _initialize_model(
115
- vllm_config: VllmConfig,
116
- *,
117
- prefix: str = "",
118
- model_class: Optional[type[nn.Module]] = None,
119
-) -> nn.Module:
120
- """Initialize a model with the given configurations."""
121
- model_config = vllm_config.model_config
122
- if model_class is None:
123
- model_class, _ = get_model_architecture(model_config)
124
-
125
- if vllm_config.quant_config is not None:
126
- configure_quant_config(vllm_config.quant_config, model_class)
127
-
128
- signatures = inspect.signature(model_class.__init__)
129
- all_params = [param.name for param in signatures.parameters.values()]
130
- if "vllm_config" in all_params and "prefix" in all_params:
131
- # new-style model class
132
- with set_current_vllm_config(vllm_config, check_compile=True):
133
- return model_class(vllm_config=vllm_config, prefix=prefix)
134
-
135
- msg = ("vLLM model class should accept `vllm_config` and `prefix` as "
136
- "input arguments. Possibly you have an old-style model class"
137
- " registered from out of tree and it is used for new vLLM version. "
138
- "Check https://docs.vllm.ai/en/latest/design/arch_overview.html "
139
- "for the design and update the model class accordingly.")
140
- warnings.warn(msg, DeprecationWarning, stacklevel=2)
141
-
142
- logger.warning(
143
- "Trying to guess the arguments for old-style model class %s",
144
- model_class,
145
- )
146
- # try to be compatible with old-style model class
147
- kwargs = {}
148
- if "prefix" in all_params:
149
- kwargs["prefix"] = prefix
150
- if "config" in all_params:
151
- kwargs["config"] = model_config.hf_config
152
- if "cache_config" in all_params:
153
- kwargs["cache_config"] = vllm_config.cache_config
154
- if "quant_config" in all_params:
155
- kwargs["quant_config"] = vllm_config.quant_config
156
- if "lora_config" in all_params:
157
- kwargs["lora_config"] = vllm_config.lora_config
158
- if "scheduler_config" in all_params:
159
- kwargs["scheduler_config"] = vllm_config.scheduler_config
160
- with set_current_vllm_config(vllm_config, check_compile=True):
161
- return model_class(**kwargs)
162
-
163
-
164
-def _process_weights_after_loading(model: nn.Module, model_config: ModelConfig,
165
- target_device: torch.device) -> None:
166
- for _, module in model.named_modules():
167
- if isinstance(module, QKVCrossParallelLinear):
168
- # NOTE(Isotr0py): special case for cross QKV layer because
169
- # q and kv proj aren't registered as submodules intentionally
170
- module.process_weights_after_loading()
171
- continue
172
- quant_method = getattr(module, "quant_method", None)
173
- if isinstance(quant_method, QuantizeMethodBase):
174
- # When quant methods need to process weights after loading
175
- # (for repacking, quantizing, etc), they expect parameters
176
- # to be on the global target device. This scope is for the
177
- # case where cpu offloading is used, where we will move the
178
- # parameters onto device for processing and back off after.
179
- with device_loading_context(module, target_device):
180
- quant_method.process_weights_after_loading(module)
181
-
182
- # Currently only used by MLA.
183
- # NOTE: This intentionally happens after other modules so we can easily
184
- # decompress the weights for MLA.
185
- for _, module in model.named_modules():
186
- if isinstance(module, Attention) and \
187
- hasattr(module, "process_weights_after_loading"):
188
- # TODO(lucas): see if there is a way to unify the signatures
189
- # of process_weights_after_loading
190
- module.process_weights_after_loading(model_config.dtype)
191
-
192
-
193
-class BaseModelLoader(ABC):
194
- """Base class for model loaders."""
195
-
196
- def __init__(self, load_config: LoadConfig):
197
- self.load_config = load_config
198
-
199
- @abstractmethod
200
- def download_model(self, model_config: ModelConfig) -> None:
201
- """Download a model so that it can be immediately loaded."""
202
- raise NotImplementedError
203
-
204
- @abstractmethod
205
- def load_model(self, *, vllm_config: VllmConfig) -> nn.Module:
206
- """Load a model with the given configurations."""
207
- raise NotImplementedError
208
-
209
-
210
-class DefaultModelLoader(BaseModelLoader):
211
- """Model loader that can load different file types from disk."""
212
-
213
- @dataclasses.dataclass
214
- class Source:
215
- """A source for weights."""
216
-
217
- model_or_path: str
218
- """The model ID or path."""
219
-
220
- revision: Optional[str]
221
- """The optional model revision."""
222
-
223
- prefix: str = ""
224
- """A prefix to prepend to all weights."""
225
-
226
- fall_back_to_pt: bool = True
227
- """Whether .pt weights can be used."""
228
-
229
- allow_patterns_overrides: Optional[list[str]] = None
230
- """If defined, weights will load exclusively using these patterns."""
231
-
232
- counter_before_loading_weights: float = 0.0
233
- counter_after_loading_weights: float = 0.0
234
-
235
- def __init__(self, load_config: LoadConfig):
236
- super().__init__(load_config)
237
- if load_config.model_loader_extra_config:
238
- raise ValueError(f"Model loader extra config is not supported for "
239
- f"load format {load_config.load_format}")
240
-
241
- def _maybe_download_from_modelscope(
242
- self, model: str, revision: Optional[str]) -> Optional[str]:
243
- """Download model from ModelScope hub if VLLM_USE_MODELSCOPE is True.
244
-
245
- Returns the path to the downloaded model, or None if the model is not
246
- downloaded from ModelScope."""
247
- if VLLM_USE_MODELSCOPE:
248
- # download model from ModelScope hub,
249
- # lazy import so that modelscope is not required for normal use.
250
- # pylint: disable=C.
251
- from modelscope.hub.snapshot_download import snapshot_download
252
-
253
- if not os.path.exists(model):
254
- # Use file lock to prevent multiple processes from
255
- # downloading the same model weights at the same time.
256
- with get_lock(model, self.load_config.download_dir):
257
- model_path = snapshot_download(
258
- model_id=model,
259
- cache_dir=self.load_config.download_dir,
260
- local_files_only=huggingface_hub.constants.
261
- HF_HUB_OFFLINE,
262
- revision=revision,
263
- ignore_file_pattern=self.load_config.ignore_patterns,
264
- )
265
- else:
266
- model_path = model
267
- return model_path
268
- return None
269
-
270
- def _prepare_weights(
271
- self,
272
- model_name_or_path: str,
273
- revision: Optional[str],
274
- fall_back_to_pt: bool,
275
- allow_patterns_overrides: Optional[list[str]],
276
- ) -> Tuple[str, List[str], bool]:
277
- """Prepare weights for the model.
278
-
279
- If the model is not local, it will be downloaded."""
280
- model_name_or_path = (self._maybe_download_from_modelscope(
281
- model_name_or_path, revision) or model_name_or_path)
282
-
283
- is_local = os.path.isdir(model_name_or_path)
284
- load_format = self.load_config.load_format
285
- use_safetensors = False
286
- index_file = SAFE_WEIGHTS_INDEX_NAME
287
- # Some quantized models use .pt files for storing the weights.
288
- if load_format == LoadFormat.AUTO:
289
- allow_patterns = ["*.safetensors", "*.bin"]
290
- elif (load_format == LoadFormat.SAFETENSORS
291
- or load_format == LoadFormat.FASTSAFETENSORS):
292
- use_safetensors = True
293
- allow_patterns = ["*.safetensors"]
294
- elif load_format == LoadFormat.MISTRAL:
295
- use_safetensors = True
296
- allow_patterns = ["consolidated*.safetensors"]
297
- index_file = "consolidated.safetensors.index.json"
298
- elif load_format == LoadFormat.PT:
299
- allow_patterns = ["*.pt"]
300
- elif load_format == LoadFormat.NPCACHE:
301
- allow_patterns = ["*.bin"]
302
- else:
303
- raise ValueError(f"Unknown load_format: {load_format}")
304
-
305
- if fall_back_to_pt:
306
- allow_patterns += ["*.pt"]
307
-
308
- if allow_patterns_overrides is not None:
309
- allow_patterns = allow_patterns_overrides
310
-
311
- if not is_local:
312
- hf_folder = download_weights_from_hf(
313
- model_name_or_path,
314
- self.load_config.download_dir,
315
- allow_patterns,
316
- revision,
317
- ignore_patterns=self.load_config.ignore_patterns,
318
- )
319
- else:
320
- hf_folder = model_name_or_path
321
-
322
- hf_weights_files: List[str] = []
323
- for pattern in allow_patterns:
324
- hf_weights_files += glob.glob(os.path.join(hf_folder, pattern))
325
- if len(hf_weights_files) > 0:
326
- if pattern == "*.safetensors":
327
- use_safetensors = True
328
- break
329
-
330
- if use_safetensors:
331
- # For models like Mistral-7B-Instruct-v0.3
332
- # there are both sharded safetensors files and a consolidated
333
- # safetensors file. Using both breaks.
334
- # Here, we download the `model.safetensors.index.json` and filter
335
- # any files not found in the index.
336
- if not is_local:
337
- download_safetensors_index_file_from_hf(
338
- model_name_or_path,
339
- index_file,
340
- self.load_config.download_dir,
341
- revision,
342
- )
343
- hf_weights_files = filter_duplicate_safetensors_files(
344
- hf_weights_files, hf_folder, index_file)
345
- else:
346
- hf_weights_files = filter_files_not_needed_for_inference(
347
- hf_weights_files)
348
-
349
- if len(hf_weights_files) == 0:
350
- raise RuntimeError(
351
- f"Cannot find any model weights with `{model_name_or_path}`")
352
-
353
- return hf_folder, hf_weights_files, use_safetensors
354
-
355
- def _get_weights_iterator(
356
- self, source: "Source"
357
- ) -> Generator[Tuple[str, torch.Tensor], None, None]:
358
- """Get an iterator for the model weights based on the load format."""
359
- hf_folder, hf_weights_files, use_safetensors = self._prepare_weights(
360
- source.model_or_path, source.revision, source.fall_back_to_pt,
361
- source.allow_patterns_overrides)
362
- if self.load_config.load_format == LoadFormat.NPCACHE:
363
- # Currently np_cache only support *.bin checkpoints
364
- assert use_safetensors is False
365
- weights_iterator = np_cache_weights_iterator(
366
- source.model_or_path,
367
- self.load_config.download_dir,
368
- hf_folder,
369
- hf_weights_files,
370
- self.load_config.use_tqdm_on_load,
371
- )
372
- elif use_safetensors:
373
- if self.load_config.load_format == LoadFormat.FASTSAFETENSORS:
374
- weights_iterator = fastsafetensors_weights_iterator(
375
- hf_weights_files,
376
- self.load_config.use_tqdm_on_load,
377
- )
378
- else:
379
- weights_iterator = safetensors_weights_iterator(
380
- hf_weights_files,
381
- self.load_config.use_tqdm_on_load,
382
- )
383
- else:
384
- weights_iterator = pt_weights_iterator(
385
- hf_weights_files,
386
- self.load_config.use_tqdm_on_load,
387
- self.load_config.pt_load_map_location,
388
- )
389
-
390
- if current_platform.is_tpu():
391
- # In PyTorch XLA, we should call `xm.mark_step` frequently so that
392
- # not too many ops are accumulated in the XLA program.
393
- import torch_xla.core.xla_model as xm
394
-
395
- def _xla_weights_iterator(iterator: Generator):
396
- for weights in iterator:
397
- yield weights
398
- xm.mark_step()
399
-
400
- weights_iterator = _xla_weights_iterator(weights_iterator)
401
-
402
- elif current_platform.is_hpu():
403
- import habana_frameworks.torch.core as htcore
404
-
405
- def _hpu_weights_iterator(iterator: Generator):
406
- for weights in iterator:
407
- yield weights
408
- htcore.mark_step()
409
-
410
- weights_iterator = _hpu_weights_iterator(weights_iterator)
411
-
412
- if self.counter_before_loading_weights == 0.0:
413
- self.counter_before_loading_weights = time.perf_counter()
414
- # Apply the prefix.
415
- return ((source.prefix + name, tensor)
416
- for (name, tensor) in weights_iterator)
417
-
418
- def get_all_weights(
419
- self,
420
- model_config: ModelConfig,
421
- model: nn.Module,
422
- ) -> Generator[Tuple[str, torch.Tensor], None, None]:
423
- primary_weights = DefaultModelLoader.Source(
424
- model_config.model,
425
- model_config.revision,
426
- prefix="",
427
- fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load",
428
- True),
429
- allow_patterns_overrides=getattr(model, "allow_patterns_overrides",
430
- None),
431
- )
432
- yield from self._get_weights_iterator(primary_weights)
433
-
434
- secondary_weights = cast(
435
- Iterable[DefaultModelLoader.Source],
436
- getattr(model, "secondary_weights", ()),
437
- )
438
- for source in secondary_weights:
439
- yield from self._get_weights_iterator(source)
440
-
441
- def download_model(self, model_config: ModelConfig) -> None:
442
- self._prepare_weights(model_config.model,
443
- model_config.revision,
444
- fall_back_to_pt=True,
445
- allow_patterns_overrides=None)
446
-
447
- def load_model(self, vllm_config: VllmConfig) -> nn.Module:
448
- device_config = vllm_config.device_config
449
- model_config = vllm_config.model_config
450
- target_device = torch.device(device_config.device)
451
- with set_default_torch_dtype(model_config.dtype):
452
- with target_device:
453
- model = _initialize_model(vllm_config=vllm_config)
454
-
455
- weights_to_load = {name for name, _ in model.named_parameters()}
456
- loaded_weights = model.load_weights(
457
- self.get_all_weights(model_config, model))
458
- self.counter_after_loading_weights = time.perf_counter()
459
- logger.info(
460
- "Loading weights took %.2f seconds",
461
- self.counter_after_loading_weights -
462
- self.counter_before_loading_weights)
463
- # We only enable strict check for non-quantized models
464
- # that have loaded weights tracking currently.
465
- if model_config.quantization is None and loaded_weights is not None:
466
- weights_not_loaded = weights_to_load - loaded_weights
467
- if weights_not_loaded:
468
- raise ValueError(
469
- "Following weights were not initialized from "
470
- f"checkpoint: {weights_not_loaded}")
471
-
472
- _process_weights_after_loading(model, model_config, target_device)
473
-
474
- return model.eval()
475
-
476
-
477
-class DummyModelLoader(BaseModelLoader):
478
- """Model loader that will set model weights to random values."""
479
-
480
- def __init__(self, load_config: LoadConfig):
481
- super().__init__(load_config)
482
- if load_config.model_loader_extra_config:
483
- raise ValueError(f"Model loader extra config is not supported for "
484
- f"load format {load_config.load_format}")
485
-
486
- def download_model(self, model_config: ModelConfig) -> None:
487
- pass # Nothing to download
488
-
489
- def load_model(self, vllm_config: VllmConfig) -> nn.Module:
490
- device_config = vllm_config.device_config
491
- model_config = vllm_config.model_config
492
- target_device = torch.device(device_config.device)
493
- with set_default_torch_dtype(model_config.dtype):
494
- with target_device:
495
- model = _initialize_model(vllm_config=vllm_config)
496
- # NOTE(woosuk): For accurate performance evaluation, we assign
497
- # random values to the weights.
498
- initialize_dummy_weights(model)
499
-
500
- _process_weights_after_loading(model, model_config, target_device)
501
- return model.eval()
502
-
503
-
504
-class TensorizerLoader(BaseModelLoader):
505
- """Model loader using CoreWeave's tensorizer library."""
506
-
507
- def __init__(self, load_config: LoadConfig):
508
- super().__init__(load_config)
509
- if isinstance(load_config.model_loader_extra_config, TensorizerConfig):
510
- self.tensorizer_config = load_config.model_loader_extra_config
511
- else:
512
- self.tensorizer_config = TensorizerConfig(
513
- **load_config.model_loader_extra_config)
514
-
515
- def _verify_config(self, model_config: ModelConfig,
516
- parallel_config: ParallelConfig):
517
- self.tensorizer_config.verify_with_model_config(model_config)
518
- self.tensorizer_config.verify_with_parallel_config(parallel_config)
519
-
520
- def _get_weights_iterator(
521
- self, ) -> Generator[Tuple[str, torch.Tensor], None, None]:
522
- tensorizer_args = self.tensorizer_config._construct_tensorizer_args()
523
- return tensorizer_weights_iterator(tensorizer_args)
524
-
525
- def _load_model_serialized_cpu(
526
- self,
527
- vllm_config: VllmConfig,
528
- ) -> nn.Module:
529
- """Load a serialized model with tensorizer to the CPU.
530
-
531
- This is only necessary when the model isn't vLLM-tensorized (see
532
- examples/other/tensorize_vllm_model.py) This should still
533
- be faster than default HuggingFace loading, but will be slower than
534
- loading a vLLM-tensorized model.
535
- """
536
- device_config = vllm_config.device_config
537
- model_config = vllm_config.model_config
538
- with set_default_torch_dtype(model_config.dtype):
539
- with torch.device(device_config.device):
540
- model = _initialize_model(vllm_config=vllm_config)
541
-
542
- model.load_weights(self._get_weights_iterator())
543
- return model.eval()
544
-
545
- def _load_model_serialized(
546
- self,
547
- vllm_config: VllmConfig,
548
- ) -> nn.Module:
549
- """Load a serialized model with tensorizer.
550
-
551
- Expects a vLLM-tensorized model. See the
552
- examples/other/tensorize_vllm_model.py example script
553
- for serializing vLLM models."""
554
-
555
- device_config = vllm_config.device_config
556
- model_config = vllm_config.model_config
557
-
558
- with set_default_torch_dtype(model_config.dtype):
559
- with torch.device(device_config.device):
560
- model_class = get_model_architecture(model_config)[0]
561
-
562
- tensorizer_config = copy.copy(self.tensorizer_config)
563
- tensorizer_config.model_class = model_class
564
- tensorizer_config.hf_config = model_config.hf_config
565
- tensorizer_config.dtype = model_config.dtype
566
-
567
- model = load_with_tensorizer(tensorizer_config,
568
- vllm_config=vllm_config)
569
- return model.eval()
570
-
571
- def download_model(self, model_config: ModelConfig) -> None:
572
- self.tensorizer_config.verify_with_model_config(model_config)
573
-
574
- with self.tensorizer_config.open_stream():
575
- pass
576
-
577
- def load_model(self, vllm_config: VllmConfig) -> nn.Module:
578
- model_config = vllm_config.model_config
579
- parallel_config = vllm_config.parallel_config
580
- self._verify_config(model_config, parallel_config)
581
-
582
- if parallel_config.tensor_parallel_size > 1:
583
- from vllm.distributed import get_tensor_model_parallel_rank
584
-
585
- self.tensorizer_config.tensorizer_uri = (
586
- self.tensorizer_config.tensorizer_uri %
587
- get_tensor_model_parallel_rank())
588
-
589
- if is_vllm_tensorized(self.tensorizer_config):
590
- return self._load_model_serialized(vllm_config=vllm_config)
591
- return self._load_model_serialized_cpu(vllm_config=vllm_config)
592
-
593
- @staticmethod
594
- def save_model(
595
- model: torch.nn.Module,
596
- tensorizer_config: TensorizerConfig,
597
- ) -> None:
598
- serialize_vllm_model(
599
- model=model,
600
- tensorizer_config=tensorizer_config,
601
- )
602
-
603
-
604
-class ShardedStateLoader(BaseModelLoader):
605
- """
606
- Model loader that directly loads each worker's model state dict, which
607
- enables a fast load path for large tensor-parallel models where each worker
608
- only needs to read its own shard rather than the entire checkpoint. See
609
- `examples/offline_inference/save_sharded_state.py` for creating a sharded
610
- checkpoint.
611
- """
612
-
613
- DEFAULT_PATTERN = "model-rank-{rank}-part-{part}.safetensors"
614
-
615
- def __init__(self,
616
- load_config: LoadConfig,
617
- runai_model_streamer: bool = False):
618
- super().__init__(load_config)
619
-
620
- self.runai_model_streamer = runai_model_streamer
621
- extra_config = ({} if load_config.model_loader_extra_config is None
622
- else load_config.model_loader_extra_config.copy())
623
- self.pattern = extra_config.pop("pattern", self.DEFAULT_PATTERN)
624
- if extra_config:
625
- raise ValueError(f"Unexpected extra config keys for load format "
626
- f"{load_config.load_format}: "
627
- f"{load_config.model_loader_extra_config.keys()}")
628
-
629
- @staticmethod
630
- def _filter_subtensors(
631
- tensors: Dict[str, torch.Tensor], ) -> Dict[str, torch.Tensor]:
632
- """
633
- Filter out all tensors that share the same memory or a subset of the
634
- memory of another tensor.
635
- """
636
- same_storage_groups: Dict[Any, List[Tuple[str, torch.Tensor]]] = (
637
- collections.defaultdict(list))
638
- for key, tensor in tensors.items():
639
- if tensor.numel():
640
- ptr = tensor.untyped_storage().data_ptr()
641
- same_storage_groups[tensor.device, ptr].append((key, tensor))
642
-
643
- def get_end_ptr(tensor: torch.Tensor) -> int:
644
- return tensor.view(-1)[-1].data_ptr() + tensor.element_size()
645
-
646
- result: Dict[str, torch.Tensor] = {}
647
- for group in same_storage_groups.values():
648
- for k, t in group:
649
- a, b = t.data_ptr(), get_end_ptr(t)
650
- for k2, t2 in group:
651
- if not t2.is_contiguous():
652
- continue
653
- a2, b2 = t2.data_ptr(), get_end_ptr(t2)
654
- if a < a2 or b2 < b:
655
- continue
656
- if a2 < a or b < b2 or not t.is_contiguous():
657
- break # t2 covers strictly more memory than t.
658
- if k2 < k:
659
- # Same tensors, keep the one with the smaller key.
660
- break
661
- else:
662
- result[k] = t
663
- return result
664
-
665
- def _prepare_weights(self, model_name_or_path: str,
666
- revision: Optional[str]):
667
- if is_s3(model_name_or_path) or os.path.isdir(model_name_or_path):
668
- return model_name_or_path
669
- else:
670
- allow_patterns = ["*.safetensors"]
671
- return download_weights_from_hf(
672
- model_name_or_path,
673
- self.load_config.download_dir,
674
- allow_patterns,
675
- revision,
676
- ignore_patterns=self.load_config.ignore_patterns,
677
- )
678
-
679
- def download_model(self, model_config: ModelConfig) -> None:
680
- self._prepare_weights(model_config.model, model_config.revision)
681
-
682
- def load_model(self, vllm_config: VllmConfig) -> nn.Module:
683
- device_config = vllm_config.device_config
684
- model_config = vllm_config.model_config
685
- target_device = torch.device(device_config.device)
686
-
687
- from vllm.distributed import get_tensor_model_parallel_rank
688
-
689
- model_weights = model_config.model
690
- if hasattr(model_config, "model_weights"):
691
- model_weights = model_config.model_weights
692
- local_model_path = model_weights
693
-
694
- with set_default_torch_dtype(model_config.dtype):
695
- with target_device:
696
- model = _initialize_model(vllm_config=vllm_config)
697
- _process_weights_after_loading(model, model_config,
698
- target_device)
699
- rank = get_tensor_model_parallel_rank()
700
- pattern = os.path.join(
701
- local_model_path,
702
- self.pattern.format(rank=rank, part="*"),
703
- )
704
-
705
- filepaths = []
706
- if is_s3(local_model_path):
707
- file_pattern = f"*{self.pattern.format(rank=rank, part=' * ')}"
708
- filepaths = s3_glob(path=local_model_path,
709
- allow_pattern=[file_pattern])
710
- else:
711
- filepaths = glob.glob(pattern)
712
- if not filepaths:
713
- # TODO: support un-sharded checkpoints too
714
- raise ValueError(
715
- f"Could not find checkpoint files '{pattern}', only "
716
- f"pre-sharded checkpoints are currently supported!")
717
- state_dict = self._filter_subtensors(model.state_dict())
718
- for key, tensor in self.iterate_over_files(filepaths):
719
- # If loading with LoRA enabled, additional padding may
720
- # be added to certain parameters. We only load into a
721
- # narrowed view of the parameter data.
722
- param_data = state_dict[key].data
723
- param_shape = state_dict[key].shape
724
- for dim, size in enumerate(tensor.shape):
725
- if size < param_shape[dim]:
726
- param_data = param_data.narrow(dim, 0, size)
727
- if tensor.shape != param_shape:
728
- logger.warning(
729
- "loading tensor of shape %s into "
730
- "parameter '%s' of shape %s",
731
- tensor.shape,
732
- key,
733
- param_shape,
734
- )
735
- param_data.copy_(tensor)
736
- state_dict.pop(key)
737
- if state_dict:
738
- raise ValueError(
739
- f"Missing keys {tuple(state_dict)} in loaded state!")
740
- return model.eval()
741
-
742
- def iterate_over_files(
743
- self, paths) -> Generator[Tuple[str, torch.Tensor], None, None]:
744
- if self.runai_model_streamer:
745
- yield from runai_safetensors_weights_iterator(paths, True)
746
- else:
747
- from safetensors.torch import safe_open
748
- for path in paths:
749
- with safe_open(path, framework="pt") as f:
750
- for key in f.keys(): # noqa: SIM118
751
- tensor = f.get_tensor(key)
752
- yield key, tensor
753
-
754
- @staticmethod
755
- def save_model(
756
- model: torch.nn.Module,
757
- path: str,
758
- pattern: Optional[str] = None,
759
- max_size: Optional[int] = None,
760
- ) -> None:
761
- from safetensors.torch import save_file
762
-
763
- from vllm.distributed import get_tensor_model_parallel_rank
764
-
765
- if pattern is None:
766
- pattern = ShardedStateLoader.DEFAULT_PATTERN
767
- rank = get_tensor_model_parallel_rank()
768
- part_idx = 0
769
- total_size = 0
770
- state_dict = ShardedStateLoader._filter_subtensors(model.state_dict())
771
- state_dict_part: Dict[str, torch.Tensor] = {}
772
- for key, tensor in state_dict.items():
773
- param_size = tensor.nelement() * tensor.element_size()
774
- if max_size is not None and total_size + param_size > max_size:
775
- filename = pattern.format(rank=rank, part=part_idx)
776
- save_file(
777
- state_dict_part,
778
- os.path.join(path, filename),
779
- )
780
- part_idx += 1
781
- total_size = 0
782
- state_dict_part = {}
783
- state_dict_part[key] = tensor
784
- total_size += param_size
785
- if len(state_dict_part) > 0:
786
- filename = pattern.format(rank=rank, part=part_idx)
787
- save_file(
788
- state_dict_part,
789
- os.path.join(path, filename),
790
- )
791
-
792
-
793
-class BitsAndBytesModelLoader(BaseModelLoader):
794
- """Model loader to load model weights with BitAndBytes quantization."""
795
-
796
- possible_config_file_names = ["adapter_config.json"]
797
-
798
- def __init__(self, load_config: LoadConfig):
799
- super().__init__(load_config)
800
-
801
- # Save the module names without sharding.
802
- self.unsharded_weights_modules: List[str] = []
803
- # Save the module names that are sharded by column.
804
- self.column_sharded_weights_modules: List[str] = []
805
- # Store all module names (from transformers) that support
806
- # BNB quantization.
807
- self.target_modules: List[str] = []
808
- # mapping weight names from transformers to vllm.
809
- self.weight_mapper: Callable = lambda name: name
810
-
811
- def _get_weight_files(
812
- self,
813
- model_name_or_path: str,
814
- allowed_patterns: List[str],
815
- revision: Optional[str] = None,
816
- ) -> Tuple[str, List[str], str]:
817
- """Retrieve weight files. Download the files if necessary.
818
-
819
- Return the weight files and the file pattern."""
820
- is_local = os.path.isdir(model_name_or_path)
821
-
822
- if is_local:
823
- for pattern in allowed_patterns:
824
- weight_files = glob.glob(
825
- os.path.join(model_name_or_path, pattern))
826
- if weight_files:
827
- return model_name_or_path, weight_files, pattern
828
- else:
829
- hf_api = HfApi()
830
- repo_files = hf_api.list_repo_files(repo_id=model_name_or_path)
831
- for pattern in allowed_patterns:
832
- matching_files = fnmatch.filter(repo_files, pattern)
833
- if matching_files:
834
- hf_folder = download_weights_from_hf(
835
- model_name_or_path,
836
- self.load_config.download_dir,
837
- [pattern],
838
- revision,
839
- ignore_patterns=self.load_config.ignore_patterns,
840
- )
841
- return hf_folder, glob.glob(
842
- os.path.join(hf_folder, pattern)), pattern
843
-
844
- raise RuntimeError(
845
- f"No model weights found in: `{model_name_or_path}`")
846
-
847
- def _prepare_weights(self, model_name_or_path: str,
848
- revision: Optional[str]) -> Tuple[List[str], bool]:
849
- """Prepare weight files for the model."""
850
-
851
- allowed_patterns = ["*.safetensors", "*.bin", "*.pt"]
852
-
853
- hf_folder, hf_weights_files, matched_pattern = self._get_weight_files(
854
- model_name_or_path, allowed_patterns, revision)
855
-
856
- use_safetensors = matched_pattern == "*.safetensors"
857
- is_local = os.path.isdir(model_name_or_path)
858
- index_file = SAFE_WEIGHTS_INDEX_NAME
859
- if use_safetensors:
860
- # For models like Mistral-7B-Instruct-v0.3
861
- # there are both sharded safetensors files and a consolidated
862
- # safetensors file. Using both breaks.
863
- # Here, we download the `model.safetensors.index.json` and filter
864
- # any files not found in the index.
865
- if not is_local:
866
- download_safetensors_index_file_from_hf(
867
- model_name_or_path,
868
- index_file,
869
- self.load_config.download_dir,
870
- revision,
871
- )
872
- hf_weights_files = filter_duplicate_safetensors_files(
873
- hf_weights_files, hf_folder, index_file)
874
- else:
875
- hf_weights_files = filter_files_not_needed_for_inference(
876
- hf_weights_files)
877
-
878
- if len(hf_weights_files) == 0:
879
- raise RuntimeError(
880
- f"Cannot find any model weights with `{model_name_or_path}`")
881
-
882
- return hf_weights_files, use_safetensors
883
-
884
- def _hf_weight_iter(self, hf_weights_files, use_safetensors: bool):
885
- if use_safetensors:
886
- iterator = safetensors_weights_iterator(
887
- hf_weights_files,
888
- self.load_config.use_tqdm_on_load,
889
- )
890
- else:
891
- iterator = pt_weights_iterator(
892
- hf_weights_files,
893
- self.load_config.use_tqdm_on_load,
894
- self.load_config.pt_load_map_location,
895
- )
896
- for org_name, param in iterator:
897
- # mapping weight names from transformers to vllm while preserving
898
- # original names.
899
- mapped_name = self.weight_mapper(org_name)
900
- yield org_name, mapped_name, param
901
-
902
- def _get_quantized_weights_iterator(
903
- self,
904
- model_name_or_path: str,
905
- revision: Optional[str],
906
- pre_quant: bool,
907
- load_8bit: bool,
908
- ) -> Tuple[Generator[Tuple[str, torch.Tensor], None, None], Dict[str,
909
- Any]]:
910
- """Get an iterator to the model weights with bitsandbytes quantization,
911
- as well as the quantization state dictionary."""
912
-
913
- # only load the bitsandbytes module when needed
914
- try:
915
- import bitsandbytes
916
-
917
- if bitsandbytes.__version__ < "0.45.3":
918
- raise ImportError("bitsandbytes version is wrong. Please "
919
- "install bitsandbytes>=0.45.3.")
920
- except ImportError as err:
921
- raise ImportError("Please install bitsandbytes>=0.45.3 via "
922
- "`pip install bitsandbytes>=0.45.3` to use "
923
- "bitsandbytes quantizer.") from err
924
-
925
- hf_weights_files, use_safetensors = self._prepare_weights(
926
- model_name_or_path, revision)
927
-
928
- quant_state_dict: Dict[str, Any] = {}
929
-
930
- if pre_quant:
931
- if load_8bit:
932
- return self._quantized_8bit_generator(
933
- hf_weights_files, use_safetensors,
934
- quant_state_dict), quant_state_dict
935
- else:
936
- return self._quantized_4bit_generator(
937
- hf_weights_files, use_safetensors,
938
- quant_state_dict), quant_state_dict
939
-
940
- return self._unquantized_generator(hf_weights_files, use_safetensors,
941
- quant_state_dict), quant_state_dict
942
-
943
- def _is_8bit_weight_name(self, weight_name: str):
944
- quantized_suffix = {".scb", ".weight_format"}
945
- return any(weight_name.lower().endswith(suffix)
946
- for suffix in quantized_suffix)
947
-
948
- def _is_4bit_weight_name(self, weight_name: str):
949
- quantized_suffix = {
950
- "absmax",
951
- "quant_map",
952
- "nested_absmax",
953
- "nested_quant_map",
954
- "bitsandbytes",
955
- }
956
- suffix = weight_name.split(".")[-1]
957
- return any(q_suffix in suffix for q_suffix in quantized_suffix)
958
-
959
- def _quantized_8bit_generator(self, hf_weights_files, use_safetensors,
960
- quant_state_dict) -> Generator:
961
- for (
962
- org_weight_name,
963
- mapped_weight_name,
964
- weight_tensor,
965
- ) in self._hf_weight_iter(hf_weights_files, use_safetensors):
966
- if not mapped_weight_name.lower().endswith(".scb"):
967
- continue
968
-
969
- weight_key = mapped_weight_name.lower().replace(".scb", ".weight")
970
- quant_state_dict[weight_key] = weight_tensor
971
-
972
- for (
973
- org_weight_name,
974
- mapped_weight_name,
975
- weight_tensor,
976
- ) in self._hf_weight_iter(hf_weights_files, use_safetensors):
977
- if self._is_8bit_weight_name(mapped_weight_name):
978
- continue
979
-
980
- if mapped_weight_name in quant_state_dict:
981
- set_weight_attrs(weight_tensor, {"load_in_8bit": True})
982
- yield org_weight_name, weight_tensor
983
- else:
984
- yield org_weight_name, weight_tensor
985
-
986
- def _quantized_4bit_generator(self, hf_weights_files, use_safetensors,
987
- quant_state_dict) -> Generator:
988
- from bitsandbytes.functional import QuantState
989
-
990
- # First iterate over all quant state weights
991
- weight_iterator = self._hf_weight_iter(hf_weights_files,
992
- use_safetensors)
993
- temp_state_dict = {}
994
- for (
995
- org_weight_name,
996
- mapped_weight_name,
997
- weight_tensor,
998
- ) in weight_iterator:
999
- if not self._is_4bit_weight_name(mapped_weight_name):
1000
- continue
1001
- # bitsandbytes library requires
1002
- # weight.quant_state.bitsandbytes__* in CPU
1003
- if "quant_state.bitsandbytes" in mapped_weight_name:
1004
- temp_state_dict[mapped_weight_name] = weight_tensor.cpu().data
1005
- else:
1006
- temp_state_dict[mapped_weight_name] = weight_tensor
1007
-
1008
- # Closure to parse quant_state for each prequant weight
1009
- def _parse_quant_state(param_name: str,
1010
- temp_state_dict: Dict) -> QuantState:
1011
- quant_state = {}
1012
- for k in temp_state_dict:
1013
- if param_name + "." in k:
1014
- quant_state[k] = temp_state_dict[k]
1015
-
1016
- return QuantState.from_dict(quant_state,
1017
- device=current_platform.device_type)
1018
-
1019
- # Second iterate over all prequant and normal weights
1020
- # pre quantized weights would have a quant_state
1021
- for (
1022
- org_weight_name,
1023
- mapped_weight_name,
1024
- weight_tensor,
1025
- ) in self._hf_weight_iter(hf_weights_files, use_safetensors):
1026
- if self._is_4bit_weight_name(mapped_weight_name):
1027
- continue
1028
-
1029
- if (f"{mapped_weight_name}.quant_state.bitsandbytes__nf4"
1030
- in temp_state_dict) or (
1031
- f"{mapped_weight_name}.quant_state.bitsandbytes__fp4"
1032
- in temp_state_dict):
1033
- quant_state = _parse_quant_state(mapped_weight_name,
1034
- temp_state_dict)
1035
- quant_state_dict[mapped_weight_name] = quant_state
1036
- yield org_weight_name, weight_tensor
1037
- else:
1038
- yield org_weight_name, weight_tensor
1039
-
1040
- def _unquantized_generator(self, hf_weights_files, use_safetensors,
1041
- quant_state_dict) -> Generator:
1042
- from bitsandbytes.functional import quantize_4bit
1043
-
1044
- tp_size = get_tensor_model_parallel_world_size()
1045
- tp_rank = get_tensor_model_parallel_rank()
1046
-
1047
- for (
1048
- org_weight_name,
1049
- mapped_weight_name,
1050
- weight_tensor,
1051
- ) in self._hf_weight_iter(hf_weights_files, use_safetensors):
1052
- if any(target_module in mapped_weight_name
1053
- for target_module in self.target_modules
1054
- ) and mapped_weight_name.endswith(".weight"):
1055
- # Without sharding
1056
- if any(
1057
- mapped_weight_name.startswith(module)
1058
- for module in self.unsharded_weights_modules):
1059
- weight_sub_tensor = weight_tensor
1060
- # Shard by column
1061
- elif any(
1062
- mapped_weight_name.startswith(module)
1063
- for module in self.column_sharded_weights_modules):
1064
- total_size = weight_tensor.size(-1)
1065
- start_index = total_size // tp_size * tp_rank
1066
- end_index = total_size // tp_size * (tp_rank + 1)
1067
- weight_sub_tensor = weight_tensor[...,
1068
- start_index:end_index]
1069
- # Weights have fused on disk. In this case, we assume that the
1070
- # weight and module use same name.
1071
- elif any(
1072
- mapped_weight_name.startswith(module)
1073
- for module in self.maybe_fused_weights_modules):
1074
- # special case for fused weights
1075
- # get the size of each shard weight tensor
1076
- total_shard_sizes = next(
1077
- (sizes for module, sizes in
1078
- self.maybe_fused_weights_modules.items()
1079
- if mapped_weight_name.startswith(module)))
1080
- total_size = weight_tensor.size(0)
1081
- assert total_size == sum(total_shard_sizes)
1082
- # get the start/end index of each shard weight tensor
1083
- total_start_index = list(
1084
- itertools.accumulate([0] + total_shard_sizes))[:-1]
1085
- shard_weights_index = [(
1086
- idx + size // tp_size * tp_rank,
1087
- idx + size // tp_size * (tp_rank + 1),
1088
- ) for idx, size in zip(total_start_index,
1089
- total_shard_sizes)]
1090
- # slice and reorder the weight tensor
1091
- weight_tensor = [
1092
- weight_tensor[start_index:end_index, ...]
1093
- for start_index, end_index in shard_weights_index
1094
- ]
1095
- weight_sub_tensor = torch.cat(weight_tensor, dim=0)
1096
- # Shard by row
1097
- else:
1098
- total_size = weight_tensor.size(0)
1099
- start_index = total_size // tp_size * tp_rank
1100
- end_index = total_size // tp_size * (tp_rank + 1)
1101
- weight_sub_tensor = weight_tensor[start_index:end_index,
1102
- ...]
1103
-
1104
- # bitsandbytes requires data in GPU
1105
- if weight_sub_tensor.is_cuda:
1106
- loaded_weight = weight_sub_tensor
1107
- else:
1108
- loaded_weight = weight_sub_tensor.cuda()
1109
-
1110
- # remove the following after the issue is fixed:
1111
- # https://github.com/bitsandbytes-foundation/bitsandbytes/issues/1342
1112
- if loaded_weight.is_contiguous() is False:
1113
- loaded_weight = loaded_weight.contiguous()
1114
-
1115
- with set_default_torch_dtype(torch.float32):
1116
- processed_weight, quant_state = quantize_4bit(
1117
- loaded_weight,
1118
- compress_statistics=True,
1119
- quant_type="nf4",
1120
- )
1121
-
1122
- quant_state_dict[mapped_weight_name] = quant_state
1123
- else:
1124
- processed_weight = weight_tensor
1125
- yield org_weight_name, processed_weight
1126
-
1127
- def _get_bnb_target_modules(self, model: nn.Module) -> None:
1128
-
1129
- for name, module in model.named_modules():
1130
- if isinstance(module, (LinearBase, )):
1131
- if modules_info := self.modules_mapping.get_sub_modules(name):
1132
- # Map vllm's names to transformers's names.
1133
- rep_name, sub_modules = modules_info
1134
- for sub_name in sub_modules:
1135
- self.target_modules.append(
1136
- name.replace(rep_name, sub_name))
1137
- # Add original module name even if the module has stacked map,
1138
- # in case model has a mixture of disk-merged and disk-splitted
1139
- # weights with same last name.
1140
- self.target_modules.append(name)
1141
-
1142
- assert (self.target_modules
1143
- ), "vllm currently does not support BNB quantization for"
1144
- f" {type(model).__name__}"
1145
-
1146
- def _load_weights(self, model_config: ModelConfig,
1147
- model: nn.Module) -> None:
1148
- if not hasattr(model, "load_weights"):
1149
- raise AttributeError(
1150
- "The required method 'load_weights' is not defined in class"
1151
- f" {type(model).__name__}.")
1152
-
1153
- if not hasattr(model, "packed_modules_mapping"):
1154
- raise AttributeError(
1155
- f"Model {type(model).__name__} does not support BitsAndBytes "
1156
- "quantization yet. No 'packed_modules_mapping' found.")
1157
-
1158
- self.modules_mapping = ParamMapping(
1159
- copy.deepcopy(model.packed_modules_mapping))
1160
-
1161
- # For some models like Molmo, we need to use hf_to_vllm_mapper
1162
- # to ensure correct loading of weights.
1163
- if hf_to_vllm_mapper := getattr(model, "hf_to_vllm_mapper", None):
1164
- self.weight_mapper = lambda name: hf_to_vllm_mapper._map_name(name)
1165
-
1166
- # Modules whose weights might have fused on disk
1167
- # we need their output_sizes to make shard in flight correctly with TP
1168
- self.maybe_fused_weights_modules: Dict[str, List[int]] = {}
1169
- self._get_bnb_target_modules(model)
1170
- for name, module in model.named_modules():
1171
- # Some modules like `ReplicatedLinear` should not have their weights
1172
- # sharded. The reason for implementing it this way is to avoid new
1173
- # static variable in the model implementation.
1174
- if isinstance(module, (ReplicatedLinear, )):
1175
- self.unsharded_weights_modules.append(name)
1176
- # `QKVParallelLinear` and `MergedColumnParallelLinear` might have
1177
- # fused weights on disk. We need to use the output sizes of these
1178
- # modules to shard the weights correctly.
1179
- elif isinstance(module,
1180
- (QKVParallelLinear, MergedColumnParallelLinear)):
1181
- self.maybe_fused_weights_modules[name] = module.output_sizes
1182
- # In TP, these weights are partitioned along the column
1183
- # dimension (dim=-1)
1184
- elif isinstance(module, (RowParallelLinear, )):
1185
- self.column_sharded_weights_modules.append(name)
1186
-
1187
- self.model_type = type(model).__name__
1188
-
1189
- logger.info("Loading weights with BitsAndBytes quantization. "
1190
- "May take a while ...")
1191
-
1192
- quant_config = getattr(model_config.hf_config, "quantization_config",
1193
- None)
1194
-
1195
- pre_quant = False
1196
- if quant_config is not None:
1197
- quant_method = quant_config.get("quant_method")
1198
- if quant_method == "bitsandbytes":
1199
- pre_quant = True
1200
- else:
1201
- raise ValueError(
1202
- f"BitsAndBytes loader does not support {quant_method} "
1203
- "quantization")
1204
-
1205
- # The quant_states in pre_quantized models cannot work with a split
1206
- # weight tensor. So TP does not work with pre_quantized bnb models.
1207
- if pre_quant and get_tensor_model_parallel_world_size() > 1:
1208
- raise ValueError(
1209
- "Prequant BitsAndBytes models with tensor parallelism is not "
1210
- "supported. Please try with pipeline parallelism.")
1211
-
1212
- load_8bit = False
1213
- if pre_quant:
1214
- load_8bit = quant_config.get("load_in_8bit", False)
1215
-
1216
- qweight_iterator, quant_state_dict = (
1217
- self._get_quantized_weights_iterator(model_config.model,
1218
- model_config.revision,
1219
- pre_quant, load_8bit))
1220
-
1221
- weights_to_load = {name for name, _ in model.named_parameters()}
1222
- loaded_weights = model.load_weights(qweight_iterator)
1223
- # Some models may have weights loading tracker unimplemented.
1224
- if loaded_weights is not None:
1225
- weights_not_loaded = weights_to_load - loaded_weights
1226
- if weights_not_loaded:
1227
- raise ValueError("Following weights were not initialized from "
1228
- f"checkpoint: {weights_not_loaded}")
1229
-
1230
- torch.cuda.empty_cache()
1231
-
1232
- param_dict = dict(model.named_parameters())
1233
- stacked_quant_state_dict: Dict[str, Dict[int, Any]] = {}
1234
- # TODO: Change this lazy import to normal import
1235
- # after the checks are updated to run on a new version
1236
- from vllm.model_executor.models.utils import is_pp_missing_parameter
1237
-
1238
- for quant_param_name in quant_state_dict:
1239
- if is_pp_missing_parameter(quant_param_name, model):
1240
- continue
1241
-
1242
- non_stacked_param_name = quant_param_name
1243
-
1244
- shard_index = 0
1245
- for shard_name, (
1246
- weight_name,
1247
- index,
1248
- ) in self.modules_mapping.inverse_packed_mapping.items():
1249
- # Some models, such as MiniCPM V2.5/2.6, contain both
1250
- # module names 'kv_proj' and 'qkv_proj'. To prevent 'kv_proj'
1251
- # from being incorrectly identified as being present in
1252
- # 'vpm.encoder.layers.0.self_attn.qkv_proj.weight
1253
- shard_pos = quant_param_name.find(shard_name)
1254
- can_correct_rename = (shard_pos
1255
- > 0) and (quant_param_name[shard_pos - 1]
1256
- == ".")
1257
- # If the quant_param_name is packed, it won't occur in the
1258
- # param_dict before renaming.
1259
- new_quant_param_name = quant_param_name.replace(
1260
- shard_name, weight_name)
1261
- need_rename = (quant_param_name not in param_dict) \
1262
- and (new_quant_param_name in param_dict)
1263
- if can_correct_rename and need_rename:
1264
- shard_index = index
1265
- quant_param_name = new_quant_param_name
1266
- break
1267
-
1268
- # Models like Clip/Siglip may skip some layers in initialization,
1269
- # causing unused quant_param_name in state_dict.
1270
- if quant_param_name not in param_dict:
1271
- continue
1272
-
1273
- if quant_param_name not in stacked_quant_state_dict:
1274
- stacked_quant_state_dict[quant_param_name] = {}
1275
-
1276
- stacked_quant_state_dict[quant_param_name][shard_index] = (
1277
- quant_state_dict[non_stacked_param_name])
1278
-
1279
- # save quant_states and offsets as the attributes of the parameters
1280
- for param_name, param in param_dict.items():
1281
- if param_name in stacked_quant_state_dict:
1282
- quant_states = stacked_quant_state_dict[param_name]
1283
- set_weight_attrs(param, {"bnb_quant_state": quant_states})
1284
-
1285
- pack_ratio = getattr(param, "pack_factor", -1)
1286
- if pack_ratio == -1:
1287
- raise ValueError(
1288
- f"pack_factor not set for parameter {param_name}.")
1289
-
1290
- num_elements = [0] * len(quant_states)
1291
- for seq, quant_state in quant_states.items():
1292
- num_elements[seq] = (math.prod(quant_state.shape) //
1293
- pack_ratio)
1294
-
1295
- offsets = np.concatenate(([0], np.cumsum(num_elements)))
1296
- # Make torch infer_schema happy
1297
- offsets = torch.tensor(offsets).cpu()
1298
- set_weight_attrs(param, {"bnb_shard_offsets": offsets})
1299
-
1300
- if load_8bit:
1301
- set_weight_attrs(
1302
- param, {"matmul_state": [None] * len(quant_states)})
1303
-
1304
- def download_model(self, model_config: ModelConfig) -> None:
1305
- self._prepare_weights(model_config.model, model_config.revision)
1306
-
1307
- def load_model(self, vllm_config: VllmConfig) -> nn.Module:
1308
- device_config = vllm_config.device_config
1309
- model_config = vllm_config.model_config
1310
- with set_default_torch_dtype(model_config.dtype):
1311
- with torch.device(device_config.device):
1312
- model = _initialize_model(vllm_config=vllm_config)
1313
-
1314
- self._load_weights(model_config, model)
1315
-
1316
- return model.eval()
1317
-
1318
-
1319
-class GGUFModelLoader(BaseModelLoader):
1320
- """
1321
- Model loader that can load GGUF files. This is useful for loading models
1322
- that are quantized with GGUF and saved in the GGUF format. This loader
1323
- supports loading both full models and sharded models.
1324
- """
1325
-
1326
- def __init__(self, load_config: LoadConfig):
1327
- super().__init__(load_config)
1328
- if load_config.model_loader_extra_config:
1329
- raise ValueError(f"Model loader extra config is not supported for "
1330
- f"load format {load_config.load_format}")
1331
-
1332
- def _prepare_weights(self, model_name_or_path: str):
1333
- if os.path.isfile(model_name_or_path):
1334
- return model_name_or_path
1335
- else:
1336
- raise ValueError(f"{model_name_or_path} is not a file.")
1337
-
1338
- def _get_gguf_weights_map(self, model_config: ModelConfig):
1339
- """
1340
- GGUF uses this naming convention for their tensors from HF checkpoint:
1341
- `blk.N.BB.weight` and `blk.N.BB.bias`
1342
- where N signifies the block number of a layer, and BB signifies the
1343
- attention/mlp layer components.
1344
- See "Standardized tensor names" in
1345
- https://github.com/ggerganov/ggml/blob/master/docs/gguf.md for details.
1346
- """
1347
- config = model_config.hf_config
1348
- model_type = config.model_type
1349
- gguf_to_hf_name_map = {}
1350
- # hack: ggufs have a different name than transformers
1351
- if model_type == "cohere":
1352
- model_type = "command-r"
1353
- if model_type in ("deepseek_v3", "deepseek_v2"):
1354
- model_type = "deepseek2"
1355
- # GGUF layer map assumes that we will have a merged expert weights
1356
- # so we need to map them manually
1357
- for idx in range(config.num_hidden_layers):
1358
- gguf_to_hf_name_map[f"blk.{idx}.exp_probs_b.bias"] = \
1359
- f"model.layers.{idx}.mlp.gate.e_score_correction_bias"
1360
- gguf_to_hf_name_map[f"blk.{idx}.ffn_down_exps.weight"] = \
1361
- f"model.layers.{idx}.mlp.experts.0.down_proj.weight"
1362
- gguf_to_hf_name_map[f"blk.{idx}.ffn_gate_exps.weight"] = \
1363
- f"model.layers.{idx}.mlp.experts.0.gate_proj.weight"
1364
- gguf_to_hf_name_map[f"blk.{idx}.ffn_up_exps.weight"] = \
1365
- f"model.layers.{idx}.mlp.experts.0.up_proj.weight"
1366
-
1367
- arch = None
1368
- for key, value in gguf.MODEL_ARCH_NAMES.items():
1369
- if value == model_type:
1370
- arch = key
1371
- break
1372
- if arch is None:
1373
- raise RuntimeError(f"Unknown gguf model_type: {model_type}")
1374
- num_layers = config.num_hidden_layers
1375
- name_map = gguf.get_tensor_name_map(arch, num_layers)
1376
- with torch.device("meta"):
1377
- dummy_model = AutoModelForCausalLM.from_config(
1378
- config, trust_remote_code=model_config.trust_remote_code)
1379
- state_dict = dummy_model.state_dict()
1380
-
1381
- for hf_name in state_dict:
1382
- name, suffix = hf_name.rsplit(".", 1)
1383
- gguf_name = name_map.get_name(name)
1384
- gguf_to_hf_name_map[f"{gguf_name}.{suffix}"] = hf_name
1385
- return gguf_to_hf_name_map
1386
-
1387
- def _get_weights_iterator(
1388
- self, model_name_or_path: str, gguf_to_hf_name_map: Dict[str, str]
1389
- ) -> Generator[Tuple[str, torch.Tensor], None, None]:
1390
- return gguf_quant_weights_iterator(model_name_or_path,
1391
- gguf_to_hf_name_map)
1392
-
1393
- def download_model(self, model_config: ModelConfig) -> None:
1394
- self._prepare_weights(model_config.model)
1395
-
1396
- def load_model(self, vllm_config: VllmConfig) -> nn.Module:
1397
- device_config = vllm_config.device_config
1398
- model_config = vllm_config.model_config
1399
- local_model_path = self._prepare_weights(model_config.model)
1400
- gguf_weights_map = self._get_gguf_weights_map(model_config)
1401
- # we can only know if tie word embeddings after mapping weights
1402
- if "lm_head.weight" in get_gguf_extra_tensor_names(
1403
- local_model_path, gguf_weights_map):
1404
- model_config.hf_config.update({"tie_word_embeddings": True})
1405
-
1406
- target_device = torch.device(device_config.device)
1407
- with set_default_torch_dtype(model_config.dtype):
1408
- with target_device:
1409
- model = _initialize_model(vllm_config=vllm_config)
1410
- model.load_weights(
1411
- self._get_weights_iterator(local_model_path, gguf_weights_map))
1412
-
1413
- _process_weights_after_loading(model, model_config, target_device)
1414
- return model
1415
-
1416
-
1417
-class RunaiModelStreamerLoader(BaseModelLoader):
1418
- """
1419
- Model loader that can load safetensors
1420
- files from local FS or S3 bucket.
1421
- """
1422
-
1423
- def __init__(self, load_config: LoadConfig):
1424
- super().__init__(load_config)
1425
- if load_config.model_loader_extra_config:
1426
- extra_config = load_config.model_loader_extra_config
1427
-
1428
- if ("concurrency" in extra_config
1429
- and isinstance(extra_config.get("concurrency"), int)):
1430
- os.environ["RUNAI_STREAMER_CONCURRENCY"] = str(
1431
- extra_config.get("concurrency"))
1432
-
1433
- if ("memory_limit" in extra_config
1434
- and isinstance(extra_config.get("memory_limit"), int)):
1435
- os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] = str(
1436
- extra_config.get("memory_limit"))
1437
-
1438
- runai_streamer_s3_endpoint = os.getenv(
1439
- 'RUNAI_STREAMER_S3_ENDPOINT')
1440
- aws_endpoint_url = os.getenv('AWS_ENDPOINT_URL')
1441
- if (runai_streamer_s3_endpoint is None
1442
- and aws_endpoint_url is not None):
1443
- os.environ["RUNAI_STREAMER_S3_ENDPOINT"] = aws_endpoint_url
1444
-
1445
- def _prepare_weights(self, model_name_or_path: str,
1446
- revision: Optional[str]) -> List[str]:
1447
- """Prepare weights for the model.
1448
-
1449
- If the model is not local, it will be downloaded."""
1450
-
1451
- is_s3_path = is_s3(model_name_or_path)
1452
- is_local = os.path.isdir(model_name_or_path)
1453
- safetensors_pattern = "*.safetensors"
1454
- index_file = SAFE_WEIGHTS_INDEX_NAME
1455
-
1456
- hf_folder = (model_name_or_path if
1457
- (is_local or is_s3_path) else download_weights_from_hf(
1458
- model_name_or_path,
1459
- self.load_config.download_dir,
1460
- [safetensors_pattern],
1461
- revision,
1462
- ignore_patterns=self.load_config.ignore_patterns,
1463
- ))
1464
- if is_s3_path:
1465
- hf_weights_files = s3_glob(path=hf_folder,
1466
- allow_pattern=[safetensors_pattern])
1467
- else:
1468
- hf_weights_files = glob.glob(
1469
- os.path.join(hf_folder, safetensors_pattern))
1470
-
1471
- if not is_local and not is_s3_path:
1472
- download_safetensors_index_file_from_hf(
1473
- model_name_or_path, index_file, self.load_config.download_dir,
1474
- revision)
1475
-
1476
- if not hf_weights_files:
1477
- raise RuntimeError(
1478
- f"Cannot find any safetensors model weights with "
1479
- f"`{model_name_or_path}`")
1480
-
1481
- return hf_weights_files
1482
-
1483
- def _get_weights_iterator(
1484
- self, model_or_path: str,
1485
- revision: str) -> Generator[Tuple[str, torch.Tensor], None, None]:
1486
- """Get an iterator for the model weights based on the load format."""
1487
- hf_weights_files = self._prepare_weights(model_or_path, revision)
1488
- return runai_safetensors_weights_iterator(
1489
- hf_weights_files,
1490
- self.load_config.use_tqdm_on_load,
1491
- )
1492
-
1493
- def download_model(self, model_config: ModelConfig) -> None:
1494
- """Download model if necessary"""
1495
- self._prepare_weights(model_config.model, model_config.revision)
1496
-
1497
- def load_model(self, vllm_config: VllmConfig) -> nn.Module:
1498
- """Perform streaming of the model to destination"""
1499
- device_config = vllm_config.device_config
1500
- model_config = vllm_config.model_config
1501
-
1502
- target_device = torch.device(device_config.device)
1503
- with set_default_torch_dtype(model_config.dtype):
1504
- with target_device:
1505
- model = _initialize_model(vllm_config=vllm_config)
1506
-
1507
- model_weights = model_config.model
1508
- if hasattr(model_config, "model_weights"):
1509
- model_weights = model_config.model_weights
1510
- model.load_weights(
1511
- self._get_weights_iterator(model_weights,
1512
- model_config.revision))
1513
-
1514
- _process_weights_after_loading(model, model_config, target_device)
1515
- return model.eval()
1516
-
1517
-
1518
-def get_model_loader(load_config: LoadConfig) -> BaseModelLoader:
1519
- """Get a model loader based on the load format."""
1520
- if isinstance(load_config.load_format, type):
1521
- return load_config.load_format(load_config)
1522
-
1523
- if load_config.load_format == LoadFormat.DUMMY:
1524
- return DummyModelLoader(load_config)
1525
-
1526
- if load_config.load_format == LoadFormat.TENSORIZER:
1527
- return TensorizerLoader(load_config)
1528
-
1529
- if load_config.load_format == LoadFormat.SHARDED_STATE:
1530
- return ShardedStateLoader(load_config)
1531
-
1532
- if load_config.load_format == LoadFormat.BITSANDBYTES:
1533
- return BitsAndBytesModelLoader(load_config)
1534
-
1535
- if load_config.load_format == LoadFormat.GGUF:
1536
- return GGUFModelLoader(load_config)
1537
-
1538
- if load_config.load_format == LoadFormat.RUNAI_STREAMER:
1539
- return RunaiModelStreamerLoader(load_config)
1540
-
1541
- if load_config.load_format == LoadFormat.RUNAI_STREAMER_SHARDED:
1542
- return ShardedStateLoader(load_config, runai_model_streamer=True)
1543
-
1544
- return DefaultModelLoader(load_config)
vllm/model_executor/model_loader/runai_streamer_loader.py
new
+120
@@ -0,0 +1,120 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# ruff: noqa: SIM117
3
+import glob
4
+import os
5
+from typing import Generator, List, Optional, Tuple
6
+
7
+import torch
8
+from torch import nn
9
+from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
10
+
11
+from vllm.config import LoadConfig, ModelConfig, VllmConfig
12
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
13
+from vllm.model_executor.model_loader.utils import (
14
+ initialize_model, process_weights_after_loading, set_default_torch_dtype)
15
+from vllm.model_executor.model_loader.weight_utils import (
16
+ download_safetensors_index_file_from_hf, download_weights_from_hf,
17
+ runai_safetensors_weights_iterator)
18
+from vllm.transformers_utils.s3_utils import glob as s3_glob
19
+from vllm.transformers_utils.utils import is_s3
20
+
21
+
22
+class RunaiModelStreamerLoader(BaseModelLoader):
23
+ """
24
+ Model loader that can load safetensors
25
+ files from local FS or S3 bucket.
26
+ """
27
+
28
+ def __init__(self, load_config: LoadConfig):
29
+ super().__init__(load_config)
30
+ if load_config.model_loader_extra_config:
31
+ extra_config = load_config.model_loader_extra_config
32
+
33
+ if ("concurrency" in extra_config
34
+ and isinstance(extra_config.get("concurrency"), int)):
35
+ os.environ["RUNAI_STREAMER_CONCURRENCY"] = str(
36
+ extra_config.get("concurrency"))
37
+
38
+ if ("memory_limit" in extra_config
39
+ and isinstance(extra_config.get("memory_limit"), int)):
40
+ os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] = str(
41
+ extra_config.get("memory_limit"))
42
+
43
+ runai_streamer_s3_endpoint = os.getenv(
44
+ 'RUNAI_STREAMER_S3_ENDPOINT')
45
+ aws_endpoint_url = os.getenv('AWS_ENDPOINT_URL')
46
+ if (runai_streamer_s3_endpoint is None
47
+ and aws_endpoint_url is not None):
48
+ os.environ["RUNAI_STREAMER_S3_ENDPOINT"] = aws_endpoint_url
49
+
50
+ def _prepare_weights(self, model_name_or_path: str,
51
+ revision: Optional[str]) -> List[str]:
52
+ """Prepare weights for the model.
53
+
54
+ If the model is not local, it will be downloaded."""
55
+
56
+ is_s3_path = is_s3(model_name_or_path)
57
+ is_local = os.path.isdir(model_name_or_path)
58
+ safetensors_pattern = "*.safetensors"
59
+ index_file = SAFE_WEIGHTS_INDEX_NAME
60
+
61
+ hf_folder = (model_name_or_path if
62
+ (is_local or is_s3_path) else download_weights_from_hf(
63
+ model_name_or_path,
64
+ self.load_config.download_dir,
65
+ [safetensors_pattern],
66
+ revision,
67
+ ignore_patterns=self.load_config.ignore_patterns,
68
+ ))
69
+ if is_s3_path:
70
+ hf_weights_files = s3_glob(path=hf_folder,
71
+ allow_pattern=[safetensors_pattern])
72
+ else:
73
+ hf_weights_files = glob.glob(
74
+ os.path.join(hf_folder, safetensors_pattern))
75
+
76
+ if not is_local and not is_s3_path:
77
+ download_safetensors_index_file_from_hf(
78
+ model_name_or_path, index_file, self.load_config.download_dir,
79
+ revision)
80
+
81
+ if not hf_weights_files:
82
+ raise RuntimeError(
83
+ f"Cannot find any safetensors model weights with "
84
+ f"`{model_name_or_path}`")
85
+
86
+ return hf_weights_files
87
+
88
+ def _get_weights_iterator(
89
+ self, model_or_path: str,
90
+ revision: str) -> Generator[Tuple[str, torch.Tensor], None, None]:
91
+ """Get an iterator for the model weights based on the load format."""
92
+ hf_weights_files = self._prepare_weights(model_or_path, revision)
93
+ return runai_safetensors_weights_iterator(
94
+ hf_weights_files,
95
+ self.load_config.use_tqdm_on_load,
96
+ )
97
+
98
+ def download_model(self, model_config: ModelConfig) -> None:
99
+ """Download model if necessary"""
100
+ self._prepare_weights(model_config.model, model_config.revision)
101
+
102
+ def load_model(self, vllm_config: VllmConfig) -> nn.Module:
103
+ """Perform streaming of the model to destination"""
104
+ device_config = vllm_config.device_config
105
+ model_config = vllm_config.model_config
106
+
107
+ target_device = torch.device(device_config.device)
108
+ with set_default_torch_dtype(model_config.dtype):
109
+ with target_device:
110
+ model = initialize_model(vllm_config=vllm_config)
111
+
112
+ model_weights = model_config.model
113
+ if hasattr(model_config, "model_weights"):
114
+ model_weights = model_config.model_weights
115
+ model.load_weights(
116
+ self._get_weights_iterator(model_weights,
117
+ model_config.revision))
118
+
119
+ process_weights_after_loading(model, model_config, target_device)
120
+ return model.eval()
vllm/model_executor/model_loader/sharded_state_loader.py
new
+210
@@ -0,0 +1,210 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import collections
4
+import glob
5
+import os
6
+from typing import Any, Dict, Generator, List, Optional, Tuple
7
+
8
+import torch
9
+from torch import nn
10
+
11
+from vllm.config import LoadConfig, ModelConfig, VllmConfig
12
+from vllm.logger import init_logger
13
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
14
+from vllm.model_executor.model_loader.utils import (
15
+ initialize_model, process_weights_after_loading, set_default_torch_dtype)
16
+from vllm.model_executor.model_loader.weight_utils import (
17
+ download_weights_from_hf, runai_safetensors_weights_iterator)
18
+from vllm.transformers_utils.s3_utils import glob as s3_glob
19
+from vllm.transformers_utils.utils import is_s3
20
+
21
+logger = init_logger(__name__)
22
+
23
+
24
+class ShardedStateLoader(BaseModelLoader):
25
+ """
26
+ Model loader that directly loads each worker's model state dict, which
27
+ enables a fast load path for large tensor-parallel models where each worker
28
+ only needs to read its own shard rather than the entire checkpoint. See
29
+ `examples/offline_inference/save_sharded_state.py` for creating a sharded
30
+ checkpoint.
31
+ """
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):
38
+ super().__init__(load_config)
39
+
40
+ self.runai_model_streamer = runai_model_streamer
41
+ extra_config = ({} if load_config.model_loader_extra_config is None
42
+ else load_config.model_loader_extra_config.copy())
43
+ self.pattern = extra_config.pop("pattern", self.DEFAULT_PATTERN)
44
+ if extra_config:
45
+ raise ValueError(f"Unexpected extra config keys for load format "
46
+ f"{load_config.load_format}: "
47
+ f"{load_config.model_loader_extra_config.keys()}")
48
+
49
+ @staticmethod
50
+ def _filter_subtensors(
51
+ tensors: Dict[str, torch.Tensor], ) -> Dict[str, torch.Tensor]:
52
+ """
53
+ Filter out all tensors that share the same memory or a subset of the
54
+ memory of another tensor.
55
+ """
56
+ same_storage_groups: Dict[Any, List[Tuple[str, torch.Tensor]]] = (
57
+ collections.defaultdict(list))
58
+ for key, tensor in tensors.items():
59
+ if tensor.numel():
60
+ ptr = tensor.untyped_storage().data_ptr()
61
+ same_storage_groups[tensor.device, ptr].append((key, tensor))
62
+
63
+ def get_end_ptr(tensor: torch.Tensor) -> int:
64
+ return tensor.view(-1)[-1].data_ptr() + tensor.element_size()
65
+
66
+ result: Dict[str, torch.Tensor] = {}
67
+ for group in same_storage_groups.values():
68
+ for k, t in group:
69
+ a, b = t.data_ptr(), get_end_ptr(t)
70
+ for k2, t2 in group:
71
+ if not t2.is_contiguous():
72
+ continue
73
+ a2, b2 = t2.data_ptr(), get_end_ptr(t2)
74
+ if a < a2 or b2 < b:
75
+ continue
76
+ if a2 < a or b < b2 or not t.is_contiguous():
77
+ break # t2 covers strictly more memory than t.
78
+ if k2 < k:
79
+ # Same tensors, keep the one with the smaller key.
80
+ break
81
+ else:
82
+ result[k] = t
83
+ return result
84
+
85
+ def _prepare_weights(self, model_name_or_path: str,
86
+ revision: Optional[str]):
87
+ if is_s3(model_name_or_path) or os.path.isdir(model_name_or_path):
88
+ return model_name_or_path
89
+ else:
90
+ allow_patterns = ["*.safetensors"]
91
+ return download_weights_from_hf(
92
+ model_name_or_path,
93
+ self.load_config.download_dir,
94
+ allow_patterns,
95
+ revision,
96
+ ignore_patterns=self.load_config.ignore_patterns,
97
+ )
98
+
99
+ def download_model(self, model_config: ModelConfig) -> None:
100
+ self._prepare_weights(model_config.model, model_config.revision)
101
+
102
+ def load_model(self, vllm_config: VllmConfig) -> nn.Module:
103
+ device_config = vllm_config.device_config
104
+ model_config = vllm_config.model_config
105
+ target_device = torch.device(device_config.device)
106
+
107
+ from vllm.distributed import get_tensor_model_parallel_rank
108
+
109
+ model_weights = model_config.model
110
+ if hasattr(model_config, "model_weights"):
111
+ model_weights = model_config.model_weights
112
+ local_model_path = model_weights
113
+
114
+ with set_default_torch_dtype(model_config.dtype):
115
+ with target_device:
116
+ model = initialize_model(vllm_config=vllm_config)
117
+ process_weights_after_loading(model, model_config,
118
+ target_device)
119
+ rank = get_tensor_model_parallel_rank()
120
+ pattern = os.path.join(
121
+ local_model_path,
122
+ self.pattern.format(rank=rank, part="*"),
123
+ )
124
+
125
+ filepaths = []
126
+ if is_s3(local_model_path):
127
+ file_pattern = f"*{self.pattern.format(rank=rank, part=' * ')}"
128
+ filepaths = s3_glob(path=local_model_path,
129
+ allow_pattern=[file_pattern])
130
+ else:
131
+ filepaths = glob.glob(pattern)
132
+ if not filepaths:
133
+ # TODO: support un-sharded checkpoints too
134
+ raise ValueError(
135
+ f"Could not find checkpoint files '{pattern}', only "
136
+ f"pre-sharded checkpoints are currently supported!")
137
+ state_dict = self._filter_subtensors(model.state_dict())
138
+ for key, tensor in self.iterate_over_files(filepaths):
139
+ # If loading with LoRA enabled, additional padding may
140
+ # be added to certain parameters. We only load into a
141
+ # narrowed view of the parameter data.
142
+ param_data = state_dict[key].data
143
+ param_shape = state_dict[key].shape
144
+ for dim, size in enumerate(tensor.shape):
145
+ if size < param_shape[dim]:
146
+ param_data = param_data.narrow(dim, 0, size)
147
+ if tensor.shape != param_shape:
148
+ logger.warning(
149
+ "loading tensor of shape %s into "
150
+ "parameter '%s' of shape %s",
151
+ tensor.shape,
152
+ key,
153
+ param_shape,
154
+ )
155
+ param_data.copy_(tensor)
156
+ state_dict.pop(key)
157
+ if state_dict:
158
+ raise ValueError(
159
+ f"Missing keys {tuple(state_dict)} in loaded state!")
160
+ return model.eval()
161
+
162
+ def iterate_over_files(
163
+ self, paths) -> Generator[Tuple[str, torch.Tensor], None, None]:
164
+ if self.runai_model_streamer:
165
+ yield from runai_safetensors_weights_iterator(paths, True)
166
+ else:
167
+ from safetensors.torch import safe_open
168
+ for path in paths:
169
+ with safe_open(path, framework="pt") as f:
170
+ for key in f.keys(): # noqa: SIM118
171
+ tensor = f.get_tensor(key)
172
+ yield key, tensor
173
+
174
+ @staticmethod
175
+ def save_model(
176
+ model: torch.nn.Module,
177
+ path: str,
178
+ pattern: Optional[str] = None,
179
+ max_size: Optional[int] = None,
180
+ ) -> None:
181
+ from safetensors.torch import save_file
182
+
183
+ from vllm.distributed import get_tensor_model_parallel_rank
184
+
185
+ if pattern is None:
186
+ pattern = ShardedStateLoader.DEFAULT_PATTERN
187
+ rank = get_tensor_model_parallel_rank()
188
+ part_idx = 0
189
+ total_size = 0
190
+ state_dict = ShardedStateLoader._filter_subtensors(model.state_dict())
191
+ state_dict_part: Dict[str, torch.Tensor] = {}
192
+ for key, tensor in state_dict.items():
193
+ param_size = tensor.nelement() * tensor.element_size()
194
+ if max_size is not None and total_size + param_size > max_size:
195
+ filename = pattern.format(rank=rank, part=part_idx)
196
+ save_file(
197
+ state_dict_part,
198
+ os.path.join(path, filename),
199
+ )
200
+ part_idx += 1
201
+ total_size = 0
202
+ state_dict_part = {}
203
+ state_dict_part[key] = tensor
204
+ total_size += param_size
205
+ if len(state_dict_part) > 0:
206
+ filename = pattern.format(rank=rank, part=part_idx)
207
+ save_file(
208
+ state_dict_part,
209
+ os.path.join(path, filename),
210
+ )
vllm/model_executor/model_loader/tensorizer_loader.py
new
+119
@@ -0,0 +1,119 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# ruff: noqa: SIM117
3
+import copy
4
+from typing import Generator, Tuple
5
+
6
+import torch
7
+from torch import nn
8
+
9
+from vllm.config import LoadConfig, ModelConfig, ParallelConfig, VllmConfig
10
+from vllm.logger import init_logger
11
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
12
+from vllm.model_executor.model_loader.tensorizer import (
13
+ TensorizerConfig, is_vllm_tensorized, load_with_tensorizer,
14
+ serialize_vllm_model, tensorizer_weights_iterator)
15
+from vllm.model_executor.model_loader.utils import (get_model_architecture,
16
+ initialize_model,
17
+ set_default_torch_dtype)
18
+
19
+logger = init_logger(__name__)
20
+
21
+
22
+class TensorizerLoader(BaseModelLoader):
23
+ """Model loader using CoreWeave's tensorizer library."""
24
+
25
+ def __init__(self, load_config: LoadConfig):
26
+ super().__init__(load_config)
27
+ if isinstance(load_config.model_loader_extra_config, TensorizerConfig):
28
+ self.tensorizer_config = load_config.model_loader_extra_config
29
+ else:
30
+ self.tensorizer_config = TensorizerConfig(
31
+ **load_config.model_loader_extra_config)
32
+
33
+ def _verify_config(self, model_config: ModelConfig,
34
+ parallel_config: ParallelConfig):
35
+ self.tensorizer_config.verify_with_model_config(model_config)
36
+ self.tensorizer_config.verify_with_parallel_config(parallel_config)
37
+
38
+ def _get_weights_iterator(
39
+ self, ) -> Generator[Tuple[str, torch.Tensor], None, None]:
40
+ tensorizer_args = self.tensorizer_config._construct_tensorizer_args()
41
+ return tensorizer_weights_iterator(tensorizer_args)
42
+
43
+ def _load_model_serialized_cpu(
44
+ self,
45
+ vllm_config: VllmConfig,
46
+ ) -> nn.Module:
47
+ """Load a serialized model with tensorizer to the CPU.
48
+
49
+ This is only necessary when the model isn't vLLM-tensorized (see
50
+ examples/other/tensorize_vllm_model.py) This should still
51
+ be faster than default HuggingFace loading, but will be slower than
52
+ loading a vLLM-tensorized model.
53
+ """
54
+ device_config = vllm_config.device_config
55
+ model_config = vllm_config.model_config
56
+ with set_default_torch_dtype(model_config.dtype):
57
+ with torch.device(device_config.device):
58
+ model = initialize_model(vllm_config=vllm_config)
59
+
60
+ model.load_weights(self._get_weights_iterator())
61
+ return model.eval()
62
+
63
+ def _load_model_serialized(
64
+ self,
65
+ vllm_config: VllmConfig,
66
+ ) -> nn.Module:
67
+ """Load a serialized model with tensorizer.
68
+
69
+ Expects a vLLM-tensorized model. See the
70
+ examples/other/tensorize_vllm_model.py example script
71
+ for serializing vLLM models."""
72
+
73
+ device_config = vllm_config.device_config
74
+ model_config = vllm_config.model_config
75
+
76
+ with set_default_torch_dtype(model_config.dtype):
77
+ with torch.device(device_config.device):
78
+ model_class = get_model_architecture(model_config)[0]
79
+
80
+ tensorizer_config = copy.copy(self.tensorizer_config)
81
+ tensorizer_config.model_class = model_class
82
+ tensorizer_config.hf_config = model_config.hf_config
83
+ tensorizer_config.dtype = model_config.dtype
84
+
85
+ model = load_with_tensorizer(tensorizer_config,
86
+ vllm_config=vllm_config)
87
+ return model.eval()
88
+
89
+ def download_model(self, model_config: ModelConfig) -> None:
90
+ self.tensorizer_config.verify_with_model_config(model_config)
91
+
92
+ with self.tensorizer_config.open_stream():
93
+ pass
94
+
95
+ def load_model(self, vllm_config: VllmConfig) -> nn.Module:
96
+ model_config = vllm_config.model_config
97
+ parallel_config = vllm_config.parallel_config
98
+ self._verify_config(model_config, parallel_config)
99
+
100
+ if parallel_config.tensor_parallel_size > 1:
101
+ from vllm.distributed import get_tensor_model_parallel_rank
102
+
103
+ self.tensorizer_config.tensorizer_uri = (
104
+ self.tensorizer_config.tensorizer_uri %
105
+ get_tensor_model_parallel_rank())
106
+
107
+ if is_vllm_tensorized(self.tensorizer_config):
108
+ return self._load_model_serialized(vllm_config=vllm_config)
109
+ return self._load_model_serialized_cpu(vllm_config=vllm_config)
110
+
111
+ @staticmethod
112
+ def save_model(
113
+ model: torch.nn.Module,
114
+ tensorizer_config: TensorizerConfig,
115
+ ) -> None:
116
+ serialize_vllm_model(
117
+ model=model,
118
+ tensorizer_config=tensorizer_config,
119
+ )
vllm/model_executor/model_loader/utils.py
+131
-2
@@ -1,6 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
"""Utilities for selecting and loading models."""
3
import contextlib
4
+import inspect
5
+import warnings
6
+from contextlib import contextmanager
7
from dataclasses import dataclass, field
8
from typing import Dict, List, Optional, Tuple, Type
9
@@ -9,14 +12,18 @@ import transformers
12
from torch import nn
13
from transformers.dynamic_module_utils import get_class_from_dynamic_module
14
12
-from vllm.config import ModelConfig, ModelImpl
15
+from vllm.attention import Attention
16
+from vllm.config import (ModelConfig, ModelImpl, VllmConfig,
17
+ set_current_vllm_config)
18
from vllm.logger import init_logger
19
+from vllm.model_executor.layers.linear import QKVCrossParallelLinear
20
from vllm.model_executor.layers.quantization.base_config import (
15
- QuantizationConfig)
21
+ QuantizationConfig, QuantizeMethodBase)
22
from vllm.model_executor.models import ModelRegistry
23
from vllm.model_executor.models.adapters import (as_classification_model,
24
as_embedding_model,
25
as_reward_model)
26
+from vllm.utils import is_pin_memory_available
27
28
logger = init_logger(__name__)
29
@@ -30,6 +37,128 @@ def set_default_torch_dtype(dtype: torch.dtype):
37
torch.set_default_dtype(old_dtype)
38
39
40
+def initialize_model(
41
+ vllm_config: VllmConfig,
42
+ *,
43
+ prefix: str = "",
44
+ model_class: Optional[type[nn.Module]] = None,
45
+) -> nn.Module:
46
+ """Initialize a model with the given configurations."""
47
+ model_config = vllm_config.model_config
48
+ if model_class is None:
49
+ model_class, _ = get_model_architecture(model_config)
50
+
51
+ if vllm_config.quant_config is not None:
52
+ configure_quant_config(vllm_config.quant_config, model_class)
53
+
54
+ signatures = inspect.signature(model_class.__init__)
55
+ all_params = [param.name for param in signatures.parameters.values()]
56
+ if "vllm_config" in all_params and "prefix" in all_params:
57
+ # new-style model class
58
+ with set_current_vllm_config(vllm_config, check_compile=True):
59
+ return model_class(vllm_config=vllm_config, prefix=prefix)
60
+
61
+ msg = ("vLLM model class should accept `vllm_config` and `prefix` as "
62
+ "input arguments. Possibly you have an old-style model class"
63
+ " registered from out of tree and it is used for new vLLM version. "
64
+ "Check https://docs.vllm.ai/en/latest/design/arch_overview.html "
65
+ "for the design and update the model class accordingly.")
66
+ warnings.warn(msg, DeprecationWarning, stacklevel=2)
67
+
68
+ logger.warning(
69
+ "Trying to guess the arguments for old-style model class %s",
70
+ model_class,
71
+ )
72
+ # try to be compatible with old-style model class
73
+ kwargs = {}
74
+ if "prefix" in all_params:
75
+ kwargs["prefix"] = prefix
76
+ if "config" in all_params:
77
+ kwargs["config"] = model_config.hf_config
78
+ if "cache_config" in all_params:
79
+ kwargs["cache_config"] = vllm_config.cache_config
80
+ if "quant_config" in all_params:
81
+ kwargs["quant_config"] = vllm_config.quant_config
82
+ if "lora_config" in all_params:
83
+ kwargs["lora_config"] = vllm_config.lora_config
84
+ if "scheduler_config" in all_params:
85
+ kwargs["scheduler_config"] = vllm_config.scheduler_config
86
+ with set_current_vllm_config(vllm_config, check_compile=True):
87
+ return model_class(**kwargs)
88
+
89
+
90
+def process_weights_after_loading(model: nn.Module, model_config: ModelConfig,
91
+ target_device: torch.device) -> None:
92
+ for _, module in model.named_modules():
93
+ if isinstance(module, QKVCrossParallelLinear):
94
+ # NOTE(Isotr0py): special case for cross QKV layer because
95
+ # q and kv proj aren't registered as submodules intentionally
96
+ module.process_weights_after_loading()
97
+ continue
98
+ quant_method = getattr(module, "quant_method", None)
99
+ if isinstance(quant_method, QuantizeMethodBase):
100
+ # When quant methods need to process weights after loading
101
+ # (for repacking, quantizing, etc), they expect parameters
102
+ # to be on the global target device. This scope is for the
103
+ # case where cpu offloading is used, where we will move the
104
+ # parameters onto device for processing and back off after.
105
+ with device_loading_context(module, target_device):
106
+ quant_method.process_weights_after_loading(module)
107
+
108
+ # Currently only used by MLA.
109
+ # NOTE: This intentionally happens after other modules so we can easily
110
+ # decompress the weights for MLA.
111
+ for _, module in model.named_modules():
112
+ if isinstance(module, Attention) and \
113
+ hasattr(module, "process_weights_after_loading"):
114
+ # TODO(lucas): see if there is a way to unify the signatures
115
+ # of process_weights_after_loading
116
+ module.process_weights_after_loading(model_config.dtype)
117
+
118
+
119
+@contextmanager
120
+def device_loading_context(module: torch.nn.Module,
121
+ target_device: torch.device):
122
+ if target_device.type == "cpu":
123
+ # If target is CPU, no need to move anything
124
+ yield module
125
+ return
126
+
127
+ original_device_states: Dict[str, torch.device] = {}
128
+
129
+ # Store original device states and move parameters to GPU if they're on CPU
130
+ for name, p in module.named_parameters():
131
+ if p.device.type == "cpu":
132
+ original_device_states[name] = p.device
133
+ p.data = p.data.to(target_device)
134
+ # Parameters already on target device are not touched
135
+
136
+ try:
137
+ yield module
138
+
139
+ finally:
140
+ # Restore parameters to their original devices, ignoring new parameters
141
+ pin_memory = is_pin_memory_available()
142
+ for name, p in module.named_parameters():
143
+ if name in original_device_states:
144
+ original_device: torch.device = original_device_states[name]
145
+ if original_device.type == "cpu":
146
+ # `torch.empty_like` does not support `pin_memory` argument
147
+ cpu_data = torch.empty_strided(
148
+ size=p.data.size(),
149
+ stride=p.data.stride(),
150
+ dtype=p.data.dtype,
151
+ layout=p.data.layout,
152
+ device="cpu",
153
+ pin_memory=pin_memory,
154
+ )
155
+ cpu_data.copy_(p.data)
156
+ p.data = cpu_data
157
+ else:
158
+ p.data = p.data.to(original_device)
159
+ # New parameters or parameters already on target device are untouched
160
+
161
+
162
def resolve_transformers_arch(model_config: ModelConfig,
163
architectures: list[str]):
164
for i, arch in enumerate(architectures):
vllm/model_executor/models/mllama4.py
+2
-2
@@ -37,7 +37,7 @@ from vllm.model_executor.layers.linear import (ColumnParallelLinear,
37
RowParallelLinear)
38
from vllm.model_executor.layers.quantization import QuantizationConfig
39
from vllm.model_executor.layers.rotary_embedding import get_rope
40
-from vllm.model_executor.model_loader.loader import _initialize_model
40
+from vllm.model_executor.model_loader.utils import initialize_model
41
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
42
from vllm.model_executor.sampling_metadata import SamplingMetadata
43
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -670,7 +670,7 @@ class Llama4ForConditionalGeneration(nn.Module, SupportsMultiModal,
670
self.config,
671
None,
672
prefix=maybe_prefix(prefix, "multi_modal_projector"))
673
- self.language_model = _initialize_model(
673
+ self.language_model = initialize_model(
674
vllm_config=vllm_config.with_hf_config(config.text_config,
675
["LlamaForCausalLM"]),
676
prefix=maybe_prefix(prefix, "language_model"),
vllm/model_executor/models/ultravox.py
+1
-1
@@ -17,7 +17,7 @@ from vllm.config import VllmConfig
17
from vllm.forward_context import get_forward_context
18
from vllm.model_executor.layers.activation import MulAndSilu, get_act_fn
19
from vllm.model_executor.layers.layernorm import RMSNorm
20
-from vllm.model_executor.model_loader.loader import DefaultModelLoader
20
+from vllm.model_executor.model_loader import DefaultModelLoader
21
from vllm.model_executor.models.module_mapping import MultiModelKeys
22
from vllm.model_executor.sampling_metadata import SamplingMetadata
23
from vllm.multimodal import MULTIMODAL_REGISTRY
vllm/model_executor/models/utils.py
+2
-2
@@ -273,7 +273,7 @@ def init_vllm_registered_model(
273
Helper function to initialize an inner model registered to vLLM,
274
based on the arguments passed to the outer vLLM model.
275
"""
276
- from vllm.model_executor.model_loader.loader import _initialize_model
276
+ from vllm.model_executor.model_loader.utils import initialize_model
277
278
if hf_config is None and architectures is not None:
279
# So that the architectures field is overridden
@@ -283,7 +283,7 @@ def init_vllm_registered_model(
283
vllm_config = vllm_config.with_hf_config(hf_config,
284
architectures=architectures)
285
286
- return _initialize_model(vllm_config=vllm_config, prefix=prefix)
286
+ return initialize_model(vllm_config=vllm_config, prefix=prefix)
287
288
289
@overload
vllm/v1/spec_decode/eagle.py
+1
-1
@@ -7,7 +7,7 @@ from vllm.config import (CompilationLevel, VllmConfig,
7
get_layers_from_vllm_config, set_current_vllm_config)
8
from vllm.forward_context import set_forward_context
9
from vllm.logger import init_logger
10
-from vllm.model_executor.model_loader.loader import get_model_loader
10
+from vllm.model_executor.model_loader import get_model_loader
11
from vllm.model_executor.model_loader.utils import set_default_torch_dtype
12
from vllm.model_executor.models import ModelRegistry
13
from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
vllm/v1/worker/gpu_worker.py
+1
-1
@@ -318,7 +318,7 @@ class Worker(WorkerBase):
318
pattern: Optional[str] = None,
319
max_size: Optional[int] = None,
320
) -> None:
321
- from vllm.model_executor.model_loader.loader import ShardedStateLoader
321
+ from vllm.model_executor.model_loader import ShardedStateLoader
322
ShardedStateLoader.save_model(
323
self.model_runner.model,
324
path,
vllm/worker/model_runner.py
+2
-2
@@ -1220,7 +1220,7 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
1220
pattern: Optional[str] = None,
1221
max_size: Optional[int] = None,
1222
) -> None:
1223
- from vllm.model_executor.model_loader.loader import ShardedStateLoader
1223
+ from vllm.model_executor.model_loader import ShardedStateLoader
1224
ShardedStateLoader.save_model(
1225
self.model,
1226
path,
@@ -1232,7 +1232,7 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
1232
self,
1233
tensorizer_config: TensorizerConfig,
1234
) -> None:
1235
- from vllm.model_executor.model_loader.loader import TensorizerLoader
1235
+ from vllm.model_executor.model_loader import TensorizerLoader
1236
TensorizerLoader.save_model(
1237
self.model,
1238
tensorizer_config=tensorizer_config,