游凯超:我与vLLM的2024,很Passion!

科技   2024-12-30 00:01   吉林  


MLNLP社区是国内外知名的机器学习与自然语言处理社区,受众覆盖国内外NLP硕博生、高校老师以及企业研究人员。
社区的愿景是促进国内外自然语言处理,机器学习学术界、产业界和广大爱好者之间的交流和进步,特别是初学者同学们的进步。
    转载自 | 知乎
作者 | 游凯超

楔子

我与 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等重量级人物。更令人惊叹的是,他们都非常技术导向,乐于探讨具体的技术细节。那种思想碰撞的火花,让人深刻体会到硅谷之所以成为创新沃土,绝非偶然。

与Chris Lattner谈笑风生

除了PyTorch Conference这样的年度盛会,vLLM社区也会定期组织线下Meetup,类似开发者见面会。在这些活动中,我有机会与诸多技术专家探讨前沿技术,获得了不少宝贵的经验。广受欢迎的社区课程CUDA Mode也举办了首次线下Meetup,我更是亲眼见到了Andrej Karparthy、CUDA编程入门课的主讲老师——UIUC胡文美教授等人,成功实现线下追星。

如果你来湾区,千万不要错过各种Meetup!比如通过lu.ma网站,你可以轻松订阅大量AI相关的技术meetup,感受硅谷的开放与活力。

vLLM:智能时代的Linux雏形

自2023年6月开源以来,vLLM已经走过了一年半的历程。我很荣幸能够参与到这个项目的发展中,为其成长贡献我的一份力量。

随着智能时代的浪潮席卷而来,我们正见证一场深刻的技术革命。从云端服务器到端侧智能设备,再到自动驾驶汽车,各个领域都在积极探索大模型的应用场景,而vLLM已经在这些场景中实现了广泛的生产部署。我相信,vLLM将逐步发展成为智能时代的“Linux”——一个高效、稳定且开源的系统软件,支撑着智能时代的基础架构。

为更好地践行开源与开放的精神,vLLM正在加入Linux基金会。我们希望借助这一平台,进一步壮大社区,与更多开发者和企业携手,共同建设和完善智能时代的生态系统。正如Linux曾经推动了操作系统的普及,vLLM也有望在智能时代留下浓墨重彩的一笔。

硬件彩票:与其抽奖,不如和庄家合作

访问快要结束时,我斗胆去请教David Patterson教授,问他AI硬件的未来会怎么发展。这位计算机体系结构的泰斗只是淡淡一笑,没有直接回答,而是推荐我去读一篇论文:《The Hardware Lottery》。

与David Patterson教授(以及他那满墙的荣誉证书)远程合影

论文中提出了一个发人深省的观点:在摩尔定律仍然有效的时代,软件和硬件的发展基本是各自为战。写软件的人几乎不用考虑硬件的特性,因为硬件性能每隔两年翻一倍,软件的性能自然也会跟着变快。然而,摩尔定律的黄金年代已经结束。虽然NVIDIA的“黄氏定律”仍在推动硬件性能增长,但这种增长更多体现在专用领域。如果算法没有利用那些特化的硬件特性,就像没抽中“硬件彩票”——注定很难跑赢时代的红利。例如,Hinton提出的capsule network,由于对GPU不友好,目前尚未得到广泛应用。

这让我联想到一个有趣的传闻:NVIDIA在P100上首次推出FP16数值格式时,芯片量产后却发现训练无法收敛,算法研究人员拒绝使用P100,这几乎让他们的数值格式征途“出师未捷身先死”。后来,是混合精度训练让P100化险为夷,从而开启了V100、A100等后续芯片的辉煌。再后来,NVIDIA在开发FP8数值格式时,更是未雨绸缪,先通过大量的软件仿真验证了可行性,才敢造芯片。

这一系列故事让我感到深受启发:硬件亲和性,已经成为决定算法成功的关键。作为一名算法研究人员,与其天马行空地研究算法(抽奖),期待着未来的硬件会对算法亲和,不如直接学习理解硬件,设计对当前硬件亲和的算法,直接与庄家合作,岂不是必然抽中彩票?或许,最理想的情况,就像那个影响了FP8的男人那样,必须算法、软件、硬件协同设计。

展望:泡沫与奇迹

vLLM能受到如此广泛的关注,离不开大模型这两年飞速发展的浪潮。发展得快,泡沫自然也就多。泡沫会不会破灭?什么时候破灭?这些问题没人能预测。但当我向许多经验丰富的“过来人”请教时,他们最常提到的参考是互联网。

回望互联网历史,确实在2000年左右经历了泡沫的破灭。但二十年后再看,即便是当年泡沫顶峰时期那些天马行空的设想,都远远不及互联网如今对世界的深远影响。很多人二十年前吹过的“牛”,不仅已经实现,而且还远超当时的想象。

硅谷先锋Roy Amara曾言道:“对突破性技术,人们往往在短期内高估其影响,但在长期内低估其潜力。” 历史总是在不断重复着螺旋式上升,AI或许就走在类似当年互联网的道路上。也许二十年后再回首,我们会发现,现在我们就站在下一个“互联网级奇迹”的起点上。


谨以此文,纪念即将过去的 2024 年,并提前祝大家 2025 新年快乐!


技术交流群邀请函

△长按添加小助手

扫描二维码添加小助手微信

请备注:姓名-学校/公司-研究方向
(如:小张-哈工大-对话系统)
即可申请加入自然语言处理/Pytorch等技术交流群

关于我们

MLNLP 社区是由国内外机器学习与自然语言处理学者联合构建的民间学术社区,目前已经发展为国内外知名的机器学习与自然语言处理社区,旨在促进机器学习,自然语言处理学术界、产业界和广大爱好者之间的进步。
社区可以为相关从业者的深造、就业及研究等方面提供开放交流平台。欢迎大家关注和加入我们。

机器学习算法与自然语言处理
关注AI前沿技术,助力AI学者进步
 最新文章