618ZXW

2024年のvLLMへの旅:清華大学卒業生のvLLM開発パス

データホエール・ユー・カイチャオ

Datawhaleの推奨事項

著者: ユー・カイチャオ、清華大学ソフトウェア学院博士課程学生

くさび

vLLM との私のつながりは、5 年前の夏に遡ります。

2019年、カリフォルニア大学バークレー校のRISELabで、マイケル・ジョーダン教授の指導のもと、夏季研究プログラムに参加しました。ある日、偶然、新しく入学した博士課程の学生と出会い、少し勇気を出してWeChatで彼を追加してしまいました。この「ソーシャルアドベンチャー」が5年後、私の人生を大きく変えることになるとは、当時は知る由もありませんでした。

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年末、マイケル・ジョーダン教授が再び清華大学を訪れました。私は機械学習研究の現状について戸惑い、機械学習システムの研究に重点を移したいと伝えました。偶然にも、ジョーダン教授の旧友であり、機械学習システムのトップエキスパートであるイオン・ストイカ氏も来訪していました。それを聞いたジョーダン教授はすぐに私をイオン氏に推薦してくれました。イオン氏との会話の中で、私は「自分が学んだこと」を披露し、ドメイン特化型ハードウェアと機械学習コンパイラの将来について詳しく説明しました。イオン氏は非常に同意して話を聞いてくれ、MI300XやInferentiaといった複数のチップ(正直に言うと、当時はこれらのチップの名前すら聞いたことがありませんでした)を含むプロジェクトを主導していることを教えてくれました。結局、私たちはすぐに意気投合し、研究室と指導教官の強力なサポートのおかげで、カリフォルニア大学バークレー校を再び訪れる機会を得ることができました。

2024年、私はスーツケースを手に、再びバークレーの土を踏んだ。5年が経っていたが、何もかもが変わっていないように思えた。イオン・ストイカの研究室は5年前と変わらず馴染みのある場所で、研究室の管理者もKatttのままだった。唯一の変化は、研究室の名前がRISELabからSkyLabに変わったことだった。そして、イオンが言っていた複数のチップを処理するプロジェクトは、5年前に出会った博士課程の学生ともう一人の韓国人が共同設立したvLLMプロジェクトだった。まるで運命づけられていたかのようだった。vLLMとの私の物語が、正式に始まったのだ。

Spark、Ray、vLLM という名前が付けられた SkyLab の歴史的な看板は、この研究所が大規模なオープンソース ソフトウェアを大量に生み出してきたことを反映しています。

入門:格闘技の世界へ

2024年3月までに、vLLMはコミュニティ内で既に一定の認知度を得ていましたが、プロジェクト管理はまだ未熟な状態でした。プロジェクトに参加した後、私はまず、PyTorchから成熟したオープンソース管理エクスペリエンスをvLLMに移植しました。例えば、課題テンプレート(課題を提起する各人に完全なランタイム環境情報を提供することを義務付け、問題の特定を容易にする)やPRテンプレートなどです。

プロジェクトを素早く理解するために、vLLMのGitHubアップデートをすべて購読しました。毎朝、目が覚めて最初にやることは、新しい問題やプルリクエストを確認し、プロジェクトの進捗状況を追跡し、できる限りの質問に答えることでした。この「vLLMオンコール」モードは半年以上続き、プロジェクトを大まかに把握できた時点で、ついに「オンコール」状態を終えました。

頭への一撃:PyTorch 2.2 の影

最初の仕事は、一見すると信じられないほど簡単なものに見えました。vLLMが依存するPyTorchをバージョン2.1から2.2にアップグレードするのです。しかし、この一見「初心者レベル」の作業は、私の人生の選択を疑わせるほどでした。

数週間のデバッグを経て、最終的に、次の 4 つの条件が同時に満たされると、迂回アプローチが機能することを発見しました。

  1. L4 マシン (NVLink なし) を使用します。
  2. マルチ GPU 並列推論を有効にします。
  3. cudagraph を有効にします。
  4. NCCL バージョン 2.19 以上。

vLLM のメモリ使用量が不思議なことに 2 GiB 増加します。

この問題を解決するために、私はNVIDIAのCUDAプログラミング、ドライバ、ランタイムに関するドキュメントを隅々まで読み、コードを一つ一つ丹念にチェックしました。最終的に、問題はNCCLの実験的な機能にあり、解決策は環境変数を設定してその機能を無効にすることだけで済む(たった1行のコード)ことがわかりました。しかし、その「たった1行のコード」を見つけるのに丸3ヶ月もかかりました。

