はじめに
いつものようにいろいろとGoogleくんに聴いていたら、Google TPUのドキュメントが更新されていたので、記録のためにブログに残します。
Google TPUの説明
— Vengineer＠ (@Vengineer) 2023年1月22日
v4の説明が追加されただけでなく、Tensorflow, Pytorch, JAXでどうやって動くかの説明も追加されていますねhttps://t.co/zl6CxsYAGv
Google TPU
Google TPUの v4/v3/v2 について、詳しく書かれています。
構成
- v4 : 2 x TensorCore ( 4 x MXU、1 x ベクトル ユニット、1 x スカラー ユニット)
- v3 : 2 x TensorCore ( 2 x MXU、1 x ベクトル ユニット、1 x スカラー ユニット)
- v3 : 1 x TensorCore ( 1 x MXU、1 x ベクトル ユニット、1 x スカラー ユニット)
チップ当たりの性能
- v4 : 275 TFLOPS（bf16 または int8）(v3 に対して、2.2357倍、動作周波数で 1.117倍、MXUが2倍 = 1.117x2 = 2.234でだいたいあっている)
- v3 : 123 TFLOPS（bf16）
- v2 : 45 TFLOPS
メモリ帯域
- v4 : 32 GiB、1200 GBps
- v3 : 32 GiB、900 GBps
- v2 : 16 GiB
上記のサイトにはなかった動作周波数 (Wikipedia)
- v4 : 1050MHz (v3 に対して、1.117倍)
- v3 : 940MHz
- v2 : 700MHz
Pod Size
- v4 : 4096 chips
- v3 : 1024 chips
- v2 : 256 chips
ここまでは今まで公開されていたのでいいんですが、これ以外に色々と載っています。
- v4 : MXU 重み付け読み込みの高速化と 8 ビットモードのサポートにより、バッチサイズを小さくし、推論のレイテンシを改善しました
- v4 : チップあたり 6 個の相互接続リンクにより、ネットワークダイヤグラムがより小さいネットワーク トポロジを実現
- v3 : v2からメモリが倍 (16GiB => 32GiB)
TPU は、次の構成で使用できます。
- シングル TPU ボード (4個のTPU chipを搭載したボード)
- TPU Pod : TPU ボードごとに TPU ホストがある
- TPU Pod スライス : 32、128、512、1024、2048
v2とv3では、TPUノードという構成がある
TPU ノード アーキテクチャは、gRPC を介して TPU ホストと通信するユーザー VM で構成されています。このアーキテクチャを使用する場合、TPU ホストに直接アクセスできないため、トレーニングと TPU のエラーをデバッグすることが困難になります。
下図を説明のために引用します。
これに対して、TPU VM アーキテクチャを使用すると、TPU デバイスに物理的に接続されている VM に SSH 接続できます。VM への root アクセス権があるため、任意のコードを実行できます。コンパイラとランタイム デバッグログ、エラー メッセージにアクセスできます。
下図を説明のために引用します。
ここまででなんで、TPUノードとTPU VMの説明をしているのかな？と思いましたが、この後の「サポートされているフレームワーク」のための前振りだったようです。
TensorFlow と PyTorch/JAX では、TPUノードでの動作が違うようです。
TPUノードのTPUホストでは、TensorFlowのruntimeが動いていて、gRPC 経由で TensorFlow の User VM/PyTorch の User VM/JAX の User VM と通信しています。
おわりに
このブログでも Google TPU については何度か書いていますが、TPUノードやTPU VMについては、ちょっと情報が増えてよかったです。
