作者：清華軟院 博士生 游凱超

楔子

我與 vLLM 的緣分，還得從五年前的那個(gè)暑假說起。

2019 年，我在UC Berkeley 的RISELab跟隨Michael Jordan 教授進(jìn)行暑期研修。某天，我偶然遇到一位新入學(xué)的博士生，厚著臉皮加了他的微信。當(dāng)時(shí)的我怎么也不會(huì)想到，這一“社交冒險(xiǎn)”會(huì)在五年后改變我的人生軌跡。

ChatGPT 之變

時(shí)間快進(jìn)到 2022 年年底，ChatGPT 橫空出世。曾經(jīng)和我一起玩泥巴的青蘋果同學(xué)已經(jīng)成為ChatGPT訓(xùn)練師，而我還在AI頂會(huì)與隨機(jī)分配的審稿人展開“雞同鴨講”式的爭(zhēng)論。頂會(huì)的內(nèi)卷氣氛讓我無比焦慮，也讓我意識(shí)到：我需要做出改變。

彼時(shí)，我正專注于研究一種能加速Conv-BN模塊微調(diào)效率的優(yōu)化方法。算法雖然不復(fù)雜，但需要深入修改模型的計(jì)算圖。一番調(diào)研后，我接觸到了PyTorch 的torch.fx模塊，順勢(shì)了解了PyTorch 2.剛發(fā)布的torch .compile。我對(duì)這些內(nèi)容越看越感興趣，一口氣看完了2022年P(guān)yTorch大會(huì)的全部視頻。

學(xué)習(xí)過程中，我結(jié)識(shí)了PyTorch編譯器團(tuán)隊(duì)的Jason Ansel。他不但教我 Dynamo 的原理，還讓我對(duì)機(jī)器學(xué)習(xí)系統(tǒng)（MLSys）這個(gè)領(lǐng)域有了初步了解。就這樣，我從優(yōu)化Conv-BN的訓(xùn)練效率，一路探索到機(jī)器學(xué)習(xí)編譯器的核心原理。與此同時(shí)，我拜讀了計(jì)算機(jī)體系結(jié)構(gòu)的泰斗David Patterson和編譯器大師Chris Lattner的經(jīng)典講座，他們都提到：在后摩爾定律 時(shí)代，算力增長的唯一途徑是專用芯片，而機(jī)器學(xué)習(xí)編譯器會(huì)成為重要的研究方向。

2023 年年底，Michael Jordan教授再次訪問清華。我向他訴說了對(duì)機(jī)器學(xué)習(xí)研究現(xiàn)狀的困惑，表達(dá)了轉(zhuǎn)向機(jī)器學(xué)習(xí)系統(tǒng)研究的興趣。巧合的是，他的老朋友Ion Stoica ——機(jī)器學(xué)習(xí)系統(tǒng)領(lǐng)域的頂級(jí)專家——也一同前來訪問。Jordan教授一聽，直接就把我推薦給了Ion。在和Ion的交流中，我“現(xiàn)學(xué)現(xiàn)賣”，詳細(xì)闡述領(lǐng)域?qū)Ｓ糜布蜋C(jī)器學(xué)習(xí)編譯器的未來。他聽完非常認(rèn)同，還告訴我，他正在主導(dǎo)一個(gè)項(xiàng)目，涉及處理多種芯片，比如MI300X 、Inferentia等（老實(shí)說，當(dāng)時(shí)這些芯片名字我連聽都沒聽過）。聊到最后，我們一拍即合，在實(shí)驗(yàn)室和導(dǎo)師的鼎力支持下，我幸運(yùn)地拿到了再次訪問UC Berkeley的機(jī)會(huì)。

