专栏名称: 吃果冻不吐果冻皮

专注于AI工程化（LLM、MLOps、LLMOps、RAG、Agent）落地。

游凯超：我与vLLM的2024

吃果冻不吐果冻皮 · 公众号 · · 2024-12-26 20:30

正文

原文： https://zhuanlan.zhihu.com/p/14430956145

楔子

我与 vLLM 的缘分，还得从五年前的那个暑假说起。

2019 年，我在UC Berkeley的RISELab跟随Michael Jordan教授进行暑期研修。某天，我偶然遇到一位新入学的博士生，厚着脸皮加了他的微信。当时的我怎么也不会想到，这一“社交冒险”会在五年后改变我的人生轨迹。

ChatGPT之变

时间快进到 2022 年年底，ChatGPT 横空出世。曾经和我一起玩泥巴的青苹果同学已经成为ChatGPT训练师，而我还在AI顶会与随机分配的审稿人展开“鸡同鸭讲”式的争论。顶会的内卷气氛让我无比焦虑，也让我意识到：我需要做出改变。

彼时，我正专注于研究一种能加速Conv-BN模块微调效率的优化方法。算法虽然不复杂，但需要深入修改模型的计算图。一番调研后，我接触到了PyTorch的torch.fx模块，顺势了解了PyTorch 2.刚发布的torch.compile。我对这些内容越看越感兴趣，一口气看完了2022年PyTorch大会的全部视频。

学习过程中，我结识了PyTorch编译器团队的Jason Ansel。他不但教我 Dynamo 的原理，还让我对机器学习系统（MLSys）这个领域有了初步了解。就这样，我从优化Conv-BN的训练效率，一路探索到机器学习编译器的核心原理。与此同时，我拜读了计算机体系结构的泰斗David Patterson和编译器大师Chris Lattner的经典讲座，他们都提到：在后摩尔定律时代，算力增长的唯一途径是专用芯片，而机器学习编译器会成为重要的研究方向。

2023 年年底，Michael Jordan教授再次访问清华。我向他诉说了对机器学习研究现状的困惑，表达了转向机器学习系统研究的兴趣。巧合的是，他的老朋友Ion Stoica——机器学习系统领域的顶级专家——也一同前来访问。Jordan教授一听，直接就把我推荐给了Ion。在和Ion的交流中，我“现学现卖”，详细阐述领域专用硬件和机器学习编译器的未来。他听完非常认同，还告诉我，他正在主导一个项目，涉及处理多种芯片，比如MI300X、Inferentia等（老实说，当时这些芯片名字我连听都没听过）。聊到最后，我们一拍即合，在实验室和导师的鼎力支持下，我幸运地拿到了再次访问UC Berkeley的机会。

2024 年，我又一次拖着行李箱，站在了伯克利的土地上。五年过去了，这里的一切似乎都没变。Ion Stoica的实验室，还是五年前熟悉的地点，连实验室管理员都还是当年的Kattt。唯一的变化，是实验室的名字从 RISELab 改成了 SkyLab。而Ion提到的那个需要处理多种芯片的项目，正是由五年前我遇到的那位博士生和另一位韩国小哥共同创立的vLLM项目。一切仿佛命中注定，我和vLLM的故事，正式拉开了序幕。

SkyLab充满历史沧桑感的门牌，从spark，到ray，再到vLLM，这个实验室盛产大型开源软件

上手：初入江湖

2024年3月，vLLM已经在社区内小有名气，但项目管理还略显原始。我加入项目后，首先将PyTorch的一些成熟开源管理经验移植到vLLM，例如issue模板（要求每个提 issue 的人必须提供完整的运行环境信息，从而方便我们定位问题）、PR模板等等。

为了快速了解项目，我订阅了vLLM的所有GitHub消息。每天早晨，睁眼后的第一件事就是查看新增的issue和PR，了解项目动态，并解答我能解答的问题。就这样，我的“vLLM oncall”模式持续了大半年，直到基本掌握了项目的全貌，才结束了“oncall”状态。

下马威：PyTorch 2.2 的阴影

