はじめに

Googleは、TPU v5e を発表したのはこのブログでも以下のように紹介しました

vengineer.hatenablog.com

今日は、v5e ではなく、v5 について、探っていきます。

Google AI生成に聞いてみた

以前(10月9日)、Google検索にて、AI生成を使って、Google TPU v5e について調べたときの結果は、下記のようになっていました。

Googleさんの生成AIが教えてくれました。

TPU v5 というのがあって、Tensor Coreは1個ではないと pic.twitter.com/Z8aD2JFzYn
— Vengineer＠ (@Vengineer) 2023年10月9日

質問が「Google TPU v5e」だったので、Google 生成AIでは、v5 について説明してくれています。

TPU v5e は、TPU v5 とはことなり、Tensor Core が 1つしかありません。

とあり、v5e とは別に v5 があり、v5e が Tensor Core が 1 つに対して、v5 は Tensor Core が 2つ以上あるということが分かります。

Google が答えてくれたのでこれは確かなことです。

が、これを検索したのは 10月9日です。今、Google検索の会話で検索すると、Google TPU v5e の検索結果しかでてきません。

あー、Google さん、やらかしていたんですね。。。

PCI DeviceID

Google JAX のgithubを調べていたら、TPUのPCI DeviceIDを下記のように見つけました。

Google TPU の PCI DEVICE ID
- v2/v3 (何故か同じ) => 0x0027
- v4 => 0x005e
- v5e => 0x0063

- 0x0056 : v4より前なので、v4i かな？
- 0x0062 : v5eより前なので、v5 かな？https://t.co/qyljBjK5M4 pic.twitter.com/Q68kLIDjvr
— Vengineer＠ (@Vengineer) 2023年10月8日

TPU v2 : 0x0027
TPU v3 : 0x0027
TPU v4 : 0x005e
TPU v5e : 0x0064
Testing only : 0x0056
Testing only : 0x0062

v2 と v3は、同じなんですね。。。

Testing only のところに、0x0056 と 0x0062 がありますね。。。

0x0056 は、Tensorflow/tpu の github によると、以下のように v4i のようです。

0x0062は、v5e より小さいのでこれが v5 なんでしょうか？

おわりに

下記のように、v5e には、TensorCore が1個入っていて、4つのMXUが入っています。

TPU v4 では、Tensor Coreが2個で、各Tensor Coreには MXUが4つ入っています。

となると、v5 は、Tensor Core が 2個か4個なんでしょうね。

そして、V5e の冷却とのための装置が下記の写真(説明のために、ここ/Cloud Watchから引用します)を見ると、かなり多きはヒートシンクが載っています。Tensor Coreが1個のv5e でこれだけのヒートシンクが載るということは 2個以上の Tensor Coreが載っていると思われる v5 は空冷じゃ無理で、v4のように水冷にする必要があるんじゃないのでは？と思っています。