2024 年，我又一次拖著行李箱，站在了伯克利 的土地上。五年過去了，這里的一切似乎都沒變。Ion Stoica的實(shí)驗(yàn)室，還是五年前熟悉的地點(diǎn)，連實(shí)驗(yàn)室管理員都還是當(dāng)年的Kattt。唯一的變化，是實(shí)驗(yàn)室的名字從 RISELab 改成了 SkyLab。而Ion提到的那個(gè)需要處理多種芯片的項(xiàng)目，正是由五年前我遇到的那位博士生和另一位韓國小哥共同創(chuàng)立的vLLM項(xiàng)目。一切仿佛命中注定，我和vLLM的故事，正式拉開了序幕。

SkyLab充滿歷史滄桑感的門牌，從spark，到ray，再到vLLM，這個(gè)實(shí)驗(yàn)室盛產(chǎn)大型開源軟件 上手：初入江湖

2024年3月，vLLM已經(jīng)在社區(qū)內(nèi)小有名氣，但項(xiàng)目管理還略顯原始。我加入項(xiàng)目后，首先將PyTorch的一些成熟開源管理經(jīng)驗(yàn)移植到vLLM，例如issue模板（要求每個(gè)提 issue 的人必須提供完整的運(yùn)行環(huán)境信息，從而方便我們定位問題）、PR 模板等等。

為了快速了解項(xiàng)目，我訂閱了vLLM的所有GitHub 消息。每天早晨，睜眼后的第一件事就是查看新增的issue和PR，了解項(xiàng)目動(dòng)態(tài)，并解答我能解答的問題。就這樣，我的“vLLM oncall”模式持續(xù)了大半年，直到基本掌握了項(xiàng)目的全貌，才結(jié)束了“oncall”狀態(tài)。

下馬威：PyTorch 2.2 的陰影

我接手的第一個(gè)任務(wù)，乍一看再簡單不過：將vLLM依賴的PyTorch從2.1升級(jí)到2.2。然而，這個(gè)看似“入門級(jí)”的活兒，差點(diǎn)讓我懷疑人生。

經(jīng)過連續(xù)幾周的調(diào)試，我最終發(fā)現(xiàn)，當(dāng)以下四個(gè)條件同時(shí)滿足時(shí)： 曲線救國

1. 使用 L4 機(jī)器（沒有 NVLink ）。

2. 開啟 多卡并行推理 。

3. 啟用 cudagraph。

4. NCCL 版本為 2.19 或以上。

vLLM 的顯存占用會(huì)神秘地增加 2 GiB。

為了解決這個(gè)問題，我硬著頭皮一行行代碼排查，一頭扎進(jìn)NVIDIA的CUDA 編程文檔、驅(qū)動(dòng)文檔和運(yùn)行時(shí)文檔中。最終，我發(fā)現(xiàn)問題出在NCCL的一個(gè)實(shí)驗(yàn)特性上，解決方案僅需設(shè)置一個(gè)環(huán)境變量 以禁用該功能（一行代碼即可搞定）。但為了找到這“一行代碼”，我整整折騰了三個(gè)月。

在徹底解決問題之前，我們只能采取權(quán)宜之計(jì)，先強(qiáng)制安裝一個(gè)獨(dú)立于PyTorch的NCCL 2.18版本。這又讓我不得不參與vLLM的CI 流程、發(fā)布流程等一系列開源項(xiàng)目 管理事務(wù)，順便接手了 vLLM 的分布式推理，并創(chuàng)建了vllm.distributed子模塊。

沒想到，半年后，幾乎相同的問題在vLLM的RLHF 流程中再次出現(xiàn)，并且驚動(dòng)了John Schulman 。是的，就是那個(gè)OpenAI的John Schulman，發(fā)明了PPO 的那個(gè)男人。得知他也用vLLM，我非常激動(dòng)，以最高優(yōu)先級(jí)解決了問題，把RLHF的權(quán)重更新時(shí)間從3分鐘壓縮到了4秒鐘，還和他一起完成了一個(gè)PR。

苦日子：從 GPU Poor 到 GPU Rich

四五月間，我們陸續(xù)收到一些反饋，vLLM在H100上的性能表現(xiàn)不盡人意。然而問題來了——我們自己連一臺(tái)H100都沒有！當(dāng)時(shí)，我還在用實(shí)驗(yàn)室的V100開發(fā)代碼，我們的CI流程也捉襟見肘，只能驗(yàn)證代碼的基本正確性，卻完全無法跟蹤性能變化。

