Googleの主力TPU Trilliumが一般利用可能に！パフォーマンスの大幅向上で、AIモデルのトレーニング効率が飛躍的に向上

今年初めに、Googleはこれまでで最もパワフルな第6世代TPUであるTrilliumを発表しました。本日、TrilliumがGoogle Cloudのお客様向けに正式にリリースされました。

Googleは、これまでで最も強力なAIモデルである最新のGemini 2.0のトレーニングにTrillium TPUを使用しました。今、企業やスタートアップ企業も、同じ強力で効率的で持続可能なインフラストラクチャを活用できます。

AIスーパーコンピューターの中核：Trillium TPU

Trillium TPUは、Google Cloud AI Hypercomputerの重要な構成要素です。AI Hypercomputerは、性能最適化されたハードウェア、オープンソースソフトウェア、最先端のMLフレームワーク、柔軟な消費モデル統合システムを採用した画期的なスーパーコンピューターアーキテクチャです。Trillium TPUの正式リリースに伴い、GoogleはAI Hypercomputerのオープンソースソフトウェア層にも重要な強化を加え、XLAコンパイラやJAX、PyTorch、TensorFlowなどの一般的なフレームワークを最適化することで、AIトレーニング、調整、サービスにおいて比類のないコストパフォーマンスを実現しました。

さらに、大規模なホストDRAM（高帯域幅メモリまたはHBMの補完）を使用したホストオフロードなどの機能により、より高いレベルの効率性が実現します。AI Hypercomputerを使用すると、13Petabits/秒の双方向帯域幅を持つJupiterネットワークアーキテクチャから、最大10万個のTrilliumチップの展開からこれまで以上に最大限の価値を引き出すことができます。このアーキテクチャは、単一の分散型トレーニングジョブを数十万個のアクセラレータに拡張できます。

AI21Labsなどの顧客はすでにTrilliumを使用しており、顧客への有益なAIソリューションの提供を迅速化しています。

AI21LabsのCTOであるBarak Lenz氏は次のように述べています。「AI21では、MambaとJamba言語モデルのパフォーマンスと効率性を向上させるために継続的に努力しています。TPU v4の長期ユーザーとして、Google CloudのTrilliumの機能に感銘を受けています。規模、速度、コスト効率の向上は非常に顕著です。Trilliumは、次世代の複雑な言語モデルの開発を加速し、顧客により強力でアクセスしやすいAIソリューションを提供する上で重要な役割を果たすと確信しています。」

Trilliumのパフォーマンスが大幅に向上し、複数の指標で記録を更新

前世代と比較して、Trilliumは次の点で大幅な改善が図られています。

トレーニング性能が4倍以上に向上

推論スループットが3倍向上

エネルギー効率が67％向上

チップあたりのピーク計算性能が4.7倍向上

高帯域幅メモリ（HBM）容量が2倍に

チップ間インターコネクト（ICI）帯域幅が2倍に

単一のJupiterネットワークアーキテクチャに10万個のTrilliumチップが含まれる

ドルあたりのトレーニング性能が2.5倍向上、ドルあたりの推論性能が1.4倍向上

これらの強化により、Trilliumは、次のようなさまざまなAIワークロードで優れたパフォーマンスを発揮できます。

拡張AIトレーニングワークロード

LLMのトレーニング（大規模モデルや混合専門家（MoE）モデルを含む）

推論性能と集合スケジューリング

埋め込み密度の高いモデル

トレーニングと推論のコストパフォーマンスを提供

Trilliumはさまざまなワークロードでどのように優れたパフォーマンスを発揮するのか？

拡張AIトレーニングワークロード

Gemini 2.0のような大規模モデルのトレーニングには、大量のデータと計算が必要です。Trilliumのほぼ線形な拡張能力により、高速なチップ間インターコネクトで接続された256チップのpodと最先端のJupiterデータセンターネットワークにおいて、複数のTrilliumホスト間でワークロードを効率的かつ効果的に分散することで、これらのモデルのトレーニング速度を大幅に高速化できます。これは、TPUマルチチップ、大規模トレーニング向けのフルスタック技術によって実現され、Titaniumによってさらに最適化されています。Titaniumは、ホストアダプターからネットワークアーキテクチャまでを網羅する動的なデータセンターレベルのオフロードシステムです。

Trilliumは、3072個のチップで構成される12個のpodの展開で99％の拡張効率を実現し、6144個のチップを持つ24個のpodでは94％の拡張効率を示し、gpt3-175bを事前にトレーニングしました。これは、データセンターネットワークを跨いで実行する場合でも同様です。

LLMのトレーニング（大規模モデルや混合専門家（MoE）モデルを含む）

GeminiのようなLLMは、数十億のパラメーターを持つ非常に強力で複雑なものです。このような大規模なLLMのトレーニングには、膨大な計算能力と協調設計されたソフトウェアの最適化が必要です。Trilliumは、Llama-2-70bやgpt3-175bなどの大規模なLLMのトレーニングにおいて、前世代のCloud TPU v5eよりも4倍高速です。

