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

性能/消費電力比が優秀なGoogle TPU　AIプロセッサーの昨今

初代Google TPUの2倍の性能を誇る第2世代

　以下、この2つのソースをもとに解説する。まずGoogle TPU v2の基本構成が下の画像だ。1つのチップに2つのダイが搭載され、おのおののダイに8GBのHBMが装着される構成になっている。

初代Google TPUの反省は外部メモリーの帯域が低すぎたことだそうで、それもあってHBMを使って大幅に高速化された

　それぞれのコアの詳細は下の画像だ。MAC Unitは128×128と1/4のサイズになったが、その代わり一度に全ユニットの計算が可能になっている。

ダイあたりで言えば、PCIe Gen3 x32でホストと接続される構造となる。VPUはMXUの10分の1程度の性能とされる

　また16MBのVmemも搭載されている。加えて言えば、データ型も初代のINT 8からTPU v2ではBfloat16/FP32に切り替わった。推論だけならINT 8のままでいいのだろうが、学習に向けてはやはりもう少し精度というか桁数が欲しい、というニーズに応えたものである。

　ダイのフロアプランは下の画像がわかりやすい。MXUよりもVPU+Vmemの方がはるかに大きなエリアを占めているのがわかる。

PCIe Linkそのものはダイに1ヵ所か所で、このあたりだけ対称性が崩れている。製造プロセスやダイサイズなどは未公表だが、プロセスは時期的なものから考えるとTSMCの20nmではないかとみられている。動作速度は700MHz

　またこのGoogle TPU v2の設計の最中に、Bach normalizationという論文が発表されている。この仕組みを、Google TPU v2ではハードウェアとソフトウェアでサポートすることで、学習速度を最大14倍にできたとしている。

具体的にはベクトルユニットのスループットを初期デザインの8倍にしたほか、逆平方根の計算を行なうハードウェアを追加したとしている

　ちなみにCloud TPU v2が64台のPodで11.5TFlopsとされているので、1台(つまり4チップ)あたり180TFlops、ダイ1個あたり45TFlopsという計算になる。

こちらはHot Chipsの論文より

　Google TPU v2では先にも書いたがBfloat16/FP32で計算しているので、処理速度はチップ1つあたり45TOP/秒という計算になり、これはおおむね初代Google TPUの2倍の性能に相当する。