我接手的第一个任务，乍一看再简单不过：将vLLM依赖的PyTorch从2.1升级到2.2。然而，这个看似“入门级”的活儿，差点让我怀疑人生。

经过连续几周的调试，我最终发现，当以下四个条件同时满足时：曲线救国

使用 L4 机器（没有 NVLink）。
开启多卡并行推理。
启用 cudagraph。
NCCL 版本为 2.19 或以上。

vLLM 的显存占用会神秘地增加 2 GiB。

为了解决这个问题，我硬着头皮一行行代码排查，一头扎进NVIDIA的CUDA编程文档、驱动文档和运行时文档中。最终，我发现问题出在NCCL的一个实验特性上，解决方案仅需设置一个环境变量以禁用该功能（一行代码即可搞定）。但为了找到这“一行代码”，我整整折腾了三个月。

在彻底解决问题之前，我们只能采取权宜之计，先强制安装一个独立于PyTorch的NCCL 2.18版本。这又让我不得不参与vLLM的CI流程、发布流程等一系列开源项目管理事务，顺便接手了 vLLM 的分布式推理，并创建了vllm.distributed子模块。

没想到，半年后，几乎相同的问题在vLLM的RLHF流程中再次出现，并且惊动了John Schulman。是的，就是那个OpenAI的John Schulman，发明了PPO的那个男人。得知他也用vLLM，我非常激动，以最高优先级解决了问题，把RLHF的权重更新时间从3分钟压缩到了4秒钟，还和他一起完成了一个PR。

苦日子：从 GPU Poor 到 GPU Rich

四五月间，我们陆续收到一些反馈，vLLM在H100上的性能表现不尽人意。然而问题来了——我们自己连一台H100都没有！当时，我还在用实验室的V100开发代码，我们的CI流程也捉襟见肘，只能验证代码的基本正确性，却完全无法跟踪性能变化。

有一次，一个贡献者提交了一个看似无害的改动，我审了一眼，觉得没问题，就痛快地同意合并了。谁知道第二天就有人反馈，这个PR把整体速度拖慢了好几倍！当时听了，真是哭笑不得——身为项目开发者，我们居然连代码的性能都搞不清楚，反倒是社区的一些有钱用户，在长期跟踪测试每个commit的性能。

究其根本，大模型推理项目的性能测试需要高端GPU，而我们CI的资金早已捉襟见肘。整个项目陷入了青黄不接的窘境，我们一度怀疑：这个项目怎么维护下去？还要不要维护？毕竟，vLLM是个“烧钱机器”，它消耗的资源远远超过我们几个学生的生活费，而未来还需要更多投入。我们要从哪里找到这些钱？

幸运的是，社区的热情与支持在关键时刻给予了我们极大帮助。在我们四处求援之后，NVIDIA 送来了一台满血H100和一台H200；AWS和Google Cloud等云厂商，捐赠了大量计算资源；真格、红杉等创投机构，也慷慨解囊，帮助我们解决了燃眉之急。

与NVIDIA送的H200机器合影

可以说，vLLM虽然诞生于伯克利，但它的成长靠的是“百家饭”。我们深深感激这一份支持，也因此坚定了一个信念：vLLM的今天离不开社区的帮助，未来我们一定要更好地回馈社区。

迎战 LLaMA 3.1 405B

vLLM 支持了许多开源模型，且通常在模型发布当天就能推出相关支持，我们管这叫“day 1 support”。但这种“神速”，背后隐藏着大量看不见的 “day -1 support”——在模型发布前，我们已经悄悄做了大量的准备工作。这其中，令我印象最深刻的，就是迎战LlaMA 3.1 405B。

今年四月，Meta发布了LLaMA 3系列模型。乍看之下，它与LLaMA和LLaMA 2在架构上并没有太大变化，也不需要vLLM提供额外支持。然而，Meta的发布博客中提到还有一个400B+的模型正在训练。这消息一下子让我们坐不住了。

要知道，405B模型仅权重就需要800GiB的显存。即便是最顶级的H100机器也撑不住这种规模。于是，我们立刻着手开发多机分布式推理功能，包括针对非RDMA机器的流水线并行推理、单机测试的CPU offloading等等，最终成功支持了LLaMA 3.1 405B模型的推理。