有一次，一個(gè)貢獻(xiàn)者提交了一個(gè)看似無害的改動(dòng)，我審了一眼，覺得沒問題，就痛快地同意合并了。誰知道第二天就有人反饋，這個(gè)PR把整體速度拖慢了好幾倍！當(dāng)時(shí)聽了，真是哭笑不得——身為項(xiàng)目開發(fā)者，我們居然連代碼的性能都搞不清楚，反倒是社區(qū)的一些有錢用戶，在長期跟蹤測(cè)試每個(gè)commit的性能。

究其根本，大模型推理 項(xiàng)目的性能測(cè)試需要高端GPU，而我們CI的資金早已捉襟見肘。整個(gè)項(xiàng)目陷入了青黃不接的窘境，我們一度懷疑：這個(gè)項(xiàng)目怎么維護(hù)下去？還要不要維護(hù)？畢竟，vLLM是個(gè)“燒錢機(jī)器”，它消耗的資源遠(yuǎn)遠(yuǎn)超過我們幾個(gè)學(xué)生的生活費(fèi)，而未來還需要更多投入。我們要從哪里找到這些錢？

幸運(yùn)的是，社區(qū)的熱情與支持在關(guān)鍵時(shí)刻給予了我們極大幫助。在我們四處求援之后，NVIDIA 送來了一臺(tái)滿血H100和一臺(tái)H200；AWS 和Google Cloud等云廠商，捐贈(zèng)了大量計(jì)算資源；真格、紅杉 等創(chuàng)投機(jī)構(gòu)，也慷慨解囊，幫助我們解決了燃眉之急。

與NVIDIA送的H200機(jī)器合影

可以說，vLLM雖然誕生于伯克利，但它的成長靠的是“百家飯 ”。我們深深感激這一份支持，也因此堅(jiān)定了一個(gè)信念：vLLM的今天離不開社區(qū)的幫助，未來我們一定要更好地回饋社區(qū)。

迎戰(zhàn) LLaMA 3.1 405B

vLLM 支持了許多開源模型，且通常在模型發(fā)布當(dāng)天就能推出相關(guān)支持，我們管這叫“day 1 support”。但這種“神速”，背后隱藏著大量看不見的 “day -1 support”——在模型發(fā)布前，我們已經(jīng)悄悄做了大量的準(zhǔn)備工作。這其中，令我印象最深刻的，就是迎戰(zhàn)LlaMA 3.1 405B。

今年四月，Meta 發(fā)布了LLaMA 3 系列模型。乍看之下，它與LLaMA和LLaMA 2 在架構(gòu)上并沒有太大變化，也不需要vLLM提供額外支持。然而，Meta的發(fā)布博客中提到還有一個(gè)400B+的模型正在訓(xùn)練。這消息一下子讓我們坐不住了。

要知道，405B模型僅權(quán)重就需要800GiB的顯存。即便是最頂級(jí)的H100機(jī)器也撐不住這種規(guī)模。于是，我們立刻著手開發(fā)多機(jī)分布式推理功能，包括針對(duì)非RDMA 機(jī)器的流水線并行推理、單機(jī)測(cè)試的CPU offloading等等，最終成功支持了LLaMA 3.1 405B模型的推理。

有趣的是，后來 Meta 告訴我們，他們?cè)诤鸵恍┖献骰锇闇y(cè)試405B模型時(shí)，發(fā)現(xiàn)這些廠商根本不知道如何部署。無奈之下，Meta只能緊急為405B模型開發(fā)FP8量化版本。vLLM對(duì)滿血非量化版405B模型的多機(jī)部署解決方案，使得Meta的十個(gè)官方發(fā)布合作伙伴中，有八個(gè)選擇了 vLLM。

模型架構(gòu)的微創(chuàng)新一直在路上，vLLM團(tuán)隊(duì)也一直在努力，增加對(duì)各種模型的支持。毫不夸張地說，vLLM是支持開源模型類型最廣泛的推理框架。

