feat: Allow passing arbitrary Tensorizer serialization and deserialization kwargs; update docs

docs/models/extensions/tensorizer.md

@@ -8,9 +8,97 @@ vLLM model tensors that have been serialized to disk, an HTTP/HTTPS endpoint, or
 at runtime extremely quickly directly to the GPU, resulting in significantly
 shorter Pod startup times and CPU memory usage. Tensor encryption is also supported.
 
-For more information on CoreWeave's Tensorizer, please refer to
-[CoreWeave's Tensorizer documentation](https://github.com/coreweave/tensorizer). For more information on serializing a vLLM model, as well a general usage guide to using Tensorizer with vLLM, see
-the [vLLM example script](https://docs.vllm.ai/en/latest/examples/tensorize_vllm_model.html).
+vLLM fully integrates Tensorizer in to its model loading machinery. The
+following will give a brief overview on how to get started with using
+Tensorizer on vLLM.
+
+## The basics
+To load a model using Tensorizer, it first needs to be serialized by Tensorizer.
+The example script in [examples/others/tensorize_vllm_model.py](https://github.com/vllm-project/vllm/blob/main/examples/others/tensorize_vllm_model.py)
+takes care of this process.
+
+The core frontend object of note integrating Tensorizer is 
+`TensorizerConfig`, defined [here](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/model_loader/tensorizer.py#L135-L214)
+It's a config object holding important state and passed to any serialization 
+or deserialization operation. When loading with Tensorizer using the vLLM 
+library rather than through a model-serving entrypoint, it gets passed to 
+the `LLM` entrypoint class directly. Here's an example of loading a model
+saved at `"s3://my-bucket/vllm/facebook/opt-125m/v1/model.tensors"`.
+
+```python
+from vllm import LLM
+from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
+
+path_to_tensors = "s3://my-bucket/vllm/facebook/opt-125m/v1/model.tensors"
+
+model_ref = "facebook/opt-125m"
+tensorizer_config = TensorizerConfig(
+        tensorizer_uri=path_to_tensors,
+    )
+
+llm = LLM(
+    model_ref,
+    load_format="tensorizer",
+    model_loader_extra_config=tensorizer_config,
+)
+```
+
+But that code won't work unless you actually have your serialized model 
+tensors `model.tensors`, so let's walk through a basic example of serializing 
+`facebook/opt-125m` using the example script, and then loading it for inference.
+
+## Saving a vLLM model with Tensorizer
+To save a model with Tensorizer, call the example script with the necessary
+CLI arguments. The docstring for the script itself explains the CLI args
+and how to use it properly in great detail, and we'll use one of the
+examples from the docstring directly, assuming we want to save our model at
+our S3 bucket example `s3://my-bucket`:
+
+```bash
+python examples/others/tensorize_vllm_model.py \
+   --model facebook/opt-125m \
+   serialize \
+   --serialized-directory s3://my-bucket \
+   --suffix v1
+```
+
+This saves the model tensors at 
+`s3://my-bucket/vllm/facebook/opt-125m/v1/model.tensors`.
+
+## Serving the model using Tensorizer
+Once the model is serialized where you want it, you need to pass
+`--load-format=tensorizer` as well as a JSON string to the
+`--model-loader-extra-config` CLI arg for `vllm serve`, specifying
+all the keyword arguments you'd normally pass to the `TensorizerConfig`
+initializer, as if you were instantiating it like this, for instance:
+`TensorizerConfig(**{"tensorizer_uri": "foo"})`. Here's an example bash
+script that neatly lays out the JSON string before passing it to
+`--model-loader-extra-config`:
+
+```bash
+#!/bin/bash
+
+read -r -d '' JSON << EOF
+{
+  "tensorizer_uri": "s3://my-bucket/vllm/facebook/opt-125m/v1/model.tensors",
+  "stream_kwargs": {"force_http": "False"},
+  "deserialization_kwargs": {"verify_hash": "True", "num_readers": 8}
+}
+EOF
+
+MODEL_LOADER_EXTRA_CONFIG=$(echo "$JSON" | tr -d '\n')
+
+vllm serve facebook/opt-125m \
+  --load-format=tensorizer \
+  --model-loader-extra-config="$MODEL_LOADER_EXTRA_CONFIG"
+```
+
+Note in this case, if the directory to the model artifacts at
+`s3://my-bucket/vllm/facebook/opt-125m/v1/` doesn't at least have a `config.
+json` file, you'll want to pass `facebook/opt-125m` as the model tag like
+it was done in the example script above. In our example, we just added a
+`model.tensors` file to that directory. In this case, vLLM will take care of
+resolving the other model artifacts using HF Hub.
 
 !!! note
     Note that to use this feature you will need to install `tensorizer` by running `pip install vllm[tensorizer]`.

examples/others/tensorize_vllm_model.py

@@ -170,6 +170,14 @@ def parse_args():
         "where `suffix` is given by `--suffix` or a random UUID if not "
         "provided.")
 
+    serialize_parser.add_argument(
+        "--serialization-kwargs",
+        type=str,
+        required=False,
+        help=("A JSON string containing additional keyword arguments that "
+              "will be passed to Tensorizer's `TensorSerializer` during "
+              "serialization."))
+
     serialize_parser.add_argument(
         "--keyfile",
         type=str,
@@ -295,6 +303,12 @@ def deserialize():
             encryption_keyfile=keyfile,
             **credentials)
 
+        if args.serialization_kwargs:
+            serialization_kwargs = json.loads(args.serialization_kwargs)
+            tensorizer_config.serialization_kwargs = serialization_kwargs
+            print("Found serialization kwargs: ", serialization_kwargs)
+
+
         if args.lora_path:
             tensorizer_config.lora_dir = tensorizer_config.tensorizer_dir
             tensorize_lora_adapter(args.lora_path, tensorizer_config)

tests/entrypoints/openai/test_tensorizer_entrypoint.py

@@ -59,12 +59,18 @@ def tensorize_model_and_lora(tmp_dir, model_uri):
 def server(model_uri, tensorize_model_and_lora):
     model_loader_extra_config = {
         "tensorizer_uri": model_uri,
+        "stream_kwargs": {
+            "force_http": False,
+        },
+        "deserialization_kwargs": {
+            "verify_hash": True,
+            "num_readers": 8,
+        }
     }
 
     ## Start OpenAI API server
     args = [
-        "--load-format", "tensorizer", "--device", "cuda",
-        "--model-loader-extra-config",
+        "--load-format", "tensorizer", "--model-loader-extra-config",
         json.dumps(model_loader_extra_config), "--enable-lora"
     ]
 

tests/lora/test_llama_tp.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 import subprocess
 import sys
+import json
 from typing import Union
 
 import pytest
@@ -210,7 +211,8 @@ def test_tp2_serialize_and_deserialize_lora(tmp_path, sql_lora_files,
             f"{VLLM_PATH}/examples/others/tensorize_vllm_model.py", "--model",
             MODEL_PATH, "--lora-path", lora_path, "--tensor-parallel-size",
             str(tp_size), "serialize", "--serialized-directory",
-            str(tmp_path), "--suffix", suffix
+            str(tmp_path), "--suffix", suffix, "--serialization-kwargs",
+            json.dumps({"limit_cpu_concurrency": 4})
         ],
                                 check=True,
                                 capture_output=True,

tests/tensorizer_loader/conftest.py

@@ -1,10 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
+from typing import Callable
+
 import pytest
+import os
 
+from vllm import LLM
 from vllm.distributed import cleanup_dist_env_and_memory
 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
 
 
+@pytest.fixture(autouse=True)
+def allow_insecure_serialization():
+    os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
+
 @pytest.fixture(autouse=True)
 def cleanup():
     cleanup_dist_env_and_memory(shutdown_ray=True)
@@ -14,3 +22,10 @@ def cleanup():
 def tensorizer_config():
     config = TensorizerConfig(tensorizer_uri="vllm")
     return config
+
+
+def assert_from_collective_rpc(engine: LLM,
+                               closure: Callable,
+                               closure_kwargs: dict):
+    res = engine.collective_rpc(method=closure, kwargs=closure_kwargs)
+    return all(res)