「万能チップは1つで」なんて忘れよう:TPU 8でGoogleはAI覇権レースに挑む
x86は追放:GoogleはTPUをArmベースのAxionコアと組み合わせる
Googleは水曜(現地時間)ラスベガスで開催した年次イベント「Cloud Next」で、社内製の新しいAIアクセラレータを2種類発表した。1つは学習を高速化することを目的に、もう1つはモデル提供(推論)のコストを引き下げることを狙っている。
「チョコレート・ファクトリー」では、第8世代のテンソル処理ユニットが、学習で最大2.8倍高速であるという。また、LLM推論においては、昨年の「Ironwood TPU」と比べて、1ドルあたりの性能が80%高いという。
それを実現するため、Googleはアクセラレータ開発を二つの並行ラインで進めている。TPU 8tは学習向け、TPU 8iは推論向けだ。これらのチップは同様の基盤の上に作られているが、それぞれのワークロードにおけるボトルネックを個別に解消することを目的としている。
Googleがこの道を最初に歩んだわけではない。AIチップ開発の初期段階で、Amazon Web Servicesは推論と学習の両方に最適化されたアクセラレータが必要だと認識していた。
Nvidiaもこの種の「特化」を試みてはいるが、同じ程度ではない。GPUベンダーのBlackwell Ultra世代は、AI推論向けに特別に最適化されており、高精度の“力仕事”を、Blackwellの兄弟製品に比べてメモリとFP4計算で50%増の伸びと引き換えにしている。
デュアル・トラッキングに加えて、GoogleはTPUホスト向けに、x86プロセッサをやめて自社開発のArmベースのAxion CPUへ切り替えています。今年の初めにはAmazonも、GravitonとTrainium 3で同様のことを行いました。
専門化へのGoogleの取り組みは、TPUそのものよりもはるかに深いところにあります。検索と広告の巨大企業は、推論と学習の両方におけるスケーリング損失を最小化するために、異なるネットワークトポロジーを備えた新しいクラスターを開発してきました。
現代のAIワークロードは、いまや単一のアクセラレータで動くことはほとんどありません。そのため、各チップがどれだけ速いかよりも、それらのワークロードを複数チップに効率よくスケールできるかのほうが重要になることが多いのです。
TPU 8tは学習向けに調整
過去のアクセラレータと同様、GoogleのTPU 8tは大規模な学習を想定して設計されています。同社はOpenAIやAnthropicほどの派手さはないかもしれませんが、Googleは依然として精力的にモデルを作り続けています。
内部では、同社はIronwoodチップで導入されたベクトル、行列乗算、SparseCore埋め込みアクセラレータの配合を微調整し、有効な浮動小数点スループットを最大化しています。
それぞれのアクセラレータは、高帯域幅メモリ(HBM)216 GBを搭載し、6.5 TB/sの帯域幅に対応可能。さらにオンチップSRAM 128 MB、4-bit浮動小数点の演算で最大12.6 petaFLOPS、チップ間(chip-to-chip)の帯域幅で最大19.2 Tbpsを備えます。
Googleがクラウド基盤に展開する計画も同時に発表していた、NvidiaのRubin GPUと比べると、新しいTPUは少しおとなしく見えるかもしれません。
RubinはFP4学習で最大35 petaFLOPS、帯域幅22 TB/sに相当するHBM4 288 GBをうたっています。つまり個別にはNvidiaのGPUのほうが速いものの、実際にはそれほど重要ではありません。新しいフロンティア型モデルを学習するとき、GPUは1枚ではなく何千枚も使うからです。そしてスケールの面では、Googleが優位にあります。
Nvidiaの最新GPUは、EthernetまたはInfiniBandへスケールアウトする必要が出るまで、単一のNVLinkドメインで最大576のアクセラレータに対応しています。
一方TPU 8tは、光学回路スイッチ(オプトメカニカルなスイッチング技術)を使います。これはパケットスイッチよりも、電話交換機のように動作するもので、単一の統合ポッドで最大9,600のアクセラレータを接続します。
その後、複数のポッドを同社の新しいVirgo Networkでつなぎ、さらに大きな計算ドメインを支えます。
単にチップ間のメッシュを大きくしただけではないようです。Googleは、ポート密度が非常に高いパケットスイッチを使っているように見えます。
理解している範囲では、これらはフラット寄りの2階層・オールツーオール構成として配置されており、1データセンターあたり最大134,000台のTPU、複数サイト接続時には最大100万台のTPUまで接続できる能力があるとのことです。Googleはそう主張しています。
新しいネットワーク基盤に加えて、Googleは、合計10 TB/sのデータをアクセラレータのメモリへ直接投入できるManaged Lustreストレージシステムも開発しています。
Googleによれば、これらの技術に加えて、信頼性・可用性・保守性(RAS)機能が改善されたことで、学習最適化されたTPUが「goodput(実効スループット)97%」で稼働できるようになります。
この場合のgoodputとは、TPUが実際に学習に費やす時間のことを指します。
同社はブログ記事の中で、「あらゆるハードウェア障害、ネットワークの停止、チェックポイントの再起動は、クラスターが学習していない時間です。そしてフロンティア規模の学習では、1パーセントポイントごとが、実際のアクティブ学習時間の何日分にもつながり得ます」と説明しています。
TPU 8i:推論の転機
推論は自己回帰型ワークロードです。つまり生成される各トークンごとに、モデルのアクティブな重み全体をメモリ経由でストリーミングする必要があります。計算も重要ではありますが、主なボトルネックになりやすいのはメモリ帯域幅です。
推論に特化したTPU 8iでは、GoogleはFLOPSを一部犠牲にして、より大きなSRAMキャッシュと、より高速で高容量のメモリプールを得ています。
紙の上では、このチップはNvidiaのBlackwellアクセラレータとほぼ同等の位置づけです。TPU 8iは、オンチップSRAM 384 MBに供給されるFP4演算で10.1 petaFLOPSを備え、帯域幅8.6 TB/sに対応するHBM 288 GBも搭載します。いっぽう、インターコネクトの帯域幅は、学習向けの兄弟製品と変わっていません。
Googleによれば、このようなSRAMの活用により、TPU 8iは重要なキー・バリューキャッシュ(モデルの短期記憶)をより多くオンチップに保持し、コアがデータ待ちになる時間を減らせるとのことです。
またTPU 8iは、GoogleのSparseCoresをやめて、集合(collective)アクセラレーションエンジン(CAE)を採用しています。名前のとおり、このアクセラレータブロックは、all-reduceやall-gatherのような集合通信をオフロードすることで推論を高速化し、結果として生じる同期の待ち(stall)を短縮して、そうであればチップのテンソルコアがアイドル状態になっていた時間を減らします。
実際、TPU 8iでGoogleが注力しているのは、推論スタック全体で可能な限りレイテンシを削ることのように見えます。
混合専門家(MoE)アーキテクチャが一般的になってくると、集合通信はかなり問題になってきました。これらのモデルは複数のサブモデル、適切に「専門家(experts)」と呼ばれるものから構成され、そのうち各トークン生成時に有効化されるのは一部です。
メリットは、同じ水準の性能を維持するのに必要なメモリ帯域幅を必ずしも追加せずに、モデルをより大きくできることです。デメリットは、トークンごとに次のトークンで使われる専門家の集合が変わり得て、その専門家が別々のアクセラレータ上に存在する可能性があるため、チップ間通信の予測可能性が低くなる点です。
Googleは、CAEが集合通信のレイテンシを5分の1に抑えると言っています。これにより、同じハードウェアにより多くのユーザーを詰め込めるようになり、経済性(コスト効率)が向上する、というわけです。
- No-Nvidia interconnect club delivers 2.0 spec before v1.0 silicon ships
- How Nvidia learned to embrace the light in its quest for scale
- Alibaba has made 470,000 AI chips, admits they're inferior and may always be
- Decoding Nvidia's Groq-powered LPX and the rest of its new rack systems
推論向けに最適化したチップに加えて、GoogleはBoardflyと呼ばれるトポロジーも開発しています。ネットワークの配置は、HPCクラスターで一般的に用いられるDragonflyトポロジーに、どこか似たところがあります。そして光学回路スイッチを使って、1,152チップ(いずれかの時点で稼働しているのは1,024)を接続できます。
Boardflyの主要な利点は、3Dトーラスで最大16ホップだったチップ間レイテンシを、わずか7ホップにまで削減できることです。これにより、MoEや推論モデルを動かしているときのチップ間レイテンシもさらに低減されます。
推論クラスターの作り方を新たなアーキテクチャに合わせて変えているのはGoogleだけではありません。12月にAWSは、同様の理由から推論における3Dトーラスを捨てましたが、その代わりに、より一般的なパケットスイッチ方式のファブリックを選んでいます。
TPU 8アクセラレータ2基はいずれも、今年後半にGoogle Cloud Platformでインスタンスとして、またはクラウド事業者のフルスタックAIハイパーコンピュータープラットフォームの一部として、一般提供される予定です。このプラットフォームには、LLMを大規模に導入または学習するために必要なネットワーク、ストレージ、コンピュート、ソフトウェアがすべてまとめて含まれています。 ®
関連記事
より絞り込んだトピック
より広いトピック
詳細
関連するより狭いトピック
- AIOps
- Android
- アプリストア
- Chrome
- Chromium
- DeepSeek
- 災害復旧
- Gemini
- Google AI
- Google I/O
- Google Nest
- Google Project Zero
- GPT-3
- GPT-4
- G Suite
- Kubernetes
- 大規模言語モデル
- 機械学習
- MCubed
- ニューラルネットワーク
- NLP
- Open Compute Project
- Pixel
- プライバシーサンドボックス
- PUE
- 生成AIの検索拡張(Retrieval Augmented Generation)
- ソフトウェア定義データセンター
- スター・ウォーズ
- Tavis Ormandy
- TOPS
