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, 08:30 UTC
[Frontend] [Core] Add Tensorizer support for V1, LoRA adapter serialization and deserialization (#17926)
Signed-off-by: Sanger Steel <sangersteel@gmail.com>
Sanger Steel committed
May 22, 2025 at 21:44 UTC
c32e249a23169353ccc02d7c6099a8c90ca4bbf6
16 files changed
+608
-199
.buildkite/test-pipeline.yaml
+1
@@ -128,6 +128,7 @@ steps:
128
- VLLM_USE_V1=0 pytest -v -s entrypoints/llm/test_guided_generate.py # it needs a clean process
129
- pytest -v -s entrypoints/openai --ignore=entrypoints/openai/test_oot_registration.py --ignore=entrypoints/openai/test_chat_with_tool_reasoning.py --ignore=entrypoints/openai/correctness/
130
- pytest -v -s entrypoints/test_chat_utils.py
131
+ - pytest -v -s entrypoints/openai/test_tensorizer_entrypoint.py
132
- VLLM_USE_V1=0 pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests
133
134
- label: Distributed Tests (4 GPUs) # 10min
examples/other/tensorize_vllm_model.py
+82
-14
@@ -6,11 +6,12 @@ import json
6
import os
7
import uuid
8
9
-from vllm import LLM
9
+from vllm import LLM, SamplingParams
10
from vllm.engine.arg_utils import EngineArgs
11
-from vllm.model_executor.model_loader.tensorizer import (TensorizerArgs,
12
- TensorizerConfig,
13
- tensorize_vllm_model)
11
+from vllm.lora.request import LoRARequest
12
+from vllm.model_executor.model_loader.tensorizer import (
13
+ TensorizerArgs, TensorizerConfig, tensorize_lora_adapter,
14
+ tensorize_vllm_model)
15
from vllm.utils import FlexibleArgumentParser
16
17
# yapf conflicts with isort for this docstring
@@ -27,7 +28,7 @@ https://github.com/coreweave/tensorizer
28
To serialize a model, install vLLM from source, then run something
29
like this from the root level of this repository:
30
30
-python -m examples.other.tensorize_vllm_model \
31
+python examples/other/tensorize_vllm_model.py \
32
--model facebook/opt-125m \
33
serialize \
34
--serialized-directory s3://my-bucket \
@@ -47,7 +48,7 @@ providing a `--keyfile` argument.
48
To deserialize a model, you can run something like this from the root
49
level of this repository:
50
50
-python -m examples.other.tensorize_vllm_model \
51
+python examples/other/tensorize_vllm_model.py \
52
--model EleutherAI/gpt-j-6B \
53
--dtype float16 \
54
deserialize \
@@ -69,7 +70,7 @@ For more information on the available arguments for serializing, run
70
71
Or for deserializing:
72
72
-`python -m examples.other.tensorize_vllm_model deserialize --help`.
73
+`python examples/other/tensorize_vllm_model.py deserialize --help`.
74
75
Once a model is serialized, tensorizer can be invoked with the `LLM` class
76
directly to load models:
@@ -90,11 +91,27 @@ TensorizerConfig arguments desired.
91
In order to see all of the available arguments usable to configure
92
loading with tensorizer that are given to `TensorizerConfig`, run:
93
93
-`python -m examples.other.tensorize_vllm_model deserialize --help`
94
+`python examples/other/tensorize_vllm_model.py deserialize --help`
95
96
under the `tensorizer options` section. These can also be used for
97
deserialization in this example script, although `--tensorizer-uri` and
98
`--path-to-tensors` are functionally the same in this case.
99
+
100
+Tensorizer can also be used to save and load LoRA adapters. A LoRA adapter
101
+can be serialized directly with the path to the LoRA adapter on HF Hub and
102
+a TensorizerConfig object. In this script, passing a HF id to a LoRA adapter
103
+will serialize the LoRA adapter artifacts to `--serialized-directory`.
104
+
105
+You can then use the LoRA adapter with `vllm serve`, for instance, by ensuring
106
+the LoRA artifacts are in your model artifacts directory and specifying
107
+`--enable-lora`. For instance:
108
+
109
+```
110
+vllm serve <model_path> \
111
+ --load-format tensorizer \
112
+ --model-loader-extra-config '{"tensorizer_uri": "<model_path>.tensors"}' \
113
+ --enable-lora
114
+```
115
"""
116
117
@@ -107,6 +124,19 @@ def parse_args():
124
"also supported, although libsodium must be installed to "
125
"use it.")
126
parser = EngineArgs.add_cli_args(parser)
127
+
128
+ parser.add_argument(
129
+ "--lora-path",
130
+ type=str,
131
+ required=False,
132
+ help="Path to a LoRA adapter to "
133
+ "serialize along with model tensors. This can then be deserialized "
134
+ "along with the model by passing a tensorizer_config kwarg to "
135
+ "LoRARequest with type TensorizerConfig. See the docstring for this "
136
+ "for a usage example."
137
+
138
+ )
139
+
140
subparsers = parser.add_subparsers(dest='command')
141
142
serialize_parser = subparsers.add_parser(
@@ -169,11 +199,42 @@ def parse_args():
199
200
201
def deserialize():
172
- llm = LLM(model=args.model,
173
- load_format="tensorizer",
174
- tensor_parallel_size=args.tensor_parallel_size,
175
- model_loader_extra_config=tensorizer_config
176
- )
202
+ if args.lora_path:
203
+ tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir
204
+ llm = LLM(model=args.model,
205
+ load_format="tensorizer",
206
+ tensor_parallel_size=args.tensor_parallel_size,
207
+ model_loader_extra_config=tensorizer_config,
208
+ enable_lora=True,
209
+ )
210
+ sampling_params = SamplingParams(
211
+ temperature=0,
212
+ max_tokens=256,
213
+ stop=["[/assistant]"]
214
+ )
215
+
216
+ # Truncating this as the extra text isn't necessary
217
+ prompts = [
218
+ "[user] Write a SQL query to answer the question based on ..."
219
+ ]
220
+
221
+ # Test LoRA load
222
+ print(
223
+ llm.generate(
224
+ prompts,
225
+ sampling_params,
226
+ lora_request=LoRARequest("sql-lora",
227
+ 1,
228
+ args.lora_path,
229
+ tensorizer_config = tensorizer_config)
230
+ )
231
+ )
232
+ else:
233
+ llm = LLM(model=args.model,
234
+ load_format="tensorizer",
235
+ tensor_parallel_size=args.tensor_parallel_size,
236
+ model_loader_extra_config=tensorizer_config
237
+ )
238
return llm
239
240
@@ -197,7 +258,10 @@ if __name__ == '__main__':
258
259
model_name = model_ref.split("/")[1]
260
200
- keyfile = args.keyfile if args.keyfile else None
261
+ if args.command == "serialize" or args.command == "deserialize":
262
+ keyfile = args.keyfile
263
+ else:
264
+ keyfile = None
265
266
if args.model_loader_extra_config:
267
config = json.loads(args.model_loader_extra_config)
@@ -228,6 +292,10 @@ if __name__ == '__main__':
292
encryption_keyfile=keyfile,
293
**credentials)
294
295
+ if args.lora_path:
296
+ tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir
297
+ tensorize_lora_adapter(args.lora_path, tensorizer_config)
298
+
299
tensorize_vllm_model(engine_args, tensorizer_config)
300
301
elif args.command == "deserialize":
tests/entrypoints/openai/test_tensorizer_entrypoint.py
new
+97
@@ -0,0 +1,97 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import gc
3
+import json
4
+import tempfile
5
+
6
+import openai
7
+import pytest
8
+import pytest_asyncio
9
+import torch.cuda
10
+
11
+from vllm.engine.arg_utils import EngineArgs
12
+from vllm.model_executor.model_loader.tensorizer import (
13
+ TensorizerConfig, tensorize_lora_adapter, tensorize_vllm_model)
14
+
15
+from ...utils import RemoteOpenAIServer
16
+
17
+MODEL_NAME = "unsloth/llama-3.2-1b-Instruct"
18
+LORA_PATH = "davzoku/finqa_adapter_1b"
19
+
20
+
21
+def _cleanup():
22
+ gc.collect()
23
+ torch.cuda.empty_cache()
24
+
25
+
26
+@pytest.fixture(autouse=True)
27
+def cleanup():
28
+ _cleanup()
29
+
30
+
31
+@pytest.fixture(scope='module')
32
+def tmp_dir():
33
+ with tempfile.TemporaryDirectory() as path:
34
+ yield path
35
+
36
+
37
+@pytest.fixture(scope='module')
38
+def model_uri(tmp_dir):
39
+ yield f"{tmp_dir}/model.tensors"
40
+
41
+
42
+@pytest.fixture(scope="module")
43
+def tensorize_model_and_lora(tmp_dir, model_uri):
44
+ tensorizer_config = TensorizerConfig(tensorizer_uri=model_uri,
45
+ lora_dir=tmp_dir)
46
+ args = EngineArgs(model=MODEL_NAME, device="cuda")
47
+
48
+ tensorize_lora_adapter(LORA_PATH, tensorizer_config)
49
+ tensorize_vllm_model(args, tensorizer_config)
50
+
51
+ # Manually invoke a _cleanup() here, as the cleanup()
52
+ # fixture won't be guaranteed to be called after this
53
+ # when this fixture is used for a test
54
+ _cleanup()
55
+ yield
56
+
57
+
58
+@pytest.fixture(scope="module")
59
+def server(model_uri, tensorize_model_and_lora):
60
+ model_loader_extra_config = {
61
+ "tensorizer_uri": model_uri,
62
+ }
63
+
64
+ ## Start OpenAI API server
65
+ args = [
66
+ "--load-format", "tensorizer", "--device", "cuda",
67
+ "--model-loader-extra-config",
68
+ json.dumps(model_loader_extra_config), "--enable-lora"
69
+ ]
70
+
71
+ with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
72
+ yield remote_server
73
+
74
+
75
+@pytest_asyncio.fixture
76
+async def client(server):
77
+ async with server.get_async_client() as async_client:
78
+ yield async_client
79
+
80
+
81
+@pytest.mark.asyncio
82
+@pytest.mark.parametrize("model_name", [MODEL_NAME])
83
+async def test_single_completion(client: openai.AsyncOpenAI, model_name: str):
84
+ _cleanup()
85
+ completion = await client.completions.create(model=model_name,
86
+ prompt="Hello, my name is",
87
+ max_tokens=5,
88
+ temperature=0.0)
89
+
90
+ assert completion.id is not None
91
+ assert completion.choices is not None and len(completion.choices) == 1
92
+ assert completion.model == MODEL_NAME
93
+ assert len(completion.choices) == 1
94
+ assert len(completion.choices[0].text) >= 5
95
+ assert completion.choices[0].finish_reason == "length"
96
+ assert completion.usage == openai.types.CompletionUsage(
97
+ completion_tokens=5, prompt_tokens=6, total_tokens=11)
tests/lora/test_llama_tp.py
+108
-12
@@ -1,12 +1,17 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+import subprocess
3
+import sys
4
+from typing import Union
5
6
import pytest
7
import ray
8
9
import vllm
10
+from vllm import LLM
11
from vllm.lora.request import LoRARequest
12
+from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
13
9
-from ..utils import create_new_process_for_each_test, multi_gpu_test
14
+from ..utils import VLLM_PATH, create_new_process_for_each_test, multi_gpu_test
15
16
MODEL_PATH = "meta-llama/Llama-2-7b-hf"
17
@@ -36,7 +41,10 @@ def v1(run_with_both_engines_lora):
41
pass
42
43
39
-def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
44
+def do_sample(llm: vllm.LLM,
45
+ lora_path: str,
46
+ lora_id: int,
47
+ tensorizer_config_dict: Union[dict, None] = None) -> list[str]:
48
prompts = [
49
"[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_74 (icao VARCHAR, airport VARCHAR)\n\n question: Name the ICAO for lilongwe international airport [/user] [assistant]", # noqa: E501
50
"[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_11 (nationality VARCHAR, elector VARCHAR)\n\n question: When Anchero Pantaleone was the elector what is under nationality? [/user] [assistant]", # noqa: E501
@@ -45,15 +53,28 @@ def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
53
"[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_name_60 (pick INTEGER, former_wnba_team VARCHAR)\n\n question: What pick was a player that previously played for the Minnesota Lynx? [/user] [assistant]", # noqa: E501
54
"[user] Write a SQL query to answer the question based on the table schema.\n\n context: CREATE TABLE table_28138035_4 (womens_doubles VARCHAR, mens_singles VARCHAR)\n\n question: Name the women's doubles for werner schlager [/user] [assistant]" # noqa: E501
55
]
56
+
57
sampling_params = vllm.SamplingParams(temperature=0,
58
max_tokens=256,
59
skip_special_tokens=False,
60
stop=["[/assistant]"])
52
- outputs = llm.generate(
53
- prompts,
54
- sampling_params,
55
- lora_request=LoRARequest(str(lora_id), lora_id, lora_path)
56
- if lora_id else None)
61
+
62
+ if tensorizer_config_dict is not None:
63
+ outputs = llm.generate(
64
+ prompts,
65
+ sampling_params,
66
+ lora_request=LoRARequest(
67
+ str(lora_id),
68
+ lora_id,
69
+ lora_path,
70
+ tensorizer_config_dict=tensorizer_config_dict)
71
+ if lora_id else None)
72
+ else:
73
+ outputs = llm.generate(
74
+ prompts,
75
+ sampling_params,
76
+ lora_request=LoRARequest(str(lora_id), lora_id, lora_path)
77
+ if lora_id else None)
78
# Print the outputs.
79
generated_texts: list[str] = []
80
for output in outputs:
@@ -64,18 +85,32 @@ def do_sample(llm: vllm.LLM, lora_path: str, lora_id: int) -> list[str]:
85
return generated_texts
86
87
67
-def generate_and_test(llm, sql_lora_files):
88
+def generate_and_test(llm,
89
+ sql_lora_files,
90
+ tensorizer_config_dict: Union[dict, None] = None):
91
print("lora adapter created")
69
- assert do_sample(llm, sql_lora_files, lora_id=0) == EXPECTED_NO_LORA_OUTPUT
92
+ assert do_sample(llm,
93
+ sql_lora_files,
94
+ tensorizer_config_dict=tensorizer_config_dict,
95
+ lora_id=0) == EXPECTED_NO_LORA_OUTPUT
96
97
print("lora 1")
72
- assert do_sample(llm, sql_lora_files, lora_id=1) == EXPECTED_LORA_OUTPUT
98
+ assert do_sample(llm,
99
+ sql_lora_files,
100
+ tensorizer_config_dict=tensorizer_config_dict,
101
+ lora_id=1) == EXPECTED_LORA_OUTPUT
102
103
print("no lora")
75
- assert do_sample(llm, sql_lora_files, lora_id=0) == EXPECTED_NO_LORA_OUTPUT
104
+ assert do_sample(llm,
105
+ sql_lora_files,
106
+ tensorizer_config_dict=tensorizer_config_dict,
107
+ lora_id=0) == EXPECTED_NO_LORA_OUTPUT
108
109
print("lora 2")
78
- assert do_sample(llm, sql_lora_files, lora_id=2) == EXPECTED_LORA_OUTPUT
110
+ assert do_sample(llm,
111
+ sql_lora_files,
112
+ tensorizer_config_dict=tensorizer_config_dict,
113
+ lora_id=2) == EXPECTED_LORA_OUTPUT
114
115
print("removing lora")
116
@@ -153,3 +188,64 @@ def test_llama_lora_tp4_fully_sharded_loras(sql_lora_files):
188
enable_chunked_prefill=True,
189
)
190
generate_and_test(llm, sql_lora_files)
191
+
192
+
193
+@multi_gpu_test(num_gpus=2)
194
+@create_new_process_for_each_test()
195
+def test_tp2_serialize_and_deserialize_lora(tmp_path, sql_lora_files,
196
+ sql_lora_huggingface_id):
197
+
198
+ # Run the tensorizing of the LoRA adapter and the model in a subprocess
199
+ # to guarantee cleanup
200
+
201
+ tp_size = 2
202
+ model_name = "model-rank-%03d.tensors"
203
+
204
+ model_ref = MODEL_PATH
205
+ lora_path = sql_lora_huggingface_id
206
+ suffix = "test"
207
+ try:
208
+ result = subprocess.run([
209
+ sys.executable,
210
+ f"{VLLM_PATH}/examples/other/tensorize_vllm_model.py", "--model",
211
+ MODEL_PATH, "--lora-path", lora_path, "--tensor-parallel-size",
212
+ str(tp_size), "serialize", "--serialized-directory",
213
+ str(tmp_path), "--suffix", suffix
214
+ ],
215
+ check=True,
216
+ capture_output=True,
217
+ text=True)
218
+ except subprocess.CalledProcessError as e:
219
+ print("Tensorizing failed.")
220
+ print("STDOUT:\n", e.stdout)
221
+ print("STDERR:\n", e.stderr)
222
+ raise
223
+
224
+ print("STDOUT:\n", result.stdout)
225
+
226
+ model_uri = tmp_path / "vllm" / model_ref / suffix / model_name
227
+ tensorizer_config = TensorizerConfig(tensorizer_uri=str(model_uri))
228
+ tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir
229
+
230
+ loaded_vllm_model = LLM(model=model_ref,
231
+ load_format="tensorizer",
232
+ enable_lora=True,
233
+ enforce_eager=True,
234
+ model_loader_extra_config=tensorizer_config,
235
+ max_num_seqs=13,
236
+ tensor_parallel_size=2,
237
+ max_loras=2)
238
+
239
+ tensorizer_config_dict = tensorizer_config.to_dict()
240
+
241
+ print("lora adapter created")
242
+ assert do_sample(loaded_vllm_model,
243
+ sql_lora_files,
244
+ tensorizer_config_dict=tensorizer_config_dict,
245
+ lora_id=0) == EXPECTED_NO_LORA_OUTPUT
246
+
247
+ print("lora 1")
248
+ assert do_sample(loaded_vllm_model,
249
+ sql_lora_files,
250
+ tensorizer_config_dict=tensorizer_config_dict,
251
+ lora_id=1) == EXPECTED_LORA_OUTPUT
tests/tensorizer_loader/conftest.py
-8
@@ -5,14 +5,6 @@ from vllm.distributed import cleanup_dist_env_and_memory
5
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
6
7
8
-@pytest.fixture(scope="function", autouse=True)
9
-def use_v0_only(monkeypatch):
10
- """
11
- Tensorizer only tested on V0 so far.
12
- """
13
- monkeypatch.setenv('VLLM_USE_V1', '0')
14
-
15
-
8
@pytest.fixture(autouse=True)
9
def cleanup():
10
cleanup_dist_env_and_memory(shutdown_ray=True)
tests/tensorizer_loader/test_tensorizer.py
+47
-115
@@ -1,17 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
import gc
4
-import json
4
import os
5
import pathlib
6
import subprocess
8
-from functools import partial
7
from unittest.mock import MagicMock, patch
8
11
-import openai
9
import pytest
10
import torch
14
-from huggingface_hub import snapshot_download
11
12
from vllm import SamplingParams
13
from vllm.engine.arg_utils import EngineArgs
@@ -22,12 +18,11 @@ from vllm.model_executor.model_loader.tensorizer import (TensorizerConfig,
18
is_vllm_tensorized,
19
load_with_tensorizer,
20
open_stream,
25
- serialize_vllm_model,
21
tensorize_vllm_model)
22
# yapf: enable
28
-from vllm.utils import PlaceholderModule, import_from_path
23
+from vllm.utils import PlaceholderModule
24
30
-from ..utils import VLLM_PATH, RemoteOpenAIServer
25
+from ..utils import VLLM_PATH
26
27
try:
28
from tensorizer import EncryptionParams
@@ -103,6 +98,7 @@ def test_can_deserialize_s3(vllm_runner):
98
@pytest.mark.skipif(not is_curl_installed(), reason="cURL is not installed")
99
def test_deserialized_encrypted_vllm_model_has_same_outputs(
100
vllm_runner, tmp_path):
101
+ args = EngineArgs(model=model_ref)
102
with vllm_runner(model_ref) as vllm_model:
103
model_path = tmp_path / (model_ref + ".tensors")
104
key_path = tmp_path / (model_ref + ".key")
@@ -110,15 +106,13 @@ def test_deserialized_encrypted_vllm_model_has_same_outputs(
106
107
outputs = vllm_model.generate(prompts, sampling_params)
108
113
- config_for_serializing = TensorizerConfig(tensorizer_uri=model_path,
114
- encryption_keyfile=key_path)
109
+ config_for_serializing = TensorizerConfig(tensorizer_uri=str(model_path),
110
+ encryption_keyfile=str(key_path))
111
116
- vllm_model.apply_model(
117
- partial(serialize_vllm_model,
118
- tensorizer_config=config_for_serializing))
112
+ tensorize_vllm_model(args, config_for_serializing)
113
120
- config_for_deserializing = TensorizerConfig(tensorizer_uri=model_path,
121
- encryption_keyfile=key_path)
114
+ config_for_deserializing = TensorizerConfig(
115
+ tensorizer_uri=str(model_path), encryption_keyfile=str(key_path))
116
117
with vllm_runner(model_ref,
118
load_format="tensorizer",
@@ -154,113 +148,46 @@ def test_deserialized_hf_model_has_same_outputs(hf_runner, vllm_runner,
148
assert outputs == deserialized_outputs
149
150
157
-def test_vllm_model_can_load_with_lora(vllm_runner, tmp_path):
158
- multilora_inference = import_from_path(
159
- "examples.offline_inference.multilora_inference",
160
- EXAMPLES_PATH / "offline_inference/multilora_inference.py",
161
- )
162
-
163
- model_ref = "meta-llama/Llama-2-7b-hf"
164
- lora_path = snapshot_download(repo_id="yard1/llama-2-7b-sql-lora-test")
165
- test_prompts = multilora_inference.create_test_prompts(lora_path)
166
-
167
- # Serialize model before deserializing and binding LoRA adapters
168
- with vllm_runner(model_ref) as vllm_model:
169
- model_path = tmp_path / (model_ref + ".tensors")
170
-
171
- vllm_model.apply_model(
172
- partial(
173
- serialize_vllm_model,
174
- tensorizer_config=TensorizerConfig(tensorizer_uri=model_path)))
175
-
176
- with vllm_runner(
177
- model_ref,
178
- load_format="tensorizer",
179
- model_loader_extra_config=TensorizerConfig(
180
- tensorizer_uri=model_path,
181
- num_readers=1,
182
- ),
183
- enable_lora=True,
184
- max_loras=1,
185
- max_lora_rank=8,
186
- max_cpu_loras=2,
187
- max_num_seqs=50,
188
- max_model_len=1000,
189
- ) as loaded_vllm_model:
190
- multilora_inference.process_requests(
191
- loaded_vllm_model.model.llm_engine, test_prompts)
192
-
193
- assert loaded_vllm_model
194
-
195
-
196
-def test_load_without_tensorizer_load_format(vllm_runner):
151
+def test_load_without_tensorizer_load_format(vllm_runner, capfd):
152
model = None
198
- with pytest.raises(ValueError):
153
+ try:
154
model = vllm_runner(
155
model_ref,
156
model_loader_extra_config=TensorizerConfig(tensorizer_uri="test"))
202
- del model
203
- gc.collect()
204
- torch.cuda.empty_cache()
205
-
206
-
207
-@pytest.mark.skipif(not is_curl_installed(), reason="cURL is not installed")
208
-def test_openai_apiserver_with_tensorizer(vllm_runner, tmp_path):
209
- ## Serialize model
210
- with vllm_runner(model_ref) as vllm_model:
211
- model_path = tmp_path / (model_ref + ".tensors")
212
-
213
- vllm_model.apply_model(
214
- partial(
215
- serialize_vllm_model,
216
- tensorizer_config=TensorizerConfig(tensorizer_uri=model_path)))
217
-
218
- model_loader_extra_config = {
219
- "tensorizer_uri": str(model_path),
220
- }
221
-
222
- ## Start OpenAI API server
223
- openai_args = [
224
- "--dtype",
225
- "float16",
226
- "--load-format",
227
- "tensorizer",
228
- "--model-loader-extra-config",
229
- json.dumps(model_loader_extra_config),
230
- ]
231
-
232
- with RemoteOpenAIServer(model_ref, openai_args) as server:
233
- print("Server ready.")
234
-
235
- client = server.get_client()
236
- completion = client.completions.create(model=model_ref,
237
- prompt="Hello, my name is",
238
- max_tokens=5,
239
- temperature=0.0)
240
-
241
- assert completion.id is not None
242
- assert len(completion.choices) == 1
243
- assert len(completion.choices[0].text) >= 5
244
- assert completion.choices[0].finish_reason == "length"
245
- assert completion.usage == openai.types.CompletionUsage(
246
- completion_tokens=5, prompt_tokens=6, total_tokens=11)
247
-
248
-
249
-def test_raise_value_error_on_invalid_load_format(vllm_runner):
157
+ except RuntimeError:
158
+ out, err = capfd.readouterr()
159
+ combined_output = out + err
160
+ assert ("ValueError: Model loader extra config "
161
+ "is not supported for load "
162
+ "format LoadFormat.AUTO") in combined_output
163
+ finally:
164
+ del model
165
+ gc.collect()
166
+ torch.cuda.empty_cache()
167
+
168
+
169
+def test_raise_value_error_on_invalid_load_format(vllm_runner, capfd):
170
model = None
251
- with pytest.raises(ValueError):
171
+ try:
172
model = vllm_runner(
173
model_ref,
174
load_format="safetensors",
175
model_loader_extra_config=TensorizerConfig(tensorizer_uri="test"))
256
- del model
257
- gc.collect()
258
- torch.cuda.empty_cache()
176
+ except RuntimeError:
177
+ out, err = capfd.readouterr()
178
+
179
+ combined_output = out + err
180
+ assert ("ValueError: Model loader extra config is not supported "
181
+ "for load format LoadFormat.SAFETENSORS") in combined_output
182
+ finally:
183
+ del model
184
+ gc.collect()
185
+ torch.cuda.empty_cache()
186
187
188
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires 2 GPUs")
262
-def test_tensorizer_with_tp_path_without_template(vllm_runner):
263
- with pytest.raises(ValueError):
189
+def test_tensorizer_with_tp_path_without_template(vllm_runner, capfd):
190
+ try:
191
model_ref = "EleutherAI/pythia-1.4b"
192
tensorized_path = f"s3://tensorized/{model_ref}/fp16/model.tensors"
193
@@ -275,6 +202,13 @@ def test_tensorizer_with_tp_path_without_template(vllm_runner):
202
tensor_parallel_size=2,
203
disable_custom_all_reduce=True,
204
)
205
+ except RuntimeError:
206
+ out, err = capfd.readouterr()
207
+ combined_output = out + err
208
+ assert ("ValueError: For a sharded model, tensorizer_uri "
209
+ "should include a string format template like '%04d' "
210
+ "to be formatted with the rank "
211
+ "of the shard") in combined_output
212
213
214
@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires 2 GPUs")
@@ -288,7 +222,6 @@ def test_deserialized_encrypted_vllm_model_with_tp_has_same_outputs(
222
enforce_eager=True,
223
) as base_model:
224
outputs = base_model.generate(prompts, sampling_params)
291
- base_model.model.llm_engine.model_executor.shutdown()
225
226
# load model with two shards and serialize with encryption
227
model_path = str(tmp_path / (model_ref + "-%02d.tensors"))
@@ -296,7 +229,7 @@ def test_deserialized_encrypted_vllm_model_with_tp_has_same_outputs(
229
230
tensorizer_config = TensorizerConfig(
231
tensorizer_uri=model_path,
299
- encryption_keyfile=key_path,
232
+ encryption_keyfile=str(key_path),
233
)
234
235
tensorize_vllm_model(
@@ -331,14 +264,13 @@ def test_vllm_tensorized_model_has_same_outputs(vllm_runner, tmp_path):
264
model_ref = "facebook/opt-125m"
265
model_path = tmp_path / (model_ref + ".tensors")
266
config = TensorizerConfig(tensorizer_uri=str(model_path))
267
+ args = EngineArgs(model=model_ref, device="cuda")
268
269
with vllm_runner(model_ref) as vllm_model:
270
outputs = vllm_model.generate(prompts, sampling_params)
271
338
- vllm_model.apply_model(
339
- partial(serialize_vllm_model, tensorizer_config=config))
340
-
341
- assert is_vllm_tensorized(config)
272
+ tensorize_vllm_model(args, config)
273
+ assert is_vllm_tensorized(config)
274
275
with vllm_runner(model_ref,
276
load_format="tensorizer",
vllm/engine/arg_utils.py
+1
-2
@@ -1195,8 +1195,7 @@ class EngineArgs:
1195
#############################################################
1196
# Unsupported Feature Flags on V1.
1197
1198
- if (self.load_format == LoadFormat.TENSORIZER.value
1199
- or self.load_format == LoadFormat.SHARDED_STATE.value):
1198
+ if self.load_format == LoadFormat.SHARDED_STATE.value:
1199
_raise_or_fallback(
1200
feature_name=f"--load_format {self.load_format}",
1201
recommend_to_remove=False)
vllm/lora/models.py
+44
-29
@@ -29,6 +29,7 @@ from vllm.lora.utils import (from_layer, from_layer_logits_processor,
29
get_supported_lora_modules,
30
is_regex_target_modules,
31
parse_fine_tuned_lora_name, replace_submodule)
32
+from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
33
from vllm.model_executor.models import SupportsLoRA, supports_multimodal
34
from vllm.model_executor.models.interfaces import is_pooling_model
35
from vllm.model_executor.models.module_mapping import MultiModelKeys
@@ -185,19 +186,19 @@ class LoRAModel(AdapterModel):
186
187
@classmethod
188
def from_local_checkpoint(
188
- cls,
189
- lora_dir: str,
190
- expected_lora_modules: list[str],
191
- peft_helper: PEFTHelper,
192
- *,
193
- lora_model_id: Optional[int] = None,
194
- device: str = "cuda",
195
- dtype: Optional[torch.dtype] = None,
196
- target_embedding_padding: Optional[int] = None,
197
- embedding_modules: Optional[dict[str, str]] = None,
198
- embedding_padding_modules: Optional[list[str]] = None,
199
- weights_mapper: Optional[WeightsMapper] = None,
200
- ) -> "LoRAModel":
189
+ cls,
190
+ lora_dir: str,
191
+ expected_lora_modules: list[str],
192
+ peft_helper: PEFTHelper,
193
+ *,
194
+ lora_model_id: Optional[int] = None,
195
+ device: str = "cuda",
196
+ dtype: Optional[torch.dtype] = None,
197
+ target_embedding_padding: Optional[int] = None,
198
+ embedding_modules: Optional[dict[str, str]] = None,
199
+ embedding_padding_modules: Optional[list[str]] = None,
200
+ weights_mapper: Optional[WeightsMapper] = None,
201
+ tensorizer_config_dict: Optional[dict] = None) -> "LoRAModel":
202
"""Create a LoRAModel from a local checkpoint.
203
204
Args:
@@ -219,10 +220,36 @@ class LoRAModel(AdapterModel):
220
lora_dir, "new_embeddings.safetensors")
221
new_embeddings_bin_file_path = os.path.join(lora_dir,
222
"new_embeddings.bin")
223
+ tensors: dict[str, torch.Tensor] = {}
224
+ unexpected_modules: list[Union[list[str], str]] = []
225
+
226
+ def check_unexpected_modules(modules: dict):
227
+ for lora_module in modules.keys(): # noqa
228
+ module_name, _, _ = parse_fine_tuned_lora_name(
229
+ lora_module, weights_mapper)
230
+ part_name = module_name.split(".")[-1]
231
+ if part_name not in expected_lora_modules:
232
+ unexpected_modules.append(module_name)
233
+ if unexpected_modules:
234
+ raise ValueError(
235
+ f"While loading {lora_dir}, expected"
236
+ f" target modules in {expected_lora_modules}"
237
+ f" but received {unexpected_modules}."
238
+ f" Please verify that the loaded LoRA module is correct")
239
223
- unexpected_modules: list[Union[list[str], str]]
224
- if os.path.isfile(lora_tensor_path):
225
- tensors: dict[str, torch.Tensor] = {}
240
+ if tensorizer_config_dict:
241
+ from tensorizer import TensorDeserializer
242
+
243
+ tensorizer_config = TensorizerConfig(**tensorizer_config_dict)
244
+ lora_tensor_path = os.path.join(tensorizer_config.tensorizer_dir,
245
+ "adapter_model.tensors")
246
+ tensorizer_args = tensorizer_config._construct_tensorizer_args()
247
+ tensors = TensorDeserializer(lora_tensor_path,
248
+ dtype=tensorizer_config.dtype,
249
+ **tensorizer_args.deserializer_params)
250
+ check_unexpected_modules(tensors)
251
+
252
+ elif os.path.isfile(lora_tensor_path):
253
# Find unexpected modules.
254
# Use safetensor key as a source of truth to find expected modules.
255
# in peft if you have target_modules A, B, C and C does not exist
@@ -232,20 +259,8 @@ class LoRAModel(AdapterModel):
259
unexpected_modules = []
260
with safetensors.safe_open(lora_tensor_path,
261
framework="pt") as f: # type: ignore
235
- for lora_module in f.keys(): # noqa
236
- module_name, _, _ = parse_fine_tuned_lora_name(
237
- lora_module, weights_mapper)
238
- part_name = module_name.split(".")[-1]
239
- if part_name not in expected_lora_modules:
240
- unexpected_modules.append(module_name)
241
- if unexpected_modules:
242
- raise ValueError(
243
- f"While loading {lora_dir}, expected"
244
- f" target modules in {expected_lora_modules}"
245
- f" but received {unexpected_modules}."
246
- f" Please verify that the loaded LoRA module is correct"
247
- )
262
# Load tensors if there are only expected modules.
263
+ check_unexpected_modules(f)
264
for module in f.keys(): # noqa
265
tensors[module] = f.get_tensor(module)
266
elif os.path.isfile(lora_bin_file_path):
vllm/lora/peft_helper.py
+24
-4
@@ -10,6 +10,7 @@ from typing import Literal, Optional, Union
10
11
from vllm.config import LoRAConfig
12
from vllm.logger import init_logger
13
+from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
14
15
logger = init_logger(__name__)
16
@@ -89,12 +90,31 @@ class PEFTHelper:
90
return cls(**filtered_dict)
91
92
@classmethod
92
- def from_local_dir(cls, lora_path: str,
93
- max_position_embeddings: Optional[int]) -> "PEFTHelper":
93
+ def from_local_dir(
94
+ cls,
95
+ lora_path: str,
96
+ max_position_embeddings: Optional[int],
97
+ tensorizer_config_dict: Optional[dict] = None) -> "PEFTHelper":
98
lora_config_path = os.path.join(lora_path, "adapter_config.json")
99
96
- with open(lora_config_path) as f:
97
- config = json.load(f)
100
+ if tensorizer_config_dict:
101
+ tensorizer_config = TensorizerConfig(**tensorizer_config_dict)
102
+ tensorizer_args = tensorizer_config._construct_tensorizer_args()
103
+ from tensorizer.stream_io import open_stream
104
+ lora_config_path = os.path.join(tensorizer_config.lora_dir,
105
+ "adapter_config.json")
106
+ with open_stream(lora_config_path,
107
+ mode="rb",
108
+ **tensorizer_args.stream_params) as f:
109
+ config = json.load(f)
110
+
111
+ logger.info("Successfully deserialized LoRA config from %s",
112
+ tensorizer_config.lora_dir)
113
+
114
+ else:
115
+ with open(lora_config_path) as f:
116
+ config = json.load(f)
117
+
118
config["vllm_max_position_embeddings"] = max_position_embeddings
119
return cls.from_dict(config)
120
vllm/lora/request.py
+1
@@ -31,6 +31,7 @@ class LoRARequest(
31
lora_local_path: Optional[str] = msgspec.field(default=None)
32
long_lora_max_len: Optional[int] = None
33
base_model_name: Optional[str] = msgspec.field(default=None)
34
+ tensorizer_config_dict: Optional[dict] = None
35
36
def __post_init__(self):
37
if self.lora_local_path:
vllm/lora/worker_manager.py
+3
-1
@@ -100,7 +100,8 @@ class WorkerLoRAManager(AbstractWorkerManager):
100
lora_path = get_adapter_absolute_path(lora_request.lora_path)
101
102
peft_helper = PEFTHelper.from_local_dir(
103
- lora_path, self.max_position_embeddings)
103
+ lora_path, self.max_position_embeddings,
104
+ lora_request.tensorizer_config_dict)
105
106
# Validates the LoRA configuration against requirements before
107
# loading weights, throwing an exception if validation fails.
@@ -125,6 +126,7 @@ class WorkerLoRAManager(AbstractWorkerManager):
126
self.lora_config.lora_extra_vocab_size,
127
embedding_modules=self.embedding_modules,
128
embedding_padding_modules=self.embedding_padding_modules,
129
+ tensorizer_config_dict=lora_request.tensorizer_config_dict,
130
weights_mapper=hf_to_vllm_mapper)
131
132
except FileNotFoundError as e:
vllm/model_executor/model_loader/tensorizer.py
+170
-12
@@ -1,24 +1,28 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
import argparse
4
+import contextlib
5
+import contextvars
6
import dataclasses
7
import io
8
+import json
9
import os
10
import re
11
+import threading
12
import time
13
from collections.abc import Generator
14
from dataclasses import dataclass
15
from functools import partial
12
-from typing import BinaryIO, Optional, Union
16
+from typing import Any, BinaryIO, Optional, Union
17
18
import torch
19
from torch import nn
20
+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
25
from vllm.engine.arg_utils import EngineArgs
21
-from vllm.engine.llm_engine import LLMEngine
26
from vllm.logger import init_logger
27
from vllm.model_executor.layers.vocab_parallel_embedding import (
28
VocabParallelEmbedding)
@@ -58,9 +62,79 @@ __all__ = [
62
logger = init_logger(__name__)
63
64
65
+class MetaTensorMode(TorchDispatchMode):
66
+
67
+ def __torch_dispatch__(self, func, types, args=(), kwargs=None):
68
+ kwargs = kwargs or {}
69
+
70
+ if func._schema.name == "aten::empty" and "device" not in kwargs:
71
+ kwargs["device"] = "meta"
72
+
73
+ return func(*args, **kwargs)
74
+
75
+
76
+def meta_tensor_mode(loading_code=None, ):
77
+
78
+ if loading_code is None:
79
+ return _NoInitOrTensorImpl.context_manager()
80
+ elif callable(loading_code):
81
+ with _NoInitOrTensorImpl.context_manager():
82
+ return loading_code()
83
+ else:
84
+ raise TypeError(
85
+ "expected a callable to evaluate,"
86
+ " or None if being used as a context manager;"
87
+ f' got an object of type "{type(loading_code).__name__}" instead.')
88
+
89
+
90
+class _NoInitOrTensorImpl:
91
+ _MODULES = (torch.nn.Linear, torch.nn.Embedding, torch.nn.LayerNorm)
92
+ _MODULE_ORIGINALS = tuple((m, m.reset_parameters) for m in _MODULES)
93
+
94
+ is_active = contextvars.ContextVar("_NoInitOrTensorImpl.is_active",
95
+ default=False)
96
+ _count_active: int = 0
97
+ _count_active_lock = threading.Lock()
98
+
99
+ @classmethod
100
+ @contextlib.contextmanager
101
+ def context_manager(cls):
102
+ if cls.is_active.get():
103
+ yield
104
+ return
105
+
106
+ with cls._count_active_lock:
107
+ cls._count_active += 1
108
+ if cls._count_active == 1:
109
+ for mod in cls._MODULES:
110
+ mod.reset_parameters = cls._disable(mod.reset_parameters)
111
+
112
+ reset_token = cls.is_active.set(True)
113
+
114
+ try:
115
+ with MetaTensorMode():
116
+ yield
117
+ finally:
118
+ cls.is_active.reset(reset_token)
119
+ with cls._count_active_lock:
120
+ cls._count_active -= 1
121
+ if cls._count_active == 0:
122
+ for mod, original in cls._MODULE_ORIGINALS:
123
+ mod.reset_parameters = original
124
+
125
+ @staticmethod
126
+ def _disable(func):
127
+
128
+ def wrapper(*args, **kwargs):
129
+ if not _NoInitOrTensorImpl.is_active.get():
130
+ return func(*args, **kwargs)
131
+
132
+ return wrapper
133
+
134
+
135
@dataclass
136
class TensorizerConfig:
63
- tensorizer_uri: str
137
+ tensorizer_uri: Union[str, None] = None
138
vllm_tensorized: Optional[bool] = False
139
verify_hash: Optional[bool] = False
140
num_readers: Optional[int] = None
@@ -71,12 +145,29 @@ class TensorizerConfig:
145
model_class: Optional[type[torch.nn.Module]] = None
146
hf_config: Optional[PretrainedConfig] = None
147
dtype: Optional[Union[str, torch.dtype]] = None
148
+ lora_dir: Optional[str] = None
149
_is_sharded: bool = False
150
151
def __post_init__(self):
152
# check if the configuration is for a sharded vLLM model
153
self._is_sharded = isinstance(self.tensorizer_uri, str) \
154
and re.search(r'%0\dd', self.tensorizer_uri) is not None
155
+ if not self.tensorizer_uri and not self.lora_dir:
156
+ raise ValueError("tensorizer_uri must be provided.")
157
+ if not self.tensorizer_uri and self.lora_dir:
158
+ self.tensorizer_uri = f"{self.lora_dir}/adapter_model.tensors"
159
+ assert self.tensorizer_uri is not None, ("tensorizer_uri must be "
160
+ "provided.")
161
+ self.tensorizer_dir = os.path.dirname(self.tensorizer_uri)
162
+ self.lora_dir = self.tensorizer_dir
163
+
164
+ @classmethod
165
+ def as_dict(cls, *args, **kwargs) -> dict[str, Any]:
166
+ cfg = TensorizerConfig(*args, **kwargs)
167
+ return dataclasses.asdict(cfg)
168
+
169
+ def to_dict(self) -> dict[str, Any]:
170
+ return dataclasses.asdict(self)
171
172
def _construct_tensorizer_args(self) -> "TensorizerArgs":
173
tensorizer_args = {
@@ -140,7 +231,9 @@ class TensorizerArgs:
231
232
Args:
233
tensorizer_uri: Path to serialized model tensors. Can be a local file
143
- path or a S3 URI.
234
+ path or a S3 URI. This is a required field unless lora_dir is
235
+ provided and the config is meant to be used for the
236
+ `tensorize_lora_adapter` function.
237
vllm_tensorized: If True, indicates that the serialized model is a
238
vLLM model. This is used to determine the behavior of the
239
TensorDeserializer when loading tensors from a serialized model.
@@ -296,10 +389,10 @@ class TensorizerAgent:
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?
299
- with no_init_or_tensor(), set_current_vllm_config(self.vllm_config,
300
- check_compile=True):
392
+ with meta_tensor_mode(), set_current_vllm_config(self.vllm_config,
393
+ check_compile=True):
394
return self.tensorizer_config.model_class(
302
- vllm_config=self.vllm_config, )
395
+ vllm_config=self.vllm_config)
396
397
def _resize_lora_embeddings(self):
398
"""Modify LoRA embedding layers to use bigger tensors
@@ -467,8 +560,73 @@ def tensorize_vllm_model(engine_args: EngineArgs,
560
) as stream:
561
stream.write(encryption_params.key)
562
470
- engine = LLMEngine.from_engine_args(engine_args)
471
- engine.model_executor.collective_rpc(
472
- "save_tensorized_model",
473
- kwargs=dict(tensorizer_config=tensorizer_config),
474
- )
563
+ from vllm import LLMEngine
564
+ from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
565
+
566
+ if not envs.VLLM_USE_V1:
567
+ engine = LLMEngine.from_engine_args(engine_args)
568
+ engine.model_executor.collective_rpc(
569
+ "save_tensorized_model",
570
+ kwargs=dict(tensorizer_config=tensorizer_config),
571
+ )
572
+ else:
573
+ engine = V1LLMEngine.from_vllm_config(engine_config)
574
+ engine.collective_rpc(
575
+ "save_tensorized_model",
576
+ kwargs=dict(tensorizer_config=tensorizer_config),
577
+ )
578
+
579
+
580
+def tensorize_lora_adapter(lora_path: str,
581
+ tensorizer_config: TensorizerConfig):
582
+ """
583
+ Uses tensorizer to serialize a LoRA adapter. Assumes that the files
584
+ needed to load a LoRA adapter are a safetensors-format file called
585
+ adapter_model.safetensors and a json config file called adapter_config.json.
586
+
587
+ Serializes the files in the tensorizer_config.lora_dir
588
+ """
589
+ import safetensors
590
+
591
+ from vllm.lora.utils import get_adapter_absolute_path
592
+
593
+ lora_dir = get_adapter_absolute_path(lora_path)
594
+
595
+ tensor_path = config_path = ""
596
+
597
+ for file in os.listdir(lora_dir):
598
+ if file.startswith("adapter_model"):
599
+ tensor_path = lora_dir + "/" + file
600
+ if file.startswith("adapter_config"):
601
+ config_path = lora_dir + "/" + file
602
+ if tensor_path and config_path:
603
+ break
604
+
605
+ if tensor_path.endswith(".safetensors"):
606
+ tensors = safetensors.torch.load_file(tensor_path)
607
+ elif tensor_path.endswith(".bin"):
608
+ tensors = torch.load(tensor_path)
609
+ else:
610
+ raise ValueError("Unsupported file: %s", tensor_path)
611
+
612
+ with open(config_path) as f:
613
+ config = json.load(f)
614
+
615
+ tensorizer_args = tensorizer_config._construct_tensorizer_args()
616
+
617
+ with open_stream(f"{tensorizer_config.lora_dir}/adapter_config.json",
618
+ mode="wb+",
619
+ **tensorizer_args.stream_params) as f:
620
+
621
+ f.write(json.dumps(config).encode("utf-8"))
622
+
623
+ lora_uri = (f"{tensorizer_config.lora_dir}"
624
+ f"/adapter_model.tensors")
625
+ with open_stream(lora_uri, mode="wb+",
626
+ **tensorizer_args.stream_params) as f:
627
+ serializer = TensorSerializer(f)
628
+ serializer.write_state_dict(tensors)
629
+ serializer.close()
630
+
631
+ logger.info("Successfully serialized LoRA files to %s",
632
+ str(tensorizer_config.lora_dir))
vllm/model_executor/model_loader/tensorizer_loader.py
+4
-1
@@ -2,6 +2,7 @@
2
# ruff: noqa: SIM117
3
import copy
4
from collections.abc import Generator
5
+from typing import Union
6
7
import torch
8
from torch import nn
@@ -111,8 +112,10 @@ class TensorizerLoader(BaseModelLoader):
112
@staticmethod
113
def save_model(
114
model: torch.nn.Module,
114
- tensorizer_config: TensorizerConfig,
115
+ tensorizer_config: Union[TensorizerConfig, dict],
116
) -> None:
117
+ if isinstance(tensorizer_config, dict):
118
+ tensorizer_config = TensorizerConfig(**tensorizer_config)
119
serialize_vllm_model(
120
model=model,
121
tensorizer_config=tensorizer_config,
vllm/v1/engine/core.py
+7
@@ -340,6 +340,13 @@ class EngineCore:
340
return self.model_executor.collective_rpc(method, timeout, args,
341
kwargs)
342
343
+ def save_tensorized_model(
344
+ self,
345
+ tensorizer_config,
346
+ ) -> None:
347
+ self.model_executor.save_tensorized_model(
348
+ tensorizer_config=tensorizer_config, )
349
+
350
351
class EngineCoreProc(EngineCore):
352
"""ZMQ-wrapper for running EngineCore in background process."""
vllm/v1/worker/gpu_model_runner.py
+11
-1
@@ -25,7 +25,7 @@ from vllm.distributed.parallel_state import (
25
from vllm.forward_context import get_forward_context, set_forward_context
26
from vllm.logger import init_logger
27
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
28
-from vllm.model_executor.model_loader import get_model
28
+from vllm.model_executor.model_loader import TensorizerLoader, get_model
29
from vllm.multimodal import MULTIMODAL_REGISTRY
30
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
31
from vllm.multimodal.utils import group_mm_inputs_by_modality
@@ -60,6 +60,7 @@ from .utils import (gather_mm_placeholders, sanity_check_mm_encoder_outputs,
60
if TYPE_CHECKING:
61
import xgrammar as xgr
62
63
+ from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
64
from vllm.v1.core.sched.output import SchedulerOutput
65
else:
66
xgr = LazyLoader("xgr", globals(), "xgrammar")
@@ -1534,6 +1535,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1535
time_after_load - time_before_load)
1536
prepare_communication_buffer_for_model(self.model)
1537
1538
+ def save_tensorized_model(
1539
+ self,
1540
+ tensorizer_config: "TensorizerConfig",
1541
+ ) -> None:
1542
+ TensorizerLoader.save_model(
1543
+ self.model,
1544
+ tensorizer_config=tensorizer_config,
1545
+ )
1546
+
1547
def _get_prompt_logprobs_dict(
1548
self,
1549
hidden_states: torch.Tensor,
vllm/v1/worker/gpu_worker.py
+8
@@ -31,6 +31,7 @@ from vllm.v1.worker.worker_base import WorkerBase
31
logger = init_logger(__name__)
32
33
if TYPE_CHECKING:
34
+ from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
35
from vllm.v1.core.sched.output import SchedulerOutput
36
37
@@ -326,6 +327,13 @@ class Worker(WorkerBase):
327
max_size=max_size,
328
)
329
330
+ def save_tensorized_model(
331
+ self,
332
+ tensorizer_config: "TensorizerConfig",
333
+ ) -> None:
334
+ self.model_runner.save_tensorized_model(
335
+ tensorizer_config=tensorizer_config, )
336
+
337
338
def init_worker_distributed_environment(
339
vllm_config: VllmConfig,