重構(gòu)與優(yōu)化：永恒的話題

vLLM 一經(jīng)推出，就憑借數(shù)十倍于HuggingFace Transformers 的推理速度吸引了廣泛關(guān)注。然而，隨著社區(qū)越來越大、功能越來越多，早期缺乏性能跟蹤機(jī)制的問題逐漸顯現(xiàn)，尤其是在高端H100 GPU上，小模型的推理性能不佳。

從四五月份開始，隨著性能測(cè)試逐漸完善，我們便一直在討論重構(gòu)的必要性。通過參考類似框架如LMDeploy 、LightLLM和TRT-LLM 的經(jīng)驗(yàn)，我們?yōu)関LLM增加了基于ZMQ的API服務(wù)器、多步調(diào)度（multi-step scheduling）等大幅提升性能的特性。然而，由于 vLLM 的功能非常多，這些優(yōu)化措施有時(shí)會(huì)與某些小眾功能發(fā)生沖突，導(dǎo)致代碼中出現(xiàn)了不少分支邏輯 。

為了徹底解決這一問題，我們正在準(zhǔn)備一次大版本重構(gòu)。這次重構(gòu)將以性能優(yōu)化為核心，優(yōu)先支持常用功能，然后逐步改造那些小眾功能，最終實(shí)現(xiàn)整個(gè)框架的全面升級(jí)。一些早期用戶已經(jīng)部署了新版本的嘗鮮版，獲得了 2~3 倍的性能提升。

此外，多樣化的硬件支持也是需要重構(gòu)的重點(diǎn)。盡管 NVIDIA 是市場(chǎng)上的頭號(hào)玩家，但其他巨頭公司也紛紛推出了自家的芯片。如何兼容多種加速硬件，是我們必須面對(duì)的挑戰(zhàn)。為此，我創(chuàng)建了vllm.platforms子模塊，將硬件相關(guān)的細(xì)節(jié)集中管理，減少主干代碼中的分支邏輯。有趣的是，我發(fā)現(xiàn)PyTorch在硬件支持上也面臨類似的挑戰(zhàn)。vLLM與PyTorch，在這方面可以說是殊途同歸。正因如此，后來我推動(dòng)vLLM加入PyTorch生態(tài)系統(tǒng)就顯得順理成章。通過更緊密地融入PyTorch，我們能夠從其發(fā)展過程中吸取更多經(jīng)驗(yàn)與教訓(xùn)，同時(shí)為PyTorch社區(qū)作出我們的貢獻(xiàn)。

torch.compile集成

在出發(fā)前往伯克利之前，我給 Ion Stoica 畫的大餅，是利用torch.compile來支持多種硬件。然而vLLM的開源項(xiàng)目事情很多，我不得不將與開源相關(guān)的工作置于優(yōu)先位置，只能在閑暇時(shí)間“兼職”探索torch.compile 的集成。

一次偶然的機(jī)會(huì)，我在為Command-R模型增加支持時(shí)，發(fā)現(xiàn)torch.compile的guard系統(tǒng)存在缺陷，會(huì)導(dǎo)致重復(fù)編譯。我向Jason Ansel報(bào)告了這個(gè)問題，他建議我直接聯(lián)系PyTorch Compiler團(tuán)隊(duì)的經(jīng)理。結(jié)果，這個(gè)問題竟讓我有機(jī)會(huì)在PyTorch團(tuán)隊(duì)的例會(huì)上做了一次報(bào)告，深入分析了torch.compile在大模型推理中遇到的挑戰(zhàn)和潛在解決方案。

這次報(bào)告直接促使PyTorch團(tuán)隊(duì)將vLLM的torch.compile集成列為下半年的重點(diǎn)目標(biāo)。經(jīng)過長達(dá)半年的協(xié)作，我們?cè)赑yTorch Compiler核心成員的幫助下，基于PyTorch提供的底層能力，開發(fā)了vLLM專屬的推理優(yōu)化torch.compile技術(shù)棧 ，也將在不久后正式發(fā)布。

