第 6 世代の Google Cloud TPU「Trillium」の発表
Amin Vahdat
VP/GM, ML, Systems, and Cloud AI
※この投稿は米国時間 2024 年 5 月 15 日に、Google Cloud blog に投稿されたものの抄訳です。
生成 AI は人々がテクノロジーと関わる方法を変えると同時に、ビジネスに影響を与えるとてつもない効率化の機会をもたらします。しかしこのような進歩に伴い、非常に高性能なモデルのトレーニングとファインチューニングを行い、世界中のユーザーにインタラクティブにサービングを提供するために、コンピューティング、メモリ、通信がこれまで以上に必要となります。Google はこれまで 10 年以上にわたり、AI に特化したカスタムのハードウェアである Tensor Processing Unit(TPU)を開発し、スケーラビリティと効率性の新境地を開拓してきました。
このハードウェアが、Gemini 1.5 Flash、Imagen 3、Gemma 2 などの新しいモデルを含め、今回の Google I/O で発表した多くのイノベーションの支えとなりました。これらのモデルはすべて TPU でトレーニングされており、TPU を使用して提供されます。モデルの新たな可能性を引き出し、お客様にも同様に取り組んでいただけるよう、このたび、これまでで最も高性能かつエネルギー効率に優れた第 6 世代の TPU「Trillium」を発表しました。
Trillium TPU は、TPU v5e と比較してチップあたりのピーク コンピューティング パフォーマンスが 4.7 倍と、驚異的な性能向上を達成しています。高帯域幅メモリ(HBM)の容量と帯域幅が 2 倍になり、チップ間相互接続(ICI)の帯域幅も TPU v5e の 2 倍になりました。さらに、Trillium には第 3 世代の SparseCore が搭載されています。これは、高度なランキングやレコメンデーションのワークロードでよくある極めて大規模なエンベディングを処理するための、特別なアクセラレータです。Trillium TPU により、新世代の基盤モデルを素早くトレーニングし、レイテンシと費用を低く抑えてモデルをサービングすることが可能となります。重要なのは第 6 世代の TPU が非常にサステナブルだということであり、Trillium TPU のエネルギー効率は TPU v5e より 67% 以上向上しています。
Trillium は、高帯域幅かつ低レイテンシの単一の Pod 内で 256 TPU までスケールアップできます。Trillium TPU はマルチスライス テクノロジーと Titanium Intelligence Processing Unit(IPU)により、この Pod レベルのスケーラビリティを超えて数百 Pod にスケーリングできます。その際、マルチペタビット/秒のデータセンター ネットワークで相互接続された、建物規模のスーパーコンピュータの数万に及ぶチップを接続します。
Trillium で AI のイノベーションが次のフェーズへ
10 年以上前、Google はそれまでになかった ML 用チップの必要性を認識しました。2013 年に世界初の専用 AI アクセラレータ TPU v1 の開発に着手し、2017 年には最初の Cloud TPU を発表しました。Gemini、Imagen、Gemma などの最先端の基盤モデルや、リアルタイムの音声検索、写真の物体認識、インタラクティブな言語翻訳といった広く利用されている Google サービスの多くは、TPU がなければ実現していなかったでしょう。実際、Google Research の Transformer(最新の生成 AI の基礎となるアルゴリズム)に関する基礎的な取り組みが可能となったのは、TPU のスケーラビリティと効率性のおかげです。
Trillium でチップあたりのコンピューティング パフォーマンスが 4.7 倍に向上
TPU はニューラル ネットワークのために一から設計されました。Google は常に、AI ワークロードのトレーニングとサービングにかかる時間の短縮に取り組んでいます。Trillium では、チップあたりのピーク コンピューティング パフォーマンスが TPU v5e の 4.7 倍になりました。このレベルのパフォーマンスを達成するために、Google は行列乗算ユニット(MXU)のサイズを拡大し、クロック速度を向上させました。さらに SparseCore は、TensorCore からのランダムな細かいアクセスを戦略的にオフロードすることで、エンベディングの多いワークロードを加速します。
ICI や高帯域幅メモリ(HBM)の容量と帯域幅が 2 倍に
HBM の容量と帯域幅が倍増したことで、Trillium はより大きな重み付けと Key-Value キャッシュのさらに大規模なモデルを処理できるようになりました。次世代の HBM ではメモリ帯域幅が高くなり、電力効率が向上するほか、柔軟なチャネル アーキテクチャでメモリ スループットが向上します。そのため、大規模モデルのトレーニング時間とサービング レイテンシが改善します。モデルの重み付けと Key-Value キャッシュが 2 倍になったことで、ML ワークロードを加速するために、より多くのコンピューティング容量でより速くアクセスできるようになりました。ICI 帯域幅が倍増したため、1 つの Pod に 256 個のチップがあるカスタムの光 ICI 相互接続と、スケーラビリティをクラスタ内の数百 Pod まで拡張する Google Jupiter ネットワーキングを戦略的に組み合わせることで、トレーニングと推論のジョブを数万チップまでスケーリングできます。
Trillium が次世代の AI モデルの原動力に
Trillium TPU は次世代の AI モデルやエージェントの原動力となります。Google は、このような高度な機能でお客様のお役に立ちたいと考えています。たとえば自律走行車メーカーの Nuro は、Cloud TPU を使用してモデルをトレーニングすることで、ロボット工学を通じた日常生活の改善に取り組んでいます。AI で創薬の未来を切り拓いている Deep Genomics は、Trillium による新しい基盤モデルが患者の生活をどのように変えるかに期待を寄せています。AI の Google Cloud パートナー オブ ザ イヤーの Deloitte は、生成 AI でビジネスを変革するために Trillium を提供する予定です。また、Trillium TPU でコンテキストの長いマルチモーダル モデルのトレーニングとサービングがサポートされるため、Google DeepMind では、将来の世代の Gemini モデルのトレーニングとサービングをこれまでより速く効率的に低レイテンシで行えるようになります。
「Gemini 1.5 Pro は何万もの TPU アクセラレータを使用してトレーニングされた、Google の最大かつ最高性能の AI モデルです。私たちのチームは第 6 世代の TPU が発表されたことに沸き立っており、Gemini モデルのスケールでトレーニングおよび推論のパフォーマンスと効率が向上することを楽しみにしています。」- Google DeepMind および Google Research、チーフ サイエンティスト、Jeff Dean
「Google Cloud TPU アクセラレータは Nuro に大きな変革をもたらしました。Cloud TPU v5e を活用することで、他と比べてトレーニング ワークロードの費用対効果が 180% 向上したのです。当社は TPU を活用して AI ファーストの自律性で路上の安全性を高めるために、行動と認識のモデルを進化させました。今後は次世代の Trillium TPU を活用してさらに発展させていこうと考えています。AI インフラストラクチャ テクノロジーの最前線に立っている Google Cloud と、次世代 AI モデルで連携できることを誇りに思っています。」- Nuro、CTO、Andrew Clare 氏
「Deep Genomics では、AI と生物学の科学者からなる学際的なチームがゲノム研究開発の新たな道を切り拓いています。私たちの使命は、高度な AI モデルを作成して適用し、RNA 治療薬の分野における可能性の限界を押し広げることです。第 6 世代の Trillium TPU が発表されたことを嬉しく思うと同時に、Google Cloud と協力して AI 基盤モデルのトレーニングと推論を 100 倍以上に拡大する機会が得られたことに喜んでいます。これによって、患者さんの生活を変える当社の力が強化されるでしょう。」- Deep Genomics、設立者 / CIO、Brendan Frey 氏
「Deloitte は、クライアントがテクノロジーを活用できるように支援することを主な使命としています。Google Cloud の次世代 TPU の Trillium を導入すれば、パフォーマンスと費用を最適化して真のビジネス価値を実現できるため、大きな飛躍となります。Google とともに、当社はこれからも業界や分野を問わずクライアントと協力し、生成 AI を活用して可能性を引き出していきます。」- Deloitte UK、Alphabet Google アライアンス最高商務責任者、Matt Lacey 氏
Trillium と AI ハイパーコンピュータ
Trillium TPU は、Google Cloud の AI ハイパーコンピュータ(最先端の AI ワークロード専用に設計された画期的なスーパーコンピューティング アーキテクチャ)の一部です。このハイパーコンピュータには、パフォーマンスが最適化されたインフラストラクチャ(Trillium TPU など)、オープンソース ソフトウェア フレームワーク、柔軟な消費モデルが統合されています。Google は、JAX、PyTorch/XLA、Keras 3 などのオープンソース ライブラリへの取り組みでデベロッパーを支援します。JAX と XLA をサポートするということは、前の世代の TPU 向けに書かれた宣言型モデルの説明が、Trillium TPU の新しいハードウェアとネットワークの機能に直接対応するということです。また Google は Optimum-TPU に関して Hugging Face と提携し、モデルのトレーニングとサービングを効率化しています。
「Google Cloud とのパートナーシップにより、Hugging Face のユーザーは、TPU などの Google Cloud AI インフラストラクチャでオープンモデルを簡単にファインチューニングして実行できるようになりました。今後の第 6 世代の Trillium TPU でオープンソース AI をさらに加速できることを嬉しく思います。前の世代と比較してチップあたりのパフォーマンスが 4.7 倍向上しているため、オープンモデルが最適なパフォーマンスを継続的に提供できると期待しています。当社は新しい Optimum-TPU ライブラリを通じて、すべての AI ビルダーが Trillium のパフォーマンスを簡単に利用できるようにします。」- Hugging Face、プロダクト責任者、Jeff Boudier 氏
SADA は 2017 年以来、パートナー オブ ザ イヤーに毎年選ばれており、最大限の効果をもたらす Google Cloud サービスを提供しています。
「誇りある Google Cloud プレミア パートナーとして、SADA は世界が認める AI のパイオニアである Google と 20 年の歴史を築いてきました。当社は数千もの多様なお客様のために、AI を急速に統合しています。当社の豊富な経験と AI ハイパーコンピュータ アーキテクチャにより、Trillium で生成 AI モデルの新境地を切り拓くお客様のお手伝いができることを楽しみにしています。」- SADA、CTO、Miles Ward 氏
AI ハイパーコンピュータには、AI / ML ワークロードに必要となる柔軟な消費モデルも用意されています。また、Dynamic Workload Scheduler(DWS)を使用すると、AI / ML リソースへのアクセスが容易になるうえ、費用も最適化できます。Flex Start モードは、エントリ ポイント(Vertex AI Training、Google Kubernetes Engine(GKE)、Google Cloud Engine)に関係なく、同時に必要となるすべてのアクセラレータをスケジューリングすることで、トレーニング、ファインチューニング、バッチジョブなどの急増するワークロードのエクスペリエンスを向上させることができます。
Lightricks は、AI ハイパーコンピュータによる効率の向上に加えて、パフォーマンスの向上を通じて価値を取り戻すことに期待を寄せています。
「Cloud TPU v4 以来、TPU を Text-to-Image モデルと Text-to-Video モデルに使用してきました。TPU v5p と AI ハイパーコンピュータの効率化により、トレーニング速度が 2.5 倍も向上しました。第 6 世代の Trillium TPU は、前の世代に比べてチップあたりのコンピューティング パフォーマンスが 4.7 倍、HBM の容量と帯域幅が 2 倍に向上しています。Text-to-Video モデルを拡大しようしていた私たちにとって、ちょうどよいタイミングでした。Dynamic Workload Scheduler の Flex Start モードを使用して、バッチ推論ジョブや将来の TPU 予約を管理することも計画しています。」- Lightricks、コア生成 AI リサーチ チーム リーダー、Yoav HaCohen 博士
Google Cloud Trillium TPU の詳細
Google Cloud TPU は、大規模 AI モデルを強化するためにカスタム設計されて最適化された AI アクセラレーションの最先端です。Google Cloud でのみ利用できる TPU として、AI ソリューションのトレーニングとサービングにおいて比類ないパフォーマンスと費用対効果を実現します。大規模言語モデルの複雑さであれ、画像生成の創造的な可能性であれ、デベロッパーや研究者は TPU で AI 分野における可能性の限界を押し広げることができます。
第 6 世代の Trillium TPU は 10 年以上にわたる研究とイノベーションの集大成であり、今年中に提供が開始される予定です。Trillium TPU と AI ハイパーコンピュータについて詳しい説明をご希望の場合は、こちらのフォームにご記入ください。セールスチームからご連絡いたします。
- ML、システム、クラウド AI 担当バイス プレジデント兼ゼネラル マネージャー、Amin Vahdat