TPU v3の鍵となる改良点

TPU v2のところで掲げた鍵となる目標を振り返ってみると、TPU v3では、ハードとソフトのコデザインでハードウェアを簡素化し、ソフトウェアの動作の予測性を高めた。TPUでは、スーパスカラやキャッシュのように状況によって動作が変わる機構を使わず、VLIWやスクラッチパッドメモリを使って同一シーケンスで処理が行なわれる構造を使っている。

  • TPU

    短期開発という点では、TPU v3ではハード、ソフトのコデザインで簡素な設計を行った。動作タイミングが変わるキャッシュやOut-of-Order実行は使わず、一定のタイミングで動作するVLIWやスクラッチパッドを使った

この設計目標の達成に大きく貢献したのは、第1にbfloat16を使ったシストリックアレイの採用でコンピュート密度を高めたこと。第2にHBMによって計算のためのデータを高速で供給できるようにしたこと。第3にXLAコンパイラの最適化でハードウェアをうまく利用できるようにしたことである。

  • TPU

    高性能の実現ではbfloat16のシストリックアレイを使って演算密度を高めたことが効いている。そして、HBMの採用でメモリからのデータ供給が速くなったことが大きい。さらに、XLAコンパイラによる最適化で性能が上がっている

さらに、短期間で開発できたのは、システムファーストのアプローチをとったことと、なじみのあるインタフェースを使うインタコネクトを使ったことが効いている。

  • TPU

    スケーラビリティの向上では、システムファーストのアプローチで使い慣れたインタコネクトを使ったことで、短期にインタコネクトが完成した

新規の問題への対応という点では、線形代数問題を解くフレキシブルなデータコアの実現と、XLAコンパイラ、HBMの大容量化が威力を発揮した。

  • TPU

    汎用の線形代数演算が行えるデータコアとXLAコンパイラ、大容量のHBMメモリの組み合わせはどのような問題にも対応できる柔軟性を持っている

そして、設計面では効率の高いマトリクスユニットを使ったこと。設計を単純化したこと。高性能の設計で性能/＄を高めたことが効果をあげている。

  • TPU

    コスト効果の高い設計という点では、高効率のマトリクスユニットと簡素な設計により、性能/＄の高い設計を実現した

TPU v2、v3はマシンラーニング用のスパコン

システム設計の面では、TPU v1はシングルチップのコプロセサであった。また、実際のアプリケーションに使っているモデルの学習を行うには、1個のTPU v2チップでは60～400日かかってしまう。このような膨大な学習を実用的な時間で行うには、多数のTPU v2、v3チップを使う大規模なスパコンを作ることが求められた。

  • TPU

    TPU v1はシングルチップシステムで、推論には適したコプロセサであった。しかし、実用的なモデルの学習を行うにはTPU v2チップ1個のシステムでは60～400日かかってしまうので、マルチチップのスパコンが必要である

TPUでの学習ポッドのアーキテクチャは、次の図のように、1枚のボードに4個のTPUを搭載し、そのボードを2次元トーラス接続して大規模アレイを作っている。そして、スパコン側はTop of the Rack(TOR)のスイッチにホストCPUを接続し、各ホストからPCI Express経由でTPUボードを接続している。このホストとTPUボードを接続するネットワークはそれぞれ独立で、同期式の並列トレーニングができるようになっている。

  • TPU

    TPU学習ポッドは2次元トーラスで多数のTPUのネットワークを作り、同期型の並列学習を行って性能を上げる。制御側はTPUボードごとにホストを置き、クラスタネットワークでホストにプログラムやデータを配る

次の図の左側の2本のラックでできているのはTPU v2の256チップのスパコンである。このシステムは64枚のTPU v2ボードで作られている。

写真の右側の8ラックのシステムはTPU v3の1024チップのスパコンで、ボード当たりのTPU個数は4個で、256枚のTPU v3ボードで作られている。右下の写真に見られるように、TPU v3では4個のチップには水冷のヒートシンクが付けられている。

ラック当たりのTPUの個数はv2もv3も128個で変わりはないが、v3はピーク演算性能が2.6倍に上がっており、電力密度が高くなっているので水冷としたのかもしれない。

  • TPU

    TPU v3を1024個使うスパコンのピーク性能は＞100PFlopsとなっている。HBMのメモリ容量は32TBである

TPU v2×256個のスパコンはピーク演算性能は11.5PFlops、HBMの総容量は4TBとなっている。一方、TPU v3×1024個のスパコンのピーク性能は>100PFlopsでHBMの容量は32TBとなっている。これは同じ2ラック構成で比べると、>25PFlopsの演算性能となり、11.5PFlops×2.6倍＝29.9PFlopsとそれほど矛盾しない性能記述である。

(次回は9月11日に掲載します)