有趣的是，torch.compile集成過程中用到的一個(gè)關(guān)鍵功能，正是我去年研究PyTorch Compiler時(shí)為其新增的bytecode hook。

群英薈萃的 PyTorch Conference 與 Meetup

九月中旬，我因獲得社區(qū)創(chuàng)新獎(jiǎng)，受邀參加PyTorch 2024大會(huì)。初次踏入會(huì)場(chǎng)，我就被現(xiàn)場(chǎng)驚人的“人才密度”震撼到了。在隨機(jī)游走 的過程中，我偶遇了Flash Attention的作者Tri Dao、LLVM 的作者Chris Lattner，以及PyTorch的創(chuàng)始人Soumith等重量級(jí)人物。更令人驚嘆的是，他們都非常技術(shù)導(dǎo)向，樂于探討具體的技術(shù)細(xì)節(jié)。那種思想碰撞的火花，讓人深刻體會(huì)到硅谷之所以成為創(chuàng)新沃土 ，絕非偶然。

與Chris Lattner談笑風(fēng)生

除了PyTorch Conference這樣的年度盛會(huì)，vLLM社區(qū)也會(huì)定期組織線下Meetup，類似開發(fā)者見面會(huì)。在這些活動(dòng)中，我有機(jī)會(huì)與諸多技術(shù)專家探討前沿技術(shù)，獲得了不少寶貴的經(jīng)驗(yàn)。廣受歡迎的社區(qū)課程CUDA Mode也舉辦了首次線下Meetup，我更是親眼見到了Andrej Karparthy 、CUDA編程入門課的主講老師——UIUC胡文美教授等人，成功實(shí)現(xiàn)線下追星。

如果你來灣區(qū)，千萬不要錯(cuò)過各種Meetup！比如通過lu.ma網(wǎng)站，你可以輕松訂閱大量AI相關(guān)的技術(shù)meetup，感受硅谷的開放與活力。

vLLM:智能時(shí)代的Linux雛形

自2023年6月開源以來，vLLM已經(jīng)走過了一年半的歷程。我很榮幸能夠參與到這個(gè)項(xiàng)目的發(fā)展中，為其成長貢獻(xiàn)我的一份力量。

隨著智能時(shí)代的浪潮席卷而來，我們正見證一場(chǎng)深刻的技術(shù)革命。從云端服務(wù)器 到端側(cè)智能設(shè)備，再到自動(dòng)駕駛汽車，各個(gè)領(lǐng)域都在積極探索大模型的應(yīng)用場(chǎng)景，而vLLM已經(jīng)在這些場(chǎng)景中實(shí)現(xiàn)了廣泛的生產(chǎn)部署。我相信，vLLM將逐步發(fā)展成為智能時(shí)代的“Linux”——一個(gè)高效、穩(wěn)定且開源的系統(tǒng)軟件，支撐著智能時(shí)代的基礎(chǔ)架構(gòu)。

為更好地踐行開源與開放的精神，vLLM正在加入Linux基金會(huì) 。我們希望借助這一平臺(tái)，進(jìn)一步壯大社區(qū)，與更多開發(fā)者和企業(yè)攜手，共同建設(shè)和完善智能時(shí)代的生態(tài)系統(tǒng)。正如Linux曾經(jīng)推動(dòng)了操作系統(tǒng)的普及，vLLM也有望在智能時(shí)代留下濃墨重彩的一筆。

硬件彩票：與其抽獎(jiǎng)，不如和莊家合作

訪問快要結(jié)束時(shí)，我斗膽去請(qǐng)教David Patterson教授，問他AI硬件 的未來會(huì)怎么發(fā)展。這位計(jì)算機(jī)體系結(jié)構(gòu)的泰斗只是淡淡一笑，沒有直接回答，而是推薦我去讀一篇論文：《The Hardware Lottery》。

與David Patterson教授（以及他那滿墻的榮譽(yù)證書）遠(yuǎn)程合影