大規模なLLMに加えて、混合専門家（MoE）アーキテクチャを使用してLLMをトレーニングすることはますます一般的な方法になっており、これは、それぞれがAIタスクの異なる側面を専門とする複数の「専門家」ニューラルネットワークを組み合わせたものです。単一の統合モデルのトレーニングと比較して、トレーニング中にこれらの専門家を管理および調整することは複雑さを増します。Trilliumは、MoEモデルのトレーニングにおいて、前世代のCloud TPU v5eよりも3.8倍高速です。

さらに、Cloud TPU v5eと比較して、Trillium TPUは3倍のホスト動的ランダムアクセスメモリ（DRAM）を提供します。これにより、一部の計算をホストにオフロードし、大規模なパフォーマンスと良好なスループットを最大限に高めるのに役立ちます。Trilliumのホストオフロード機能は、Llama-3.1-405Bモデルのトレーニングにおいて、モデルFLOP利用率（MFU）に関して50％以上の性能向上をもたらしました。

推論性能と集合スケジューリング

推論においては、複数ステップの推論の重要性が高まっており、アクセラレータが増加した計算要件を効果的に処理できる必要があります。Trilliumは推論ワークロードに大きな進歩をもたらし、AIモデルをより迅速かつ効率的に展開できます。実際、Trilliumは画像拡散と大規模なLLMに最適なTPU推論性能を提供します。私たちのテストでは、Cloud TPU v5eと比較して、Stable Diffusion XLの相対推論スループット（1秒あたりの画像数）が3倍以上、Llama2-70Bの相対推論スループット（1秒あたりのトークン数）がほぼ2倍向上することが示されました。

Trilliumは、オフラインとサーバーの推論ユースケース向けの最高性能のTPUです。下図に示すように、Cloud TPU v5eと比較して、Stable Diffusion XLのオフライン推論の相対スループット（1秒あたりの画像数）は3.1倍、サーバー推論の相対スループットは2.9倍向上しています。

優れたパフォーマンスに加えて、Trilliumは新しい集合スケジューリング機能も導入しています。この機能により、Googleのスケジューリングシステムは、複数のレプリカが存在する場合に、推論ワークロードの全体的な可用性と効率性を向上させるためのインテリジェントなジョブスケジューリングの意思決定を行うことができます。これは、Google Kubernetes Engine（GKE）を通じて、単一ホストまたは複数ホストの推論ワークロードを実行する複数のTPUスライスを管理する方法を提供します。これらのスライスを1つの集合にグループ化することで、ニーズに合わせてレプリカの数などを簡単に調整できます。

埋め込み密度の高いモデル

第3世代のSparseCoreを追加することで、Trilliumは埋め込み密度の高いモデルのパフォーマンスを2倍、DLRM DCNv2のパフォーマンスを5倍向上させました。

SparseCoreはデータフロープロセッサであり、埋め込み密度の高いワークロードにより適応性の高いアーキテクチャ基盤を提供します。Trilliumの第3世代SparseCoreは、分散収集、スパースセグメント合計、パーティショニングなど、動的かつデータ依存性の操作の高速化に優れています。

トレーニングと推論のコストパフォーマンスを提供

世界で最も大規模なAIワークロードのトレーニングに必要な絶対的な性能と規模に加えて、Trilliumはドルあたりの性能の最適化も目指しています。これまで、Trilliumは、Llama2-70bやLlama3.1-405bなどの大規模なLLMのトレーニングにおいて、ドルあたりの性能がCloud TPU v5eよりも2.1倍、Cloud TPU v5pよりも2.5倍向上しています。

Trilliumは大規模なモデルを経済的に効率的な方法で並列処理することに優れています。これは、研究者や開発者が、これまでよりもはるかに低いコストで強力で効率的な画像モデルを提供できるように設計されています。Trilliumでの1000枚の画像生成コストは、オフライン推論ではCloud TPU v5eよりも27％低く、SDXLのサーバー推論ではCloud TPU v5eよりも22％低くなっています。

AIイノベーションを新たなレベルへ

Trilliumは、Google Cloud AIインフラストラクチャの大きな飛躍を表しており、さまざまなAIワークロードに信じられないほどの性能、拡張性、効率性を提供します。世界クラスの協調設計ソフトウェアを使用して数十万個のチップに拡張できる能力により、Trilliumは、より迅速なブレークスルーを実現し、優れたAIソリューションを提供できます。さらに、Trilliumの優れたコストパフォーマンスにより、AI投資の価値を最大化したい組織にとって経済的に効率的な選択肢となっています。AIの状況が進化し続ける中、Trilliumは、企業がAIの可能性を最大限に引き出すために、Google Cloudが最先端のインフラストラクチャを提供することに尽力していることを証明しています。

公式紹介：https://cloud.google.com/blog/products/compute/trillium-tpu-is-ga