跳到內容

使用 CoreWeave 的 Tensorizer 載入模型

vLLM 支援使用 CoreWeave 的 Tensorizer 載入模型。已序列化到磁碟、HTTP/HTTPS 端點或 S3 端點的 vLLM 模型張量可以在執行時極快地直接反序列化到 GPU,從而顯著縮短 Pod 啟動時間和降低 CPU 記憶體使用。還支援張量加密。

vLLM 將 Tensorizer 完全整合到其模型載入機制中。以下將簡要概述如何在 vLLM 上開始使用 Tensorizer。

安裝 Tensorizer

要安裝 tensorizer,請執行 pip install vllm[tensorizer]

基礎知識

要使用 Tensorizer 載入模型,模型首先需要由 Tensorizer 序列化。示例指令碼負責此過程。

讓我們透過使用指令碼序列化 facebook/opt-125m,然後載入它進行推理,來演示一個基本示例。

使用 Tensorizer 序列化 vLLM 模型

要使用 Tensorizer 序列化模型,請使用必要的 CLI 引數呼叫示例指令碼。指令碼本身的文件字串詳細解釋了 CLI 引數以及如何正確使用它,我們將直接使用文件字串中的一個示例,假設我們想將模型序列化並儲存到我們的 S3 儲存桶示例 s3://my-bucket

python examples/others/tensorize_vllm_model.py \
   --model facebook/opt-125m \
   serialize \
   --serialized-directory s3://my-bucket \
   --suffix v1

這會將模型張量儲存到 s3://my-bucket/vllm/facebook/opt-125m/v1。如果您打算將 LoRA 介面卡應用於已張量化的模型,您可以在上述命令中傳遞 LoRA 介面卡的 HF ID,相關檔案也將儲存到那裡

python examples/others/tensorize_vllm_model.py \
   --model facebook/opt-125m \
   --lora-path <lora_id> \
   serialize \
   --serialized-directory s3://my-bucket \
   --suffix v1

使用 Tensorizer 提供模型服務

模型序列化到所需位置後,您可以使用 vllm serveLLM 入口點載入模型。您可以將儲存模型的目錄作為 LLM()vllm servemodel 引數傳遞。例如,要提供之前儲存的帶有 LoRA 介面卡的已張量化模型服務,您可以這樣做

vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \
    --load-format tensorizer \
    --enable-lora 

或者,使用 LLM()

from vllm import LLM
llm = LLM(
    "s3://my-bucket/vllm/facebook/opt-125m/v1", 
    load_format="tensorizer",
    enable_lora=True
)

配置 Tensorizer 的選項

tensorizer 用於序列化和反序列化模型的核心物件分別是 TensorSerializerTensorDeserializer。為了向這些物件傳遞任意 kwargs 以配置序列化和反序列化過程,您可以將它們作為鍵提供給 model_loader_extra_config,分別使用 serialization_kwargsdeserialization_kwargs。上述物件的所有引數的完整文件字串可以在 tensorizerserialization.py 檔案中找到。

例如,在使用 tensorizer 進行序列化時,可以透過 TensorSerializer 初始化器中的 limit_cpu_concurrency 引數限制 CPU 併發。要將 limit_cpu_concurrency 設定為任意值,您可以在序列化時這樣做

python examples/others/tensorize_vllm_model.py \
   --model facebook/opt-125m \
   --lora-path <lora_id> \
   serialize \
   --serialized-directory s3://my-bucket \
   --serialization-kwargs '{"limit_cpu_concurrency": 2}' \
   --suffix v1

作為透過 TensorDeserializer 自定義載入過程的示例,您可以透過初始化器中的 num_readers 引數,透過 model_loader_extra_config 限制反序列化期間的併發讀取器數量,如下所示

vllm serve s3://my-bucket/vllm/facebook/opt-125m/v1 \
    --load-format tensorizer \
    --enable-lora \
    --model-loader-extra-config '{"deserialization_kwargs": {"num_readers": 2}}'

或者使用 LLM()

from vllm import LLM
llm = LLM(
    "s3://my-bucket/vllm/facebook/opt-125m/v1", 
    load_format="tensorizer",
    enable_lora=True,
    model_loader_extra_config={"deserialization_kwargs": {"num_readers": 2}}
)