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, 07:15 UTC
[Core] Support inplace model weights loading (#18745)
Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com>
22quinn committed
Jun 2, 2025 at 02:38 UTC
9760fd8f6acc752168333a25f5da80a6d2097608
13 files changed
+249
-297
tests/tensorizer_loader/test_tensorizer.py
-17
@@ -4,7 +4,6 @@ import gc
4
import os
5
import pathlib
6
import subprocess
7
-from unittest.mock import MagicMock, patch
7
8
import pytest
9
import torch
@@ -16,7 +15,6 @@ from vllm.engine.arg_utils import EngineArgs
15
from vllm.model_executor.model_loader.tensorizer import (TensorizerConfig,
16
TensorSerializer,
17
is_vllm_tensorized,
19
- load_with_tensorizer,
18
open_stream,
19
tensorize_vllm_model)
20
# yapf: enable
@@ -61,21 +59,6 @@ def write_keyfile(keyfile_path: str):
59
f.write(encryption_params.key)
60
61
64
-@patch('vllm.model_executor.model_loader.tensorizer.TensorizerAgent')
65
-def test_load_with_tensorizer(mock_agent, tensorizer_config):
66
- mock_linear_method = MagicMock()
67
- mock_agent_instance = mock_agent.return_value
68
- mock_agent_instance.deserialize.return_value = MagicMock()
69
-
70
- result = load_with_tensorizer(tensorizer_config,
71
- quant_method=mock_linear_method)
72
-
73
- mock_agent.assert_called_once_with(tensorizer_config,
74
- quant_method=mock_linear_method)
75
- mock_agent_instance.deserialize.assert_called_once()
76
- assert result == mock_agent_instance.deserialize.return_value
77
-
78
-
62
@pytest.mark.skipif(not is_curl_installed(), reason="cURL is not installed")
63
def test_can_deserialize_s3(vllm_runner):
64
model_ref = "EleutherAI/pythia-1.4b"
tests/v1/worker/test_gpu_model_runner.py
+18
@@ -94,6 +94,9 @@ def model_runner():
94
return runner
95
96
97
+model_runner_2 = model_runner
98
+
99
+
100
def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
101
new_reqs = []
102
num_scheduled_tokens = {}
@@ -366,3 +369,18 @@ def test_kv_cache_stride_order(monkeypatch, model_runner):
369
assert all(kv.is_contiguous() for kv in model_runner.kv_caches)
370
else:
371
assert all(not kv.is_contiguous() for kv in model_runner.kv_caches)
372
+
373
+
374
+def test_load_model_weights_inplace(dist_init, model_runner, model_runner_2):
375
+ # In this test, model_runner loads model + weights in one go, while
376
+ # model_runner_2 loads dummy weights first then load real weights inplace
377
+ model_runner.load_model()
378
+ original_load_format = model_runner_2.load_config.load_format
379
+ model_runner_2.load_config.load_format = "dummy"
380
+ model_runner_2.load_model() # Initial model loading with dummy weights
381
+ assert str(model_runner.get_model().state_dict()) != str(
382
+ model_runner_2.get_model().state_dict())
383
+ model_runner_2.load_config.load_format = original_load_format
384
+ model_runner_2.load_model() # Load real weights inplace
385
+ assert str(model_runner.get_model().state_dict()) == str(
386
+ model_runner_2.get_model().state_dict())
vllm/model_executor/model_loader/base_loader.py
+20
-2
@@ -1,9 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
from abc import ABC, abstractmethod
3
4
+import torch
5
import torch.nn as nn
6
7
from vllm.config import LoadConfig, ModelConfig, VllmConfig
8
+from vllm.model_executor.model_loader.utils import (
9
+ initialize_model, process_weights_after_loading, set_default_torch_dtype)
10
11
12
class BaseModelLoader(ABC):
@@ -18,7 +21,22 @@ class BaseModelLoader(ABC):
21
raise NotImplementedError
22
23
@abstractmethod
21
- def load_model(self, *, vllm_config: VllmConfig,
24
+ def load_weights(self, model: nn.Module,
25
+ model_config: ModelConfig) -> None:
26
+ """Load weights into a model. This standalone API allows
27
+ inplace weights loading for an already-initialized model"""
28
+ raise NotImplementedError
29
+
30
+ def load_model(self, vllm_config: VllmConfig,
31
model_config: ModelConfig) -> nn.Module:
32
"""Load a model with the given configurations."""
24
- raise NotImplementedError
33
+ device_config = vllm_config.device_config
34
+ target_device = torch.device(device_config.device)
35
+ with set_default_torch_dtype(model_config.dtype):
36
+ with target_device:
37
+ model = initialize_model(vllm_config=vllm_config,
38
+ model_config=model_config)
39
+ # Quantization does not happen in `load_weights` but after it
40
+ self.load_weights(model, model_config)
41
+ process_weights_after_loading(model, model_config, target_device)
42
+ return model.eval()
vllm/model_executor/model_loader/bitsandbytes_loader.py
+2
-16
@@ -14,7 +14,7 @@ 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
17
+from vllm.config import LoadConfig, ModelConfig
18
from vllm.distributed import (get_tensor_model_parallel_rank,
19
get_tensor_model_parallel_world_size)
20
# yapf: enable
@@ -28,7 +28,6 @@ from vllm.model_executor.layers.linear import (LinearBase,
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,
31
set_default_torch_dtype)
32
from vllm.model_executor.model_loader.weight_utils import (
33
download_safetensors_index_file_from_hf, download_weights_from_hf,
@@ -408,8 +407,7 @@ class BitsAndBytesModelLoader(BaseModelLoader):
407
), "vllm currently does not support BNB quantization for"
408
f" {type(model).__name__}"
409
411
- def _load_weights(self, model_config: ModelConfig,
412
- model: nn.Module) -> None:
410
+ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
411
if not hasattr(model, "load_weights"):
412
raise AttributeError(
413
"The required method 'load_weights' is not defined in class"
@@ -568,15 +566,3 @@ class BitsAndBytesModelLoader(BaseModelLoader):
566
567
def download_model(self, model_config: ModelConfig) -> None:
568
self._prepare_weights(model_config.model, model_config.revision)
571
-
572
- def load_model(self, vllm_config: VllmConfig,
573
- model_config: ModelConfig) -> nn.Module:
574
- device_config = vllm_config.device_config
575
- with set_default_torch_dtype(model_config.dtype):
576
- with torch.device(device_config.device):
577
-
578
- model = initialize_model(vllm_config=vllm_config)
579
-
580
- self._load_weights(model_config, model)
581
-
582
- return model.eval()
vllm/model_executor/model_loader/default_loader.py
+18
-32
@@ -12,11 +12,9 @@ from torch import nn
12
from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
13
14
from vllm import envs
15
-from vllm.config import LoadConfig, LoadFormat, ModelConfig, VllmConfig
15
+from vllm.config import LoadConfig, LoadFormat, ModelConfig
16
from vllm.logger import init_logger
17
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
18
-from vllm.model_executor.model_loader.utils import (
19
- initialize_model, process_weights_after_loading, set_default_torch_dtype)
18
from vllm.model_executor.model_loader.weight_utils import (
19
download_safetensors_index_file_from_hf, download_weights_from_hf,
20
fastsafetensors_weights_iterator, filter_duplicate_safetensors_files,
@@ -264,32 +262,20 @@ class DefaultModelLoader(BaseModelLoader):
262
fall_back_to_pt=True,
263
allow_patterns_overrides=None)
264
267
- def load_model(self, vllm_config: VllmConfig,
268
- model_config: ModelConfig) -> nn.Module:
269
- device_config = vllm_config.device_config
270
- target_device = torch.device(device_config.device)
271
- with set_default_torch_dtype(model_config.dtype):
272
- with target_device:
273
- model = initialize_model(vllm_config=vllm_config,
274
- model_config=model_config)
275
-
276
- weights_to_load = {name for name, _ in model.named_parameters()}
277
- loaded_weights = model.load_weights(
278
- self.get_all_weights(model_config, model))
279
- self.counter_after_loading_weights = time.perf_counter()
280
- logger.info(
281
- "Loading weights took %.2f seconds",
282
- self.counter_after_loading_weights -
283
- self.counter_before_loading_weights)
284
- # We only enable strict check for non-quantized models
285
- # that have loaded weights tracking currently.
286
- if model_config.quantization is None and loaded_weights is not None:
287
- weights_not_loaded = weights_to_load - loaded_weights
288
- if weights_not_loaded:
289
- raise ValueError(
290
- "Following weights were not initialized from "
291
- f"checkpoint: {weights_not_loaded}")
292
-
293
- process_weights_after_loading(model, model_config, target_device)
294
-
295
- return model.eval()
265
+ def load_weights(self, model: nn.Module,
266
+ model_config: ModelConfig) -> None:
267
+ weights_to_load = {name for name, _ in model.named_parameters()}
268
+ loaded_weights = model.load_weights(
269
+ self.get_all_weights(model_config, model))
270
+ self.counter_after_loading_weights = time.perf_counter()
271
+ logger.info(
272
+ "Loading weights took %.2f seconds",
273
+ self.counter_after_loading_weights -
274
+ self.counter_before_loading_weights)
275
+ # We only enable strict check for non-quantized models
276
+ # that have loaded weights tracking currently.
277
+ if model_config.quantization is None and loaded_weights is not None:
278
+ weights_not_loaded = weights_to_load - loaded_weights
279
+ if weights_not_loaded:
280
+ raise ValueError("Following weights were not initialized from "
281
+ f"checkpoint: {weights_not_loaded}")
vllm/model_executor/model_loader/dummy_loader.py
+6
-17
@@ -1,11 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
-import torch
2
import torch.nn as nn
3
5
-from vllm.config import LoadConfig, ModelConfig, VllmConfig
4
+from vllm.config import LoadConfig, ModelConfig
5
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)
6
from vllm.model_executor.model_loader.weight_utils import (
7
initialize_dummy_weights)
8
@@ -22,16 +19,8 @@ class DummyModelLoader(BaseModelLoader):
19
def download_model(self, model_config: ModelConfig) -> None:
20
pass # Nothing to download
21
25
- def load_model(self, vllm_config: VllmConfig,
26
- model_config: ModelConfig) -> nn.Module:
27
- device_config = vllm_config.device_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()
22
+ def load_weights(self, model: nn.Module,
23
+ model_config: ModelConfig) -> None:
24
+ # NOTE(woosuk): For accurate performance evaluation, we assign
25
+ # random values to the weights.
26
+ initialize_dummy_weights(model)
vllm/model_executor/model_loader/gguf_loader.py
+8
-2
@@ -92,6 +92,13 @@ class GGUFModelLoader(BaseModelLoader):
92
def download_model(self, model_config: ModelConfig) -> None:
93
self._prepare_weights(model_config.model)
94
95
+ def load_weights(self, model: nn.Module,
96
+ model_config: ModelConfig) -> None:
97
+ local_model_path = self._prepare_weights(model_config.model)
98
+ gguf_weights_map = self._get_gguf_weights_map(model_config)
99
+ model.load_weights(
100
+ self._get_weights_iterator(local_model_path, gguf_weights_map))
101
+
102
def load_model(self, vllm_config: VllmConfig,
103
model_config: ModelConfig) -> nn.Module:
104
device_config = vllm_config.device_config
@@ -106,8 +113,7 @@ class GGUFModelLoader(BaseModelLoader):
113
with set_default_torch_dtype(model_config.dtype):
114
with target_device:
115
model = initialize_model(vllm_config=vllm_config)
109
- model.load_weights(
110
- self._get_weights_iterator(local_model_path, gguf_weights_map))
116
+ self.load_weights(model, model_config)
117
118
process_weights_after_loading(model, model_config, target_device)
119
return model
vllm/model_executor/model_loader/runai_streamer_loader.py
+9
-21
@@ -9,10 +9,8 @@ import torch
9
from torch import nn
10
from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
11
12
-from vllm.config import LoadConfig, ModelConfig, VllmConfig
12
+from vllm.config import LoadConfig, ModelConfig
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)
14
from vllm.model_executor.model_loader.weight_utils import (
15
download_safetensors_index_file_from_hf, download_weights_from_hf,
16
runai_safetensors_weights_iterator)
@@ -100,21 +98,11 @@ class RunaiModelStreamerLoader(BaseModelLoader):
98
"""Download model if necessary"""
99
self._prepare_weights(model_config.model, model_config.revision)
100
103
- def load_model(self, vllm_config: VllmConfig,
104
- model_config: ModelConfig) -> nn.Module:
105
- """Perform streaming of the model to destination"""
106
- device_config = vllm_config.device_config
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()
101
+ def load_weights(self, model: nn.Module,
102
+ model_config: ModelConfig) -> None:
103
+ """Load weights into a model."""
104
+ model_weights = model_config.model
105
+ if hasattr(model_config, "model_weights"):
106
+ model_weights = model_config.model_weights
107
+ model.load_weights(
108
+ self._get_weights_iterator(model_weights, model_config.revision))
vllm/model_executor/model_loader/sharded_state_loader.py
+44
-55
@@ -9,11 +9,9 @@ from typing import Any, Optional
9
import torch
10
from torch import nn
11
12
-from vllm.config import LoadConfig, ModelConfig, VllmConfig
12
+from vllm.config import LoadConfig, ModelConfig
13
from vllm.logger import init_logger
14
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
15
-from vllm.model_executor.model_loader.utils import (
16
- initialize_model, process_weights_after_loading, set_default_torch_dtype)
15
from vllm.model_executor.model_loader.weight_utils import (
16
download_weights_from_hf, runai_safetensors_weights_iterator)
17
from vllm.transformers_utils.s3_utils import glob as s3_glob
@@ -100,11 +98,8 @@ class ShardedStateLoader(BaseModelLoader):
98
def download_model(self, model_config: ModelConfig) -> None:
99
self._prepare_weights(model_config.model, model_config.revision)
100
103
- def load_model(self, vllm_config: VllmConfig,
104
- model_config: ModelConfig) -> nn.Module:
105
- device_config = vllm_config.device_config
106
- target_device = torch.device(device_config.device)
107
-
101
+ def load_weights(self, model: nn.Module,
102
+ model_config: ModelConfig) -> None:
103
from vllm.distributed import get_tensor_model_parallel_rank
104
105
model_weights = model_config.model
@@ -112,53 +107,47 @@ class ShardedStateLoader(BaseModelLoader):
107
model_weights = model_config.model_weights
108
local_model_path = model_weights
109
115
- with set_default_torch_dtype(model_config.dtype):
116
- with target_device:
117
- model = initialize_model(vllm_config=vllm_config)
118
- process_weights_after_loading(model, model_config,
119
- target_device)
120
- rank = get_tensor_model_parallel_rank()
121
- pattern = os.path.join(
122
- local_model_path,
123
- self.pattern.format(rank=rank, part="*"),
124
- )
125
-
126
- filepaths = []
127
- if is_s3(local_model_path):
128
- file_pattern = f"*{self.pattern.format(rank=rank, part=' * ')}"
129
- filepaths = s3_glob(path=local_model_path,
130
- allow_pattern=[file_pattern])
131
- else:
132
- filepaths = glob.glob(pattern)
133
- if not filepaths:
134
- # TODO: support un-sharded checkpoints too
135
- raise ValueError(
136
- f"Could not find checkpoint files '{pattern}', only "
137
- f"pre-sharded checkpoints are currently supported!")
138
- state_dict = self._filter_subtensors(model.state_dict())
139
- for key, tensor in self.iterate_over_files(filepaths):
140
- # If loading with LoRA enabled, additional padding may
141
- # be added to certain parameters. We only load into a
142
- # narrowed view of the parameter data.
143
- param_data = state_dict[key].data
144
- param_shape = state_dict[key].shape
145
- for dim, size in enumerate(tensor.shape):
146
- if size < param_shape[dim]:
147
- param_data = param_data.narrow(dim, 0, size)
148
- if tensor.shape != param_shape:
149
- logger.warning(
150
- "loading tensor of shape %s into "
151
- "parameter '%s' of shape %s",
152
- tensor.shape,
153
- key,
154
- param_shape,
155
- )
156
- param_data.copy_(tensor)
157
- state_dict.pop(key)
158
- if state_dict:
159
- raise ValueError(
160
- f"Missing keys {tuple(state_dict)} in loaded state!")
161
- return model.eval()
110
+ rank = get_tensor_model_parallel_rank()
111
+ pattern = os.path.join(
112
+ local_model_path,
113
+ self.pattern.format(rank=rank, part="*"),
114
+ )
115
+
116
+ filepaths = []
117
+ if is_s3(local_model_path):
118
+ file_pattern = f"*{self.pattern.format(rank=rank, part=' * ')}"
119
+ filepaths = s3_glob(path=local_model_path,
120
+ allow_pattern=[file_pattern])
121
+ else:
122
+ filepaths = glob.glob(pattern)
123
+ if not filepaths:
124
+ # TODO: support un-sharded checkpoints too
125
+ raise ValueError(
126
+ f"Could not find checkpoint files '{pattern}', only "
127
+ f"pre-sharded checkpoints are currently supported!")
128
+ state_dict = self._filter_subtensors(model.state_dict())
129
+ for key, tensor in self.iterate_over_files(filepaths):
130
+ # If loading with LoRA enabled, additional padding may
131
+ # be added to certain parameters. We only load into a
132
+ # narrowed view of the parameter data.
133
+ param_data = state_dict[key].data
134
+ param_shape = state_dict[key].shape
135
+ for dim, size in enumerate(tensor.shape):
136
+ if size < param_shape[dim]:
137
+ param_data = param_data.narrow(dim, 0, size)
138
+ if tensor.shape != param_shape:
139
+ logger.warning(
140
+ "loading tensor of shape %s into "
141
+ "parameter '%s' of shape %s",
142
+ tensor.shape,
143
+ key,
144
+ param_shape,
145
+ )
146
+ param_data.copy_(tensor)
147
+ state_dict.pop(key)
148
+ if state_dict:
149
+ raise ValueError(
150
+ f"Missing keys {tuple(state_dict)} in loaded state!")
151
152
def iterate_over_files(
153
self, paths) -> Generator[tuple[str, torch.Tensor], None, None]:
vllm/model_executor/model_loader/tensorizer.py
+66
-99
@@ -21,7 +21,8 @@ from torch.utils._python_dispatch import TorchDispatchMode
21
from transformers import PretrainedConfig
22
23
import vllm.envs as envs
24
-from vllm.config import ModelConfig, ParallelConfig, set_current_vllm_config
24
+from vllm.config import (ModelConfig, ParallelConfig, VllmConfig,
25
+ set_current_vllm_config)
26
from vllm.engine.arg_utils import EngineArgs
27
from vllm.logger import init_logger
28
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -208,12 +209,6 @@ class TensorizerConfig:
209
**tensorizer_args.stream_params)
210
211
211
-def load_with_tensorizer(tensorizer_config: TensorizerConfig,
212
- **extra_kwargs) -> nn.Module:
213
- tensorizer = TensorizerAgent(tensorizer_config, **extra_kwargs)
214
- return tensorizer.deserialize()
215
-
216
-
212
@dataclass
213
class TensorizerArgs:
214
tensorizer_uri: Union[io.BufferedIOBase, io.RawIOBase, BinaryIO, str,
@@ -366,100 +361,72 @@ class TensorizerArgs:
361
return tensorizer_args
362
363
369
-class TensorizerAgent:
370
- """
371
- A class for performing tensorizer deserializations specifically for
372
- vLLM models using plaid_mode. Uses TensorizerArgs to configure the
373
- behavior of the TensorDeserializer when loading tensors from a serialized
374
- model. For deserializations of HuggingFace models, TensorDeserializer is
375
- instead used as an iterator directly in the func hf_model_weights_iterator
376
- in vllm/model_executor/model_loader/weight_utils.py
377
- """
378
-
379
- def __init__(self, tensorizer_config: TensorizerConfig, vllm_config):
380
- self.tensorizer_config = tensorizer_config
381
- self.tensorizer_args = (
382
- self.tensorizer_config._construct_tensorizer_args())
383
- self.vllm_config = vllm_config
384
- self.model = self._init_model()
385
-
386
- def _init_model(self):
387
- assert self.tensorizer_config.hf_config is not None
388
- model_args = self.tensorizer_config.hf_config
389
- model_args.torch_dtype = self.tensorizer_config.dtype
390
- assert self.tensorizer_config.model_class is not None
391
- # TODO: Do we need to consider old-style model class?
392
- with meta_tensor_mode(), set_current_vllm_config(self.vllm_config,
393
- check_compile=True):
394
- return self.tensorizer_config.model_class(
395
- vllm_config=self.vllm_config)
396
-
397
- def _resize_lora_embeddings(self):
398
- """Modify LoRA embedding layers to use bigger tensors
399
- to allow for adapter added tokens."""
400
- for child in self.model.modules():
401
- if (isinstance(child, VocabParallelEmbedding)
402
- and child.weight.shape[0]
403
- < child.num_embeddings_per_partition):
404
- new_weight = torch.empty(child.num_embeddings_per_partition,
405
- child.embedding_dim,
406
- dtype=child.weight.dtype,
407
- device=child.weight.device)
408
- new_weight[:child.weight.shape[0]].copy_(child.weight.data)
409
- new_weight[child.weight.shape[0]:].fill_(0)
410
- child.weight.data = new_weight
411
-
412
- def _check_tensors_on_meta_device(self):
413
- for tensor in self.model.state_dict().values():
414
- if tensor.device.type == 'meta':
415
- raise ValueError(
416
- "The serialized model contains tensors on the meta device,"
417
- " indicating that some tensors were not loaded properly."
418
- " Please check that the parameters of the model being"
419
- " specified match that of the serialized model, such as"
420
- " its quantization.")
421
-
422
- def deserialize(self):
423
- """
424
- Deserialize the model using the TensorDeserializer. This method is
425
- specifically for vLLM models using tensorizer's plaid_mode.
426
-
427
- The deserializer makes use of tensorizer_args.stream_params
428
- to configure the behavior of the stream when loading tensors from a
429
- serialized model. The deserializer_params are used to configure the
430
- behavior of the TensorDeserializer when loading tensors themselves.
431
- Documentation on these params can be found in TensorizerArgs
432
-
433
- Returns:
434
- nn.Module: The deserialized model.
435
- """
436
- before_mem = get_mem_usage()
437
- start = time.perf_counter()
438
- with _read_stream(
439
- self.tensorizer_config.tensorizer_uri,
440
- **self.tensorizer_args.stream_params
441
- ) as stream, TensorDeserializer(
364
+def _check_tensors_on_meta_device(model: nn.Module) -> None:
365
+ for tensor in model.state_dict().values():
366
+ if tensor.device.type == 'meta':
367
+ raise ValueError(
368
+ "The serialized model contains tensors on the meta device,"
369
+ " indicating that some tensors were not loaded properly."
370
+ " Please check that the parameters of the model being"
371
+ " specified match that of the serialized model, such as"
372
+ " its quantization.")
373
+
374
+
375
+def _resize_lora_embeddings(model: nn.Module):
376
+ """Modify LoRA embedding layers to use bigger tensors
377
+ to allow for adapter added tokens."""
378
+ for child in model.modules():
379
+ if (isinstance(child, VocabParallelEmbedding) and child.weight.shape[0]
380
+ < child.num_embeddings_per_partition):
381
+ new_weight = torch.empty(child.num_embeddings_per_partition,
382
+ child.embedding_dim,
383
+ dtype=child.weight.dtype,
384
+ device=child.weight.device)
385
+ new_weight[:child.weight.shape[0]].copy_(child.weight.data)
386
+ new_weight[child.weight.shape[0]:].fill_(0)
387
+ child.weight.data = new_weight
388
+
389
+
390
+def init_tensorizer_model(tensorizer_config: TensorizerConfig,
391
+ vllm_config: VllmConfig) -> nn.Module:
392
+ assert tensorizer_config.hf_config is not None
393
+ model_args = tensorizer_config.hf_config
394
+ model_args.torch_dtype = tensorizer_config.dtype
395
+ assert tensorizer_config.model_class is not None
396
+ # TODO: Do we need to consider old-style model class?
397
+ with meta_tensor_mode(), set_current_vllm_config(vllm_config,
398
+ check_compile=True):
399
+ return tensorizer_config.model_class(vllm_config=vllm_config)
400
+
401
+
402
+def deserialize_tensorizer_model(model: nn.Module,
403
+ tensorizer_config: TensorizerConfig) -> None:
404
+ tensorizer_args = tensorizer_config._construct_tensorizer_args()
405
+ before_mem = get_mem_usage()
406
+ start = time.perf_counter()
407
+ with _read_stream(
408
+ tensorizer_config.tensorizer_uri,
409
+ **tensorizer_args.stream_params) as stream, TensorDeserializer(
410
stream,
443
- dtype=self.tensorizer_config.dtype,
411
+ dtype=tensorizer_config.dtype,
412
device=f'cuda:{torch.cuda.current_device()}',
445
- **self.tensorizer_args.deserializer_params) as deserializer:
446
- deserializer.load_into_module(self.model)
447
- end = time.perf_counter()
448
-
449
- total_bytes_str = convert_bytes(deserializer.total_tensor_bytes)
450
- duration = end - start
451
- per_second = convert_bytes(deserializer.total_tensor_bytes / duration)
452
- after_mem = get_mem_usage()
453
- deserializer.close()
454
- logger.info("Deserialized %s in %0.2fs, %s/s", total_bytes_str,
455
- end - start, per_second)
456
- logger.info("Memory usage before: %s", before_mem)
457
- logger.info("Memory usage after: %s", after_mem)
458
-
459
- self._check_tensors_on_meta_device()
460
- self._resize_lora_embeddings()
461
- del self.model.vllm_tensorized_marker
462
- return self.model.eval()
413
+ **tensorizer_args.deserializer_params) as deserializer:
414
+ deserializer.load_into_module(model)
415
+ end = time.perf_counter()
416
+
417
+ total_bytes_str = convert_bytes(deserializer.total_tensor_bytes)
418
+ duration = end - start
419
+ per_second = convert_bytes(deserializer.total_tensor_bytes / duration)
420
+ after_mem = get_mem_usage()
421
+ deserializer.close()
422
+ logger.info("Deserialized %s in %0.2fs, %s/s", total_bytes_str,
423
+ end - start, per_second)
424
+ logger.info("Memory usage before: %s", before_mem)
425
+ logger.info("Memory usage after: %s", after_mem)
426
+
427
+ _check_tensors_on_meta_device(model)
428
+ _resize_lora_embeddings(model)
429
+ del model.vllm_tensorized_marker
430
431
432
def tensorizer_weights_iterator(
vllm/model_executor/model_loader/tensorizer_loader.py
+29
-29
@@ -11,8 +11,8 @@ from vllm.config import LoadConfig, ModelConfig, ParallelConfig, VllmConfig
11
from vllm.logger import init_logger
12
from vllm.model_executor.model_loader.base_loader import BaseModelLoader
13
from vllm.model_executor.model_loader.tensorizer import (
14
- TensorizerConfig, is_vllm_tensorized, load_with_tensorizer,
15
- serialize_vllm_model, tensorizer_weights_iterator)
14
+ TensorizerConfig, deserialize_tensorizer_model, init_tensorizer_model,
15
+ is_vllm_tensorized, serialize_vllm_model, tensorizer_weights_iterator)
16
from vllm.model_executor.model_loader.utils import (get_model_architecture,
17
initialize_model,
18
set_default_torch_dtype)
@@ -61,38 +61,34 @@ class TensorizerLoader(BaseModelLoader):
61
model.load_weights(self._get_weights_iterator())
62
return model.eval()
63
64
- def _load_model_serialized(
65
- self,
66
- vllm_config: VllmConfig,
67
- ) -> nn.Module:
68
- """Load a serialized model with tensorizer.
69
-
70
- Expects a vLLM-tensorized model. See the
71
- examples/others/tensorize_vllm_model.py example script
72
- for serializing vLLM models."""
73
-
74
- device_config = vllm_config.device_config
75
- model_config = vllm_config.model_config
76
-
77
- with set_default_torch_dtype(model_config.dtype):
78
- with torch.device(device_config.device):
79
- model_class = get_model_architecture(model_config)[0]
80
-
81
- tensorizer_config = copy.copy(self.tensorizer_config)
82
- tensorizer_config.model_class = model_class
83
- tensorizer_config.hf_config = model_config.hf_config
84
- tensorizer_config.dtype = model_config.dtype
85
-
86
- model = load_with_tensorizer(tensorizer_config,
87
- vllm_config=vllm_config)
88
- return model.eval()
89
-
64
def download_model(self, model_config: ModelConfig) -> None:
65
self.tensorizer_config.verify_with_model_config(model_config)
66
67
with self.tensorizer_config.open_stream():
68
pass
69
70
+ def _patch_tensorizer_config(
71
+ self, model_config: ModelConfig) -> TensorizerConfig:
72
+ model_class = get_model_architecture(model_config)[0]
73
+ tensorizer_config = copy.copy(self.tensorizer_config)
74
+ tensorizer_config.model_class = model_class
75
+ tensorizer_config.hf_config = model_config.hf_config
76
+ tensorizer_config.dtype = model_config.dtype
77
+ return tensorizer_config
78
+
79
+ def load_weights(self, model: nn.Module,
80
+ model_config: ModelConfig) -> None:
81
+ """Load serialized model weights with tensorizer.
82
+
83
+ Expects a vLLM-tensorized model. See the
84
+ examples/others/tensorize_vllm_model.py example script
85
+ for serializing vLLM models."""
86
+ if is_vllm_tensorized(self.tensorizer_config):
87
+ tensorizer_config = self._patch_tensorizer_config(model_config)
88
+ deserialize_tensorizer_model(model, tensorizer_config)
89
+ else:
90
+ model.load_weights(self._get_weights_iterator())
91
+
92
def load_model(self, vllm_config: VllmConfig,
93
model_config: ModelConfig) -> nn.Module:
94
parallel_config = vllm_config.parallel_config
@@ -106,7 +102,11 @@ class TensorizerLoader(BaseModelLoader):
102
get_tensor_model_parallel_rank())
103
104
if is_vllm_tensorized(self.tensorizer_config):
109
- return self._load_model_serialized(vllm_config=vllm_config)
105
+ tensorizer_config = self._patch_tensorizer_config(model_config)
106
+ model = init_tensorizer_model(tensorizer_config=tensorizer_config,
107
+ vllm_config=vllm_config)
108
+ self.load_weights(model, model_config)
109
+ return model
110
return self._load_model_serialized_cpu(vllm_config=vllm_config)
111
112
@staticmethod
vllm/v1/worker/gpu_model_runner.py
+13
-2
@@ -28,7 +28,7 @@ from vllm.forward_context import (DPMetadata, get_forward_context,
28
set_forward_context)
29
from vllm.logger import init_logger
30
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
31
-from vllm.model_executor.model_loader import TensorizerLoader, get_model
31
+from vllm.model_executor.model_loader import TensorizerLoader, get_model_loader
32
from vllm.multimodal import MULTIMODAL_REGISTRY
33
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
34
from vllm.multimodal.utils import group_mm_inputs_by_modality
@@ -1564,7 +1564,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1564
logger.info("Starting to load model %s...", self.model_config.model)
1565
with DeviceMemoryProfiler() as m: # noqa: SIM117
1566
time_before_load = time.perf_counter()
1567
- self.model = get_model(vllm_config=self.vllm_config)
1567
+ model_loader = get_model_loader(self.load_config)
1568
+ if not hasattr(self, "model"):
1569
+ logger.info("Loading model from scratch...")
1570
+ self.model = model_loader.load_model(
1571
+ vllm_config=self.vllm_config,
1572
+ model_config=self.model_config)
1573
+ else:
1574
+ logger.info(
1575
+ "Model was already initialized. Loading weights inplace..."
1576
+ )
1577
+ model_loader.load_weights(self.model,
1578
+ model_config=self.model_config)
1579
if self.lora_config:
1580
self.model = self.load_lora_model(self.model,
1581
self.model_config,
vllm/v1/worker/tpu_model_runner.py
+16
-5
@@ -21,7 +21,7 @@ from vllm.config import VllmConfig, get_layers_from_vllm_config
21
from vllm.forward_context import set_forward_context
22
from vllm.logger import init_logger
23
from vllm.lora.layers import BaseLayerWithLoRA
24
-from vllm.model_executor.model_loader import get_model
24
+from vllm.model_executor.model_loader import get_model_loader
25
from vllm.multimodal import MULTIMODAL_REGISTRY
26
from vllm.multimodal.inputs import (BatchedTensorInputs, MultiModalKwargs,
27
PlaceholderRange)
@@ -171,7 +171,7 @@ class TPUModelRunner(LoRAModelRunnerMixin):
171
self.encoder_cache_size = encoder_cache_size
172
173
# Lazy initialization
174
- # self.model: nn.Module # Set after load_model
174
+ self.model: nn.Module # Set after load_model
175
self.kv_caches: list[torch.Tensor] = []
176
# req_id -> (input_id -> encoder_output)
177
self.encoder_cache: dict[str, dict[int, torch.Tensor]] = {}
@@ -419,7 +419,6 @@ class TPUModelRunner(LoRAModelRunnerMixin):
419
return len(unscheduled_req_ids) > 0 or len(req_ids_to_add) > 0
420
421
def get_model(self) -> nn.Module:
422
- assert self.model is not None
422
return self.model
423
424
def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
@@ -936,7 +935,18 @@ class TPUModelRunner(LoRAModelRunnerMixin):
935
"vllm.model_executor.layers.vocab_parallel_embedding."
936
"get_tensor_model_parallel_rank",
937
return_value=xm_tp_rank):
939
- model = get_model(vllm_config=self.vllm_config)
938
+ # model = get_model(vllm_config=self.vllm_config)
939
+ model_loader = get_model_loader(self.load_config)
940
+ if not hasattr(self, "model"):
941
+ logger.info("Loading model from scratch...")
942
+ model = model_loader.load_model(vllm_config=self.vllm_config,
943
+ model_config=self.model_config)
944
+ else:
945
+ logger.info(
946
+ "Model was already initialized. Loading weights inplace..."
947
+ )
948
+ model_loader.load_weights(self.model,
949
+ model_config=self.model_config)
950
if self.lora_config is not None:
951
model = self.load_lora_model(model, self.model_config,
952
self.scheduler_config,
@@ -947,7 +957,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):
957
# loading.
958
xm.mark_step()
959
xm.wait_device_ops()
950
- self.model = model
960
+ if not hasattr(self, "model"):
961
+ self.model = model
962
self.sampler = TPUSampler()
963
964
@torch.no_grad()