GKE における Cloud TPU への対応の詳細
Google Cloud Japan Team
※この投稿は米国時間 2023 年 9 月 1 日に、Google Cloud blog に投稿されたものの抄訳です。
Cloud TPU は、大規模言語モデル（LLM）、Vision Transformer ベースのアーキテクチャ、拡散ベースの生成モデル、およびレコメンデーション システムのトレーニング用に構築されています。しかしこれまでは、障害処理やロギング、モニタリングといった基礎的な操作に手作業が必要であり、Cloud TPU を使って大規模な AI ワークロードをオーケストレーションするには大変な労力がかかる可能性がありました。
こうした課題に対処するため、今週、スケーラブルなフルマネージド Kubernetes サービスである Google Kubernetes Engine（GKE）での Cloud TPU の一般提供開始を発表しました。つまり、待ち望まれていたとおり、GKE において Cloud TPU v4 と最新の第 5 世代 Cloud TPU v5e の両方への対応が開始します。この対応開始を受けて、Cloud TPU 上で大規模な AI ワークロードのオーケストレーションを管理するために GKE を活用すると、AI 開発の生産性が向上します。
GKE における Cloud TPU の利点
GKE の Cloud TPU では Google のカスタム設計された TPU のコスト パフォーマンスと GKE の機能の高さが共存しており、GKE のシンプルさを活かして AI ワークロードを表現、実行することができるようになりました。自動スケーリング、自動プロビジョニング、自動修復、自動アップグレードといった GKE の機能を活用して、あらゆる種類のワークロードを一貫した運用環境で実行することで、Day 2（2 日目からの）運用の負担を軽減し、リソースを効率的に管理できます。さらに、Weights & Biases の Launch や Google Cloud Managed Service for Prometheus などのツールを GKE で Cloud TPU と併用することで、TPU と GKE のエコシステムの利点を相互に活用できます。
GKE の TPU VM ノードプールによるオーケストレーションの簡素化
すでに Kubernetes 上でウェブ ワークロード、ステートレス ワークロード、ステートフル ワークロード、バッチ ワークロードなどのワークロードを実行している組織であれば、ワークロードをオーケストレートするプラットフォームが 1 つであることの運用上の利点を実感しているはずです。
GKE における Cloud TPU の一般提供により、お客様の AI ワークロードは、GKE が実現する一貫したオーケストレーション、利用率の向上、マルチテナント対応環境の恩恵を受けることができます。Cloud TPU ノードプールは自動プロビジョニングと自動スケーリングが可能であるため、オーケストレーションをさらに簡素化し、必要なときだけリソースを使用できます。ワークロードと予算に合わせて、スポット、オンデマンド、予約済みのいずれかの Cloud TPU インスタンスを選択できます。
写真および動画作成プラットフォームのパイオニアである Lightricks は、コンテンツ クリエイターやブランドに革新的なプラットフォームを提供する AI ファースト企業です。GKE を使用した TPU のオーケストレーションにより、Lightricks の業務は簡素化されました。
「GKE の TPU を使用することで、TPU 上で大規模モデルのトレーニングを行うために構築したカスタム スクリプトの多くが不要になりました。」 - Lightricks、コア生成 AI リサーチチーム リーダー Yoav HaCohen 博士
GKE における Cloud TPU マルチスライスを用いた大規模トレーニング
新しいマルチスライス テクノロジーにより Cloud TPU 上で水平スケーリングが可能になり、1 回のトレーニングで数万チップまでスケールアップできるため、大規模な AI モデルのトレーニングがさらに効率化されます。効率的なスケーリングとは、リソースの使用量を比例的に増加させることで、ほぼ同じ速度でワークロードを実行できることを意味します。
さらに、新しく導入された JobSet API を使用して、マルチスライス ジョブの抽象化が簡素化されています。JobSet の抽象化により、GKE には、数百の TPU Pod スライス上の数千の TPU チップに広がるトレーニング ジョブのための堅牢な障害処理メカニズムも備わっています。
TPU VM ノードプールの自動スケーリングによる、費用対効果に優れた推論
Kubernetes オートスケーラーを使用すれば、TPU ノードプールで、推論ワークロードやインタラクティブな開発ユースケースの TPU フットプリントを簡単にスケールアップできます。未使用の TPU ノードは GKE クラスタ オートスケーラーによって削除されるので、高費用なリソースを節約できます。
Kueue でリソースの共有と利用を最適化
GKE の Cloud TPU で Kueue などの API が提供するキュー機能を使うと、組織内の複数のチームが TPU リソースをタイムシェアできる、新しいリソース共有機能が実現します。ワークロードは従来どおり Kubernetes スケジューラを使用しますが、Kueue では、優先順位の低いジョブが優先順位の高いタスクの合間にアイドル状態のリソースを使用できるように、異なる優先順位を設定してワークロードを管理できるようになりました。これにより、複数のチーム間で TPU リソースをより最適に利用できるようになります。
自動修復による障害に強いトレーニングとサービング
GKE は Cloud TPU VM ノードプールの健全性を追跡し、TPU ノード（数は問いません）に障害が発生した場合に修復アクティビティを開始します。これにより、予測可能な方法でトレーニング ワークロードおよびサービング ワークロードの障害に対処できます。修復イベントがすべて終わると、ワークロードが自動的に再開されます。
GKE のモニタリングとロギングから実用的な分析情報を得る
GKE ネイティブなロギング機能とモニタリング機能を活用して、完全なトレーニング ログやサービングログを取得できます。GKE のシステム指標を有効にしている場合、デューティ サイクルや高帯域幅メモリ（HBM）利用率などの Cloud TPU インフラストラクチャ指標が Cloud Monitoring で利用でき、Google Cloud コンソールの Google Kubernetes Engine セクションに表示されるため、設定を行わなくても TPU リソースの利用状況を把握できます。
AI ワークロードの効率的かつ基盤的なオーケストレーター
GKE において Cloud TPU VM に対応したことにより、Google Cloud 上で AI / ML インフラストラクチャを構築するための基盤となるオーケストレーション レイヤーが提供されます。GKE は、TPU と GPU の両方を使用して ML モデルを運用するための高度なツールスイートをネイティブに提供します。これらのツールには、高度なロギング機能、モニタリング機能、自動修復機能、自動プロビジョニング機能が含まれます。
GKE で Cloud TPU と GPU の両方を柔軟にオーケストレーション
GKE の Cloud TPU により、お客様は、統合された環境で両方のアクセラレータ オプションの長所を活用することで、AI プラットフォームを柔軟に構築できるようになります。
Google のお客様にはすでにこのようなメリットを実感していただいています。
「GKE の TPU によって TPU と GPU がシームレスに統合されることで、ML パイプラインが強化されます。これにより特定のタスクにそれぞれの強みを活用し、レイテンシと推論費用を削減できるようになります。たとえば、TPU 上で大規模なテキスト エンコーダを使用してテキスト プロンプトをバッチで効率的に処理し、その後に GPU を使用して社内拡散モデルを実行し、テキスト埋め込みを利用して見事な画像を生成しています。」 - Lightricks、コア生成 AI リサーチチーム リーダー Yoav HaCohen 博士
「Anthropic は AI の安全性を追求し研究を行う会社です。信頼性が高く、解釈と制御が可能な AI システムの構築を専門としています。Google Cloud とは、当社のモデルを効率的にトレーニング、デプロイ、共有するために協力してきました。Google Kubernetes Engine（GKE）により、GPU と TPU のインフラストラクチャを大規模に実行し最適化できるようになります。また、Vertex AI により、Vertex AI Model Garden を通じて当社のモデルを顧客に配布できるようになります。A3 と TPU v5e を搭載した Google の次世代 AI インフラストラクチャは、AI の次の波を構築し続ける当社のワークロードにコスト パフォーマンスのメリットをもたらすでしょう。」 - Anthropic、共同設立者、Tom Brown 氏
次のステップ
GKE における Cloud TPU の詳細と利用開始方法については、こちらのドキュメントをご覧ください。
- Cloud TPU シニア プロダクト マネージャー Vaibhav Singh
- Google Kubernetes Engine グループ プロダクト マネージャー Nathan Beach
