Google が TPUv4 とそのクレイジーな光学的に再構成可能な AI ネットワークを詳細に解説

Aug 17, 2023

Hot Chips 2023 で、Google は光学的に再構成可能なクレイジーな AI ネットワークを披露しました。 同社は、AI トレーニング クラスターのパフォーマンスの向上、消費電力の削減、および柔軟性の向上を実現するために、光回線スイッチングを行っています。 さらに驚くべきことは、これを何年にもわたって運用してきたことです。

生放送なので誤字脱字はご容赦ください。

この大きな目標は、Google TPU チップを結合することです。

こちらは 7nm Google TPUv4 です。 今週、TPUv5 についてさらに多くのことが聞かれるようになると予想されます。 Google は通常、1 世代前のハードウェアに関する論文やプレゼンテーションを行うことができます。 TPU v4i は推論バージョンでしたが、これは TPUv4 に焦点を当てた話です。

Googleは、5ミリ秒のサービス時間SLAを満たすことができるように、通常の電力と比較して電力を過剰にプロビジョニングしていると述べています。 したがって、チップ上の TDP ははるかに高くなりますが、これは SLA バーストを満たすバーストを可能にするためです。

これは TPUv4 アーキテクチャ図です。 Google は、これらの TPU チップを単一のアクセラレータとしてだけでなく、スケールアウトして大規模インフラストラクチャの一部として実行できるように構築しています。

これは、これまでに見た中で最も明確な表の 1 つで、Google TPUv4 と TPUv3 の統計を示しています。

Google はピーク FLOPS を 2 倍以上に向上させましたが、TPUv3 と TPUv4 の間の電力は削減しました。

Google には TPUv4 に SparseCore アクセラレータが組み込まれています。

Google の TPUv4 SparseCore のパフォーマンスは次のとおりです。

ボード自体には 4 つの TPUv4 チップが搭載されており、水冷式です。 Googleは、液体冷却に変更するにはデータセンターと運用をやり直す必要があるが、電力削減にはそれだけの価値があると述べた。 右側のバルブは、液体冷却チューブを通る流れを制御します。 Googleによれば、これはファン速度コントローラーのようなものだが、液体用だという。

Google はまた、これは 2020 年の設計であるため、ホストに PCIe Gen3 x16 を使用しているとも述べています。

Google では、多くのデータセンターと同様にトップオブラックから電力を供給していますが、多数の相互接続を備えています。 ラック内では Google は電気 DAC を使用できますが、ラックの外では光ケーブルを使用する必要があります。

各システムには、4096 個の相互接続されたチップを備えた 64 個のラックがあります。 ある意味、256 ノードの NVIDIA の AI クラスターには半分の GPU が搭載されています。

また、ラックの端には CDU ラックがあります。 液体冷却について詳しく知りたい場合は、「Gigabyte および CoolIT での液体冷却サーバーの仕組み」を参照してください。 近々、さらに多くの液体冷却コンテンツを提供する予定です。 グーグルによれば、液体の流量は、はしご消防車のホース内の水よりも速いという。

各ラックは、TPU 間に光回線スイッチング (OCS) を備えた 4x4x4 キューブ (64 ノード) です。 ラック内の接続は DAC です。 立方体の面はすべて光学的です。

ここではOCSについて見ていきます。 電気スイッチを使用する代わりに、OCS を使用すると、チップ間が直接接続されます。 Google は 2D MEMS アレイ、レンズ、カメラなどを内部に備えています。 ネットワークのオーバーヘッドをすべて回避することで、データをより効率的に共有できるようになります。 余談ですが、これはある意味 DLP テレビに似ているように感じます。

Googleによると、スーパーポッドには1万6000以上の接続と十分な距離のファイバーがあり、ロードアイランド州を取り囲むことができるという。

ポイントツーポイント通信が非常に多いため、多くのファイバーストランドが必要になります。

さらに、各プールをより大きなプールに接続できます。

OCS は再構成可能なため、ノードの使用率が高くなります。

その後、Google は光ルーティングを調整することでトポロジを変更できます。

ここで Google はさまざまなトポロジの利点を示しています。

Google は、モデルのニーズの変化がシステムの変更を引き起こす可能性があると述べているため、これは重要です。

これは、最大 3072 チップでの線形高速化を伴う Google の対数スケールでのスケーリングです。

Google はまた、データ アクセスをローカルに保つために、オンチップ メモリを 128 MB に増加しました。

以下は、Google と NVIDIA A100 をワットあたりのパフォーマンスで比較したものです。

これは、2 つのポッドで 6144 個の TPU をトレーニングする PaLM モデルです。

それは膨大な数です！