有趣的是，后来 Meta 告诉我们，他们在和一些合作伙伴测试405B模型时，发现这些厂商根本不知道如何部署。无奈之下，Meta只能紧急为405B模型开发FP8量化版本。vLLM对满血非量化版405B模型的多机部署解决方案，使得Meta的十个官方发布合作伙伴中，有八个选择了 vLLM。

模型架构的微创新一直在路上，vLLM团队也一直在努力，增加对各种模型的支持。毫不夸张地说，vLLM是支持开源模型类型最广泛的推理框架。

重构与优化：永恒的话题

vLLM 一经推出，就凭借数十倍于HuggingFace Transformers的推理速度吸引了广泛关注。然而，随着社区越来越大、功能越来越多，早期缺乏性能跟踪机制的问题逐渐显现，尤其是在高端H100 GPU上，小模型的推理性能不佳。

从四五月份开始，随着性能测试逐渐完善，我们便一直在讨论重构的必要性。通过参考类似框架如LMDeploy、LightLLM和TRT-LLM的经验，我们为vLLM增加了基于ZMQ的API服务器、多步调度（multi-step scheduling）等大幅提升性能的特性。然而，由于 vLLM 的功能非常多，这些优化措施有时会与某些小众功能发生冲突，导致代码中出现了不少分支逻辑。

为了彻底解决这一问题，我们正在准备一次大版本重构。这次重构将以性能优化为核心，优先支持常用功能，然后逐步改造那些小众功能，最终实现整个框架的全面升级。一些早期用户已经部署了新版本的尝鲜版，获得了 2~3 倍的性能提升。

此外，多样化的硬件支持也是需要重构的重点。尽管 NVIDIA 是市场上的头号玩家，但其他巨头公司也纷纷推出了自家的芯片。如何兼容多种加速硬件，是我们必须面对的挑战。为此，我创建了vllm.platforms子模块，将硬件相关的细节集中管理，减少主干代码中的分支逻辑。有趣的是，我发现PyTorch在硬件支持上也面临类似的挑战。vLLM与PyTorch，在这方面可以说是殊途同归。正因如此，后来我推动vLLM加入PyTorch生态系统就显得顺理成章。通过更紧密地融入PyTorch，我们能够从其发展过程中吸取更多经验与教训，同时为PyTorch社区作出我们的贡献。

`torch.compile` 集成

在出发前往伯克利之前，我给 Ion Stoica 画的大饼，是利用torch.compile来支持多种硬件。然而vLLM的开源项目事情很多，我不得不将与开源相关的工作置于优先位置，只能在闲暇时间“兼职”探索torch.compile的集成。

一次偶然的机会，我在为Command-R模型增加支持时，发现torch.compile的guard系统存在缺陷，会导致重复编译。我向Jason Ansel报告了这个问题，他建议我直接联系PyTorch Compiler团队的经理。结果，这个问题竟让我有机会在PyTorch团队的例会上做了一次报告，深入分析了torch.compile在大模型推理中遇到的挑战和潜在解决方案。

这次报告直接促使PyTorch团队将vLLM的torch.compile集成列为下半年的重点目标。经过长达半年的协作，我们在PyTorch Compiler核心成员的帮助下，基于PyTorch提供的底层能力，开发了vLLM专属的推理优化torch.compile技术栈，也将在不久后正式发布。

有趣的是，torch.compile集成过程中用到的一个关键功能，正是我去年研究PyTorch Compiler时为其新增的bytecode hook。

群英荟萃的 PyTorch Conference 与 Meetup

九月中旬，我因获得社区创新奖，受邀参加PyTorch 2024大会。初次踏入会场，我就被现场惊人的“人才密度”震撼到了。在随机游走的过程中，我偶遇了Flash Attention的作者Tri Dao、LLVM的作者Chris Lattner，以及PyTorch的创始人Soumith等重量级人物。更令人惊叹的是，他们都非常技术导向，乐于探讨具体的技术细节。那种思想碰撞的火花，让人深刻体会到硅谷之所以成为创新沃土，绝非偶然。