[Frontend] [Core] feat: Add model loading using tensorizer

[sangstar]

Feature Request Issue

Tensorizer Support

This PR allows models used for the OpenAI-compatible API server to be loaded using
Coreweave's Tensorizer, enabling extremely fast (faster than cached Safetensors) model loads from HTTP/HTTPS,
Redis, and S3 endpoints.

The key changes involved are:

1. Adds TensorizerConfig to the set of configs.
2. Adds tensorizer_loader.py to vllm/model_executor that provides utility functions
   for tensorizer.
3. Adds multiple args to the vLLM's OpenAI inference service entrypoint that allows
   the user to specify the path to serialized-by-tensorizer model tensors, as well as
   arguments for tensorizer's deserializer.
4. Allows deserialization of serialized model tensors in HuggingFace's model format,
   as well as supporting deserializing serialized vLLM-formatted models, allowing the
   use of loading with plaid_mode, which can allow Llama 2 13B non-locally to start serving requests
   in as little as 10 seconds. Also supports encrypting and decrypting model tensors.
5. Adds a tensorize_vllm_model.py script to examples/ that allows vLLM models to be serialized and
   deserialized with tensorizer.
6. Adds tensorizer as an optional dependency.

Credentialing for S3 is supported by passing a user's access and secret key to
S3_ACCESS_KEY_ID and S3_SECRET_ACCESS_KEY environment variables respectively. It can also be specified as CLI args to the api server entrypoint.

Model loading benchmarks

Tensorizer can load models like Llama 2 13B in as little as 10 seconds. In order to do so, a model must be
serialized using TensorSerializer to a .tensors file located either locally or through a S3, HTTP/HTTPS, or Redis
endpoint. --tensorizer-uri must be specified with the serialized tensors location when invoking the API server.

Example usage:

python -m vllm.entrypoints.openai.api_server \
--host 0.0.0.0 \
--model EleutherAI/pythia-6.9b \
--load-format tensorizer \
--tensorizer-uri s3://tensorized/EleutherAI/pythia-6.9b/fp16/model.tensors

If a vLLM model is serialized, plaid_mode can be used, which loads much faster. The following plot demonstrates model loading time benchmarks for vLLM's OpenAI-compatible inference server on a Nvidia A40 GPU.

Tensorizer is so fast that it loads models faster than Safetensors even locally.

[sangstar]

@cadedaniel @rkooo567 Pinging for an assigned reviewer from someone on the team when possible!

[sangstar]

@cadedaniel @rkooo567 @Yard1 @WoosukKwon @zhuohan123 @ywang96

All tests are passing. Can I get eyes on this please? Cheers!

[sangstar]

Added additional tests and some fixes; checks are all passing!

[rkooo567]

LGTM if @Yard1 gives a go

[ywang96]

Thank you very much @sangster for your contribution to add this feature, and @Yard1 @rkooo567 for the detailed reviews!

I've tested this PR on different scenarios (from S3, from a local directory with serialized weights using the example script, etc) and left a few feedback on the PR. PTAL, thanks!

[sangstar]

@ywang96 @Yard1 @rkooo567

Thank you all very much for your reviews! I've implemented the changes from @ywang96 's comments. To summarize:

• An error is raised if the tensor parallel size exceeds 1 and attempting to use Tensorizer (test added)
• The serialization step in examples/tensorize_vllm_model.py now instantiates the model to serialize using LLMEngine
• Meta tensors found when deserializing will raise an error
• Removed forcing float16 from the parser for examples/tensorize_vllm_model.py
• I've also additionally added a PerformanceWarning when trying to load a tensorized model with quantization, as that is a bit unstable at the moment (I may try to look in to this in another PR) (test added).
• Added the Tensorizer testing folder for the CI suite

[sangstar]

Some minor fixes to ensure the testing suite can run the tensorizer tests. All passing! Thanks very much again for the reviews @rkooo567 @Yard1 @ywang96 let me know if anything else needed! :)

[ywang96]

LGTM! Thank you again @sangstar for all the work and test coverage on this PR to add this feature!

[bbrowning]

I was able to successfully test this in a vLLM 0.4.1 container running on OpenShift, both with models serialized with the Tensorizer library directly and for vLLM-serialized models. Once I cranked up the Pod's CPU and increased the num_readers parameter, I got about an 8x speedup in my case when loading the same model via vLLM-serialized tensorize files compared to not using tensorizer at all and just downloading the safetensors from S3 to local disk then loading with vLLM. This took my overall cold start time of this Pod from a bit over 4 minutes to 30 seconds. There may be even more performance available in my setup with additional tweaking, but this is already a great win.

INFO 05-03 11:11:38 tensorizer.py:337] Deserialized 14.5 GB in 15.21s, 953.1 MB/s

That's an awesome improvement, and thank you!

[sangstar]

I was able to successfully test this in a vLLM 0.4.1 container running on OpenShift, both with models serialized with the Tensorizer library directly and for vLLM-serialized models. Once I cranked up the Pod's CPU and increased the num_readers parameter, I got about an 8x speedup in my case when loading the same model via vLLM-serialized tensorize files compared to not using tensorizer at all and just downloading the safetensors from S3 to local disk then loading with vLLM. This took my overall cold start time of this Pod from a bit over 4 minutes to 30 seconds. There may be even more performance available in my setup with additional tweaking, but this is already a great win.

INFO 05-03 11:11:38 tensorizer.py:337] Deserialized 14.5 GB in 15.21s, 953.1 MB/s

That's an awesome improvement, and thank you!

I'm thrilled to hear that! I currently actually have a new PR up #4208 that uses the full 2.9.0 release, has better usage documentation, and automated inferring a vLLM-serialized model.