GKE 経由でマルチホスト生成 AI 大規模オープンモデルをデプロイ、サービングする方法
Cindy Xing
Software Engineering Manager, Google
Puvi Pandian
Software Engineering Manager, Google
※この投稿は米国時間 2024 年 11 月 9 日に、Google Cloud blog に投稿されたものの抄訳です。
LLM（大規模言語モデル）の進歩に支えられて生成 AI が爆発的な成長を遂げる中、デベロッパーにとってオープンモデルへのアクセスはこれまで以上に重要になっています。オープンモデルは一般に公開されている事前トレーニング済みの基礎的な LLM です。Google Cloud の Vertex AI、Kaggle、Hugging Face などのプラットフォームでは、すでにデータ サイエンティスト、ML エンジニア、アプリケーション デベロッパーが簡単にアクセスできるオープンモデルが提供されています。
これらのモデルには強力なインフラストラクチャとデプロイ機能を必要とするものもあります。それに伴い、今回、GKE（Google Kubernetes Engine）経由で Llama 3.1 405B FP16 LLM などのオープンモデルをデプロイ、サービングできる機能をリリースしたことをお知らせいたします。Meta が公開した 4,050 億のパラメータを備える Llama 3.1 は、一般知識、推論能力、コーディング能力において著しい向上を示しています。FP（浮動小数点）16 適合率で実行して 4,050 億のパラメータを保存および処理する場合、このモデルは推論に 750 GB を超える GPU メモリを必要とします。今回の記事で説明する GKE ソリューションにより、このような大規模なモデルのデプロイとサービングの課題をより簡単に達成できるようになります。
カスタマー エクスペリエンス
Google Cloud のお客様は、Vertex AI Model Garden にアクセスして Llama 3.1 モデルタイルを選択することで、Llama 3.1 LLM を表示できます。
デプロイボタンをクリックしたら、GKE を選択し、Llama 3.1 405B FP16 モデルを選択できます。
このページでは、自動生成された Kubernetes yaml と、Llama 3.1 405B FP16 のデプロイおよびサービングに関する詳細な手順を確認できます。
マルチホスト デプロイとマルチホスト サービング
Llama 3.1 405B FP16 LLM には 750 GB 以上の GPU メモリが必要であり、そのデプロイとサービングには大きな課題が生じます。モデルの重みによって消費されるメモリに加えて、KV（Key-Value）キャッシュ ストレージやより長いシーケンス長のサポートといった要素も、全体的なメモリ要件に影響します。現在、Google Cloud Platform で最も強力な GPU サービスである A3 仮想マシンには、それぞれ 80 GB の HBM（高帯域幅メモリ）を備えた 8 つの H100 Nvidia GPU が搭載されています。FP16 Llama 3.1 405B モデルのような LLM を提供する場合、マルチホスト デプロイとマルチホスト サービングが唯一の実行可能なソリューションとなります。ここでは Ray と vLLM で LeaderWorkerSet を使用し、GKE 経由でデプロイします。
LeaderWorkerSet
LeaderWorkerSet（LWS）は、マルチホスト推論のワークロード要件に対応するために特別に開発された Deployment API であり、複数ノード上の複数デバイス間でのモデルのシャーディングと実行を容易にします。Kubernetes Deployment API として構築された LWS は、クラウドにもアクセラレータにも依存せず、GPU と TPU の両方で実行できます。LWS は、以下に示すように、アップストリームの StatefulSet API を基本的な構成要素として活用しています。
LWS アーキテクチャ内では、Pod のグループは単一のエンティティとして管理されます。このグループ内の各 Pod には、0 から n-1 までの一意のインデックスが割り当てられ、インデックス 0 の Pod がグループのリーダーとして指定されます。グループ内の各 Pod の作成は同時に実行され、同一のライフサイクルを共有します。LWS はグループレベルでのロールアウトとローリング アップデートを促進します。各グループは、ローリング アップデート、スケーリング、および配置専用トポロジへのマッピングのための、単一のユニットと見なされます。各グループのアップグレード プロセスは単一の最小単位として実行され、グループ内のすべての Pod が同時に更新されます。トポロジを考慮した配置もオプションでサポートすることで、同じグループ内のすべての Pod を同じトポロジ内に共存させることが可能です。障害処理のコンテキストでも、グループは単一のエンティティとして扱われ、オール オア ナッシング型の再起動をサポートしています。有効にすると、グループ内の 1 つの Pod に障害が発生した場合や、いずれかの Pod 内の 1 つのコンテナが再起動した場合に、グループ内のすべての Pod が再作成されます。
LWS フレームワーク内では、レプリカというコンセプトには単一のリーダーと一連のワーカーで構成されるグループが含まれます。LWS はリーダー用とワーカー用の 2 つのテンプレートをサポートします。LWS は HPA のスケール エンドポイントを提供し、レプリカ数を動的にスケールできるようにします。
vLLM と LWS によるマルチホスト デプロイ
vLLM は人気のオープンソース モデル サーバーであり、テンソル並列処理とパイプライン並列処理を採用することで、マルチノード / マルチ GPU 推論をサポートします。また、Megatron-LM のテンソル並列アルゴリズムを使用して分散テンソル並列処理をサポートします。パイプラインの並列処理では、vLLM はマルチノード推論用に Ray を使用して分散ランタイムを管理します。
テンソル並列処理では、モデルを複数の GPU に水平にパーティショニングし、テンソル並列サイズが各ノード内の GPU 数に相当します。このアプローチでは、GPU 間の高速ネットワーク通信が必要になることに注意してください。
一方、パイプライン並列処理では、モデルをレイヤごとに垂直にパーティショニングするため、GPU 間で常時通信を行う必要はありません。通常、これはマルチホスト サービングに使用されるノード数に相当します。
これらの並列処理戦略の組み合わせは、Llama 3.1 405B FP16 モデル全体に対応するために不可欠です。それぞれ 8 基の H100 GPU を搭載した 2 つの A3 ノードは、合計 1,280 GB のメモリ容量を提供し、このモデルの 750 GB というメモリ要件に十分に対応します。この構成では、Key-Value（KV）キャッシュに必要なバッファメモリも提供され、長いコンテキスト長もサポートできます。この LWS デプロイでは、テンソルの並列サイズは 8 に、パイプラインの並列サイズは 2 に設定されています。
まとめ
今回のブログでは、LWS がマルチホスト サービングに必要不可欠な機能をどのように提供するかを説明しました。この技術は、Llama 3.1 405B FP8 などの小規模なモデルをより費用対効果の高いマシンで提供することもでき、価格性能比を最適化します。詳細については、モデルに適したマシンタイプの選択方法を紹介するこちらのブログ投稿をご覧ください。LWS はオープンソースであり、コミュニティの寄与が大きな役割を果たしています。GitHub にアクセスして、詳細をご確認ください。直接貢献いただける場合も歓迎しています。
Google Cloud Platform は生成 AI ワークロードの導入を支援しており、Vertex AI Model Garden を利用して、マネージド Vertex AI バックエンドまたは GKE DIY（セルフサービス）クラスタ上でオープンモデルをデプロイして提供できます。Google の目標はシームレスなカスタマー エクスペリエンスを生み出すことであり、その一例として挙げられるのがマルチホスト デプロイとマルチホスト サービングです。皆様からのフィードバックをお待ちしています。
