ロードマップでわかる！当世プロセッサー事情 第729回

電気を利用せずにスイッチングできるGoogle TPU v4　AIプロセッサーの昨今

チップの消費電力は最大でも200W程度

　話を戻すと、もう1つ大きな違いはチップの集積数だ。TPU v3の場合は4本のリンクで合計280GB/秒だが、TPU v4ではこれが6本になり帯域も300GB/秒になっている。6本、ということは3次元メッシュが構成できるという話で、4本で2次元メッシュしか構成できないTPU v3よりも多くのチップを同等のレイテンシーで接続可能であり、実際接続できる数も最大1024チップから4096チップに増えている。

　例えばTPU v3とv4でどちらも最大構成を正方形なり立方体の構成で接続するとすると、TPU v3なら32×32構成になるので、正方形の端から端まで通信するのに31×2で62hop必要になる。

　一方TPU v3では16×16×16構成になるので、立方体の端から端までの通信は15×3で45hopで済む計算になる。TPU v3よりも4倍も多い数のプロセッサー間通信が、v3よりも少ないhop数で通信できる、というのがここでのメリットというわけだ。ちなみにメッシュといっても実際はHyperCube型に両端が接続されている構造だ(論文では3Dトーラス構造としている)。

　ちなみにTPU v4のパッケージが下の画像だ。チップの消費電力は最大でも200W程度(HBMの分も含む)とそれほど大きくないのだが、空冷にするとどうしても体積を喰う(さすがに1Uサーバーに収めるには冷却能力的に厳しいだろう)ということで液冷となっている。

搭載されているのは2020年ということもあり、HBM2。ただし帯域はTPU v3が900GB/秒だったのに対し1200GB/秒に向上している。容量はTPU v3/v4ともに32GB

右側から冷却液が供給され、左に排出される格好だが、一番長い緑のパイプの用途が不明。後段にもう1枚別のカードなどがつながる格好だろうか？(ホストCPU向けの可能性もあり)

　この液冷のモジュールは1本のラックに128枚(つまりTPU v4が512個)収められ、このラックを8つ並べたのがTPU v4 Podと称されている。

この8ラック合計の演算性能は1.1EFlopsほどになる。ちょうどNVIDIAのDGX GH200と同じ程度だ

　なぜここまで高密度にするかと言えば、TPU v4同士のネットワーク接続の距離を最小に収めたいためである。2Uあるいは3Uサイズにすると、明らかにラックの数が増え、それだけネットワークケーブルの長さを伸ばす必要がある。

　おそらく、液冷にする方が高密度化によりネットワークのレイテンシーも削減できるし、到達距離が短い分ネットワークの消費電力も減り、性能も向上する。液冷システムを追加する分消費電力は増えるが、これは性能と消費電力のバーターという話であろう(ついでに言えば、空冷よりも効率的に冷却できる分、トータルでの消費電力はむしろ減っているかもしれない)。

　ネットワーク周りではもう1つ、スイッチも独特である。論文によれば、2つ上の画像にある基板には4つのPCI Expressコネクターと16個のOSFPコネクターが搭載されているとされる。

　PCI Expressの方が基板の表面に見えているもので、これはホストとの接続用と思われる(OCuLinkコネクターだろう)。ちょうどそのPCIeの裏側にOSFPコネクターが4つづつ、合計16個並んでいる格好だ。このOSFPコネクターはTPU v4同士を接続するためのものと思われる。

　OSFPというのはイーサネットのトランシーバー用モジュールの規格の1つで、下の画像のようなのモジュールを装着する格好だが、2020年という時期から考えるとまだ400Gbpsの対応は難しく、おそらく100Gbpsのモジュールを装着しているものと思われる。

この中に、最大400Gbpsでの通信が可能なイーサネットのトランシーバーが内蔵されている。OSFPそのものは、このモジュールと、このモジュールを装着できるレセプタクルに関する電気・機械的形状などを定めた規格

　ちなみにイーサネットを使っているからといって別にTCP/IPを通しているわけではなく、独自のネットワークプロトコルを通しているものと思われる(CelebrasやTenstorrentと同じやり方だ)。