TensorFlow Processor Unitアーキテクチャ

Googleは今月初め、1年前に発表したTPUに関わる詳細な情報を公開した。プロジェクトの上級アーキテクトであるNorm Jouppi氏とそのチームは、TPUを使用したニューラルネットワーク実行のパフォーマンスが、Nvidia K80やHaswell E5-2699 v3といった他の代表的なハイパフォーマンスプロセッサの数倍に到達した、と報告している。

“TPUは推論において、K80 GPUやHaswell CPUより約15倍～30倍高速です...(テストを実施した)６つのアプリ中４つは、TPU上ではメモリ帯域幅の制限を受けています。K80 GPUと同等のメモリシステムを持つように修正されれば、GPUやCPUに比較して30～50倍は高速になるでしょう ... [さらに]TPUは、効率測定値であるTOPS/Watt計測においては、標準的なプロセッサの30～80倍という値を示しています。”

カスタムASICを開発したそもそもの動機は、Googleの翻訳APIに必要な計算要件を供用することにあった。世界中の携帯電話で翻訳APIを1日1回、わずか3分間使用するためには、数十ヶ所というデータセンタの増設が必要なのだ。

アーキテクチャに関する白書には、実験的な設計やデータ収集、および一連のNNを処理するエグゼキュータコンポーネントとして、K80およびE5-2699を対象に実施された分析の内容が記録されている。今のところTPUは、ニューラルネットワークのトレーニングには使用されていない。その主な利用対象は行列乗算機である。これは行列乗算ユニットとオンボードメモリ、キャッシングシステムで構成されており、ニューラルネットワークの複数層を処理することができる。その中には、MLPとCNNで頻繁に使用される、レイヤ間のパーセプトロンの出力と状態保持動作も含まれる。

ただしTPUは、ニューラルネットワーク実装のみを目的としたものではなく、汎用性を意図しており、そのアーキテクチャは、Jouppi氏とそのチームが調査した包括的なユースケースに基づいている。そのようにした理由にはTPUの提供スケジュールも関連しているが、それだけではない。チップが実行する最も基本的なオペレーションである、行列浮動小数点演算を最適化するために必要な柔軟性とも関係がある。アプリケーション全体を実現ためにはCPUないしGPUとの統合が必要だが、TPUではPCIeバスを利用することで、これを簡易に実現している。

TPUのアーキテクチャではニューラルネットワークの実行以外に、トレーニングあるいはTensorFlowアプリケーションの部分的実行のためのCPUないしGPUが必要となる。例えば、アプリケーションがTensorFlow実行部分にインストラクションを送るためにデータのフロントロードやロジックの実行が必要な場合、それらをCPUないしGPUで実行した上でTPUに送信しなければならない。この意味から言えば、TPUはグラフィックカードやFPUに極めて近いものだ。

“TPUはCPUやGPUのようにプログラム可能です。特定のニューラルネットワークモデル用に設計されたものではなく、さまざまなネットワーク(畳み込み、LSTMモデル、完全に接続された大規模モデル)上でCISC命令を実行します。そのようにプログラム可能ではありますが、ベクタやスカラではなく、プリミティブとして行列を使用しています。”

CPUやGPUの事変的(time-varying)最適化に比較して、TPUは決定論的な性質を備えているため、電力消費量あたりのテラOPSないしテラFLOPSに関するベンチマークではTPUが上回っている。TOPS/ワットの測定値では、標準的なプロセッサの30～80倍であると報告されている。

“TPUは、K80 GPUの32ビット浮動小数点データパスに比較して、電力と面積を大幅に削減した8ビット整数シストリック(systolic)行列乗算器を活用することによって、25倍のMAC(65,536個の8bit対2,493個の32bit)、3.5倍のオンチップメモリ(28MiB対8MiB)を搭載しながら、K80の半分以下の電力消費を比較的小型のダイで実現しています。この大容量メモリがアプリケーションの演算強度を向上することで、豊富なMACをさらにフル活用できるようにしているのです ... コンピュータアーキテクチャにおいてこれほどの大きな性能差があるというのは、コマーシャルな製品としては希有のものです。このためTPUは、ドメイン固有アーキテクチャの原型になる可能性を持っています。”

実験的デザインの研究段階の一環として、Jouppie氏とそのチームは、GoogleプラットフォームにおけるNNの使用状況についての調査を行なった。その結果として氏らは、スループット上の要件よりも時間的な要件を持ったアプリケーションの数が当初考えていた以上に多いことを見出した。そこから氏らは、大規模で安価なリソースでは、たとえ利用効率が多少低くてもハイパフォーマンスを高いコスト効率で提供可能であるという、Amdhalの法則を再認識するに至ったのだ。

TPUの実験的プロファイルは、2つのMLPとCNN、LSTMという、６つのニューラルネットワークで構成されている。2つのMLPとLSTMはメモリバウンドであるため、試験時のメモリ帯域幅の順列の調整がパフォーマンスに最も大きく影響している。これはMLPとCNNが、重み付けなどいくつかの面において、過去のレイヤでのアウトプットを再利用していることが影響している可能性がある。対してLSTMでは、時間ステップに関する再利用部分はあるものの、過去のレイヤのアウトプットを選択的に維持しているだけであるため、パフォーマンスの制限要因としてはメモリバウンドより計算バウンドの部分がより大きくなっている。TPUが使用するPCIeバスのI/O帯域幅を考慮したとしても、これは事実として成立する。

現在の暫定的なTPUアーキテクチャにおいてメモリとキャッシュ処理を完全することにより、15ヶ月以上という現行の提供スケジュールを前提とするならば、K80およびES5-2699に対して30倍から50倍のパフォーマンス向上が実現可能である、とJouppi氏のチームでは主張している。