跳到内容

使用 Run:ai Model Streamer 加载模型

Run:ai Model Streamer 是一个在读取张量的同时将其流式传输到 GPU 内存的库。更多信息请参阅 Run:ai Model Streamer 文档

vLLM 支持使用 Run:ai Model Streamer 以 Safetensors 格式加载权重。您首先需要安装 vLLM RunAI 可选依赖项。

pip3 install vllm[runai]

若要将其作为兼容 OpenAI 的服务器运行,请添加 --load-format runai_streamer 标志。

vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \
    --load-format runai_streamer

若要从 AWS S3 对象存储运行模型,请执行

vllm serve s3://core-llm/Llama-3-8b \
    --load-format runai_streamer

若要从 Google Cloud Storage 运行模型,请执行

vllm serve gs://core-llm/Llama-3-8b \
    --load-format runai_streamer

若要从 Azure Blob Storage 运行模型,请执行

AZURE_STORAGE_ACCOUNT_NAME=<account> \
vllm serve az://<container>/<model-path> \
    --load-format runai_streamer

身份验证使用 DefaultAzureCredential,它支持 az login、托管标识、环境变量(AZURE_CLIENT_IDAZURE_TENANT_IDAZURE_CLIENT_SECRET)以及其他方法。

若要从 S3 兼容的对象存储运行模型,请执行

RUNAI_STREAMER_S3_USE_VIRTUAL_ADDRESSING=0 \
AWS_EC2_METADATA_DISABLED=true \
AWS_ENDPOINT_URL=https://storage.googleapis.com \
vllm serve s3://core-llm/Llama-3-8b \
    --load-format runai_streamer

可调参数

您可以使用 --model-loader-extra-config 来调整参数。

您可以调整 distributed 参数,该参数控制是否应使用分布式流式传输。目前仅在 CUDA 和 ROCM 设备上可用。这可以显著改善从对象存储或高吞吐量网络文件共享加载模型的时间。您可以在此处阅读有关分布式流式传输的更多信息。

vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \
    --load-format runai_streamer \
    --model-loader-extra-config '{"distributed":true}'

您可以调整 concurrency 参数,它控制并发级别以及将张量从文件读取到 CPU 缓冲区的操作系统线程数。对于从 S3 读取数据,它指的是主机向 S3 服务器打开的客户端实例数量。

vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \
    --load-format runai_streamer \
    --model-loader-extra-config '{"concurrency":16}'

您可以控制并限制张量从文件读取到的 CPU 内存缓冲区大小。您可以在此处阅读有关 CPU 缓冲区内存限制的更多信息。

vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \
    --load-format runai_streamer \
    --model-loader-extra-config '{"memory_limit":5368709120}'

注意

有关可调参数以及可通过环境变量配置的其他参数的进一步说明,请阅读环境变量文档

分片模型加载

vLLM 还支持使用 Run:ai Model Streamer 加载分片模型。这对于跨多个文件拆分的大型模型特别有用。要使用此功能,请使用 --load-format runai_streamer_sharded 标志。

vllm serve /path/to/sharded/model --load-format runai_streamer_sharded

分片加载器要求模型文件遵循与常规分片状态加载器相同的命名模式:model-rank-{rank}-part-{part}.safetensors。您可以使用 --model-loader-extra-config 中的 pattern 参数自定义此模式。

vllm serve /path/to/sharded/model \
    --load-format runai_streamer_sharded \
    --model-loader-extra-config '{"pattern":"custom-model-rank-{rank}-part-{part}.safetensors"}'

要创建分片模型文件,您可以使用 examples/offline_inference/save_sharded_state.py 中提供的脚本。该脚本演示了如何保存与 Run:ai Model Streamer 分片加载器兼容的分片格式模型。

分片加载器支持与常规 Run:ai Model Streamer 相同的所有可调参数,包括 concurrencymemory_limit。这些参数可以通过相同的方式进行配置。

vllm serve /path/to/sharded/model \
    --load-format runai_streamer_sharded \
    --model-loader-extra-config '{"concurrency":16, "memory_limit":5368709120}'

注意

对于张量并行或流水线并行模型,分片加载器特别高效,因为每个工作节点只需要读取自己的分片,而不是整个检查点。