vLLM 模型权重流式加载实战Run:ai Model Streamer 从对象存储加载模型与分片检查点【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm本文围绕 vLLM 对 Run:ai Model Streamer 的集成展开介绍如何通过--load-format runai_streamer将 Safetensors 权重以并发流式方式读入 GPU如何在 AWS S3、Google Cloud Storage、Azure Blob 等对象存储上直接启动 OpenAI 兼容服务以及如何用--model-loader-extra-config调整distributed、concurrency、memory_limit等参数并结合仓库源码说明加载器的工作流程、参数校验逻辑与runai_streamer_sharded分片加载路径的实现细节帮助你在大模型部署场景中显著缩短权重加载阶段。Run:ai Model Streamer 是什么Run:ai Model Streamer 是一个并发读取 tensor 并将其流式传输到 GPU 内存的库。vLLM 通过该库支持以流式方式加载 Safetensors 格式的模型权重避免传统“整文件读入 CPU 内存再上卡”的两段式路径。该功能在 vLLM 中通过两个加载格式提供runai_streamer适用于常规未分片的Safetensors 检查点权重源可以是本地文件系统、S3、GCS 或 Azure Blob Storagerunai_streamer_sharded适用于预先按 worker 分片的检查点每个 worker 只读取属于自己的 shard特别适合张量并行/流水线并行场景。在 vLLM 的加载器注册表中这两种格式分别映射到RunaiModelStreamerLoader和ShardedStateLoader见 model_loader 注册表_LOAD_FORMAT_TO_MODEL_LOADER: dict[str, type[BaseModelLoader]] { ... runai_streamer: RunaiModelStreamerLoader, runai_streamer_sharded: ShardedStateLoader, sharded_state: ShardedStateLoader, ... }安装 Run:ai 可选依赖使用 Run:ai Model Streamer 前需要先安装 vLLM 的 runai 可选依赖pip3 install vllm[runai]从 setup.py 的 extras 定义可以看到该依赖安装的是runai-model-streamer[s3,gcs,azure] 0.15.7即核心库加上 S3、GCS、Azure 三个对象存储后端。仓库的 ROCm 测试依赖 同样锁定了runai-model-streamer[s3,gcs,azure]0.15.7说明该集成在 ROCm 平台同样受支持。此外从源码结构看vLLM 对runai_model_streamer采用惰性导入策略runai_utils.py 中若import失败会退化为PlaceholderModule占位只有在真正使用对象存储路径s3://、gs://、az://前缀时才会触发缺失依赖的报错因此未安装该依赖的本地加载场景不受影响。基本用法本地目录流式加载在启动 OpenAI 兼容服务时添加--load-format runai_streamer即可vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \ --load-format runai_streamerrunai_streamer与runai_streamer_sharded这两个加载格式的完整语义定义在 LoadConfig 的文档字符串中runai_streamer使用 Run:ai Model Streamer 加载 Safetensors 权重runai_streamer_sharded则从预分片的 checkpoint 文件加载权重。源码视角加载流程与参数校验RunaiModelStreamerLoader定义在 runai_streamer_loader.py。其构造函数对--model-loader-extra-config传入的键做白名单校验只允许distributed、concurrency、memory_limit三个键且会先完成全部值校验再统一写入环境变量避免“某个非法值导致前一个配置已部分生效”的问题见 参数校验逻辑配置键校验规则落地的环境变量distributed必须是布尔值直接记录为self._is_distributedconcurrency必须是正整数RUNAI_STREAMER_CONCURRENCYmemory_limit必须是 -1的整数RUNAI_STREAMER_MEMORY_LIMIT此外该加载器还会处理 S3 端点的环境变量桥接若用户设置了AWS_ENDPOINT_URL而未显式设置RUNAI_STREAMER_S3_ENDPOINT会自动将后者指向前者见 端点映射逻辑这正是后文“S3 兼容对象存储”示例中设置AWS_ENDPOINT_URL即可生效的原因。权重准备阶段_prepare_weights按三种来源分派见 _prepare_weights 实现本地目录os.path.isdir判断为真时直接使用对象存储 URI以s3://、gs://、az://开头由 is_runai_obj_uri 判断时直接透传给流式加载器Hugging Face 仓库名走download_weights_from_hf下载*.safetensors文件并额外下载model.safetensors.index.json索引文件。无论哪种来源最终都通过list_safetensors对 Run:ai 库list_safetensors的封装枚举权重文件列表找不到任何 Safetensors 文件时会抛出RuntimeError。真正的流式读取发生在 runai_safetensors_weights_iterator它创建SafetensorsStreamer上下文调用stream_files(files, device..., is_distributed...)启动流式传输再通过streamer.get_tensors()逐个产出(name, tensor.clone())外层用 tqdm 显示 “Loading safetensors using Runai Model Streamer” 进度条。一个值得注意的实现细节是目标设备的选择is_cuda_alike current_platform.is_cuda_alike() device ( fcuda:{current_platform.current_device()} if is_distributed and is_cuda_alike else cpu )即只有当distributedtrue且平台为 CUDA 类设备含 ROCm 等 cuda-like 平台时tensor 才直接流式到对应 GPU其他情况一律先落 CPU buffer。这与文档中“distributed 目前仅在 CUDA 和 ROCM 设备上可用”的说明一致。从对象存储直接加载模型Run:ai Model Streamer 支持四种权重来源本地文件系统、AWS S3、Google Cloud Storage、Azure Blob Storage对应的 URI 方案在 runai_utils.py 中定义为SUPPORTED_SCHEMES [s3://, gs://, az://]。AWS S3vllm serve s3://core-llm/Llama-3-8b \ --load-format runai_streamerGoogle Cloud Storagevllm serve gs://core-llm/Llama-3-8b \ --load-format runai_streamerAzure Blob StorageAZURE_STORAGE_ACCOUNT_NAMEaccount \ vllm serve az://container/model-path \ --load-format runai_streamerAzure 场景的鉴权基于DefaultAzureCredential支持az login、托管标识managed identity、环境变量AZURE_CLIENT_ID、AZURE_TENANT_ID、AZURE_CLIENT_SECRET等 Azure 标准鉴权方式。S3 兼容对象存储对于 MinIO 等 S3 兼容的对象存储通过环境变量重定向端点即可例如将存储端指向 Google Storage 的 S3 兼容接口RUNAI_STREAMER_S3_USE_VIRTUAL_ADDRESSING0 \ AWS_EC2_METADATA_DISABLEDtrue \ AWS_ENDPOINT_URLhttps://storage.googleapis.com \ vllm serve s3://core-llm/Llama-3-8b \ --load-format runai_streamer其中RUNAI_STREAMER_S3_USE_VIRTUAL_ADDRESSING0关闭虚拟寻址路径风格寻址AWS_EC2_METADATA_DISABLEDtrue跳过 EC2 元数据探测AWS_ENDPOINT_URL指定兼容端点——如前所述该变量会被加载器自动桥接为RUNAI_STREAMER_S3_ENDPOINT。可调参数详解--model-loader-extra-config所有可调参数都通过--model-loader-extra-config以 JSON 形式传入对应LoadConfig中的 model_loader_extra_config 字段。distributed是否启用分布式流式加载distributed控制是否使用分布式流式加载。该能力目前仅在 CUDA 和 ROCm 设备上可用可以从对象存储或高吞吐网络文件共享中显著缩短加载时间。vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \ --load-format runai_streamer \ --model-loader-extra-config {distributed:true}从源码看开启后每个 worker 会把 tensor 直接流到本机cuda:{local_rank}设备参见上文runai_safetensors_weights_iterator的设备选择逻辑避免“CPU 中转 逐卡拷贝”的额外开销。concurrency并发度concurrency控制读取 tensor 时的并发级别与 OS 线程数文件 → CPU buffer 方向对于从 S3 读取的场景它就是宿主机向 S3 服务端打开的客户端实例数量。取值必须是正整数vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \ --load-format runai_streamer \ --model-loader-extra-config {concurrency:16}memory_limitCPU 缓冲上限memory_limit用于限制 tensor 读入的 CPU 内存 buffer 大小单位是字节允许-1不限制或 0的整数值。例如限制为 5 GiBvllm serve /home/meta-llama/Llama-3.2-3B-Instruct \ --load-format runai_streamer \ --model-loader-extra-config {memory_limit:5368709120}这三个参数最终都映射为 Run:ai Model Streamer 自身识别的环境变量RUNAI_STREAMER_CONCURRENCY/RUNAI_STREAMER_MEMORY_LIMIT。关于这些环境变量的完整清单包括本文未展开的其他可配置项可以参考 Run:ai Model Streamer 官方的环境变量文档run-ai/runai-model-streamer仓库docs/src/env-vars.md。分片模型加载runai_streamer_sharded对于拆分在多个文件中的大模型vLLM 还支持通过 Run:ai Model Streamer 加载预分片的检查点。每个并行 worker 只读取属于自己的 shard而不是完整 checkpoint因此在张量并行或流水线并行部署下效率尤为突出vllm serve /path/to/sharded/model --load-format runai_streamer_sharded文件命名模式与自定义 pattern分片加载器期望模型文件遵循model-rank-{rank}-part-{part}.safetensors命名模式定义在 ShardedStateLoader.DEFAULT_PATTERN可以通过--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}从源码结构看ShardedStateLoader同时服务sharded_state原生 safetensors 逐 key 读取与runai_streamer_sharded走runai_safetensors_weights_iterator流式读取两种加载格式分派逻辑在 iterate_over_files 中它按当前张量并行 rank 拼出文件通配模式本地目录用globS3 路径用s3_glob枚举见 文件匹配逻辑。需要注意该路径目前只支持已预先分片的 checkpoint找不到匹配文件时会直接报 “only pre-sharded checkpoints are currently supported” 错误。如何生成分片 checkpoint仓库提供了配套脚本 save_sharded_state_offline.py它启动一个 vLLM 引擎实例后调用save_sharded_state将各 worker 的 state dict 按 rank 落盘并把 tokenizer 等元数据文件一并复制到输出目录python save_sharded_state_offline.py \ --model /path/to/load \ --tensor-parallel-size 8 \ --output /path/to/save脚本提供--file-pattern默认即model-rank-{rank}-part-{part}.safetensors和--max-file-size默认 5 GiB控制单个 safetensors 分片文件的大小两个选项保存与加载两侧的模式必须保持一致。生成的 checkpoint 可用load_formatsharded_state或load_formatrunai_streamer_sharded加载后者在此基础上再叠加 Run:ai 的并发流式读取能力。分片加载器的可调参数runai_streamer_sharded支持与普通 Run:ai Model Streamer 相同的核心可调参数concurrency、memory_limit等配置方式一致vllm serve /path/to/sharded/model \ --load-format runai_streamer_sharded \ --model-loader-extra-config {concurrency:16, memory_limit:5368709120}不过从源码结构看有一个细微差别ShardedStateLoader的构造函数只显式 pop 出pattern键见 extra_config 处理其余键会原样保留在model_loader_extra_config中而concurrency/memory_limit到环境变量的映射发生在RunaiModelStreamerLoader.__init__。因此实际生效的调优参数以 Run:ai 库自身读取的环境变量为准——如果你的部署环境需要严格保证参数落地建议同时验证日志中的加载耗时ShardedStateLoader会以 “Loading weights took X seconds” 记录权重加载阶段耗时。行为边界与适用前提格式限制runai_streamer仅支持 Safetensors 权重_prepare_weights以*.safetensors为匹配模式目录中找不到任何 safetensors 文件时会抛出RuntimeError: Cannot find any safetensors model weights ...。extra-config 白名单runai_streamer只接受distributed、concurrency、memory_limit三个键传入其他键会抛出ValueError: Unexpected extra config keys for runai_streamer类型不合法如concurrency为非正整数同样会在校验阶段提前报错不会部分生效。依赖前提对象存储路径s3:///gs:///az://依赖vllm[runai]引入的runai-model-streamer[s3,gcs,azure]分布式流式加载在源码层面仅对 cuda-like 平台CUDA 及 ROCm启用 GPU 直流其他平台回退到 CPU buffer。分片加载限制runai_streamer_sharded只接受预分片 checkpoint且加载时按张量并行 rank 精确匹配分片文件rank 与 checkpoint 生成时的并行规模必须对应若检查点缺少某些 key加载会报 “Missing keys ... in loaded state!”。相关测试仓库在 tests/model_executor/model_loader/runai_streamer_loader/ 下提供了test_runai_model_streamer_loader.py与test_runai_model_streamer_s3.py可参考其用例了解该加载器在本地路径与 S3 路径下的预期行为。小结vLLM 通过runai_streamer与runai_streamer_sharded两种加载格式接入了 Run:ai Model Streamer前者把本地/HF/S3/GCS/Azure 上的 Safetensors 权重以可控并发流式读入设备配合distributed、concurrency、memory_limit三个参数适配不同的存储介质与内存预算后者叠加预分片检查点让每个并行 worker 只读取自己的分片。两者都通过统一的--load-format与--model-loader-extra-config入口配置参数在 RunaiModelStreamerLoader 和 ShardedStateLoader 中有明确的校验与落盘逻辑便于在生产部署中按需调优权重加载阶段。【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考