@Vengineerの戯言 : Twitter SystemVerilogの世界へようこそ、すべては、SystemC v0.9公開から始まった

はじめに

Google TPU、v1、v2、v3、v4 と出て、v2からは学習もできるようになりました。学習するためには、複数チップを接続するための interconnect が必要になります。

今回のブログでは、v2、v3、v4 の interconnect について、再度考えてみました。

Google TPU v2/v3 の interconnect

Google TPU v2/v3 については、マイナビの Ando さんの「TPU v1をベースに学習を可能としたTPU v2/v3」に詳しく書いてあります。

GoogleのAIアクセラレータ「TPU v2/v3」を読み解く - Hot Chips 32(1) TPU v1をベースに学習を可能としたTPU v2/v3 | TECH+（テックプラス）

interconnect については、「TPU v2とTPU v3の違い」にあります。

news.mynavi.jp

Hotchips2020 のプレゼンテーション資料 : Google’s Training Chips Revealed: TPUv2 and TPUv3

説明のために図を引用します。チップ内の interconnect router に 2つのコアと4つのLinkが接続しています。このLinkがinterconnectになります。各Linkは500Gbpsです。

v3では、このLinkが30％ upされたということです。記録のために図を引用します。

基板の写真は、Google Cloud TPUに載っています。説明のために、引用します。

v2 の interconnect は基板の右側に4個、左側に4個あります。右側の4個では、2個の間に小さなコネクタがあります。これはどうやらガイドというもので、基板を挿入するためのガイドになるものです。基板は左側から右側に挿入します。左側のコネクタも2個づつになっていて、真ん中に大きなコネクタが2つあります。これはHostと接続をするためのPCIe Gen3用のコネクタだと思います。

v3 の interconnect は、基板の右側に4個 (v2 と同じ感じに 2個のガイドがあります)、左側には4個のコネクタが1つのパッケージになったものが2つ。この2つのパッケージの間にあるのは v2 と同じでPCIe Gen3用のコネクタだと思います。

v3 のPCIe Gen3用のコネクタがケーブルの隠れているので、「グーグルのAI専用プロセッサ「TPU 3.0」、スパコン京と比較した実力」にあった写真(説明のために引用します)

xtech.nikkei.com

PCIe Gen3用のコネクタは2つであることが分かります。

v2の写真もあったので記録のために引用します。

v3 の左側のコネクタは、OSFP I/O コネクタっぽいです。

OSFP : - 200 Gbps (8 x 28G NRZ) - 400 Gbps (8 x 56G PAM-4)

v3 では、1つのLink が 650Gbpsのようなので、OSFPの 400 Gbps が2個あれば、OKそう。

interconnect がどのように繋がっているか？

The nextplatform の「TEARING APART GOOGLE’S TPU 3.0 AI COPROCESSOR」に Google TPU v2/v3 Pod の写真が載っています。

Google TPU v2 Podの拡大写真の一部を説明のために引用します。

上らか下に接続している赤の2本のケーブルは v2 の真ん中のコネクタに接続しています。ということは、ホストからのPCIe Gen3用のケーブルだとおもいます。同じ段には、4枚の v2 が入っていて、interconnect 用のケーブルは、(黄色、灰色、灰色、赤色)、(緑色、灰色、灰色、青色)、(青色、黄色、黄色、緑色)、(赤色、橙色、橙色、黄色)、になっています。GoogleのTPU v2の説明では、32コアの場合では左の2個を2段使う(2 Board x 2 段 x 4chips x 2コア)ようになっています。しかしながら上記のようなケーブルの色付けだと、両サイドの2Board と中の2Boardが1つの単位になっているっぽいです。そうすると、(黄色、灰色、灰色、赤色) + (赤色、橙色、橙色、黄色) と (緑色、灰色、灰色、青色) + (青色、黄色、黄色、緑色) となり配線が分かりやすいです。上下はそのまま上下で接続できるます。1段で4Board使う時は、赤色と緑色をSwitchで切り替えればそのままつながると思います。

GoogleのTPU v2発表時のブログ：Build and train machine learning models on our new Google Cloud TPUs に、interconnect と PCIe 用のケーブル接続している写真がありましたので、記録のために引用します。(黄色、灰色、ピンク、ピンク、灰色、赤色)のケーブルがコネクタに刺さっているのが分かります。