問題が完全に解決されるまでは、PyTorchとは独立したNCCL 2.18バージョンを強制的にインストールするという、一時的な対策しか取れません。そのため、私はvLLMのCIやリリースプロセスといったオープンソースプロジェクトの管理業務に積極的に参加する必要があり、またvLLMの分散推論も引き継いでvllm.distributedサブモジュールを作成しました。

予想外にも、6ヶ月後、vLLMのRLHFプロセスでほぼ同じ問題が再発し、ジョン・シュルマン氏の注目を集めました。そう、彼はOpenAIのジョン・シュルマン氏、PPOを発明した人物です。彼もvLLMを使っていたと知り、私は大変喜び、この問題を最優先で解決しました。RLHFの重み更新時間を3分から4秒に短縮することができました。PRでも彼と共同作業しました。

厳しい時代: GPU 貧弱から GPU 豊富へ

4月と5月に、H100上でのvLLMのパフォーマンスが不十分だというフィードバックをいただきました。しかし、問題が発生しました。私たち自身にH100がなかったのです!当時、私はまだラボのV100を使ってコード開発を行っており、CIプロセスは限界に達していました。コードの基本的な正しさは検証できるものの、パフォーマンスの変化を全く追跡することができませんでした。

ある時、ある貢献者が一見無害そうな変更を提出しました。私はそれをちらっと見て、問題ないと思い、すぐにマージを承認しました。ところが、その翌日、このPRによって全体の速度が数倍も遅くなったという報告が来るとは、誰も想像もしていませんでした! これを聞いた時、私は面白がると同時に憤慨しました。プロジェクト開発者である私たちは、コードのパフォーマンスさえ理解していなかったのに、コミュニティ内の裕福なユーザーの中には、すべてのコミットのパフォーマンスを綿密に追跡・テストしている人もいたのです。

幸いなことに、コミュニティの熱意とサポートは、まさに危機的な状況において大きな助けとなりました。各方面に支援を求めた結果、NVIDIAはフル装備のH100とH200を送付し、AWSやGoogle Cloudなどのクラウドプロバイダーは大量のコンピューティングリソースを寄付し、ZhenFundやSequoia Capitalなどのベンチャーキャピタルも惜しみない資金提供を行い、差し迫った問題の解決に尽力してくれました。

NVIDIAから提供されたH200マシンの写真

vLLMはバークレーで誕生しましたが、その成長は多くの方々のご支援によって支えられてきたと言っても過言ではありません。私たちは皆様のご支援に深く感謝しており、vLLMの今日の成功はコミュニティの皆様のご支援と切り離せないという信念を強めました。そして、将来に向けてコミュニティに恩返しをしていく必要があると考えています。

LLaMA 3.1 405Bに直面

vLLMは多くのオープンソースモデルをサポートしており、通常はモデルのリリース当日にサポートを開始します。私たちはこれを「Day 1サポート」と呼んでいます。しかし、この「電光石火のスピード」の裏には、目に見えない多くの「Day 1サポート」が隠されています。モデルのリリース前に、私たちはひっそりと多くの準備作業を行ってきました。その中でも、私が最も感銘を受けたのは、LlaMA 3.1 405Bへの準備作業です。

今年4月、MetaはLLaMA 3シリーズのモデルをリリースしました。一見すると、LLaMAやLLaMA 2からアーキテクチャに大きな変更はなく、vLLMによる追加サポートも必要ないように見えます。しかし、Metaのリリースブログには、4000億を超えるモデルも学習中であることが記載されていました。このニュースは、私たちを不安にさせました。

405Bモデルは重み付けだけで800GiBのGPUメモリを必要とすることを理解することが重要です。最上位機種のH100マシンでさえ、このスケールには対応できません。そのため、私たちはすぐにマルチマシン分散推論機能の開発に着手しました。これには、RDMA非対応マシン向けのパイプライン並列推論、単一マシンテスト向けのCPUオフロードなどが含まれます。最終的に、LLaMA 3.1 405Bモデルの推論をサポートすることに成功しました。

興味深いことに、Meta社は後に、405Bモデルを複数のパートナーとテストしていた際に、これらのベンダーがモデルの導入方法を全く理解していないことに気づいたと語っています。他に選択肢がなかったため、Meta社は405Bモデル用のFP8量子化バージョンを緊急に開発する必要がありました。vLLMの本格的な非量子化405Bモデル向けマルチマシン展開ソリューションは、Meta社の公式リリースパートナー10社のうち8社がvLLMを採用するきっかけとなりました。

