| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import json |
| 5 | import logging |
| 6 | import os |
| 7 | import uuid |
| 8 | |
| 9 | from vllm import LLM, SamplingParams |
| 10 | from vllm.engine.arg_utils import EngineArgs |
| 11 | from vllm.lora.request import LoRARequest |
| 12 | from vllm.model_executor.model_loader.tensorizer import ( |
| 13 | TensorizerArgs, |
| 14 | TensorizerConfig, |
| 15 | tensorize_lora_adapter, |
| 16 | tensorize_vllm_model, |
| 17 | tensorizer_kwargs_arg, |
| 18 | ) |
| 19 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 20 | |
| 21 | logger = logging.getLogger() |
| 22 | |
| 23 | |
| 24 | """ |
| 25 | tensorize_vllm_model.py is a script that can be used to serialize and |
| 26 | deserialize vLLM models. These models can be loaded using tensorizer |
| 27 | to the GPU extremely quickly over an HTTP/HTTPS endpoint, an S3 endpoint, |
| 28 | or locally. Tensor encryption and decryption is also supported, although |
| 29 | libsodium must be installed to use it. Install vllm with tensorizer support |
| 30 | using `pip install vllm[tensorizer]`. To learn more about tensorizer, visit |
| 31 | https://github.com/coreweave/tensorizer |
| 32 | |
| 33 | To serialize a model, install vLLM from source, then run something |
| 34 | like this from the root level of this repository: |
| 35 | |
| 36 | python examples/features/tensorize_vllm_model.py \ |
| 37 | --model facebook/opt-125m \ |
| 38 | serialize \ |
| 39 | --serialized-directory s3://my-bucket \ |
| 40 | --suffix v1 |
| 41 | |
| 42 | Which downloads the model from HuggingFace, loads it into vLLM, serializes it, |
| 43 | and saves it to your S3 bucket. A local directory can also be used. This |
| 44 | assumes your S3 credentials are specified as environment variables |
| 45 | in the form of `S3_ACCESS_KEY_ID`, `S3_SECRET_ACCESS_KEY`, and |
| 46 | `S3_ENDPOINT_URL`. To provide S3 credentials directly, you can provide |
| 47 | `--s3-access-key-id` and `--s3-secret-access-key`, as well as `--s3-endpoint` |
| 48 | as CLI args to this script. |
| 49 | |
| 50 | You can also encrypt the model weights with a randomly-generated key by |
| 51 | providing a `--keyfile` argument. |
| 52 | |
| 53 | To deserialize a model, you can run something like this from the root |
| 54 | level of this repository: |
| 55 | |
| 56 | python examples/features/tensorize_vllm_model.py \ |
| 57 | --model EleutherAI/gpt-j-6B \ |
| 58 | --dtype float16 \ |
| 59 | deserialize \ |
| 60 | --path-to-tensors s3://my-bucket/vllm/EleutherAI/gpt-j-6B/v1/model.tensors |
| 61 | |
| 62 | Which downloads the model tensors from your S3 bucket and deserializes them. |
| 63 | |
| 64 | You can also provide a `--keyfile` argument to decrypt the model weights if |
| 65 | they were serialized with encryption. |
| 66 | |
| 67 | To support distributed tensor-parallel models, each model shard will be |
| 68 | serialized to a separate file. The tensorizer_uri is then specified as a string |
| 69 | template with a format specifier such as '%03d' that will be rendered with the |
| 70 | shard's rank. Sharded models serialized with this script will be named as |
| 71 | model-rank-%03d.tensors |
| 72 | |
| 73 | For more information on the available arguments for serializing, run |
| 74 | `python -m examples.features.tensorize_vllm_model serialize --help`. |
| 75 | |
| 76 | Or for deserializing: |
| 77 | |
| 78 | `python examples/features/tensorize_vllm_model.py deserialize --help`. |
| 79 | |
| 80 | Once a model is serialized, tensorizer can be invoked with the `LLM` class |
| 81 | directly to load models: |
| 82 | |
| 83 | ```python |
| 84 | from vllm import LLM |
| 85 | llm = LLM( |
| 86 | "s3://my-bucket/vllm/facebook/opt-125m/v1", |
| 87 | load_format="tensorizer", |
| 88 | ) |
| 89 | ``` |
| 90 | |
| 91 | |
| 92 | A serialized model can be used during model loading for the vLLM OpenAI |
| 93 | inference server: |
| 94 | |
| 95 | ``` |
| 96 | vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \ |
| 97 | --load-format tensorizer |
| 98 | ``` |
| 99 | |
| 100 | In order to see all of the available arguments usable to configure |
| 101 | loading with tensorizer that are given to `TensorizerConfig`, run: |
| 102 | |
| 103 | `python examples/features/tensorize_vllm_model.py deserialize --help` |
| 104 | |
| 105 | under the `tensorizer options` section. These can also be used for |
| 106 | deserialization in this example script, although `--tensorizer-uri` and |
| 107 | `--path-to-tensors` are functionally the same in this case. |
| 108 | |
| 109 | Tensorizer can also be used to save and load LoRA adapters. A LoRA adapter |
| 110 | can be serialized directly with the path to the LoRA adapter on HF Hub and |
| 111 | a TensorizerConfig object. In this script, passing a HF id to a LoRA adapter |
| 112 | will serialize the LoRA adapter artifacts to `--serialized-directory`. |
| 113 | |
| 114 | You can then use the LoRA adapter with `vllm serve`, for instance, by ensuring |
| 115 | the LoRA artifacts are in your model artifacts directory and specifying |
| 116 | `--enable-lora`. For instance: |
| 117 | |
| 118 | ``` |
| 119 | vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \ |
| 120 | --load-format tensorizer \ |
| 121 | --enable-lora |
| 122 | ``` |
| 123 | """ |
| 124 | |
| 125 | |
| 126 | def get_parser(): |
| 127 | parser = FlexibleArgumentParser( |
| 128 | description="An example script that can be used to serialize and " |
| 129 | "deserialize vLLM models. These models " |
| 130 | "can be loaded using tensorizer directly to the GPU " |
| 131 | "extremely quickly. Tensor encryption and decryption is " |
| 132 | "also supported, although libsodium must be installed to " |
| 133 | "use it." |
| 134 | ) |
| 135 | parser = EngineArgs.add_cli_args(parser) |
| 136 | |
| 137 | parser.add_argument( |
| 138 | "--lora-path", |
| 139 | type=str, |
| 140 | required=False, |
| 141 | help="Path to a LoRA adapter to " |
| 142 | "serialize along with model tensors. This can then be deserialized " |
| 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 | |
| 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`" |
| 153 | ) |
| 154 | |
| 155 | serialize_parser.add_argument( |
| 156 | "--suffix", |
| 157 | type=str, |
| 158 | required=False, |
| 159 | help=( |
| 160 | "The suffix to append to the serialized model directory, which is " |
| 161 | "used to construct the location of the serialized model tensors, " |
| 162 | "e.g. if `--serialized-directory` is `s3://my-bucket/` and " |
| 163 | "`--suffix` is `v1`, the serialized model tensors will be " |
| 164 | "saved to " |
| 165 | "`s3://my-bucket/vllm/EleutherAI/gpt-j-6B/v1/model.tensors`. " |
| 166 | "If none is provided, a random UUID will be used." |
| 167 | ), |
| 168 | ) |
| 169 | serialize_parser.add_argument( |
| 170 | "--serialized-directory", |
| 171 | type=str, |
| 172 | required=True, |
| 173 | help="The directory to serialize the model to. " |
| 174 | "This can be a local directory or S3 URI. The path to where the " |
| 175 | "tensors are saved is a combination of the supplied `dir` and model " |
| 176 | "reference ID. For instance, if `dir` is the serialized directory, " |
| 177 | "and the model HuggingFace ID is `EleutherAI/gpt-j-6B`, tensors will " |
| 178 | "be saved to `dir/vllm/EleutherAI/gpt-j-6B/suffix/model.tensors`, " |
| 179 | "where `suffix` is given by `--suffix` or a random UUID if not " |
| 180 | "provided.", |
| 181 | ) |
| 182 | |
| 183 | serialize_parser.add_argument( |
| 184 | "--serialization-kwargs", |
| 185 | type=tensorizer_kwargs_arg, |
| 186 | required=False, |
| 187 | help=( |
| 188 | "A JSON string containing additional keyword arguments to " |
| 189 | "pass to Tensorizer's TensorSerializer during " |
| 190 | "serialization." |
| 191 | ), |
| 192 | ) |
| 193 | |
| 194 | serialize_parser.add_argument( |
| 195 | "--keyfile", |
| 196 | type=str, |
| 197 | required=False, |
| 198 | help=( |
| 199 | "Encrypt the model weights with a randomly-generated binary key," |
| 200 | " and save the key at this path" |
| 201 | ), |
| 202 | ) |
| 203 | |
| 204 | deserialize_parser = subparsers.add_parser( |
| 205 | "deserialize", |
| 206 | help=( |
| 207 | "Deserialize a model from `--path-to-tensors`" |
| 208 | " to verify it can be loaded and used." |
| 209 | ), |
| 210 | ) |
| 211 | |
| 212 | deserialize_parser.add_argument( |
| 213 | "--path-to-tensors", |
| 214 | type=str, |
| 215 | required=False, |
| 216 | help="The local path or S3 URI to the model tensors to deserialize. ", |
| 217 | ) |
| 218 | |
| 219 | deserialize_parser.add_argument( |
| 220 | "--serialized-directory", |
| 221 | type=str, |
| 222 | required=False, |
| 223 | help="Directory with model artifacts for loading. Assumes a " |
| 224 | "model.tensors file exists therein. Can supersede " |
| 225 | "--path-to-tensors.", |
| 226 | ) |
| 227 | |
| 228 | deserialize_parser.add_argument( |
| 229 | "--keyfile", |
| 230 | type=str, |
| 231 | required=False, |
| 232 | help=( |
| 233 | "Path to a binary key to use to decrypt the model weights," |
| 234 | " if the model was serialized with encryption" |
| 235 | ), |
| 236 | ) |
| 237 | |
| 238 | deserialize_parser.add_argument( |
| 239 | "--deserialization-kwargs", |
| 240 | type=tensorizer_kwargs_arg, |
| 241 | required=False, |
| 242 | help=( |
| 243 | "A JSON string containing additional keyword arguments to " |
| 244 | "pass to Tensorizer's `TensorDeserializer` during " |
| 245 | "deserialization." |
| 246 | ), |
| 247 | ) |
| 248 | |
| 249 | TensorizerArgs.add_cli_args(deserialize_parser) |
| 250 | |
| 251 | return parser |
| 252 | |
| 253 | |
| 254 | def merge_extra_config_with_tensorizer_config(extra_cfg: dict, cfg: TensorizerConfig): |
| 255 | for k, v in extra_cfg.items(): |
| 256 | if hasattr(cfg, k): |
| 257 | setattr(cfg, k, v) |
| 258 | logger.info( |
| 259 | "Updating TensorizerConfig with %s from " |
| 260 | "--model-loader-extra-config provided", |
| 261 | k, |
| 262 | ) |
| 263 | |
| 264 | |
| 265 | def deserialize(args, tensorizer_config): |
| 266 | if args.lora_path: |
| 267 | tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir |
| 268 | llm = LLM( |
| 269 | model=args.model, |
| 270 | load_format="tensorizer", |
| 271 | tensor_parallel_size=args.tensor_parallel_size, |
| 272 | model_loader_extra_config=tensorizer_config, |
| 273 | enable_lora=True, |
| 274 | ) |
| 275 | sampling_params = SamplingParams( |
| 276 | temperature=0, max_tokens=256, stop=["[/assistant]"] |
| 277 | ) |
| 278 | |
| 279 | # Truncating this as the extra text isn't necessary |
| 280 | prompts = ["[user] Write a SQL query to answer the question based on ..."] |
| 281 | |
| 282 | # Test LoRA load |
| 283 | print( |
| 284 | llm.generate( |
| 285 | prompts, |
| 286 | sampling_params, |
| 287 | lora_request=LoRARequest( |
| 288 | "sql-lora", |
| 289 | 1, |
| 290 | args.lora_path, |
| 291 | tensorizer_config_dict=tensorizer_config.to_serializable(), |
| 292 | ), |
| 293 | ) |
| 294 | ) |
| 295 | else: |
| 296 | llm = LLM( |
| 297 | model=args.model, |
| 298 | load_format="tensorizer", |
| 299 | tensor_parallel_size=args.tensor_parallel_size, |
| 300 | model_loader_extra_config=tensorizer_config, |
| 301 | ) |
| 302 | return llm |
| 303 | |
| 304 | |
| 305 | def main(): |
| 306 | parser = get_parser() |
| 307 | args = parser.parse_args() |
| 308 | |
| 309 | s3_access_key_id = getattr(args, "s3_access_key_id", None) or os.environ.get( |
| 310 | "S3_ACCESS_KEY_ID", None |
| 311 | ) |
| 312 | s3_secret_access_key = getattr( |
| 313 | args, "s3_secret_access_key", None |
| 314 | ) or os.environ.get("S3_SECRET_ACCESS_KEY", None) |
| 315 | s3_endpoint = getattr(args, "s3_endpoint", None) or os.environ.get( |
| 316 | "S3_ENDPOINT_URL", None |
| 317 | ) |
| 318 | |
| 319 | credentials = { |
| 320 | "s3_access_key_id": s3_access_key_id, |
| 321 | "s3_secret_access_key": s3_secret_access_key, |
| 322 | "s3_endpoint": s3_endpoint, |
| 323 | } |
| 324 | |
| 325 | model_ref = args.model |
| 326 | |
| 327 | if args.command == "serialize" or args.command == "deserialize": |
| 328 | keyfile = args.keyfile |
| 329 | else: |
| 330 | keyfile = None |
| 331 | |
| 332 | extra_config = {} |
| 333 | if args.model_loader_extra_config: |
| 334 | extra_config = json.loads(args.model_loader_extra_config) |
| 335 | |
| 336 | tensorizer_dir = args.serialized_directory or extra_config.get("tensorizer_dir") |
| 337 | tensorizer_uri = getattr(args, "path_to_tensors", None) or extra_config.get( |
| 338 | "tensorizer_uri" |
| 339 | ) |
| 340 | |
| 341 | if tensorizer_dir and tensorizer_uri: |
| 342 | parser.error( |
| 343 | "--serialized-directory and --path-to-tensors cannot both be provided" |
| 344 | ) |
| 345 | |
| 346 | if not tensorizer_dir and not tensorizer_uri: |
| 347 | parser.error( |
| 348 | "Either --serialized-directory or --path-to-tensors must be provided" |
| 349 | ) |
| 350 | |
| 351 | if args.command == "serialize": |
| 352 | engine_args = EngineArgs.from_cli_args(args) |
| 353 | |
| 354 | input_dir = tensorizer_dir.rstrip("/") |
| 355 | suffix = args.suffix if args.suffix else uuid.uuid4().hex |
| 356 | base_path = f"{input_dir}/vllm/{model_ref}/{suffix}" |
| 357 | if engine_args.tensor_parallel_size > 1: |
| 358 | model_path = f"{base_path}/model-rank-%03d.tensors" |
| 359 | else: |
| 360 | model_path = f"{base_path}/model.tensors" |
| 361 | |
| 362 | tensorizer_config = TensorizerConfig( |
| 363 | tensorizer_uri=model_path, |
| 364 | encryption_keyfile=keyfile, |
| 365 | serialization_kwargs=args.serialization_kwargs or {}, |
| 366 | **credentials, |
| 367 | ) |
| 368 | |
| 369 | if args.lora_path: |
| 370 | tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir |
| 371 | tensorize_lora_adapter(args.lora_path, tensorizer_config) |
| 372 | |
| 373 | merge_extra_config_with_tensorizer_config(extra_config, tensorizer_config) |
| 374 | tensorize_vllm_model(engine_args, tensorizer_config) |
| 375 | |
| 376 | elif args.command == "deserialize": |
| 377 | tensorizer_config = TensorizerConfig( |
| 378 | tensorizer_uri=args.path_to_tensors, |
| 379 | tensorizer_dir=args.serialized_directory, |
| 380 | encryption_keyfile=keyfile, |
| 381 | deserialization_kwargs=args.deserialization_kwargs or {}, |
| 382 | **credentials, |
| 383 | ) |
| 384 | |
| 385 | merge_extra_config_with_tensorizer_config(extra_config, tensorizer_config) |
| 386 | deserialize(args, tensorizer_config) |
| 387 | else: |
| 388 | raise ValueError("Either serialize or deserialize must be specified.") |
| 389 | |
| 390 | |
| 391 | if __name__ == "__main__": |
| 392 | main() |