@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 72d14d0eed
Fork 0

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, 03:30 UTC

[Frontend] [Core] Integrate Tensorizer in to S3 loading machinery, allow passing arbitrary arguments during save/load (#19619)

Signed-off-by: Sanger Steel <sangersteel@gmail.com> Co-authored-by: Eta <esyra@coreweave.com>

Sanger Steel committed Jul 8, 2025 at 01:47 UTC 72d14d0eed4b29e5827519283c085a7a674f3256
18 files changed +807 -189
examples/others/tensorize_vllm_model.py
+83 -25
@@ -4,6 +4,7 @@
4 import argparse
5 import dataclasses
6 import json
7 +import logging
8 import os
9 import uuid
10
@@ -15,9 +16,13 @@ from vllm.model_executor.model_loader.tensorizer import (
16 TensorizerConfig,
17 tensorize_lora_adapter,
18 tensorize_vllm_model,
19 + tensorizer_kwargs_arg,
20 )
21 from vllm.utils import FlexibleArgumentParser
22
23 +logger = logging.getLogger()
24 +
25 +
26 # yapf conflicts with isort for this docstring
27 # yapf: disable
28 """
@@ -119,7 +124,7 @@ vllm serve <model_path> \
124 """
125
126
122 -def parse_args():
127 +def get_parser():
128 parser = FlexibleArgumentParser(
129 description="An example script that can be used to serialize and "
130 "deserialize vLLM models. These models "
@@ -135,13 +140,13 @@ def parse_args():
140 required=False,
141 help="Path to a LoRA adapter to "
142 "serialize along with model tensors. This can then be deserialized "
138 - "along with the model by passing a tensorizer_config kwarg to "
139 - "LoRARequest with type TensorizerConfig. See the docstring for this "
140 - "for a usage example."
141 -
143 + "along with the model by instantiating a TensorizerConfig object, "
144 + "creating a dict from it with TensorizerConfig.to_serializable(), "
145 + "and passing it to LoRARequest's initializer with the kwarg "
146 + "tensorizer_config_dict."
147 )
148
144 - subparsers = parser.add_subparsers(dest='command')
149 + subparsers = parser.add_subparsers(dest='command', required=True)
150
151 serialize_parser = subparsers.add_parser(
152 'serialize', help="Serialize a model to `--serialized-directory`")
@@ -171,6 +176,14 @@ def parse_args():
176 "where `suffix` is given by `--suffix` or a random UUID if not "
177 "provided.")
178
179 + serialize_parser.add_argument(
180 + "--serialization-kwargs",
181 + type=tensorizer_kwargs_arg,
182 + required=False,
183 + help=("A JSON string containing additional keyword arguments to "
184 + "pass to Tensorizer's TensorSerializer during "
185 + "serialization."))
186 +
187 serialize_parser.add_argument(
188 "--keyfile",
189 type=str,
@@ -186,9 +199,17 @@ def parse_args():
199 deserialize_parser.add_argument(
200 "--path-to-tensors",
201 type=str,
189 - required=True,
202 + required=False,
203 help="The local path or S3 URI to the model tensors to deserialize. ")
204
205 + deserialize_parser.add_argument(
206 + "--serialized-directory",
207 + type=str,
208 + required=False,
209 + help="Directory with model artifacts for loading. Assumes a "
210 + "model.tensors file exists therein. Can supersede "
211 + "--path-to-tensors.")
212 +
213 deserialize_parser.add_argument(
214 "--keyfile",
215 type=str,
@@ -196,11 +217,27 @@ def parse_args():
217 help=("Path to a binary key to use to decrypt the model weights,"
218 " if the model was serialized with encryption"))
219
199 - TensorizerArgs.add_cli_args(deserialize_parser)
220 + deserialize_parser.add_argument(
221 + "--deserialization-kwargs",
222 + type=tensorizer_kwargs_arg,
223 + required=False,
224 + help=("A JSON string containing additional keyword arguments to "
225 + "pass to Tensorizer's `TensorDeserializer` during "
226 + "deserialization."))
227
201 - return parser.parse_args()
228 + TensorizerArgs.add_cli_args(deserialize_parser)
229
230 + return parser
231
232 +def merge_extra_config_with_tensorizer_config(extra_cfg: dict,
233 + cfg: TensorizerConfig):
234 + for k, v in extra_cfg.items():
235 + if hasattr(cfg, k):
236 + setattr(cfg, k, v)
237 + logger.info(
238 + "Updating TensorizerConfig with %s from "
239 + "--model-loader-extra-config provided", k
240 + )
241
242 def deserialize(args, tensorizer_config):
243 if args.lora_path:
@@ -230,7 +267,8 @@ def deserialize(args, tensorizer_config):
267 lora_request=LoRARequest("sql-lora",
268 1,
269 args.lora_path,
233 - tensorizer_config = tensorizer_config)
270 + tensorizer_config_dict = tensorizer_config
271 + .to_serializable())
272 )
273 )
274 else:
@@ -243,7 +281,8 @@ def deserialize(args, tensorizer_config):
281
282
283 def main():
246 - args = parse_args()
284 + parser = get_parser()
285 + args = parser.parse_args()
286
287 s3_access_key_id = (getattr(args, 's3_access_key_id', None)
288 or os.environ.get("S3_ACCESS_KEY_ID", None))
@@ -265,13 +304,24 @@ def main():
304 else:
305 keyfile = None
306
307 + extra_config = {}
308 if args.model_loader_extra_config:
269 - config = json.loads(args.model_loader_extra_config)
270 - tensorizer_args = \
271 - TensorizerConfig(**config)._construct_tensorizer_args()
272 - tensorizer_args.tensorizer_uri = args.path_to_tensors
273 - else:
274 - tensorizer_args = None
309 + extra_config = json.loads(args.model_loader_extra_config)
310 +
311 +
312 + tensorizer_dir = (args.serialized_directory or
313 + extra_config.get("tensorizer_dir"))
314 + tensorizer_uri = (getattr(args, "path_to_tensors", None)
315 + or extra_config.get("tensorizer_uri"))
316 +
317 + if tensorizer_dir and tensorizer_uri:
318 + parser.error("--serialized-directory and --path-to-tensors "
319 + "cannot both be provided")
320 +
321 + if not tensorizer_dir and not tensorizer_uri:
322 + parser.error("Either --serialized-directory or --path-to-tensors "
323 + "must be provided")
324 +
325
326 if args.command == "serialize":
327 eng_args_dict = {f.name: getattr(args, f.name) for f in
@@ -281,7 +331,7 @@ def main():
331 argparse.Namespace(**eng_args_dict)
332 )
333
284 - input_dir = args.serialized_directory.rstrip('/')
334 + input_dir = tensorizer_dir.rstrip('/')
335 suffix = args.suffix if args.suffix else uuid.uuid4().hex
336 base_path = f"{input_dir}/vllm/{model_ref}/{suffix}"
337 if engine_args.tensor_parallel_size > 1:
@@ -292,21 +342,29 @@ def main():
342 tensorizer_config = TensorizerConfig(
343 tensorizer_uri=model_path,
344 encryption_keyfile=keyfile,
295 - **credentials)
345 + serialization_kwargs=args.serialization_kwargs or {},
346 + **credentials
347 + )
348
349 if args.lora_path:
350 tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir
351 tensorize_lora_adapter(args.lora_path, tensorizer_config)
352
353 + merge_extra_config_with_tensorizer_config(extra_config,
354 + tensorizer_config)
355 tensorize_vllm_model(engine_args, tensorizer_config)
356
357 elif args.command == "deserialize":
304 - if not tensorizer_args:
305 - tensorizer_config = TensorizerConfig(
306 - tensorizer_uri=args.path_to_tensors,
307 - encryption_keyfile = keyfile,
308 - **credentials
309 - )
358 + tensorizer_config = TensorizerConfig(
359 + tensorizer_uri=args.path_to_tensors,
360 + tensorizer_dir=args.serialized_directory,
361 + encryption_keyfile=keyfile,
362 + deserialization_kwargs=args.deserialization_kwargs or {},
363 + **credentials
364 + )
365 +
366 + merge_extra_config_with_tensorizer_config(extra_config,
367 + tensorizer_config)
368 deserialize(args, tensorizer_config)
369 else:
370 raise ValueError("Either serialize or deserialize must be specified.")
requirements/nightly_torch_test.txt
+1 -1
@@ -1,6 +1,6 @@
1 # testing
2 pytest
3 -tensorizer>=2.9.0
3 +tensorizer==2.10.1
4 pytest-forked
5 pytest-asyncio
6 pytest-rerunfailures
requirements/rocm.txt
+1 -1
@@ -11,7 +11,7 @@ datasets
11 ray>=2.10.0,<2.45.0
12 peft
13 pytest-asyncio
14 -tensorizer>=2.9.0
14 +tensorizer==2.10.1
15 packaging>=24.2
16 setuptools>=77.0.3,<80.0.0
17 setuptools-scm>=8
requirements/test.in
+1 -1
@@ -1,6 +1,6 @@
1 # testing
2 pytest
3 -tensorizer>=2.9.0
3 +tensorizer==2.10.1
4 pytest-forked
5 pytest-asyncio
6 pytest-rerunfailures
requirements/test.txt
+1 -1
@@ -739,7 +739,7 @@ tenacity==9.0.0
739 # via
740 # lm-eval
741 # plotly
742 -tensorizer==2.9.0
742 +tensorizer==2.10.1
743 # via -r requirements/test.in
744 threadpoolctl==3.5.0
745 # via scikit-learn
setup.py
+1 -1
@@ -689,7 +689,7 @@ setup(
689 install_requires=get_requirements(),
690 extras_require={
691 "bench": ["pandas", "datasets"],
692 - "tensorizer": ["tensorizer>=2.9.0"],
692 + "tensorizer": ["tensorizer==2.10.1"],
693 "fastsafetensors": ["fastsafetensors >= 0.1.10"],
694 "runai": ["runai-model-streamer", "runai-model-streamer-s3", "boto3"],
695 "audio": ["librosa", "soundfile"], # Required for audio processing
tests/entrypoints/openai/test_tensorizer_entrypoint.py
+10 -8
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import gc
4 -import json
4 +import os
5 import tempfile
6
7 import openai
@@ -58,18 +58,20 @@ def tensorize_model_and_lora(tmp_dir, model_uri):
58
59 @pytest.fixture(scope="module")
60 def server(model_uri, tensorize_model_and_lora):
61 - model_loader_extra_config = {
62 - "tensorizer_uri": model_uri,
63 - }
61 + # In this case, model_uri is a directory with a model.tensors
62 + # file and all necessary model artifacts, particularly a
63 + # HF `config.json` file. In this case, Tensorizer can infer the
64 + # `TensorizerConfig` so --model-loader-extra-config can be completely
65 + # omitted.
66
67 ## Start OpenAI API server
68 args = [
67 - "--load-format", "tensorizer", "--device", "cuda",
68 - "--model-loader-extra-config",
69 - json.dumps(model_loader_extra_config), "--enable-lora"
69 + "--load-format", "tensorizer", "--served-model-name", MODEL_NAME,
70 + "--enable-lora"
71 ]
72
72 - with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
73 + model_dir = os.path.dirname(model_uri)
74 + with RemoteOpenAIServer(model_dir, args) as remote_server:
75 yield remote_server
76
77
tests/lora/test_llama_tp.py
+3 -2
@@ -169,7 +169,8 @@ def test_tp2_serialize_and_deserialize_lora(tmp_path, sql_lora_files,
169 f"{VLLM_PATH}/examples/others/tensorize_vllm_model.py", "--model",
170 MODEL_PATH, "--lora-path", lora_path, "--tensor-parallel-size",
171 str(tp_size), "serialize", "--serialized-directory",
172 - str(tmp_path), "--suffix", suffix
172 + str(tmp_path), "--suffix", suffix, "--serialization-kwargs",
173 + '{"limit_cpu_concurrency": 4}'
174 ],
175 check=True,
176 capture_output=True,
@@ -195,7 +196,7 @@ def test_tp2_serialize_and_deserialize_lora(tmp_path, sql_lora_files,
196 tensor_parallel_size=2,
197 max_loras=2)
198
198 - tensorizer_config_dict = tensorizer_config.to_dict()
199 + tensorizer_config_dict = tensorizer_config.to_serializable()
200
201 print("lora adapter created")
202 assert do_sample(loaded_vllm_model,
tests/tensorizer_loader/conftest.py
+85
@@ -1,9 +1,28 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Callable
4 +
5 import pytest
6
7 +from vllm import LLM, EngineArgs
8 from vllm.distributed import cleanup_dist_env_and_memory
9 +from vllm.model_executor.model_loader import tensorizer as tensorizer_mod
10 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
11 +from vllm.utils import get_distributed_init_method, get_ip, get_open_port
12 +from vllm.v1.executor.abstract import UniProcExecutor
13 +from vllm.worker.worker_base import WorkerWrapperBase
14 +
15 +MODEL_REF = "facebook/opt-125m"
16 +
17 +
18 +@pytest.fixture()
19 +def model_ref():
20 + return MODEL_REF
21 +
22 +
23 +@pytest.fixture(autouse=True)
24 +def allow_insecure_serialization(monkeypatch):
25 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
26
27
28 @pytest.fixture(autouse=True)
@@ -11,7 +30,73 @@ def cleanup():
30 cleanup_dist_env_and_memory(shutdown_ray=True)
31
32
33 +@pytest.fixture()
34 +def just_serialize_model_tensors(model_ref, monkeypatch, tmp_path):
35 +
36 + def noop(*args, **kwargs):
37 + return None
38 +
39 + args = EngineArgs(model=model_ref)
40 + tc = TensorizerConfig(tensorizer_uri=f"{tmp_path}/model.tensors")
41 +
42 + monkeypatch.setattr(tensorizer_mod, "serialize_extra_artifacts", noop)
43 +
44 + tensorizer_mod.tensorize_vllm_model(args, tc)
45 + yield tmp_path
46 +
47 +
48 @pytest.fixture(autouse=True)
49 def tensorizer_config():
50 config = TensorizerConfig(tensorizer_uri="vllm")
51 return config
52 +
53 +
54 +@pytest.fixture()
55 +def model_path(model_ref, tmp_path):
56 + yield tmp_path / model_ref / "model.tensors"
57 +
58 +
59 +def assert_from_collective_rpc(engine: LLM, closure: Callable,
60 + closure_kwargs: dict):
61 + res = engine.collective_rpc(method=closure, kwargs=closure_kwargs)
62 + return all(res)
63 +
64 +
65 +# This is an object pulled from tests/v1/engine/test_engine_core.py
66 +# Modified to strip the `load_model` method from its `_init_executor`
67 +# method. It's purely used as a dummy utility to run methods that test
68 +# Tensorizer functionality
69 +class DummyExecutor(UniProcExecutor):
70 +
71 + def _init_executor(self) -> None:
72 + """Initialize the worker and load the model.
73 + """
74 + self.driver_worker = WorkerWrapperBase(vllm_config=self.vllm_config,
75 + rpc_rank=0)
76 + distributed_init_method = get_distributed_init_method(
77 + get_ip(), get_open_port())
78 + local_rank = 0
79 + # set local rank as the device index if specified
80 + device_info = self.vllm_config.device_config.device.__str__().split(
81 + ":")
82 + if len(device_info) > 1:
83 + local_rank = int(device_info[1])
84 + rank = 0
85 + is_driver_worker = True
86 + kwargs = dict(
87 + vllm_config=self.vllm_config,
88 + local_rank=local_rank,
89 + rank=rank,
90 + distributed_init_method=distributed_init_method,
91 + is_driver_worker=is_driver_worker,
92 + )
93 + self.collective_rpc("init_worker", args=([kwargs], ))
94 + self.collective_rpc("init_device")
95 +
96 + @property
97 + def max_concurrent_batches(self) -> int:
98 + return 2
99 +
100 + def shutdown(self):
101 + if hasattr(self, 'thread_pool'):
102 + self.thread_pool.shutdown(wait=False)
tests/tensorizer_loader/test_tensorizer.py
+301 -19
@@ -1,36 +1,51 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import asyncio
5 import gc
6 +import json
7 import os
8 import pathlib
9 import subprocess
10 +import sys
11 +from typing import Any
12
13 import pytest
14 import torch
15
12 -from vllm import SamplingParams
16 +import vllm.model_executor.model_loader.tensorizer
17 +from vllm import LLM, SamplingParams
18 from vllm.engine.arg_utils import EngineArgs
14 -# yapf conflicts with isort for this docstring
19 # yapf: disable
20 from vllm.model_executor.model_loader.tensorizer import (TensorizerConfig,
21 TensorSerializer,
22 is_vllm_tensorized,
23 open_stream,
24 tensorize_vllm_model)
25 +from vllm.model_executor.model_loader.tensorizer_loader import (
26 + BLACKLISTED_TENSORIZER_ARGS)
27 # yapf: enable
28 from vllm.utils import PlaceholderModule
29
24 -from ..utils import VLLM_PATH
30 +from ..utils import VLLM_PATH, RemoteOpenAIServer
31 +from .conftest import DummyExecutor, assert_from_collective_rpc
32
33 try:
34 + import tensorizer
35 from tensorizer import EncryptionParams
36 except ImportError:
37 tensorizer = PlaceholderModule("tensorizer") # type: ignore[assignment]
38 EncryptionParams = tensorizer.placeholder_attr("EncryptionParams")
39
40 +
41 +class TensorizerCaughtError(Exception):
42 + pass
43 +
44 +
45 EXAMPLES_PATH = VLLM_PATH / "examples"
46
47 +pytest_plugins = "pytest_asyncio",
48 +
49 prompts = [
50 "Hello, my name is",
51 "The president of the United States is",
@@ -40,9 +55,37 @@ prompts = [
55 # Create a sampling params object.
56 sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
57
43 -model_ref = "facebook/opt-125m"
44 -tensorize_model_for_testing_script = os.path.join(
45 - os.path.dirname(__file__), "tensorize_vllm_model_for_testing.py")
58 +
59 +def patch_init_and_catch_error(self, obj, method_name,
60 + expected_error: type[Exception]):
61 + original = getattr(obj, method_name, None)
62 + if original is None:
63 + raise ValueError("Method '{}' not found.".format(method_name))
64 +
65 + def wrapper(*args, **kwargs):
66 + try:
67 + return original(*args, **kwargs)
68 + except expected_error as err:
69 + raise TensorizerCaughtError from err
70 +
71 + setattr(obj, method_name, wrapper)
72 +
73 + self.load_model()
74 +
75 +
76 +def assert_specific_tensorizer_error_is_raised(
77 + executor,
78 + obj: Any,
79 + method_name: str,
80 + expected_error: type[Exception],
81 +):
82 + with pytest.raises(TensorizerCaughtError):
83 + executor.collective_rpc(patch_init_and_catch_error,
84 + args=(
85 + obj,
86 + method_name,
87 + expected_error,
88 + ))
89
90
91 def is_curl_installed():
@@ -81,11 +124,10 @@ def test_can_deserialize_s3(vllm_runner):
124
125 @pytest.mark.skipif(not is_curl_installed(), reason="cURL is not installed")
126 def test_deserialized_encrypted_vllm_model_has_same_outputs(
84 - vllm_runner, tmp_path):
127 + model_ref, vllm_runner, tmp_path, model_path):
128 args = EngineArgs(model=model_ref)
129 with vllm_runner(model_ref) as vllm_model:
87 - model_path = tmp_path / (model_ref + ".tensors")
88 - key_path = tmp_path / (model_ref + ".key")
130 + key_path = tmp_path / model_ref / "model.key"
131 write_keyfile(key_path)
132
133 outputs = vllm_model.generate(prompts, sampling_params)
@@ -111,9 +153,9 @@ def test_deserialized_encrypted_vllm_model_has_same_outputs(
153
154
155 def test_deserialized_hf_model_has_same_outputs(hf_runner, vllm_runner,
114 - tmp_path):
156 + tmp_path, model_ref,
157 + model_path):
158 with hf_runner(model_ref) as hf_model:
116 - model_path = tmp_path / (model_ref + ".tensors")
159 max_tokens = 50
160 outputs = hf_model.generate_greedy(prompts, max_tokens=max_tokens)
161 with open_stream(model_path, "wb+") as stream:
@@ -123,7 +165,7 @@ def test_deserialized_hf_model_has_same_outputs(hf_runner, vllm_runner,
165 with vllm_runner(model_ref,
166 load_format="tensorizer",
167 model_loader_extra_config=TensorizerConfig(
126 - tensorizer_uri=model_path,
168 + tensorizer_uri=str(model_path),
169 num_readers=1,
170 )) as loaded_hf_model:
171 deserialized_outputs = loaded_hf_model.generate_greedy(
@@ -132,7 +174,7 @@ def test_deserialized_hf_model_has_same_outputs(hf_runner, vllm_runner,
174 assert outputs == deserialized_outputs
175
176
135 -def test_load_without_tensorizer_load_format(vllm_runner, capfd):
177 +def test_load_without_tensorizer_load_format(vllm_runner, capfd, model_ref):
178 model = None
179 try:
180 model = vllm_runner(
@@ -150,7 +192,8 @@ def test_load_without_tensorizer_load_format(vllm_runner, capfd):
192 torch.cuda.empty_cache()
193
194
153 -def test_raise_value_error_on_invalid_load_format(vllm_runner, capfd):
195 +def test_raise_value_error_on_invalid_load_format(vllm_runner, capfd,
196 + model_ref):
197 model = None
198 try:
199 model = vllm_runner(
@@ -208,7 +251,7 @@ def test_deserialized_encrypted_vllm_model_with_tp_has_same_outputs(
251 outputs = base_model.generate(prompts, sampling_params)
252
253 # load model with two shards and serialize with encryption
211 - model_path = str(tmp_path / (model_ref + "-%02d.tensors"))
254 + model_path = str(tmp_path / model_ref / "model-%02d.tensors")
255 key_path = tmp_path / (model_ref + ".key")
256
257 tensorizer_config = TensorizerConfig(
@@ -242,13 +285,12 @@ def test_deserialized_encrypted_vllm_model_with_tp_has_same_outputs(
285
286
287 @pytest.mark.flaky(reruns=3)
245 -def test_vllm_tensorized_model_has_same_outputs(vllm_runner, tmp_path):
288 +def test_vllm_tensorized_model_has_same_outputs(model_ref, vllm_runner,
289 + tmp_path, model_path):
290 gc.collect()
291 torch.cuda.empty_cache()
248 - model_ref = "facebook/opt-125m"
249 - model_path = tmp_path / (model_ref + ".tensors")
292 config = TensorizerConfig(tensorizer_uri=str(model_path))
251 - args = EngineArgs(model=model_ref, device="cuda")
293 + args = EngineArgs(model=model_ref)
294
295 with vllm_runner(model_ref) as vllm_model:
296 outputs = vllm_model.generate(prompts, sampling_params)
@@ -264,3 +306,243 @@ def test_vllm_tensorized_model_has_same_outputs(vllm_runner, tmp_path):
306 # noqa: E501
307
308 assert outputs == deserialized_outputs
309 +
310 +
311 +def test_load_with_just_model_tensors(just_serialize_model_tensors, model_ref):
312 + # For backwards compatibility, ensure Tensorizer can be still be loaded
313 + # for inference by passing the model reference name, not a local/S3 dir,
314 + # and the location of the model tensors
315 +
316 + model_dir = just_serialize_model_tensors
317 +
318 + extra_config = {"tensorizer_uri": f"{model_dir}/model.tensors"}
319 +
320 + ## Start OpenAI API server
321 + args = [
322 + "--load-format",
323 + "tensorizer",
324 + "--model-loader-extra-config",
325 + json.dumps(extra_config),
326 + ]
327 +
328 + with RemoteOpenAIServer(model_ref, args):
329 + # This test only concerns itself with being able to load the model
330 + # and successfully initialize the server
331 + pass
332 +
333 +
334 +def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path):
335 +
336 + serialization_params = {
337 + "limit_cpu_concurrency": 2,
338 + }
339 + model_ref = "facebook/opt-125m"
340 + model_path = tmp_path / (model_ref + ".tensors")
341 + config = TensorizerConfig(tensorizer_uri=str(model_path),
342 + serialization_kwargs=serialization_params)
343 + llm = LLM(model=model_ref, )
344 +
345 + def serialization_test(self, *args, **kwargs):
346 + # This is performed in the ephemeral worker process, so monkey-patching
347 + # will actually work, and cleanup is guaranteed so don't
348 + # need to reset things
349 +
350 + original_dict = serialization_params
351 + to_compare = {}
352 +
353 + original = tensorizer.serialization.TensorSerializer.__init__
354 +
355 + def tensorizer_serializer_wrapper(self, *args, **kwargs):
356 + nonlocal to_compare
357 + to_compare = kwargs.copy()
358 + return original(self, *args, **kwargs)
359 +
360 + tensorizer.serialization.TensorSerializer.__init__ = (
361 + tensorizer_serializer_wrapper)
362 +
363 + tensorizer_config = TensorizerConfig(**kwargs["tensorizer_config"])
364 + self.save_tensorized_model(tensorizer_config=tensorizer_config, )
365 + return to_compare | original_dict == to_compare
366 +
367 + kwargs = {"tensorizer_config": config.to_serializable()}
368 +
369 + assert assert_from_collective_rpc(llm, serialization_test, kwargs)
370 +
371 +
372 +def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(
373 + tmp_path, capfd):
374 +
375 + deserialization_kwargs = {
376 + "num_readers": "bar", # illegal value
377 + }
378 +
379 + serialization_params = {
380 + "limit_cpu_concurrency": 2,
381 + }
382 +
383 + model_ref = "facebook/opt-125m"
384 + model_path = tmp_path / (model_ref + ".tensors")
385 + config = TensorizerConfig(tensorizer_uri=str(model_path),
386 + serialization_kwargs=serialization_params)
387 +
388 + args = EngineArgs(model=model_ref)
389 + tensorize_vllm_model(args, config)
390 +
391 + loader_tc = TensorizerConfig(
392 + tensorizer_uri=str(model_path),
393 + deserialization_kwargs=deserialization_kwargs,
394 + )
395 +
396 + engine_args = EngineArgs(
397 + model="facebook/opt-125m",
398 + load_format="tensorizer",
399 + model_loader_extra_config=loader_tc.to_serializable(),
400 + )
401 +
402 + vllm_config = engine_args.create_engine_config()
403 + executor = DummyExecutor(vllm_config)
404 +
405 + assert_specific_tensorizer_error_is_raised(
406 + executor,
407 + tensorizer.serialization.TensorDeserializer,
408 + "__init__",
409 + TypeError,
410 + )
411 +
412 +
413 +def test_assert_stream_kwargs_passed_to_tensor_deserializer(tmp_path, capfd):
414 +
415 + deserialization_kwargs = {
416 + "num_readers": 1,
417 + }
418 +
419 + serialization_params = {
420 + "limit_cpu_concurrency": 2,
421 + }
422 +
423 + model_ref = "facebook/opt-125m"
424 + model_path = tmp_path / (model_ref + ".tensors")
425 + config = TensorizerConfig(tensorizer_uri=str(model_path),
426 + serialization_kwargs=serialization_params)
427 +
428 + args = EngineArgs(model=model_ref)
429 + tensorize_vllm_model(args, config)
430 +
431 + stream_kwargs = {"mode": "foo"}
432 +
433 + loader_tc = TensorizerConfig(
434 + tensorizer_uri=str(model_path),
435 + deserialization_kwargs=deserialization_kwargs,
436 + stream_kwargs=stream_kwargs,
437 + )
438 +
439 + engine_args = EngineArgs(
440 + model="facebook/opt-125m",
441 + load_format="tensorizer",
442 + model_loader_extra_config=loader_tc.to_serializable(),
443 + )
444 +
445 + vllm_config = engine_args.create_engine_config()
446 + executor = DummyExecutor(vllm_config)
447 +
448 + assert_specific_tensorizer_error_is_raised(
449 + executor,
450 + vllm.model_executor.model_loader.tensorizer,
451 + "open_stream",
452 + ValueError,
453 + )
454 +
455 +
456 +@pytest.mark.asyncio
457 +async def test_serialize_and_serve_entrypoints(tmp_path):
458 + model_ref = "facebook/opt-125m"
459 +
460 + suffix = "test"
461 + try:
462 + result = subprocess.run([
463 + sys.executable,
464 + f"{VLLM_PATH}/examples/others/tensorize_vllm_model.py", "--model",
465 + model_ref, "serialize", "--serialized-directory",
466 + str(tmp_path), "--suffix", suffix, "--serialization-kwargs",
467 + '{"limit_cpu_concurrency": 4}'
468 + ],
469 + check=True,
470 + capture_output=True,
471 + text=True)
472 + except subprocess.CalledProcessError as e:
473 + print("Tensorizing failed.")
474 + print("STDOUT:\n", e.stdout)
475 + print("STDERR:\n", e.stderr)
476 + raise
477 +
478 + assert "Successfully serialized" in result.stdout
479 +
480 + # Next, try to serve with vllm serve
481 + model_uri = tmp_path / "vllm" / model_ref / suffix / "model.tensors"
482 +
483 + model_loader_extra_config = {
484 + "tensorizer_uri": str(model_uri),
485 + "stream_kwargs": {
486 + "force_http": False,
487 + },
488 + "deserialization_kwargs": {
489 + "verify_hash": True,
490 + "num_readers": 8,
491 + }
492 + }
493 +
494 + cmd = [
495 + "-m", "vllm.entrypoints.cli.main", "serve", "--host", "localhost",
496 + "--load-format", "tensorizer", model_ref,
497 + "--model-loader-extra-config",
498 + json.dumps(model_loader_extra_config, indent=2)
499 + ]
500 +
501 + proc = await asyncio.create_subprocess_exec(
502 + sys.executable,
503 + *cmd,
504 + stdout=asyncio.subprocess.PIPE,
505 + stderr=asyncio.subprocess.STDOUT,
506 + )
507 +
508 + assert proc.stdout is not None
509 + fut = proc.stdout.readuntil(b"Application startup complete.")
510 +
511 + try:
512 + await asyncio.wait_for(fut, 180)
513 + except asyncio.TimeoutError:
514 + pytest.fail("Server did not start successfully")
515 + finally:
516 + proc.terminate()
517 + await proc.communicate()
518 +
519 +
520 +@pytest.mark.parametrize("illegal_value", BLACKLISTED_TENSORIZER_ARGS)
521 +def test_blacklisted_parameter_for_loading(tmp_path, vllm_runner, capfd,
522 + illegal_value):
523 +
524 + serialization_params = {
525 + "limit_cpu_concurrency": 2,
526 + }
527 +
528 + model_ref = "facebook/opt-125m"
529 + model_path = tmp_path / (model_ref + ".tensors")
530 + config = TensorizerConfig(tensorizer_uri=str(model_path),
531 + serialization_kwargs=serialization_params)
532 +
533 + args = EngineArgs(model=model_ref)
534 + tensorize_vllm_model(args, config)
535 +
536 + loader_tc = {"tensorizer_uri": str(model_path), illegal_value: "foo"}
537 +
538 + try:
539 + vllm_runner(
540 + model_ref,
541 + load_format="tensorizer",
542 + model_loader_extra_config=loader_tc,
543 + )
544 + except RuntimeError:
545 + out, err = capfd.readouterr()
546 + combined_output = out + err
547 + assert (f"ValueError: {illegal_value} is not an allowed "
548 + f"Tensorizer argument.") in combined_output
vllm/config.py
+7 -3
@@ -686,8 +686,11 @@ class ModelConfig:
686
687 # If tokenizer is same as model, download to same directory
688 if model == tokenizer:
689 - s3_model.pull_files(
690 - model, ignore_pattern=["*.pt", "*.safetensors", "*.bin"])
689 + s3_model.pull_files(model,
690 + ignore_pattern=[
691 + "*.pt", "*.safetensors", "*.bin",
692 + "*.tensors"
693 + ])
694 self.tokenizer = s3_model.dir
695 return
696
@@ -695,7 +698,8 @@ class ModelConfig:
698 if is_s3(tokenizer):
699 s3_tokenizer = S3Model()
700 s3_tokenizer.pull_files(
698 - model, ignore_pattern=["*.pt", "*.safetensors", "*.bin"])
701 + model,
702 + ignore_pattern=["*.pt", "*.safetensors", "*.bin", "*.tensors"])
703 self.tokenizer = s3_tokenizer.dir
704
705 def _init_multimodal_config(self) -> Optional["MultiModalConfig"]:
vllm/engine/arg_utils.py
+34 -2
@@ -58,7 +58,8 @@ def parse_type(return_type: Callable[[str], T]) -> Callable[[str], T]:
58
59 def _parse_type(val: str) -> T:
60 try:
61 - if return_type is json.loads and not re.match("^{.*}$", val):
61 + if return_type is json.loads and not re.match(
62 + r"(?s)^\s*{.*}\s*$", val):
63 return cast(T, nullable_kvs(val))
64 return return_type(val)
65 except ValueError as e:
@@ -80,7 +81,7 @@ def optional_type(
81
82
83 def union_dict_and_str(val: str) -> Optional[Union[str, dict[str, str]]]:
83 - if not re.match("^{.*}$", val):
84 + if not re.match(r"(?s)^\s*{.*}\s*$", val):
85 return str(val)
86 return optional_type(json.loads)(val)
87
@@ -1001,11 +1002,42 @@ class EngineArgs:
1002 override_attention_dtype=self.override_attention_dtype,
1003 )
1004
1005 + def valid_tensorizer_config_provided(self) -> bool:
1006 + """
1007 + Checks if a parseable TensorizerConfig was passed to
1008 + self.model_loader_extra_config. It first checks if the config passed
1009 + is a dict or a TensorizerConfig object directly, and if the latter is
1010 + true (by checking that the object has TensorizerConfig's
1011 + .to_serializable() method), converts it in to a serializable dict
1012 + format
1013 + """
1014 + if self.model_loader_extra_config:
1015 + if hasattr(self.model_loader_extra_config, "to_serializable"):
1016 + self.model_loader_extra_config = (
1017 + self.model_loader_extra_config.to_serializable())
1018 + for allowed_to_pass in ["tensorizer_uri", "tensorizer_dir"]:
1019 + try:
1020 + self.model_loader_extra_config[allowed_to_pass]
1021 + return False
1022 + except KeyError:
1023 + pass
1024 + return True
1025 +
1026 def create_load_config(self) -> LoadConfig:
1027
1028 if self.quantization == "bitsandbytes":
1029 self.load_format = "bitsandbytes"
1030
1031 + if (self.load_format == "tensorizer"
1032 + and self.valid_tensorizer_config_provided()):
1033 + logger.info("Inferring Tensorizer args from %s", self.model)
1034 + self.model_loader_extra_config = {"tensorizer_dir": self.model}
1035 + else:
1036 + logger.info(
1037 + "Using Tensorizer args from --model-loader-extra-config. "
1038 + "Note that you can now simply pass the S3 directory in the "
1039 + "model tag instead of providing the JSON string.")
1040 +
1041 return LoadConfig(
1042 load_format=self.load_format,
1043 download_dir=self.download_dir,
vllm/lora/models.py
+4 -3
@@ -245,9 +245,10 @@ class LoRAModel(AdapterModel):
245 lora_tensor_path = os.path.join(tensorizer_config.tensorizer_dir,
246 "adapter_model.tensors")
247 tensorizer_args = tensorizer_config._construct_tensorizer_args()
248 - tensors = TensorDeserializer(lora_tensor_path,
249 - dtype=tensorizer_config.dtype,
250 - **tensorizer_args.deserializer_params)
248 + tensors = TensorDeserializer(
249 + lora_tensor_path,
250 + dtype=tensorizer_config.dtype,
251 + **tensorizer_args.deserialization_kwargs)
252 check_unexpected_modules(tensors)
253
254 elif os.path.isfile(lora_tensor_path):
vllm/lora/peft_helper.py
+1 -1
@@ -106,7 +106,7 @@ class PEFTHelper:
106 "adapter_config.json")
107 with open_stream(lora_config_path,
108 mode="rb",
109 - **tensorizer_args.stream_params) as f:
109 + **tensorizer_args.stream_kwargs) as f:
110 config = json.load(f)
111
112 logger.info("Successfully deserialized LoRA config from %s",
vllm/model_executor/model_loader/tensorizer.py
+257 -121
@@ -5,18 +5,18 @@ import argparse
5 import contextlib
6 import contextvars
7 import dataclasses
8 -import io
8 import json
9 import os
10 +import tempfile
11 import threading
12 import time
13 -from collections.abc import Generator
14 -from dataclasses import dataclass
15 -from functools import partial
16 -from typing import TYPE_CHECKING, Any, BinaryIO, Optional, Union
13 +from collections.abc import Generator, MutableMapping
14 +from dataclasses import asdict, dataclass, field, fields
15 +from typing import TYPE_CHECKING, Any, ClassVar, Optional, Union
16
17 import regex as re
18 import torch
19 +from huggingface_hub import snapshot_download
20 from torch import nn
21 from torch.utils._python_dispatch import TorchDispatchMode
22 from transformers import PretrainedConfig
@@ -39,10 +39,6 @@ try:
39 from tensorizer.utils import (convert_bytes, get_mem_usage,
40 no_init_or_tensor)
41
42 - _read_stream, _write_stream = (partial(
43 - open_stream,
44 - mode=mode,
45 - ) for mode in ("rb", "wb+"))
42 except ImportError:
43 tensorizer = PlaceholderModule("tensorizer")
44 DecryptionParams = tensorizer.placeholder_attr("DecryptionParams")
@@ -54,9 +50,6 @@ except ImportError:
50 get_mem_usage = tensorizer.placeholder_attr("utils.get_mem_usage")
51 no_init_or_tensor = tensorizer.placeholder_attr("utils.no_init_or_tensor")
52
57 - _read_stream = tensorizer.placeholder_attr("_read_stream")
58 - _write_stream = tensorizer.placeholder_attr("_write_stream")
59 -
53 __all__ = [
54 'EncryptionParams', 'DecryptionParams', 'TensorDeserializer',
55 'TensorSerializer', 'open_stream', 'convert_bytes', 'get_mem_usage',
@@ -66,6 +59,23 @@ __all__ = [
59 logger = init_logger(__name__)
60
61
62 +def is_valid_deserialization_uri(uri: Optional[str]) -> bool:
63 + if uri:
64 + scheme = uri.lower().split("://")[0]
65 + return scheme in {"s3", "http", "https"} or os.path.exists(uri)
66 + return False
67 +
68 +
69 +def tensorizer_kwargs_arg(value):
70 + loaded = json.loads(value)
71 + if not isinstance(loaded, dict):
72 + raise argparse.ArgumentTypeError(
73 + f"Not deserializable to dict: {value}. serialization_kwargs and "
74 + f"deserialization_kwargs must be "
75 + f"deserializable from a JSON string to a dictionary. ")
76 + return loaded
77 +
78 +
79 class MetaTensorMode(TorchDispatchMode):
80
81 def __torch_dispatch__(self, func, types, args=(), kwargs=None):
@@ -137,54 +147,143 @@ class _NoInitOrTensorImpl:
147
148
149 @dataclass
140 -class TensorizerConfig:
141 - tensorizer_uri: Union[str, None] = None
142 - vllm_tensorized: Optional[bool] = False
143 - verify_hash: Optional[bool] = False
150 +class TensorizerConfig(MutableMapping):
151 + tensorizer_uri: Optional[str] = None
152 + tensorizer_dir: Optional[str] = None
153 + vllm_tensorized: Optional[bool] = None
154 + verify_hash: Optional[bool] = None
155 num_readers: Optional[int] = None
156 encryption_keyfile: Optional[str] = None
157 s3_access_key_id: Optional[str] = None
158 s3_secret_access_key: Optional[str] = None
159 s3_endpoint: Optional[str] = None
149 - model_class: Optional[type[torch.nn.Module]] = None
150 - hf_config: Optional[PretrainedConfig] = None
151 - dtype: Optional[Union[str, torch.dtype]] = None
160 lora_dir: Optional[str] = None
153 - _is_sharded: bool = False
161 + stream_kwargs: Optional[dict[str, Any]] = None
162 + serialization_kwargs: Optional[dict[str, Any]] = None
163 + deserialization_kwargs: Optional[dict[str, Any]] = None
164 + _extra_serialization_attrs: Optional[dict[str, Any]] = field(init=False,
165 + default=None)
166 + model_class: Optional[type[torch.nn.Module]] = field(init=False,
167 + default=None)
168 + hf_config: Optional[PretrainedConfig] = field(init=False, default=None)
169 + dtype: Optional[Union[str, torch.dtype]] = field(init=False, default=None)
170 + _is_sharded: bool = field(init=False, default=False)
171 + _fields: ClassVar[tuple[str, ...]]
172 + _keys: ClassVar[frozenset[str]]
173 + """
174 + Args for the TensorizerConfig class. These are used to configure the
175 + behavior of model serialization and deserialization using Tensorizer.
176 +
177 + Args:
178 + tensorizer_uri: Path to serialized model tensors. Can be a local file
179 + path or a S3 URI. This is a required field unless lora_dir is
180 + provided and the config is meant to be used for the
181 + `tensorize_lora_adapter` function. Unless a `tensorizer_dir` or
182 + `lora_dir` is passed to this object's initializer, this is a required
183 + argument.
184 + tensorizer_dir: Path to a directory containing serialized model tensors,
185 + and all other potential model artifacts to load the model, such as
186 + configs and tokenizer files. Can be passed instead of `tensorizer_uri`
187 + where the `model.tensors` file will be assumed to be in this
188 + directory.
189 + vllm_tensorized: If True, indicates that the serialized model is a
190 + vLLM model. This is used to determine the behavior of the
191 + TensorDeserializer when loading tensors from a serialized model.
192 + It is far faster to deserialize a vLLM model as it utilizes
193 + tensorizer's optimized GPU loading. Note that this is now
194 + deprecated, as serialized vLLM models are now automatically
195 + inferred as vLLM models.
196 + verify_hash: If True, the hashes of each tensor will be verified against
197 + the hashes stored in the metadata. A `HashMismatchError` will be
198 + raised if any of the hashes do not match.
199 + num_readers: Controls how many threads are allowed to read concurrently
200 + from the source file. Default is `None`, which will dynamically set
201 + the number of readers based on the number of available
202 + resources and model size. This greatly increases performance.
203 + encryption_keyfile: File path to a binary file containing a
204 + binary key to use for decryption. `None` (the default) means
205 + no decryption. See the example script in
206 + examples/others/tensorize_vllm_model.py.
207 + s3_access_key_id: The access key for the S3 bucket. Can also be set via
208 + the S3_ACCESS_KEY_ID environment variable.
209 + s3_secret_access_key: The secret access key for the S3 bucket. Can also
210 + be set via the S3_SECRET_ACCESS_KEY environment variable.
211 + s3_endpoint: The endpoint for the S3 bucket. Can also be set via the
212 + S3_ENDPOINT_URL environment variable.
213 + lora_dir: Path to a directory containing LoRA adapter artifacts for
214 + serialization or deserialization. When serializing LoRA adapters
215 + this is the only necessary parameter to pass to this object's
216 + initializer.
217 + """
218
219 def __post_init__(self):
220 # check if the configuration is for a sharded vLLM model
221 self._is_sharded = isinstance(self.tensorizer_uri, str) \
222 and re.search(r'%0\dd', self.tensorizer_uri) is not None
159 - if not self.tensorizer_uri and not self.lora_dir:
160 - raise ValueError("tensorizer_uri must be provided.")
161 - if not self.tensorizer_uri and self.lora_dir:
162 - self.tensorizer_uri = f"{self.lora_dir}/adapter_model.tensors"
163 - assert self.tensorizer_uri is not None, ("tensorizer_uri must be "
164 - "provided.")
165 - self.tensorizer_dir = os.path.dirname(self.tensorizer_uri)
166 - self.lora_dir = self.tensorizer_dir
167 -
168 - @classmethod
169 - def as_dict(cls, *args, **kwargs) -> dict[str, Any]:
170 - cfg = TensorizerConfig(*args, **kwargs)
171 - return dataclasses.asdict(cfg)
223
173 - def to_dict(self) -> dict[str, Any]:
174 - return dataclasses.asdict(self)
224 + if self.tensorizer_dir and self.tensorizer_uri:
225 + raise ValueError(
226 + "Either tensorizer_dir or tensorizer_uri must be provided, "
227 + "not both.")
228 + if self.tensorizer_dir and self.lora_dir:
229 + raise ValueError(
230 + "Only one of tensorizer_dir or lora_dir may be specified. "
231 + "Use lora_dir exclusively when serializing LoRA adapters, "
232 + "and tensorizer_dir or tensorizer_uri otherwise.")
233 + if not self.tensorizer_uri:
234 + if self.lora_dir:
235 + self.tensorizer_uri = f"{self.lora_dir}/adapter_model.tensors"
236 + elif self.tensorizer_dir:
237 + self.tensorizer_uri = f"{self.tensorizer_dir}/model.tensors"
238 + else:
239 + raise ValueError("Unable to resolve tensorizer_uri. "
240 + "A valid tensorizer_uri or tensorizer_dir "
241 + "must be provided for deserialization, and a "
242 + "valid tensorizer_uri, tensorizer_uri, or "
243 + "lora_dir for serialization.")
244 + else:
245 + self.tensorizer_dir = os.path.dirname(self.tensorizer_uri)
246 +
247 + if not self.serialization_kwargs:
248 + self.serialization_kwargs = {}
249 + if not self.deserialization_kwargs:
250 + self.deserialization_kwargs = {}
251 +
252 + def to_serializable(self) -> dict[str, Any]:
253 + # Due to TensorizerConfig needing to be msgpack-serializable, it needs
254 + # support for morphing back and forth between itself and its dict
255 + # representation
256 +
257 + # TensorizerConfig's representation as a dictionary is meant to be
258 + # linked to TensorizerConfig in such a way that the following is
259 + # technically initializable:
260 + # TensorizerConfig(**my_tensorizer_cfg.to_serializable())
261 +
262 + # This means the dict must not retain non-initializable parameters
263 + # and post-init attribute states
264 +
265 + # Also don't want to retain private and unset parameters, so only retain
266 + # not None values and public attributes
267 +
268 + raw_tc_dict = asdict(self)
269 + blacklisted = []
270 +
271 + if "tensorizer_uri" in raw_tc_dict and "tensorizer_dir" in raw_tc_dict:
272 + blacklisted.append("tensorizer_dir")
273 +
274 + if "tensorizer_dir" in raw_tc_dict and "lora_dir" in raw_tc_dict:
275 + blacklisted.append("tensorizer_dir")
276 +
277 + tc_dict = {}
278 + for k, v in raw_tc_dict.items():
279 + if (k not in blacklisted and k not in tc_dict
280 + and not k.startswith("_") and v is not None):
281 + tc_dict[k] = v
282 +
283 + return tc_dict
284
285 def _construct_tensorizer_args(self) -> "TensorizerArgs":
177 - tensorizer_args = {
178 - "tensorizer_uri": self.tensorizer_uri,
179 - "vllm_tensorized": self.vllm_tensorized,
180 - "verify_hash": self.verify_hash,
181 - "num_readers": self.num_readers,
182 - "encryption_keyfile": self.encryption_keyfile,
183 - "s3_access_key_id": self.s3_access_key_id,
184 - "s3_secret_access_key": self.s3_secret_access_key,
185 - "s3_endpoint": self.s3_endpoint,
186 - }
187 - return TensorizerArgs(**tensorizer_args) # type: ignore
286 + return TensorizerArgs(self) # type: ignore
287
288 def verify_with_parallel_config(
289 self,
@@ -209,81 +308,76 @@ class TensorizerConfig:
308 tensorizer_args = self._construct_tensorizer_args()
309
310 return open_stream(self.tensorizer_uri,
212 - **tensorizer_args.stream_params)
311 + **tensorizer_args.stream_kwargs)
312 +
313 + def keys(self):
314 + return self._keys
315 +
316 + def __len__(self):
317 + return len(fields(self))
318 +
319 + def __iter__(self):
320 + return iter(self._fields)
321 +
322 + def __getitem__(self, item: str) -> Any:
323 + if item not in self.keys():
324 + raise KeyError(item)
325 + return getattr(self, item)
326 +
327 + def __setitem__(self, key: str, value: Any) -> None:
328 + if key not in self.keys():
329 + # Disallow modifying invalid keys
330 + raise KeyError(key)
331 + setattr(self, key, value)
332 +
333 + def __delitem__(self, key, /):
334 + if key not in self.keys():
335 + raise KeyError(key)
336 + delattr(self, key)
337 +
338 +
339 +TensorizerConfig._fields = tuple(f.name for f in fields(TensorizerConfig))
340 +TensorizerConfig._keys = frozenset(TensorizerConfig._fields)
341
342
343 @dataclass
344 class TensorizerArgs:
217 - tensorizer_uri: Union[io.BufferedIOBase, io.RawIOBase, BinaryIO, str,
218 - bytes, os.PathLike, int]
219 - vllm_tensorized: Optional[bool] = False
220 - verify_hash: Optional[bool] = False
221 - num_readers: Optional[int] = None
345 + tensorizer_uri: Optional[str] = None
346 + tensorizer_dir: Optional[str] = None
347 encryption_keyfile: Optional[str] = None
223 - s3_access_key_id: Optional[str] = None
224 - s3_secret_access_key: Optional[str] = None
225 - s3_endpoint: Optional[str] = None
226 - """
227 - Args for the TensorizerAgent class. These are used to configure the behavior
228 - of the TensorDeserializer when loading tensors from a serialized model.
229 -
230 - Args:
231 - tensorizer_uri: Path to serialized model tensors. Can be a local file
232 - path or a S3 URI. This is a required field unless lora_dir is
233 - provided and the config is meant to be used for the
234 - `tensorize_lora_adapter` function.
235 - vllm_tensorized: If True, indicates that the serialized model is a
236 - vLLM model. This is used to determine the behavior of the
237 - TensorDeserializer when loading tensors from a serialized model.
238 - It is far faster to deserialize a vLLM model as it utilizes
239 - tensorizer's optimized GPU loading. Note that this is now
240 - deprecated, as serialized vLLM models are now automatically
241 - inferred as vLLM models.
242 - verify_hash: If True, the hashes of each tensor will be verified against
243 - the hashes stored in the metadata. A `HashMismatchError` will be
244 - raised if any of the hashes do not match.
245 - num_readers: Controls how many threads are allowed to read concurrently
246 - from the source file. Default is `None`, which will dynamically set
247 - the number of readers based on the number of available
248 - resources and model size. This greatly increases performance.
249 - encryption_keyfile: File path to a binary file containing a
250 - binary key to use for decryption. `None` (the default) means
251 - no decryption. See the example script in
252 - examples/others/tensorize_vllm_model.py.
253 - s3_access_key_id: The access key for the S3 bucket. Can also be set via
254 - the S3_ACCESS_KEY_ID environment variable.
255 - s3_secret_access_key: The secret access key for the S3 bucket. Can also
256 - be set via the S3_SECRET_ACCESS_KEY environment variable.
257 - s3_endpoint: The endpoint for the S3 bucket. Can also be set via the
258 - S3_ENDPOINT_URL environment variable.
259 - """
348
261 - def __post_init__(self):
262 - self.file_obj = self.tensorizer_uri
263 - self.s3_access_key_id = self.s3_access_key_id or envs.S3_ACCESS_KEY_ID
264 - self.s3_secret_access_key = (self.s3_secret_access_key
349 + def __init__(self, tensorizer_config: TensorizerConfig):
350 + for k, v in tensorizer_config.items():
351 + setattr(self, k, v)
352 + self.file_obj = tensorizer_config.tensorizer_uri
353 + self.s3_access_key_id = (tensorizer_config.s3_access_key_id
354 + or envs.S3_ACCESS_KEY_ID)
355 + self.s3_secret_access_key = (tensorizer_config.s3_secret_access_key
356 or envs.S3_SECRET_ACCESS_KEY)
266 - self.s3_endpoint = self.s3_endpoint or envs.S3_ENDPOINT_URL
267 - self.stream_params = {
268 - "s3_access_key_id": self.s3_access_key_id,
269 - "s3_secret_access_key": self.s3_secret_access_key,
270 - "s3_endpoint": self.s3_endpoint,
357 + self.s3_endpoint = tensorizer_config.s3_endpoint or envs.S3_ENDPOINT_URL
358 +
359 + self.stream_kwargs = {
360 + "s3_access_key_id": tensorizer_config.s3_access_key_id,
361 + "s3_secret_access_key": tensorizer_config.s3_secret_access_key,
362 + "s3_endpoint": tensorizer_config.s3_endpoint,
363 + **(tensorizer_config.stream_kwargs or {})
364 }
365
273 - self.deserializer_params = {
274 - "verify_hash": self.verify_hash,
275 - "encryption": self.encryption_keyfile,
276 - "num_readers": self.num_readers
366 + self.deserialization_kwargs = {
367 + "verify_hash": tensorizer_config.verify_hash,
368 + "encryption": tensorizer_config.encryption_keyfile,
369 + "num_readers": tensorizer_config.num_readers,
370 + **(tensorizer_config.deserialization_kwargs or {})
371 }
372
373 if self.encryption_keyfile:
374 with open_stream(
281 - self.encryption_keyfile,
282 - **self.stream_params,
375 + tensorizer_config.encryption_keyfile,
376 + **self.stream_kwargs,
377 ) as stream:
378 key = stream.read()
379 decryption_params = DecryptionParams.from_key(key)
286 - self.deserializer_params['encryption'] = decryption_params
380 + self.deserialization_kwargs['encryption'] = decryption_params
381
382 @staticmethod
383 def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
@@ -405,15 +499,22 @@ def init_tensorizer_model(tensorizer_config: TensorizerConfig,
499 def deserialize_tensorizer_model(model: nn.Module,
500 tensorizer_config: TensorizerConfig) -> None:
501 tensorizer_args = tensorizer_config._construct_tensorizer_args()
502 + if not is_valid_deserialization_uri(tensorizer_config.tensorizer_uri):
503 + raise ValueError(
504 + f"{tensorizer_config.tensorizer_uri} is not a valid "
505 + f"tensorizer URI. Please check that the URI is correct. "
506 + f"It must either point to a local existing file, or have a "
507 + f"S3, HTTP or HTTPS scheme.")
508 before_mem = get_mem_usage()
509 start = time.perf_counter()
410 - with _read_stream(
510 + with open_stream(
511 tensorizer_config.tensorizer_uri,
412 - **tensorizer_args.stream_params) as stream, TensorDeserializer(
512 + mode="rb",
513 + **tensorizer_args.stream_kwargs) as stream, TensorDeserializer(
514 stream,
515 dtype=tensorizer_config.dtype,
415 - device=f'cuda:{torch.cuda.current_device()}',
416 - **tensorizer_args.deserializer_params) as deserializer:
516 + device=torch.device("cuda", torch.cuda.current_device()),
517 + **tensorizer_args.deserialization_kwargs) as deserializer:
518 deserializer.load_into_module(model)
519 end = time.perf_counter()
520
@@ -442,9 +543,9 @@ def tensorizer_weights_iterator(
543 "examples/others/tensorize_vllm_model.py example script "
544 "for serializing vLLM models.")
545
445 - deserializer_args = tensorizer_args.deserializer_params
446 - stream_params = tensorizer_args.stream_params
447 - stream = open_stream(tensorizer_args.tensorizer_uri, **stream_params)
546 + deserializer_args = tensorizer_args.deserialization_kwargs
547 + stream_kwargs = tensorizer_args.stream_kwargs
548 + stream = open_stream(tensorizer_args.tensorizer_uri, **stream_kwargs)
549 with TensorDeserializer(stream, **deserializer_args,
550 device="cpu") as state:
551 yield from state.items()
@@ -465,8 +566,8 @@ def is_vllm_tensorized(tensorizer_config: "TensorizerConfig") -> bool:
566 """
567 tensorizer_args = tensorizer_config._construct_tensorizer_args()
568 deserializer = TensorDeserializer(open_stream(
468 - tensorizer_args.tensorizer_uri, **tensorizer_args.stream_params),
469 - **tensorizer_args.deserializer_params,
569 + tensorizer_args.tensorizer_uri, **tensorizer_args.stream_kwargs),
570 + **tensorizer_args.deserialization_kwargs,
571 lazy_load=True)
572 if tensorizer_config.vllm_tensorized:
573 logger.warning(
@@ -477,13 +578,41 @@ def is_vllm_tensorized(tensorizer_config: "TensorizerConfig") -> bool:
578 return ".vllm_tensorized_marker" in deserializer
579
580
581 +def serialize_extra_artifacts(
582 + tensorizer_args: TensorizerArgs,
583 + served_model_name: Union[str, list[str], None]) -> None:
584 + if not isinstance(served_model_name, str):
585 + raise ValueError(
586 + f"served_model_name must be a str for serialize_extra_artifacts, "
587 + f"not {type(served_model_name)}.")
588 +
589 + with tempfile.TemporaryDirectory() as tmpdir:
590 + snapshot_download(served_model_name,
591 + local_dir=tmpdir,
592 + ignore_patterns=[
593 + "*.pt", "*.safetensors", "*.bin", "*.cache",
594 + "*.gitattributes", "*.md"
595 + ])
596 + for artifact in os.scandir(tmpdir):
597 + if not artifact.is_file():
598 + continue
599 + with open(artifact.path, "rb") as f, open_stream(
600 + f"{tensorizer_args.tensorizer_dir}/{artifact.name}",
601 + mode="wb+",
602 + **tensorizer_args.stream_kwargs) as stream:
603 + logger.info("Writing artifact %s", artifact.name)
604 + stream.write(f.read())
605 +
606 +
607 def serialize_vllm_model(
608 model: nn.Module,
609 tensorizer_config: TensorizerConfig,
610 + model_config: "ModelConfig",
611 ) -> nn.Module:
612 model.register_parameter(
613 "vllm_tensorized_marker",
614 nn.Parameter(torch.tensor((1, ), device="meta"), requires_grad=False))
615 +
616 tensorizer_args = tensorizer_config._construct_tensorizer_args()
617
618 encryption_params = None
@@ -497,10 +626,16 @@ def serialize_vllm_model(
626 from vllm.distributed import get_tensor_model_parallel_rank
627 output_file = output_file % get_tensor_model_parallel_rank()
628
500 - with _write_stream(output_file, **tensorizer_args.stream_params) as stream:
501 - serializer = TensorSerializer(stream, encryption=encryption_params)
629 + with open_stream(output_file, mode="wb+",
630 + **tensorizer_args.stream_kwargs) as stream:
631 + serializer = TensorSerializer(stream,
632 + encryption=encryption_params,
633 + **tensorizer_config.serialization_kwargs)
634 serializer.write_module(model)
635 serializer.close()
636 +
637 + serialize_extra_artifacts(tensorizer_args, model_config.served_model_name)
638 +
639 logger.info("Successfully serialized model to %s", str(output_file))
640 return model
641
@@ -522,8 +657,9 @@ def tensorize_vllm_model(engine_args: "EngineArgs",
657 if generate_keyfile and (keyfile :=
658 tensorizer_config.encryption_keyfile) is not None:
659 encryption_params = EncryptionParams.random()
525 - with _write_stream(
660 + with open_stream(
661 keyfile,
662 + mode="wb+",
663 s3_access_key_id=tensorizer_config.s3_access_key_id,
664 s3_secret_access_key=tensorizer_config.s3_secret_access_key,
665 s3_endpoint=tensorizer_config.s3_endpoint,
@@ -537,13 +673,13 @@ def tensorize_vllm_model(engine_args: "EngineArgs",
673 engine = LLMEngine.from_engine_args(engine_args)
674 engine.model_executor.collective_rpc(
675 "save_tensorized_model",
540 - kwargs=dict(tensorizer_config=tensorizer_config),
676 + kwargs={"tensorizer_config": tensorizer_config.to_serializable()},
677 )
678 else:
679 engine = V1LLMEngine.from_vllm_config(engine_config)
680 engine.collective_rpc(
681 "save_tensorized_model",
546 - kwargs=dict(tensorizer_config=tensorizer_config),
682 + kwargs={"tensorizer_config": tensorizer_config.to_serializable()},
683 )
684
685
@@ -586,14 +722,14 @@ def tensorize_lora_adapter(lora_path: str,
722
723 with open_stream(f"{tensorizer_config.lora_dir}/adapter_config.json",
724 mode="wb+",
589 - **tensorizer_args.stream_params) as f:
725 + **tensorizer_args.stream_kwargs) as f:
726
727 f.write(json.dumps(config).encode("utf-8"))
728
729 lora_uri = (f"{tensorizer_config.lora_dir}"
730 f"/adapter_model.tensors")
731 with open_stream(lora_uri, mode="wb+",
596 - **tensorizer_args.stream_params) as f:
732 + **tensorizer_args.stream_kwargs) as f:
733 serializer = TensorSerializer(f)
734 serializer.write_state_dict(tensors)
735 serializer.close()
vllm/model_executor/model_loader/tensorizer_loader.py
+15
@@ -20,6 +20,18 @@ from vllm.model_executor.model_loader.utils import (get_model_architecture,
20
21 logger = init_logger(__name__)
22
23 +BLACKLISTED_TENSORIZER_ARGS = {
24 + "device", # vLLM decides this
25 + "dtype", # vLLM decides this
26 + "mode", # Not meant to be configurable by the user
27 +}
28 +
29 +
30 +def validate_config(config: dict):
31 + for k, v in config.items():
32 + if v is not None and k in BLACKLISTED_TENSORIZER_ARGS:
33 + raise ValueError(f"{k} is not an allowed Tensorizer argument.")
34 +
35
36 class TensorizerLoader(BaseModelLoader):
37 """Model loader using CoreWeave's tensorizer library."""
@@ -29,6 +41,7 @@ class TensorizerLoader(BaseModelLoader):
41 if isinstance(load_config.model_loader_extra_config, TensorizerConfig):
42 self.tensorizer_config = load_config.model_loader_extra_config
43 else:
44 + validate_config(load_config.model_loader_extra_config)
45 self.tensorizer_config = TensorizerConfig(
46 **load_config.model_loader_extra_config)
47
@@ -118,10 +131,12 @@ class TensorizerLoader(BaseModelLoader):
131 def save_model(
132 model: torch.nn.Module,
133 tensorizer_config: Union[TensorizerConfig, dict],
134 + model_config: ModelConfig,
135 ) -> None:
136 if isinstance(tensorizer_config, dict):
137 tensorizer_config = TensorizerConfig(**tensorizer_config)
138 serialize_vllm_model(
139 model=model,
140 tensorizer_config=tensorizer_config,
141 + model_config=model_config,
142 )
vllm/v1/worker/gpu_model_runner.py
+1
@@ -1820,6 +1820,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1820 TensorizerLoader.save_model(
1821 self.model,
1822 tensorizer_config=tensorizer_config,
1823 + model_config=self.model_config,
1824 )
1825
1826 def _get_prompt_logprobs_dict(
vllm/worker/model_runner.py
+1
@@ -1246,6 +1246,7 @@ class GPUModelRunnerBase(ModelRunnerBase[TModelInputForGPU]):
1246 TensorizerLoader.save_model(
1247 self.model,
1248 tensorizer_config=tensorizer_config,
1249 + model_config=self.model_config,
1250 )
1251
1252 def get_max_block_per_batch(self) -> int: