使用 Run:ai Model Streamer 实现模型的高效加载
随着大语言模型（LLM）在生产环境中展现出日益复杂的动态特性，其部署过程的挑战性也随之增加。这一趋势进一步凸显了采用高效模型加载策略的必要性，从而确保模型能够快速响应并满足不断变化的业务需求。大语言模型通常需要数十到数百 GB 的内存，在延迟和资源利用方面带来了重大挑战，尤其是在服务需要根据不可预测的用户需求进行扩展时。此时，冷启动问题尤为突出：模型加载到 GPU 内存所需的时间可能会带来显著延迟，严重影响最终用户体验和机器学习系统的运行效率。
传统的模型加载方法 —— 将张量文件依次从存储传输到 CPU 内存，再从 CPU 内存传输到 GPU —— 效率低且成本高，尤其是在自动扩缩容的场景下，需要从空闲状态快速启动的能力。因此，许多企业倾向于长期保留大量空闲副本，以免影响用户体验，但这也导致计算成本的显著上升。为了解决这些问题，Run:ai 推出了 Run:ai Model Streamer 工具。该工具支持从存储中并发读取模型权重并直接流式传输到 GPU，相较现有方法显著提升了加载效率。
1 Run:ai Model Streamer 介绍
Run:ai Model Streamer 的底层采用高性能的 C++ 实现，旨在加速将模型加载到 GPU 的过程，支持来自各种存储类型（如网络文件系统、S3、本地磁盘等）的模型文件。除了性能之外，Run:ai Model Streamer 还提供了 Python SDK，方便集成到现有的推理引擎中（如 vLLM）。
Run:ai Model Streamer 通过使用多线程机制，同时从文件中读取张量数据并将其加载到 CPU 内存中的专用缓冲区中。每个张量会被分配一个标识符，使得应用程序可以在从 CPU 到 GPU 的传输过程中实现张量的并发读取与传输。这意味着，在将某些张量从 CPU 内存传输到 GPU 的同时，还能从存储继续读取其它张量，显著提升了加载效率。
此外，Run:ai Model Streamer 工具充分利用了 GPU 与 CPU 拥有独立子系统的架构特性。GPU 通过 PCIe 总线连接系统，可以直接访问 CPU 内存，无需 CPU 干预。这使得 CPU 端的存储读取和 GPU 的传输可以真正实现实时并行处理，从而在两个子系统间实现高效且快速的模型加载。
Run:ai Model Streamer 的核心特性包括：
并发加载：使用多线程并行读取模型权重文件，缓解存储瓶颈，提高 GPU 利用率。即使是单个张量，也可以被多个线程同时读取。
读取负载均衡：针对不同大小的张量进行负载分配，使得存储带宽可以更充分地被利用，无论模型规模如何都能保持良好性能。
支持多种存储类型：兼容多种存储解决方案，包括本地文件系统（如 SSD）和云端对象存储（如 S3）。
无需格式转换：原生支持 safetensors 格式，无需额外的格式转换，避免加载开销。
易于集成：Run:ai Model Streamer 提供类似 safetensors 的迭代器接口，同时在后台完成并发读取；Python API 简洁易用，便于与 vLLM、TGI 等推理引擎集成，同时享受高性能 C++ 加速层的优势。
Run:ai Model Streamer 支持通过环境变量控制并发读取的参数，例如：
RUNAI_STREAMER_CONCURRENCY：控制并发线程数，默认值为 16。RUNAI_STREAMER_BLOCK_BYTESIZE：每个线程的数据块大小，对于文件系统默认是 2 MB，对于 S3 默认是 8 MB。RUNAI_STREAMER_MEMORY_LIMIT：限制 CPU 内存使用量，默认无限制。
完整的环境变量列表可以参考 Environment Variables。
2 基准测试
在基准测试中，对 3 种主流模型加载器 —— Run:ai Model Streamer、Hugging Face (HF) Safetensors Loader 和 Tensorizer 进行了对比，评估它们在不同存储场景下的模型加载效率。详情可以参考：Run:ai Model Streamer Benchmarks。
基准测试使用了 Meta-Llama-3-8B 模型，这是一个大小为 15 GB 的大语言模型，并以单一的 .safetensors 格式存储。测试运行在 AWS g5.12xlarge 实例上，该实例配备了 4 块 NVIDIA A10G GPU（为保证测试一致性，仅使用其中一块 GPU 进行所有测试）。
软件环境包括：
CUDA 12.4
vLLM 0.5.5（Transformers 4.44.2）
Run:ai Model Streamer 0.6.0
Tensorizer 2.9.0
Transformers 4.45.0.dev0
Accelerate 0.34.2
上图展示了 vLLM 引擎使用 3 种模型加载器在不同存储类型（GP3 SSD、IO2 SSD 和 S3）下完成加载并准备好进行推理所需的总时间。深色柱表示从存储加载模型到 GPU 的耗时，浅色柱则表示 vLLM 引擎整体准备就绪的总耗时（即模型加载时间加上推理引擎预热时间）。
在本地存储（GP3 和 IO2 SSD）场景下，Run:ai Model Streamer 和 Tensorizer 均显著优于 Hugging Face (HF) Safetensors Loader，使整体就绪时间几乎缩短了一半。
而在 S3 存储场景中，仅对 Run:ai Model Streamer 和 Tensorizer 进行了测试，其中 Run:ai Model Streamer 表现出明显更快的就绪速度。需要指出的是，由于 Hugging Face 的 Safetensors Loader 不支持直接从 S3 加载模型，因此未参与该场景的对比。
在各项实验中，Run:ai Model Streamer 相较其他模型加载器均展现出了显著的性能优势。
本地存储（如 GP3 和 IO2 SSD）场景指的是模型已预先下载到本地磁盘，因此加载时间不包含下载过程。而在 S3 存储场景中，模型文件存储于 S3，加载时间则包含了从 S3 下载模型的过程。从图中可以看到，Run:ai Model Streamer 在 S3 场景下反而表现更快，这与网络带宽关系密切。基准测试中，虚拟机与 S3 Bucket 位于同一区域，带宽高达 4GB/s。结合 Run:ai Model Streamer 的并发读取能力，使其能够充分利用 S3 的读取带宽，因此在 S3 场景下展现出更优性能。
3 使用 Run:ai Model Streamer
Run:ai Model Streamer 支持通过 Python SDK 进行使用，具体可参考官方文档 Using Run:ai Model Streamer。目前，vLLM 已经集成 Run:ai Model Streamer，因此本文将基于 vLLM 演示如何使用 Run:ai Model Streamer 加载模型。为了便于实验，本文在 Google Colab 上运行 vLLM，并使用了 A100 GPU。
需要注意的是，在 Google Colab 中运行 shell 命令时需要使用
!开头，这是 Colab 提供的一种方式，用来在 Jupyter Notebook 环境中运行终端命令。
3.1 安装 vLLM 和 Run:ai Model Streamer
执行以下命令，安装含有 Run:ai Model Streamer 依赖的 vLLM：
3.2 vLLM 使用 Run:ai Model Streamer 加载模型
我们先看看在不使用 Run:ai Model Streamer 的情况下，vLLM 加载模型需要多长时间。vLLM 默认使用 Hugging Face (HF) Safetensors Loader 来加载模型。本次实验使用 deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B 模型，大小约为 3.55GB。这里已经提前将模型下载到本地。也可以先执行一次 vllm serve 命令，在默认情况下，vLLM 会自动从 Hugging Face 下载模型。
输出显示，模型加载的时间为 1.67 秒。
接下来看看 vLLM 使用 Run:ai Model Streamer 加载模型的效果。通过 --load-format runai_streamer 参数指定使用 Run:ai Model Streamer 加载模型，并将并发读取线程数设置为 32。
输出显示模型加载时间为 1.39 秒，Run:ai Model Streamer 略快于 Hugging Face（HF）Safetensors Loader 的 1.67 秒。估计在模型体积更大的情况下，Run:ai Model Streamer 的优势才能凸显出来。
最后，我们来看看使用 Run:ai Model Streamer 直接从 S3 加载模型的性能表现。关于如何将模型从 Hugging Face 下载并上传至 S3，请参考文章末尾的附录部分。
使用 vllm serve 命令从 S3 Bucket s3://cr7258/DeepSeek-R1-Distill-Qwen-1.5B 加载模型，并通过 --load-format runai_streamer 参数指定使用 Run:ai Model Streamer，同时将并发线程数设置为 32。需要注意的是，默认的下载请求超时时间为 1 秒，如果网络较慢，可能会导致超时。此时可以通过设置环境变量 RUNAI_STREAMER_S3_REQUEST_TIMEOUT_MS 来延长超时时间。另外，你还需要设置 AWS_ACCESS_KEY_ID 和 AWS_SECRET_ACCESS_KEY 环境变量以确保有权限访问 S3。
输出结果显示，从 S3 直接加载模型的耗时为 6.10 秒。由于服务器访问 S3 的平均读取带宽约为 650 MB/s，远低于基准测试中 4 GB/s 的理想带宽，因此在本次实验中，Run:ai Model Streamer 从 S3 加载模型的时间高于从本地文件系统加载的时间。
4 附录
4.1 从 Hugging Face 下载模型并上传至 S3
首先在 AWS S3 创建一个 Bucket（本文以 cr7258 为例），并在 Bucket 中创建一个文件夹用于存储模型权重文件（本文以 DeepSeek-R1-Distill-Qwen-1.5B 为例）。
然后使用 huggingface-cli 从 Hugging Face 下载模型。
模型会被下载至 /root/.cache/huggingface/hub/models--deepseek-ai--DeepSeek-R1-Distill-Qwen-1.5B 目录中。
然后使用 aws s3 sync 命令将模型上传到 S3。模型文件位于 snapshots 子目录下，具体的目录名是 Hugging Face 在某次发布或更新模型时生成的哈希值，本例中为 ad9f0ae0864d7fbcd1cd905e3c6c5b069cc8b562。
注意：你需要设置好以下环境变量，确保你有权限能够访问 S3。
上传后的效果如下：
4.2 故障排查
如果在使用 Run:ai Model Streamer 加载模型时遇到问题，可以设置以下环境变量来获取更详细的日志信息。
5 总结
本文介绍了 Run:ai Model Streamer 这一高效的模型加载工具，它通过并发读取模型权重并直接流式传输到 GPU，显著提升了大型模型的加载效率。基准测试表明，相比于 Hugging Face (HF) Safetensors Loader 和 Tensorizer，Run:ai Model Streamer 在各种存储场景下均展现出了明显的性能优势，特别是在从 S3 等远程存储加载模型时效果更为显著。文章还结合实验，演示了如何在 vLLM 中集成并使用 Run:ai Model Streamer 加载模型。
6 参考资料
Run:ai Model Streamer: Performance Benchmarks：https://pages.run.ai/hubfs/PDFs/White%20Papers/Model-Streamer-Performance-Benchmarks.pdf
Loading models with Run:ai Model Streamer：https://docs.vllm.ai/en/v0.8.2/models/extensions/runai_model_streamer.html
runai-model-streamer：https://github.com/run-ai/runai-model-streamer
Loading Llama-2 70b 20x faster with Anyscale Endpoints：https://www.anyscale.com/blog/loading-llama-2-70b-20x-faster-with-anyscale-endpoints#anyscale-model-loader
'File access error' when attempting to access s3-compatible storage in Kubernetes, works in Docker：https://github.com/run-ai/runai-model-streamer/issues/55#issuecomment-2884254292
7 欢迎关注
版权声明: 本文为 InfoQ 作者【Se7en】的原创文章。
原文链接:【http://xie.infoq.cn/article/1a190da48c37c6d4320874a6c】。文章转载请联系作者。
还未添加个人签名 2020-01-10 加入
喜欢钻研技术，熟悉 Elasticsearch、Kubernetes 。
评论