モデルアーキテクチャにおけるマイクロイノベーションは継続的に進められており、vLLMチームは様々なモデルのサポート拡大に尽力しています。vLLMは、最も幅広いオープンソースモデルタイプをサポートする推論フレームワークであると言っても過言ではありません。

再構築と最適化:永遠のテーマ

vLLMはリリース当初、HuggingFace Transformersの数十倍もの推論速度を誇ることから大きな注目を集めました。しかし、コミュニティが成長し、より多くの機能が利用可能になるにつれて、初期段階ではパフォーマンス追跡メカニズムが不足していることが徐々に明らかになり、特にハイエンドのH100 GPUでは、小規模モデルの推論性能が低かったことが顕著になりました。

4月か5月頃から、パフォーマンステストが徐々に改善していくにつれ、リファクタリングの必要性について議論が始まりました。LMDeploy、LightLLM、TRT-LLMといった類似フレームワークの経験を踏まえ、ZMQベースのAPIサーバーやマルチステップスケジューリングなど、パフォーマンスを大幅に向上させる機能をvLLMに追加しました。しかし、vLLMの機能数が多いため、これらの最適化が特定のニッチな機能と競合し、コード内に多数の分岐ロジックが存在するという問題がありました。

この問題を完全に解決するため、メジャーバージョンの全面的な見直しを準備しています。この全面的な見直しでは、パフォーマンスの最適化に重点を置き、使用頻度の高い機能のサポートを優先し、使用頻度の低い機能を段階的にリファクタリングすることで、最終的にはフレームワーク全体の包括的なアップグレードを実現します。一部のアーリーアダプターは既に新しいシステムのベータ版を導入しており、2~3倍のパフォーマンス向上を実感しています。

torch.compile統合

バークレーに向かう前は、Ion Stoicaの様々なハードウェアコンポーネントをサポートするためにtorch.compile使うことを構想していました。しかし、vLLMオープンソースプロジェクトは非常に要求が厳しく、オープンソース関連の作業を優先せざるを得なくなり、余暇にtorch.compile統合を「パートタイム」で検討することができました。

偶然にも、Command-Rモデルのサポートを追加している際に、 torch.compileのガードシステムに欠陥があり、コンパイルが重複していることを発見しました。この問題をJason Anselに報告したところ、PyTorch Compilerチームのマネージャーに直接連絡するよう勧められました。驚いたことに、これがきっかけでPyTorchチームミーティングでプレゼンテーションする機会が生まれ、大規模モデル推論においてtorch.compileが直面する課題と潜在的な解決策について詳細な分析を行うことができました。

この報告は、PyTorchチームが今年後半の主要目標としてvLLM torch.compileの統合を優先することを直接促しました。6ヶ月にわたる共同作業を経て、PyTorch Compilerチームのコアメンバーの協力を得て、PyTorchが提供する基盤機能をベースにvLLM専用の推論最適化技術スタックであるtorch.compile開発し、まもなく正式リリースする予定です。

興味深いことに、 torch.compile統合プロセスで使用される重要な機能は、昨年私が調査していたときに PyTorch コンパイラに追加したバイトコード フックです。

PyTorchカンファレンスとミートアップはトップクラスの才能が集まる場です

9月中旬、コミュニティ・イノベーション・アワードを受賞したため、PyTorch 2024カンファレンスへの参加を依頼されました。会場に入ると、その圧倒的な才能の密度に圧倒されました。会場内を歩き回る中で、Flash Attentionの作者Tri Dao氏、LLVMの作者Chris Lattner氏、PyTorchの創設者Soumith氏といった著名人に偶然出会いました。さらに印象的だったのは、彼らが共通の技術志向を持ち、具体的な技術内容について議論する姿勢を見せていたことです。こうした知的交流の火花が散る様子から、シリコンバレーがイノベーションの土壌として栄えているのは偶然ではないことを改めて実感しました。

クリス・ラトナーとおしゃべりしながら笑う

PyTorchカンファレンスのような年次イベントに加え、vLLMコミュニティは開発者の集まりのようなオフラインミートアップを定期的に開催しています。これらのイベントでは、多くの技術専門家と最先端技術について議論する機会があり、貴重な経験を得ることができました。人気のコミュニティコース「CUDA Mode」も初のオフラインミートアップを開催し、UIUCのAndrej Karparthy氏やCUDAプログラミング入門コースの講師であるWenmei Hu教授など、多くの方々と直接お会いすることができました。彼らと直接会うという私の夢が叶いました。

