65,000 ノードのさらにその先へ: Google Kubernetes Engine が数兆単位のパラメータを扱う AI モデルに対応

※この投稿は米国時間 2024 年 11 月 14 日に、Google Cloud blog に投稿されたものの抄訳です。

生成 AI が進化するに従い、その革新をもたらす可能性がさまざまな業界や私たちの日常の中に現れるようになっています。そして、大規模言語モデル（LLM）のサイズの増大に伴い、昨今のモデルのパラメータ数は数千億単位に迫る勢いです。最新のモデルであればパラメータ数は 2 兆個近くになることもあり、要求される計算能力は増加する一方です。実際現時点でも、こうした大規模なモデルを最新のアクセラレータでトレーニングする場合、10,000 を超えるノードを持つクラスタが必要になります。

Google Kubernetes Engine（GKE）は 15,000 ノード（世界最大数）のクラスタをサポートし、そのような要求の厳しいトレーニング ワークロードも処理することが可能です。現在はさらに大規模なモデルを想定し、65,000 ノードを持つクラスタのサポートの導入に取り組んでいます。

GKE は最大 65,000 ノードをサポートします。これにより、他の最大手パブリック クラウド プロバイダ 2 社と比較して 10 倍以上の規模を実現できると Google Cloud は見込んでいます。

他に類を見ない規模でのトレーニングや推論が可能に

65,000 ノードにまでスケーリングできる能力は、世界で最もリソースを消費する AI ワークロードにまさに求められている機能です。この能力をアクセラレータの革新されたコンピューティング能力と組み合わせることで、モデルのトレーニング時間を短縮したり、パラメータを数兆単位にスケーリングしたりと、さまざまなことを達成できるようになります。各ノードには、複数のアクセラレータ（例: 4 つのチップを搭載した Cloud TPU v5e ノード）が備わっており、250,000 を超えるアクセラレータを 1 つのクラスタ上で管理できます。

最先端の AI モデルを開発するには、多種多様なワークロードにコンピューティング リソースを割り当てる必要があります。これには、モデルのトレーニングだけでなく、提供、推論、アドホック調査の実施、補助タスクの管理も含まれます。コンピューティング パワーを最小限のクラスタに集中させることで、推論の提供、調査、トレーニングといった種々のワークロードに対する需要の変化に速やかに適応できる柔軟性が得られます。

65,000 ノードをサポートする GKE であれば、5 つのジョブを単一のクラスタで実行できます。そしてそのジョブ一つひとつが、Google Cloud が過去に達成した世界最大の LLM 向けトレーニング ジョブの規模に匹敵します。

最先端の AI を活用するお客様にとって、このような進歩は喜ばしいことです。AI の安全性を追求する研究会社である Anthropic は、信頼性が高く、解釈と制御が可能な AI システムの構築に取り組んでおり、GKE のスケールが拡大されることを喜ばしく感じています。

「GKE でより大規模なクラスタが新たにサポートされるようになり、私たちの AI イノベーションのペースを加速させるために必要となる規模を手に入れることができました。」 - Anthropic コンピューティング担当責任者、James Bradbury 氏

イノベーションの舞台裏

この成果は、さまざまな機能強化が合わさる形で達成されます。まずひとつとして、Google は分散型 Key-Value ストアであるオープンソースの etcd から、新しくより堅牢な Spanner（規模に事実上制限が存在しない Google の分散型データベース）をベースとした Key-Value ストアへの切り替えに取り組んでいます。この変更により、より大規模な GKE クラスタをサポートできるようになるだけでなく、GKE の信頼性が新たな水準に引き上げられます。これは、クラスタ オペレーション（例、クラスタの起動やアップグレード）のレイテンシの改善や、ステートレス クラスタ コントロール プレーンの実装につながります。Google は、Spanner ベースのストレージに etcd API を実装することで後方互換性を確保し、新しいテクノロジーの導入でコアの Kubernetes に変更を加えなくてもよいようにしています。