Google TPU v3 Podの拡大写真の一部を説明のために引用します。

v2 ではホストは別Rackでしたが、v3では同じRackの真ん中に2つのホストがあり、各ホストが1つのv3に接続しています。ホストからはピンクのケーブルが2本、v3の真ん中の2個のPCIe Gen3用コネクタに接続しています。2本の灰色のケーブルは水冷用のケーブルです。v3では、v2と違って1段に2枚のBoardなので2枚のBoardはそのままつなげて、両サイドのRackとそれぞれ水色のケーブルで繋がっている感じだと思います。

v2/v3ともに、反対側のコネクタはなんらかのバックプレーンに繋がっていると思います。

Google TPU v4

Google TPU v4 に関しては、まだまだ情報はありませんが、Google の下記のツイートの中のビデオを何度も何度も見ました。

Our new TPU v4 pods are 10 times as fast as v3. And a TPU v4 pod can deliver more than an exaflop of computing power — more than if 10 million people were on their laptops right now. It’s a historic milestone. #GoogleIO pic.twitter.com/YeEncTm7c1
— Google (@Google) 2021年5月18日

説明のためにビデオをスクリーンショットを引用します。

まずは、Google TPU v4 Boardです。v2 と v3 とは違って、右側にコネクタはありません。コネクタは左側に移りました。見えている側にはPCIe Gen4用コネクタが4つ。これは Google TPU v4 のホストは、AMD EPYC (Rome) なので、PCIe Gen4 x16 が4組になっていて、各PCIe Gen4 x16 が TPU v4 チップに接続しています。基板の裏側に v3 の左側にあったようなコネクタが4個付いています。PCIe Gen4 x16とinterconnect がペアになる感じですね。1チップに OSFP I/O コネクタが4個。OSFPは 400 Gbps (8 x 56G PAM-4) なので、2個で 800 Gbps までサポートできそう。(112 Gbbps PAM-4 もあるようなので、そうなると、1.6Tbps となる)

次の写真は、TPU v4 Pod の写真です。

右側の黒いケーブルは冷却用だと思います。白い丸が付いている金属板の隣に紫のケーブルが接続されています。これは interconnect 用ケーブルのようです。ホストとのPCIe Gen4用のケーブルが見当たりません。Podの上から下にある黄色のケーブルは両サイドの8本(片側4本)来ていますが、これはなんでしょうか？ホストのネットワーク(100G/200G)のケーブルでしょうか？各ホストに4本接続しているように見えます。

v2/v3 と違って、v4には左側にコネクタが無いのでバックボーンとは接続していないようです。

v2 では、1段4枚のBoardで、8段で合計32枚のBoard = 128チップが 1 Rack でした。 v3 では、1段2枚のBoardで、16段で合計32枚のBoard = 128チップが 1 Rack でした。 v4 では、上記のビデオ(写真)では、1段2枚のBoardで、8段で合計16枚のBard = 64チップが 1 Rack のように見えます。

v4のホストは、AMD EPYC (Rome)と考えると、1個のRome に対して4個のPCIe Gen4 x16 が接続できます。2020年MLPerf v0.7のtpu-v4-16-TF.jsonを見ると、"host_processors_per_node": "2"で"accelerators_per_node": "8"とあるので、Google TPU v4 Boardが2枚でAMD EPYC (Rome)が2個の構成かもしれません。 v4のホストは、v3とは違って TPU Board と並んでいるのではなく、TPU Board の奥にある感じがします。その場合、PCIe Gen4 x16用のコネクタへのケーブルがどうなるのが分かりません。紫色のケーブルは見えますが、ホストとの接続のケーブルが見当たらないのですよ。。。

追記)、2023.4.12

We’re Training AI Twice as Fast This Year as Last New MLPerf rankings show training times plunging

に載っている TPU v4 Pod を拡大したら、PCIe と interconnect のコネクタが分かりました。一部を下記に説明のために引用します。

上段にHost、下段にTPU v4 Board。Host と TPU v4 Board は、紫のPCIe Gen4用ケーブルで接続しています。Interconnect用のケーブル(黄色)は紫のケーブルの下の各4個のコネクタから繋がっているのがわかります。TPU v4 は6個のInterconnectを持っていて、Board 内の4個は各2個を使って、Ring構造で繋がっています(両手で両側のTPU v4と繋がっている感じ)