ベイエリアにお越しの際は、ぜひ様々なミートアップに参加してみてください!例えば、lu.maというウェブサイトでは、AI関連の技術系ミートアップに簡単に登録でき、シリコンバレーのオープンさと活気を体感できます。

vLLM: インテリジェント時代の Linux プロトタイプ

vLLMは2023年6月にオープンソースとしてリリースされてから1年半が経ちました。このプロジェクトの開発に携わり、その成長に貢献できたことを光栄に思います。

インテリジェント時代の波が押し寄せ、私たちは大きな技術革命を目の当たりにしています。クラウドサーバーからエッジインテリジェントデバイス、自動運転車まで、様々な分野で大規模なアプリケーションシナリオが積極的に模索されており、vLLMは既にこれらのシナリオにおいて広く実用化されています。vLLMは、インテリジェント時代のインフラを支える、高効率で安定したオープンソースのシステムソフトウェア、つまりインテリジェント時代の「Linux」へと徐々に進化していくと信じています。

ハードウェア宝くじ: 宝くじに参加する代わりに、ブックメーカーに協力します。

訪問も終わりに近づき、思い切ってデイビッド・パターソン教授にAIハードウェアの将来について尋ねてみた。コンピュータアーキテクチャの巨匠である彼は、かすかに微笑むだけで、直接的な答えは返さず、「ハードウェア宝くじ」と題された論文を読むことを勧めてくれた。

デビッド・パターソン教授とのリモート写真(教授の壁には表彰状がいっぱい)。

興味深い逸話を思い出しました。NVIDIAがP100で初めてFP16数値形式を導入した際、量産後にトレーニングが収束しないことが分かりました。アルゴリズム研究者たちはP100の使用を拒否し、数値形式への取り組みは開始前にほぼ失敗に終わりそうになりました。その後、混合精度トレーニングによってP100は救われ、V100やA100といった後継チップの成功へと繋がりました。さらに後、FP8数値形式を開発する際には、NVIDIAはさらに慎重な対策を講じ、チップの製造に踏み切る前に、広範なソフトウェアシミュレーションを通じてその実現可能性を検証しました。

この一連の話は私に深い感銘を与えました。ハードウェアとの親和性は、アルゴリズムの成功にとって非常に重要になっているのです。アルゴリズム研究者として、(宝くじに当たるように)想像力豊かにアルゴリズムを研究し、将来のハードウェアがそれに対応していることを願うよりも、ハードウェアを直接学び理解し、現在のハードウェアと互換性のあるアルゴリズムを設計し、システムと直接連携する方がよいのではないでしょうか。そうすれば、宝くじに当たることは間違いないでしょう。おそらく理想的なシナリオは、FP8に影響を与えた人物のように、アルゴリズム、ソフトウェア、ハードウェアの協調設計を必要とするのでしょう。

展望:バブルと奇跡

vLLMが広く注目を集めているのは、過去2年間の大規模モデルの急速な発展と切り離せない関係にあります。急速な発展は当然のことながら、さらなる誇大宣伝と憶測を招きます。このバブルはいつ崩壊するのか? これらの疑問を予測できる人は誰もいません。しかし、多くの経験豊富な専門家に相談したところ、最も頻繁に挙げられた情報源はインターネットでした。

インターネットの歴史を振り返ると、確かに2000年頃にバブル崩壊を経験しました。しかし、20年後を振り返ると、バブルのピーク時に生まれた奇想天外なアイデアでさえ、インターネットが今日の世界にもたらした計り知れない影響に比べれば、取るに足らないものとなっています。20年前に語られた「自慢」の多くは、現実のものとなっただけでなく、当時の期待をはるかに上回るものとなっています。

シリコンバレーのパイオニア、ロイ・アマラはかつてこう言いました。「人々は画期的な技術の短期的な影響を過大評価しがちだが、長期的な可能性を過小評価しがちだ。」歴史は常に上昇の螺旋を繰り返すものであり、AIも初期のインターネットと同じような道を辿っているのかもしれません。おそらく20年後、振り返ってみると、私たちは今まさに次の「インターネットレベルの奇跡」の出発点に立っていることに気づくでしょう。

この記事は、2024 年の終わりを記念し、皆様が 2025 年の幸せな新年を迎えられるよう事前に祈るために捧げられています。

この記事は知书から転載し、Jizhiliu が編集したものです。

著者: カイチャオ・ユー、清華大学ソフトウェア学院博士課程学生、研究分野: 人工知能/機械学習/転移学習

志湖ホームページ:https://www.zhihu.com/people/...

いいね (3件のいいね!)↓