跳到内容

使用 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 参数调用示例脚本。脚本本身的 docstring 详细解释了 CLI 参数及其正确用法,我们将直接使用 docstring 中的一个示例,假设我们要将模型序列化并保存到我们的 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。如果您打算为您的 tensorized 模型应用 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 适配器保存的 tensorized 模型,您可以这样做:

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。有关上述对象的全部参数的完整 docstrings 可以在 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 自定义加载过程时,可以通过 model_loader_extra_config 中的 num_readers 参数在初始化器中限制反序列化期间的并发读取器数量,如下所示:

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}},
)