| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import argparse |
| 5 | import contextlib |
| 6 | import contextvars |
| 7 | import dataclasses |
| 8 | import json |
| 9 | import os |
| 10 | import tempfile |
| 11 | import threading |
| 12 | import time |
| 13 | from collections.abc import Generator, MutableMapping |
| 14 | from dataclasses import asdict, dataclass, field, fields |
| 15 | from typing import TYPE_CHECKING, Any, ClassVar |
| 16 | |
| 17 | import regex as re |
| 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, VllmConfig, set_current_vllm_config |
| 25 | from vllm.logger import init_logger |
| 26 | from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding |
| 27 | from vllm.platforms import current_platform |
| 28 | from vllm.transformers_utils.repo_utils import hf_api |
| 29 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 30 | from vllm.utils.import_utils import PlaceholderModule |
| 31 | |
| 32 | if TYPE_CHECKING: |
| 33 | from vllm.engine.arg_utils import EngineArgs |
| 34 | |
| 35 | try: |
| 36 | from tensorizer import ( |
| 37 | DecryptionParams, |
| 38 | EncryptionParams, |
| 39 | TensorDeserializer, |
| 40 | TensorSerializer, |
| 41 | ) |
| 42 | from tensorizer.stream_io import open_stream |
| 43 | from tensorizer.utils import convert_bytes, get_mem_usage, no_init_or_tensor |
| 44 | |
| 45 | except ImportError: |
| 46 | tensorizer = PlaceholderModule("tensorizer") |
| 47 | DecryptionParams = tensorizer.placeholder_attr("DecryptionParams") |
| 48 | EncryptionParams = tensorizer.placeholder_attr("EncryptionParams") |
| 49 | TensorDeserializer = tensorizer.placeholder_attr("TensorDeserializer") |
| 50 | TensorSerializer = tensorizer.placeholder_attr("TensorSerializer") |
| 51 | open_stream = tensorizer.placeholder_attr("stream_io.open_stream") |
| 52 | convert_bytes = tensorizer.placeholder_attr("utils.convert_bytes") |
| 53 | get_mem_usage = tensorizer.placeholder_attr("utils.get_mem_usage") |
| 54 | no_init_or_tensor = tensorizer.placeholder_attr("utils.no_init_or_tensor") |
| 55 | |
| 56 | __all__ = [ |
| 57 | "EncryptionParams", |
| 58 | "DecryptionParams", |
| 59 | "TensorDeserializer", |
| 60 | "TensorSerializer", |
| 61 | "open_stream", |
| 62 | "convert_bytes", |
| 63 | "get_mem_usage", |
| 64 | "no_init_or_tensor", |
| 65 | "TensorizerConfig", |
| 66 | ] |
| 67 | |
| 68 | logger = init_logger(__name__) |
| 69 | _TENSORIZER_ENGINE_CLEANUP_GRACE_S = 10.0 |
| 70 | |
| 71 | |
| 72 | def is_valid_deserialization_uri(uri: str | None) -> bool: |
| 73 | if uri: |
| 74 | scheme = uri.lower().split("://")[0] |
| 75 | return scheme in {"s3", "http", "https"} or os.path.exists(uri) |
| 76 | return False |
| 77 | |
| 78 | |
| 79 | def tensorizer_kwargs_arg(value): |
| 80 | loaded = json.loads(value) |
| 81 | if not isinstance(loaded, dict): |
| 82 | raise argparse.ArgumentTypeError( |
| 83 | f"Not deserializable to dict: {value}. serialization_kwargs and " |
| 84 | f"deserialization_kwargs must be " |
| 85 | f"deserializable from a JSON string to a dictionary. " |
| 86 | ) |
| 87 | return loaded |
| 88 | |
| 89 | |
| 90 | class MetaTensorMode(TorchDispatchMode): |
| 91 | def __torch_dispatch__(self, func, types, args=(), kwargs=None): |
| 92 | kwargs = kwargs or {} |
| 93 | |
| 94 | if func._schema.name == "aten::empty" and "device" not in kwargs: |
| 95 | kwargs["device"] = "meta" |
| 96 | |
| 97 | return func(*args, **kwargs) |
| 98 | |
| 99 | |
| 100 | def meta_tensor_mode( |
| 101 | loading_code=None, |
| 102 | ): |
| 103 | if loading_code is None: |
| 104 | return _NoInitOrTensorImpl.context_manager() |
| 105 | elif callable(loading_code): |
| 106 | with _NoInitOrTensorImpl.context_manager(): |
| 107 | return loading_code() |
| 108 | else: |
| 109 | raise TypeError( |
| 110 | "expected a callable to evaluate," |
| 111 | " or None if being used as a context manager;" |
| 112 | f' got an object of type "{type(loading_code).__name__}" instead.' |
| 113 | ) |
| 114 | |
| 115 | |
| 116 | class _NoInitOrTensorImpl: |
| 117 | _MODULES = (torch.nn.Linear, torch.nn.Embedding, torch.nn.LayerNorm) |
| 118 | _MODULE_ORIGINALS = tuple((m, m.reset_parameters) for m in _MODULES) |
| 119 | |
| 120 | is_active = contextvars.ContextVar("_NoInitOrTensorImpl.is_active", default=False) |
| 121 | _count_active: int = 0 |
| 122 | _count_active_lock = threading.Lock() |
| 123 | |
| 124 | @classmethod |
| 125 | @contextlib.contextmanager |
| 126 | def context_manager(cls): |
| 127 | if cls.is_active.get(): |
| 128 | yield |
| 129 | return |
| 130 | |
| 131 | with cls._count_active_lock: |
| 132 | cls._count_active += 1 |
| 133 | if cls._count_active == 1: |
| 134 | for mod in cls._MODULES: |
| 135 | mod.reset_parameters = cls._disable(mod.reset_parameters) |
| 136 | |
| 137 | reset_token = cls.is_active.set(True) |
| 138 | |
| 139 | try: |
| 140 | with MetaTensorMode(): |
| 141 | yield |
| 142 | finally: |
| 143 | cls.is_active.reset(reset_token) |
| 144 | with cls._count_active_lock: |
| 145 | cls._count_active -= 1 |
| 146 | if cls._count_active == 0: |
| 147 | for mod, original in cls._MODULE_ORIGINALS: |
| 148 | mod.reset_parameters = original |
| 149 | |
| 150 | @staticmethod |
| 151 | def _disable(func): |
| 152 | def wrapper(*args, **kwargs): |
| 153 | if not _NoInitOrTensorImpl.is_active.get(): |
| 154 | return func(*args, **kwargs) |
| 155 | |
| 156 | return wrapper |
| 157 | |
| 158 | |
| 159 | @dataclass |
| 160 | class TensorizerConfig(MutableMapping): |
| 161 | tensorizer_uri: str | None = None |
| 162 | tensorizer_dir: str | None = None |
| 163 | vllm_tensorized: bool | None = None |
| 164 | verify_hash: bool | None = None |
| 165 | num_readers: int | None = None |
| 166 | encryption_keyfile: str | None = None |
| 167 | s3_access_key_id: str | None = None |
| 168 | s3_secret_access_key: str | None = None |
| 169 | s3_endpoint: str | None = None |
| 170 | lora_dir: str | None = None |
| 171 | stream_kwargs: dict[str, Any] | None = None |
| 172 | serialization_kwargs: dict[str, Any] | None = None |
| 173 | deserialization_kwargs: dict[str, Any] | None = None |
| 174 | _extra_serialization_attrs: dict[str, Any] | None = field(init=False, default=None) |
| 175 | model_class: type[torch.nn.Module] | None = field(init=False, default=None) |
| 176 | hf_config: PreTrainedConfig | None = field(init=False, default=None) |
| 177 | dtype: str | torch.dtype | None = field(init=False, default=None) |
| 178 | _is_sharded: bool = field(init=False, default=False) |
| 179 | _fields: ClassVar[tuple[str, ...]] |
| 180 | _keys: ClassVar[frozenset[str]] |
| 181 | """Configuration class for Tensorizer settings. |
| 182 | |
| 183 | These settings configure the behavior of model serialization and |
| 184 | deserialization using Tensorizer. |
| 185 | |
| 186 | Attributes: |
| 187 | tensorizer_uri: Path to serialized model tensors. Can be a local file |
| 188 | path or a S3 URI. This is a required field unless lora_dir is |
| 189 | provided and the config is meant to be used for the |
| 190 | `tensorize_lora_adapter` function. Unless a `tensorizer_dir` or |
| 191 | `lora_dir` is passed to this object's initializer, this is |
| 192 | a required argument. |
| 193 | tensorizer_dir: Path to a directory containing serialized model tensors, |
| 194 | and all other potential model artifacts to load the model, such as |
| 195 | configs and tokenizer files. Can be passed instead of |
| 196 | `tensorizer_uri` where the `model.tensors` file will be assumed |
| 197 | to be in this directory. |
| 198 | vllm_tensorized: If True, indicates that the serialized model is a |
| 199 | vLLM model. This is used to determine the behavior of the |
| 200 | TensorDeserializer when loading tensors from a serialized model. |
| 201 | It is far faster to deserialize a vLLM model as it utilizes |
| 202 | tensorizer's optimized GPU loading. Note that this is now |
| 203 | deprecated, as serialized vLLM models are now automatically |
| 204 | inferred as vLLM models. |
| 205 | verify_hash: If True, the hashes of each tensor will be verified |
| 206 | against the hashes stored in the metadata. A `HashMismatchError` |
| 207 | will be raised if any of the hashes do not match. |
| 208 | num_readers: Controls how many threads are allowed to read concurrently |
| 209 | from the source file. Default is `None`, which will dynamically set |
| 210 | the number of readers based on the number of available |
| 211 | resources and model size. This greatly increases performance. |
| 212 | encryption_keyfile: File path to a binary file containing a |
| 213 | binary key to use for decryption. `None` (the default) means |
| 214 | no decryption. See the example script in |
| 215 | examples/features/tensorize_vllm_model.py. |
| 216 | s3_access_key_id: The access key for the S3 bucket. Can also be set via |
| 217 | the S3_ACCESS_KEY_ID environment variable. |
| 218 | s3_secret_access_key: The secret access key for the S3 bucket. Can also |
| 219 | be set via the S3_SECRET_ACCESS_KEY environment variable. |
| 220 | s3_endpoint: The endpoint for the S3 bucket. Can also be set via the |
| 221 | S3_ENDPOINT_URL environment variable. |
| 222 | lora_dir: Path to a directory containing LoRA adapter artifacts for |
| 223 | serialization or deserialization. When serializing LoRA adapters |
| 224 | this is the only necessary parameter to pass to this object's |
| 225 | initializer. |
| 226 | """ |
| 227 | |
| 228 | def __post_init__(self): |
| 229 | # check if the configuration is for a sharded vLLM model |
| 230 | self._is_sharded = ( |
| 231 | isinstance(self.tensorizer_uri, str) |
| 232 | and re.search(r"%0\dd", self.tensorizer_uri) is not None |
| 233 | ) |
| 234 | |
| 235 | if self.tensorizer_dir and self.lora_dir: |
| 236 | raise ValueError( |
| 237 | "Only one of tensorizer_dir or lora_dir may be specified. " |
| 238 | "Use lora_dir exclusively when serializing LoRA adapters, " |
| 239 | "and tensorizer_dir or tensorizer_uri otherwise." |
| 240 | ) |
| 241 | if self.tensorizer_dir and self.tensorizer_uri: |
| 242 | logger.warning_once( |
| 243 | "Provided both tensorizer_dir and tensorizer_uri. " |
| 244 | "Inferring tensorizer_dir from tensorizer_uri as the " |
| 245 | "latter takes precedence." |
| 246 | ) |
| 247 | self.tensorizer_dir = os.path.dirname(self.tensorizer_uri) |
| 248 | if not self.tensorizer_uri: |
| 249 | if self.lora_dir: |
| 250 | self.tensorizer_uri = f"{self.lora_dir}/adapter_model.tensors" |
| 251 | elif self.tensorizer_dir: |
| 252 | self.tensorizer_uri = f"{self.tensorizer_dir}/model.tensors" |
| 253 | else: |
| 254 | raise ValueError( |
| 255 | "Unable to resolve tensorizer_uri. " |
| 256 | "A valid tensorizer_uri or tensorizer_dir " |
| 257 | "must be provided for deserialization, and a " |
| 258 | "valid tensorizer_uri, tensorizer_uri, or " |
| 259 | "lora_dir for serialization." |
| 260 | ) |
| 261 | else: |
| 262 | self.tensorizer_dir = os.path.dirname(self.tensorizer_uri) |
| 263 | |
| 264 | if not self.serialization_kwargs: |
| 265 | self.serialization_kwargs = {} |
| 266 | if not self.deserialization_kwargs: |
| 267 | self.deserialization_kwargs = {} |
| 268 | |
| 269 | def to_serializable(self) -> dict[str, Any]: |
| 270 | # Due to TensorizerConfig needing to be msgpack-serializable, it needs |
| 271 | # support for morphing back and forth between itself and its dict |
| 272 | # representation |
| 273 | |
| 274 | # TensorizerConfig's representation as a dictionary is meant to be |
| 275 | # linked to TensorizerConfig in such a way that the following is |
| 276 | # technically initializable: |
| 277 | # TensorizerConfig(**my_tensorizer_cfg.to_serializable()) |
| 278 | |
| 279 | # This means the dict must not retain non-initializable parameters |
| 280 | # and post-init attribute states |
| 281 | |
| 282 | # Also don't want to retain private and unset parameters, so only retain |
| 283 | # not None values and public attributes |
| 284 | |
| 285 | raw_tc_dict = asdict(self) |
| 286 | blacklisted = [] |
| 287 | |
| 288 | if "tensorizer_uri" in raw_tc_dict and "tensorizer_dir" in raw_tc_dict: |
| 289 | blacklisted.append("tensorizer_dir") |
| 290 | |
| 291 | if "tensorizer_dir" in raw_tc_dict and "lora_dir" in raw_tc_dict: |
| 292 | blacklisted.append("tensorizer_dir") |
| 293 | |
| 294 | tc_dict = {} |
| 295 | for k, v in raw_tc_dict.items(): |
| 296 | if ( |
| 297 | k not in blacklisted |
| 298 | and k not in tc_dict |
| 299 | and not k.startswith("_") |
| 300 | and v is not None |
| 301 | ): |
| 302 | tc_dict[k] = v |
| 303 | |
| 304 | return tc_dict |
| 305 | |
| 306 | def _construct_tensorizer_args(self) -> "TensorizerArgs": |
| 307 | return TensorizerArgs(self) # type: ignore |
| 308 | |
| 309 | def verify_with_parallel_config( |
| 310 | self, |
| 311 | parallel_config: "ParallelConfig", |
| 312 | ) -> None: |
| 313 | if parallel_config.tensor_parallel_size > 1 and not self._is_sharded: |
| 314 | raise ValueError( |
| 315 | "For a sharded model, tensorizer_uri should include a" |
| 316 | " string format template like '%04d' to be formatted" |
| 317 | " with the rank of the shard" |
| 318 | ) |
| 319 | |
| 320 | def verify_with_model_config(self, model_config: "ModelConfig") -> None: |
| 321 | if model_config.quantization is not None and self.tensorizer_uri is not None: |
| 322 | logger.warning( |
| 323 | "Loading a model using Tensorizer with quantization on vLLM" |
| 324 | " is unstable and may lead to errors." |
| 325 | ) |
| 326 | |
| 327 | def open_stream(self, tensorizer_args: "TensorizerArgs | None" = None): |
| 328 | if tensorizer_args is None: |
| 329 | tensorizer_args = self._construct_tensorizer_args() |
| 330 | |
| 331 | return open_stream(self.tensorizer_uri, **tensorizer_args.stream_kwargs) |
| 332 | |
| 333 | def keys(self): |
| 334 | return self._keys |
| 335 | |
| 336 | def __len__(self): |
| 337 | return len(fields(self)) |
| 338 | |
| 339 | def __iter__(self): |
| 340 | return iter(self._fields) |
| 341 | |
| 342 | def __getitem__(self, item: str) -> Any: |
| 343 | if item not in self.keys(): |
| 344 | raise KeyError(item) |
| 345 | return getattr(self, item) |
| 346 | |
| 347 | def __setitem__(self, key: str, value: Any) -> None: |
| 348 | if key not in self.keys(): |
| 349 | # Disallow modifying invalid keys |
| 350 | raise KeyError(key) |
| 351 | setattr(self, key, value) |
| 352 | |
| 353 | def __delitem__(self, key, /): |
| 354 | if key not in self.keys(): |
| 355 | raise KeyError(key) |
| 356 | delattr(self, key) |
| 357 | |
| 358 | |
| 359 | TensorizerConfig._fields = tuple(f.name for f in fields(TensorizerConfig)) |
| 360 | TensorizerConfig._keys = frozenset(TensorizerConfig._fields) |
| 361 | |
| 362 | |
| 363 | @dataclass |
| 364 | class TensorizerArgs: |
| 365 | tensorizer_uri: str | None = None |
| 366 | tensorizer_dir: str | None = None |
| 367 | encryption_keyfile: str | None = None |
| 368 | |
| 369 | def __init__(self, tensorizer_config: TensorizerConfig): |
| 370 | for k, v in tensorizer_config.items(): |
| 371 | setattr(self, k, v) |
| 372 | self.file_obj = tensorizer_config.tensorizer_uri |
| 373 | self.s3_access_key_id = ( |
| 374 | tensorizer_config.s3_access_key_id or envs.S3_ACCESS_KEY_ID |
| 375 | ) |
| 376 | self.s3_secret_access_key = ( |
| 377 | tensorizer_config.s3_secret_access_key or envs.S3_SECRET_ACCESS_KEY |
| 378 | ) |
| 379 | self.s3_endpoint = tensorizer_config.s3_endpoint or envs.S3_ENDPOINT_URL |
| 380 | |
| 381 | self.stream_kwargs = { |
| 382 | "s3_access_key_id": tensorizer_config.s3_access_key_id, |
| 383 | "s3_secret_access_key": tensorizer_config.s3_secret_access_key, |
| 384 | "s3_endpoint": tensorizer_config.s3_endpoint, |
| 385 | **(tensorizer_config.stream_kwargs or {}), |
| 386 | } |
| 387 | |
| 388 | self.deserialization_kwargs = { |
| 389 | "verify_hash": tensorizer_config.verify_hash, |
| 390 | "encryption": tensorizer_config.encryption_keyfile, |
| 391 | "num_readers": tensorizer_config.num_readers, |
| 392 | **(tensorizer_config.deserialization_kwargs or {}), |
| 393 | } |
| 394 | |
| 395 | if self.encryption_keyfile: |
| 396 | with open_stream( |
| 397 | tensorizer_config.encryption_keyfile, |
| 398 | **self.stream_kwargs, |
| 399 | ) as stream: |
| 400 | key = stream.read() |
| 401 | decryption_params = DecryptionParams.from_key(key) |
| 402 | self.deserialization_kwargs["encryption"] = decryption_params |
| 403 | |
| 404 | @staticmethod |
| 405 | def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: |
| 406 | """Tensorizer CLI arguments.""" |
| 407 | # Tensorizer options arg group |
| 408 | group = parser.add_argument_group( |
| 409 | "tensorizer options", |
| 410 | description=( |
| 411 | "Options for configuring the behavior of the" |
| 412 | " tensorizer deserializer when " |
| 413 | "load_format=tensorizer is specified when " |
| 414 | "initializing an LLMEngine, either via the CLI " |
| 415 | "when running the vLLM OpenAI inference server " |
| 416 | "with a JSON string passed to " |
| 417 | "--model-loader-extra-config or as arguments given " |
| 418 | "to TensorizerConfig when passed to " |
| 419 | "model_loader_extra_config in the constructor " |
| 420 | "for LLMEngine." |
| 421 | ), |
| 422 | ) |
| 423 | |
| 424 | group.add_argument( |
| 425 | "--tensorizer-uri", |
| 426 | type=str, |
| 427 | help="Path to serialized model tensors. Can be a local file path," |
| 428 | " or an HTTP(S) or S3 URI.", |
| 429 | ) |
| 430 | group.add_argument( |
| 431 | "--verify-hash", |
| 432 | action="store_true", |
| 433 | help="If enabled, the hashes of each tensor will be verified" |
| 434 | " against the hashes stored in the file metadata. An exception" |
| 435 | " will be raised if any of the hashes do not match.", |
| 436 | ) |
| 437 | group.add_argument( |
| 438 | "--encryption-keyfile", |
| 439 | type=str, |
| 440 | default=None, |
| 441 | help="The file path to a binary file containing a binary key to " |
| 442 | "use for decryption. Can be a file path or S3 network URI.", |
| 443 | ) |
| 444 | group.add_argument( |
| 445 | "--num-readers", |
| 446 | default=None, |
| 447 | type=int, |
| 448 | help="Controls how many threads are allowed to read concurrently " |
| 449 | "from the source file. Default is `None`, which will dynamically " |
| 450 | "set the number of readers based on the available resources " |
| 451 | "and model size. This greatly increases performance.", |
| 452 | ) |
| 453 | group.add_argument( |
| 454 | "--s3-access-key-id", |
| 455 | type=str, |
| 456 | default=None, |
| 457 | help="The access key for the S3 bucket. Can also be set via the " |
| 458 | "S3_ACCESS_KEY_ID environment variable.", |
| 459 | ) |
| 460 | group.add_argument( |
| 461 | "--s3-secret-access-key", |
| 462 | type=str, |
| 463 | default=None, |
| 464 | help="The secret access key for the S3 bucket. Can also be set via " |
| 465 | "the S3_SECRET_ACCESS_KEY environment variable.", |
| 466 | ) |
| 467 | group.add_argument( |
| 468 | "--s3-endpoint", |
| 469 | type=str, |
| 470 | default=None, |
| 471 | help="The endpoint for the S3 bucket. Can also be set via the " |
| 472 | "S3_ENDPOINT_URL environment variable.", |
| 473 | ) |
| 474 | |
| 475 | return parser |
| 476 | |
| 477 | @classmethod |
| 478 | def from_cli_args(cls, args: argparse.Namespace) -> "TensorizerArgs": |
| 479 | attrs = [attr.name for attr in dataclasses.fields(cls)] |
| 480 | tensorizer_args = cls( |
| 481 | **{attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)} |
| 482 | ) |
| 483 | return tensorizer_args |
| 484 | |
| 485 | |
| 486 | def _check_tensors_on_meta_device(model: nn.Module) -> None: |
| 487 | for tensor in model.state_dict().values(): |
| 488 | if tensor.device.type == "meta": |
| 489 | raise ValueError( |
| 490 | "The serialized model contains tensors on the meta device," |
| 491 | " indicating that some tensors were not loaded properly." |
| 492 | " Please check that the parameters of the model being" |
| 493 | " specified match that of the serialized model, such as" |
| 494 | " its quantization." |
| 495 | ) |
| 496 | |
| 497 | |
| 498 | def _resize_lora_embeddings(model: nn.Module): |
| 499 | """Modify LoRA embedding layers to use bigger tensors |
| 500 | to allow for adapter added tokens.""" |
| 501 | for child in model.modules(): |
| 502 | if ( |
| 503 | isinstance(child, VocabParallelEmbedding) |
| 504 | and child.weight.shape[0] < child.num_embeddings_per_partition |
| 505 | ): |
| 506 | new_weight = torch.empty( |
| 507 | child.num_embeddings_per_partition, |
| 508 | child.embedding_dim, |
| 509 | dtype=child.weight.dtype, |
| 510 | device=child.weight.device, |
| 511 | ) |
| 512 | new_weight[: child.weight.shape[0]].copy_(child.weight.data) |
| 513 | new_weight[child.weight.shape[0] :].fill_(0) |
| 514 | child.weight.data = new_weight |
| 515 | |
| 516 | |
| 517 | def init_tensorizer_model( |
| 518 | tensorizer_config: TensorizerConfig, vllm_config: VllmConfig |
| 519 | ) -> nn.Module: |
| 520 | assert tensorizer_config.hf_config is not None |
| 521 | model_args = tensorizer_config.hf_config |
| 522 | model_args.dtype = tensorizer_config.dtype |
| 523 | assert tensorizer_config.model_class is not None |
| 524 | # TODO: Do we need to consider old-style model class? |
| 525 | with meta_tensor_mode(), set_current_vllm_config(vllm_config, check_compile=True): |
| 526 | return tensorizer_config.model_class(vllm_config=vllm_config) |
| 527 | |
| 528 | |
| 529 | def deserialize_tensorizer_model( |
| 530 | model: nn.Module, tensorizer_config: TensorizerConfig |
| 531 | ) -> None: |
| 532 | tensorizer_args = tensorizer_config._construct_tensorizer_args() |
| 533 | if not is_valid_deserialization_uri(tensorizer_config.tensorizer_uri): |
| 534 | raise ValueError( |
| 535 | f"{tensorizer_config.tensorizer_uri} is not a valid " |
| 536 | f"tensorizer URI. Please check that the URI is correct. " |
| 537 | f"It must either point to a local existing file, or have a " |
| 538 | f"S3, HTTP or HTTPS scheme." |
| 539 | ) |
| 540 | before_mem = get_mem_usage() |
| 541 | start = time.perf_counter() |
| 542 | device_index = torch.accelerator.current_device_index() |
| 543 | device_type = current_platform.device_type |
| 544 | with ( |
| 545 | open_stream( |
| 546 | tensorizer_config.tensorizer_uri, mode="rb", **tensorizer_args.stream_kwargs |
| 547 | ) as stream, |
| 548 | TensorDeserializer( |
| 549 | stream, |
| 550 | dtype=tensorizer_config.dtype, |
| 551 | device=f"{device_type}:{device_index}", |
| 552 | **tensorizer_args.deserialization_kwargs, |
| 553 | ) as deserializer, |
| 554 | ): |
| 555 | deserializer.load_into_module(model) |
| 556 | end = time.perf_counter() |
| 557 | |
| 558 | total_bytes_str = convert_bytes(deserializer.total_tensor_bytes) |
| 559 | duration = end - start |
| 560 | per_second = convert_bytes(deserializer.total_tensor_bytes / duration) |
| 561 | after_mem = get_mem_usage() |
| 562 | deserializer.close() |
| 563 | logger.info( |
| 564 | "Deserialized %s in %0.2fs, %s/s", total_bytes_str, end - start, per_second |
| 565 | ) |
| 566 | logger.info("Memory usage before: %s", before_mem) |
| 567 | logger.info("Memory usage after: %s", after_mem) |
| 568 | |
| 569 | _check_tensors_on_meta_device(model) |
| 570 | _resize_lora_embeddings(model) |
| 571 | del model.vllm_tensorized_marker |
| 572 | |
| 573 | |
| 574 | def tensorizer_weights_iterator( |
| 575 | tensorizer_args: "TensorizerArgs", |
| 576 | ) -> Generator[tuple[str, torch.Tensor], None, None]: |
| 577 | logger.warning( |
| 578 | "Deserializing HuggingFace models is not optimized for " |
| 579 | "loading on vLLM, as tensorizer is forced to load to CPU. " |
| 580 | "Consider deserializing a vLLM model instead for faster " |
| 581 | "load times. See the " |
| 582 | "examples/features/tensorize_vllm_model.py example script " |
| 583 | "for serializing vLLM models." |
| 584 | ) |
| 585 | |
| 586 | deserializer_args = tensorizer_args.deserialization_kwargs |
| 587 | stream_kwargs = tensorizer_args.stream_kwargs |
| 588 | stream = open_stream(tensorizer_args.tensorizer_uri, **stream_kwargs) |
| 589 | with TensorDeserializer(stream, **deserializer_args, device="cpu") as state: |
| 590 | yield from state.items() |
| 591 | del state |
| 592 | |
| 593 | |
| 594 | def is_vllm_tensorized(tensorizer_config: "TensorizerConfig") -> bool: |
| 595 | """Infer if the model is a vLLM model by checking the weights for |
| 596 | a vLLM tensorized marker. |
| 597 | |
| 598 | Args: |
| 599 | tensorizer_config: The TensorizerConfig object containing the |
| 600 | tensorizer_uri to the serialized model. |
| 601 | |
| 602 | Returns: |
| 603 | bool: True if the model is a vLLM model, False otherwise. |
| 604 | |
| 605 | """ |
| 606 | tensorizer_args = tensorizer_config._construct_tensorizer_args() |
| 607 | deserializer = TensorDeserializer( |
| 608 | open_stream(tensorizer_args.tensorizer_uri, **tensorizer_args.stream_kwargs), |
| 609 | **tensorizer_args.deserialization_kwargs, |
| 610 | lazy_load=True, |
| 611 | ) |
| 612 | if tensorizer_config.vllm_tensorized: |
| 613 | logger.warning( |
| 614 | "Please note that newly serialized vLLM models are automatically " |
| 615 | "inferred as vLLM models, so setting vllm_tensorized=True is " |
| 616 | "only necessary for models serialized prior to this change." |
| 617 | ) |
| 618 | return True |
| 619 | return ".vllm_tensorized_marker" in deserializer |
| 620 | |
| 621 | |
| 622 | def serialize_extra_artifacts( |
| 623 | tensorizer_args: TensorizerArgs, |
| 624 | served_model_name: str | list[str] | None, |
| 625 | revision: str | None = None, |
| 626 | ) -> None: |
| 627 | if not isinstance(served_model_name, str): |
| 628 | raise ValueError( |
| 629 | f"served_model_name must be a str for serialize_extra_artifacts, " |
| 630 | f"not {type(served_model_name)}." |
| 631 | ) |
| 632 | |
| 633 | with tempfile.TemporaryDirectory() as tmpdir: |
| 634 | hf_api().snapshot_download( |
| 635 | served_model_name, |
| 636 | revision=revision, |
| 637 | local_dir=tmpdir, |
| 638 | ignore_patterns=[ |
| 639 | "*.pt", |
| 640 | "*.safetensors", |
| 641 | "*.bin", |
| 642 | "*.cache", |
| 643 | "*.gitattributes", |
| 644 | "*.md", |
| 645 | ], |
| 646 | ) |
| 647 | for artifact in os.scandir(tmpdir): |
| 648 | if not artifact.is_file(): |
| 649 | continue |
| 650 | with ( |
| 651 | open(artifact.path, "rb") as f, |
| 652 | open_stream( |
| 653 | f"{tensorizer_args.tensorizer_dir}/{artifact.name}", |
| 654 | mode="wb+", |
| 655 | **tensorizer_args.stream_kwargs, |
| 656 | ) as stream, |
| 657 | ): |
| 658 | logger.info("Writing artifact %s", artifact.name) |
| 659 | stream.write(f.read()) |
| 660 | |
| 661 | |
| 662 | def serialize_vllm_model( |
| 663 | model: nn.Module, |
| 664 | tensorizer_config: TensorizerConfig, |
| 665 | model_config: "ModelConfig", |
| 666 | ) -> nn.Module: |
| 667 | model.register_parameter( |
| 668 | "vllm_tensorized_marker", |
| 669 | nn.Parameter(torch.tensor((1,), device="meta"), requires_grad=False), |
| 670 | ) |
| 671 | |
| 672 | tensorizer_args = tensorizer_config._construct_tensorizer_args() |
| 673 | |
| 674 | encryption_params = None |
| 675 | if (keyfile := tensorizer_config.encryption_keyfile) is not None: |
| 676 | with open(keyfile, "rb") as f: |
| 677 | key = f.read() |
| 678 | encryption_params = EncryptionParams(key=key) |
| 679 | |
| 680 | if (output_file := tensorizer_args.tensorizer_uri) is None: |
| 681 | raise ValueError("tensorizer_uri must be specified for serialization.") |
| 682 | if tensorizer_config._is_sharded: |
| 683 | from vllm.distributed import get_tensor_model_parallel_rank |
| 684 | |
| 685 | output_file = output_file % get_tensor_model_parallel_rank() |
| 686 | |
| 687 | with open_stream( |
| 688 | output_file, mode="wb+", **tensorizer_args.stream_kwargs |
| 689 | ) as stream: |
| 690 | serializer = TensorSerializer( |
| 691 | stream, |
| 692 | encryption=encryption_params, |
| 693 | **(tensorizer_config.serialization_kwargs or {}), |
| 694 | ) |
| 695 | serializer.write_module(model) |
| 696 | serializer.close() |
| 697 | |
| 698 | serialize_extra_artifacts( |
| 699 | tensorizer_args, |
| 700 | model_config.served_model_name, |
| 701 | revision=model_config.revision, |
| 702 | ) |
| 703 | |
| 704 | logger.info("Successfully serialized model to %s", str(output_file)) |
| 705 | return model |
| 706 | |
| 707 | |
| 708 | def tensorize_vllm_model( |
| 709 | engine_args: "EngineArgs", |
| 710 | tensorizer_config: TensorizerConfig, |
| 711 | generate_keyfile: bool = True, |
| 712 | ): |
| 713 | """Utility to load a model and then serialize it with Tensorizer. |
| 714 | |
| 715 | Intended to be used separately from running a vLLM server since it |
| 716 | creates its own Engine instance. |
| 717 | """ |
| 718 | engine_config = engine_args.create_engine_config() |
| 719 | tensorizer_config.verify_with_model_config(engine_config.model_config) |
| 720 | tensorizer_config.verify_with_parallel_config(engine_config.parallel_config) |
| 721 | |
| 722 | # generate the encryption key before creating the engine to support sharding |
| 723 | if ( |
| 724 | generate_keyfile |
| 725 | and (keyfile := tensorizer_config.encryption_keyfile) is not None |
| 726 | ): |
| 727 | encryption_params = EncryptionParams.random() |
| 728 | with open_stream( |
| 729 | keyfile, |
| 730 | mode="wb+", |
| 731 | s3_access_key_id=tensorizer_config.s3_access_key_id, |
| 732 | s3_secret_access_key=tensorizer_config.s3_secret_access_key, |
| 733 | s3_endpoint=tensorizer_config.s3_endpoint, |
| 734 | ) as stream: |
| 735 | stream.write(encryption_params.key) |
| 736 | |
| 737 | from vllm.v1.engine.llm_engine import LLMEngine |
| 738 | |
| 739 | engine = LLMEngine.from_vllm_config(engine_config) |
| 740 | error: BaseException | None = None |
| 741 | try: |
| 742 | engine.collective_rpc( |
| 743 | "save_tensorized_model", |
| 744 | kwargs={"tensorizer_config": tensorizer_config.to_serializable()}, |
| 745 | ) |
| 746 | except BaseException as operation_error: |
| 747 | error = operation_error |
| 748 | |
| 749 | def shutdown_engine_core() -> None: |
| 750 | engine.engine_core.shutdown( |
| 751 | timeout=( |
| 752 | envs.VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS |
| 753 | + _TENSORIZER_ENGINE_CLEANUP_GRACE_S |
| 754 | ) |
| 755 | ) |
| 756 | |
| 757 | for name, callback in ( |
| 758 | ("renderer", engine.renderer.shutdown), |
| 759 | ("engine core", shutdown_engine_core), |
| 760 | ): |
| 761 | try: |
| 762 | callback() |
| 763 | except BaseException as shutdown_error: |
| 764 | logger.exception("Failed to shut down tensorization %s", name) |
| 765 | if error is None: |
| 766 | error = shutdown_error |
| 767 | elif hasattr(error, "add_note"): |
| 768 | error.add_note(f"{name} shutdown also failed: {shutdown_error!r}") |
| 769 | |
| 770 | if error is not None: |
| 771 | raise error |
| 772 | |
| 773 | |
| 774 | def tensorize_lora_adapter(lora_path: str, tensorizer_config: TensorizerConfig): |
| 775 | """Uses tensorizer to serialize a LoRA adapter. Assumes that the files |
| 776 | needed to load a LoRA adapter are a safetensors-format file called |
| 777 | adapter_model.safetensors and a json config file called adapter_config.json. |
| 778 | |
| 779 | Serializes the files in the tensorizer_config.tensorizer_dir |
| 780 | """ |
| 781 | import safetensors |
| 782 | |
| 783 | from vllm.lora.utils import get_adapter_absolute_path |
| 784 | |
| 785 | lora_dir = get_adapter_absolute_path(lora_path) |
| 786 | |
| 787 | tensor_path = config_path = "" |
| 788 | |
| 789 | for file in os.listdir(lora_dir): |
| 790 | if file.startswith("adapter_model"): |
| 791 | tensor_path = lora_dir + "/" + file |
| 792 | if file.startswith("adapter_config"): |
| 793 | config_path = lora_dir + "/" + file |
| 794 | if tensor_path and config_path: |
| 795 | break |
| 796 | |
| 797 | if tensor_path.endswith(".safetensors"): |
| 798 | tensors = safetensors.torch.load_file(tensor_path) |
| 799 | elif tensor_path.endswith(".bin"): |
| 800 | tensors = torch.load(tensor_path, weights_only=True) |
| 801 | else: |
| 802 | raise ValueError( |
| 803 | f"Unsupported adapter model file: {tensor_path}. " |
| 804 | f"Must be a .safetensors or .bin file." |
| 805 | ) |
| 806 | |
| 807 | with open(config_path) as f: |
| 808 | config = json.load(f) |
| 809 | |
| 810 | tensorizer_args = tensorizer_config._construct_tensorizer_args() |
| 811 | |
| 812 | with open_stream( |
| 813 | f"{tensorizer_config.tensorizer_dir}/adapter_config.json", |
| 814 | mode="wb+", |
| 815 | **tensorizer_args.stream_kwargs, |
| 816 | ) as f: |
| 817 | f.write(json.dumps(config).encode("utf-8")) |
| 818 | |
| 819 | lora_uri = f"{tensorizer_config.tensorizer_dir}/adapter_model.tensors" |
| 820 | with open_stream(lora_uri, mode="wb+", **tensorizer_args.stream_kwargs) as f: |
| 821 | serializer = TensorSerializer(f) |
| 822 | serializer.write_state_dict(tensors) |
| 823 | serializer.close() |
| 824 | |
| 825 | logger.info( |
| 826 | "Successfully serialized LoRA files to %s", |
| 827 | str(tensorizer_config.tensorizer_dir), |
| 828 | ) |