論文中提出了一個(gè)發(fā)人深省的觀點(diǎn)：在摩爾定律仍然有效的時(shí)代，軟件和硬件的發(fā)展基本是各自為戰(zhàn)。寫軟件的人幾乎不用考慮硬件的特性，因?yàn)橛布阅苊扛魞赡攴槐叮浖男阅茏匀灰矔?huì)跟著變快。然而，摩爾定律的黃金年代已經(jīng)結(jié)束。雖然NVIDIA的“黃氏定律 ”仍在推動(dòng)硬件性能增長，但這種增長更多體現(xiàn)在專用領(lǐng)域。如果算法沒有利用那些特化的硬件特性，就像沒抽中“硬件彩票”——注定很難跑贏時(shí)代的紅利。例如，Hinton提出的capsule network ，由于對(duì)GPU不友好，目前尚未得到廣泛應(yīng)用。

這讓我聯(lián)想到一個(gè)有趣的傳聞：NVIDIA在P100上首次推出FP16數(shù)值格式時(shí)，芯片量產(chǎn)后卻發(fā)現(xiàn)訓(xùn)練無法收斂，算法研究人員拒絕使用P100，這幾乎讓他們的數(shù)值格式征途“出師未捷身先死”。后來，是混合精度訓(xùn)練 讓P100化險(xiǎn)為夷，從而開啟了V100、A100等后續(xù)芯片的輝煌。再后來，NVIDIA在開發(fā)FP8數(shù)值格式時(shí)，更是未雨綢繆，先通過大量的軟件仿真驗(yàn)證了可行性，才敢造芯片。

這一系列故事讓我感到深受啟發(fā)：硬件親和性，已經(jīng)成為決定算法成功的關(guān)鍵。作為一名算法研究人員，與其天馬行空地研究算法（抽獎(jiǎng)），期待著未來的硬件會(huì)對(duì)算法親和，不如直接學(xué)習(xí)理解硬件，設(shè)計(jì)對(duì)當(dāng)前硬件親和的算法，直接與莊家合作，豈不是必然抽中彩票？或許，最理想的情況，就像那個(gè)影響了FP8的男人那樣，必須算法、軟件、硬件協(xié)同設(shè)計(jì)。

展望：泡沫與奇跡

vLLM能受到如此廣泛的關(guān)注，離不開大模型這兩年飛速發(fā)展的浪潮。發(fā)展得快，泡沫自然也就多。泡沫會(huì)不會(huì)破滅？什么時(shí)候破滅？這些問題沒人能預(yù)測(cè)。但當(dāng)我向許多經(jīng)驗(yàn)豐富的“過來人”請(qǐng)教時(shí)，他們最常提到的參考是互聯(lián)網(wǎng)。

回望互聯(lián)網(wǎng)歷史，確實(shí)在2000年左右經(jīng)歷了泡沫的破滅。但二十年后再看，即便是當(dāng)年泡沫頂峰時(shí)期那些天馬行空的設(shè)想，都遠(yuǎn)遠(yuǎn)不及互聯(lián)網(wǎng)如今對(duì)世界的深遠(yuǎn)影響。很多人二十年前吹過的“牛”，不僅已經(jīng)實(shí)現(xiàn)，而且還遠(yuǎn)超當(dāng)時(shí)的想象。

硅谷先鋒Roy Amara曾言道：“對(duì)突破性技術(shù)，人們往往在短期內(nèi)高估其影響，但在長期內(nèi)低估其潛力。” 歷史總是在不斷重復(fù)著螺旋式上升，AI或許就走在類似當(dāng)年互聯(lián)網(wǎng)的道路上。也許二十年后再回首，我們會(huì)發(fā)現(xiàn)，現(xiàn)在我們就站在下一個(gè)“互聯(lián)網(wǎng)級(jí)奇跡”的起點(diǎn)上。

謹(jǐn)以此文，紀(jì)念即將過去的 2024 年，并提前祝大家 2025 新年快樂！