また、Kubernetes コントロール プレーンを管理する GKE インフラストラクチャが抜本的に見直されたおかげで、GKE のスケーリングが大幅に高速化され、より少ない遅延でお客様のデプロイの需要に応えられるようになりました。このようなクラスタ コントロール プレーンの強化は複数の恩恵をもたらすものであり、高い整合性を維持した状態での高ボリューム オペレーションの実行などが可能になります。コントロール プレーンは、運用上のレイテンシを予測する能力はそのままに、このようなオペレーションに自動的に適応するようになりました。これは、SaaS などの大規模で動的なアプリケーションや、障害復旧とフォールバック、バッチデプロイ、テスト環境などにおいて特に重要なことであり、チャーンレートが高いときにはより重要になります。

また Google は、Google Cloud を AI ワークロード構築に最適なプラットフォームとするために、IaaS と GKE の機能のイノベーションにも継続的に取り組んでいます。以下は最近のこうしたイノベーションの一例です。

• セカンダリ ブートディスク。コンテナ イメージのキャッシュ保存でワークロードの起動が高速になります。

• フルマネージドの DCGM 指標。アクセラレータのモニタリングが改善されます。

• Hyperdisk ML。現在一般提供されているスケーラブルなアプリケーション向けの高性能ストレージ ソリューションです。

• サーバーレス GPU。現在、Cloud Run で利用可能です。

• カスタム コンピューティング クラス。コンピューティング リソースの割り当てとスケーリングの管理能力が改善されます。

• Trillium のサポート。Google の第 6 世代 の TPU である Trillium は、現時点で性能とエネルギー効率において最高水準の TPU です。

• A3 Ultra VM のサポート。A3 Ultra VM は Google の新しい Titanium ML ネットワーク アダプタと、NVIDIA H200 Tensor Core GPU を搭載。RDMA over Converged Ethernet（RoCE）による 3.2 Tbps のノンブロッキング GPU 間トラフィックを実現します。A3 Ultra VM のプレビュー版は来月利用可能になります。

オープンソースに対する継続的な取り組み

Google には長年にわたる確かなオープンソースの文化があり、Google はオープンソース コミュニティへの貢献を惜しみません。これには、Kubernetes のスケーリングに関するものも含まれます。65,000 ノード クラスタのサポートに伴い、このレベルのスケーリングに必要な機能の最適化および強化をすべてコアのオープンソース Kubernetes に組み込みました。

Kubernetes を AI プラットフォームに最も適した基盤とするために Google が投資しているのは、拡張性だけではありません。Kubernetes プロジェクトに対する過去 2 年間の貢献の一部を紹介します。

• Job API の大規模な見直しを推進

• K8S Batch Working Group のインキュベーションを通じ、研究、HPC、AI ワークロードに関するコミュニティを構築し、Kueue.sh などのツールを構築（このツールは Kubernetes でのジョブ キューイングのデファクト スタンダードとなりつつある）

• Kubeflow エコシステムに組み込まれ、異種のジョブ（例: driver-executer）の実行をサポートする JobSet 演算子を作成

• マルチホスト推論のユースケースのための Leader Worker Set コントローラを作成

• JetStream の高度に最適化された内部モデルサーバーを公開

• Kubernetes Serving Working Group のインキュベーションを展開し、モデル指標の標準化、Serving Catalog、Inference Gateway、などの複数の取り組みを推進

Google Cloud は、コンテナ化されたワークロードの実行に最適なプラットフォームを提供することにすべてを注ぎ、これからもイノベーションの限界を押し広げていきます。このような新たな進歩は、次世代の AI テクノロジーを Google がサポートするために欠かせない橋頭堡となります。詳しくは、Kubernetes ポッドキャストをお聴きください。Kaslin Fields が Maciek Rozacki と Wojtek Tyczynski を迎え、GKE の 65,000 ノードのサポートに関して話し合います。単一の GKE クラスタ上で 65,000 ノードを実行するデモは、こちらからご覧いただけます。

-Cloud ランタイム プロダクト担当シニア ディレクター Drew Bradstock
-Google Kubernetes Engine グループ プロダクト マネージャー Maciek Różacki