Computer Architecture 6th Editionの7章"Domain-Specific Architecture" を読む (7.4章 Google's Tensor Processing Unit)

Computer Architecture, Sixth Edition: A Quantitative Approach (The Morgan Kaufmann Series in Computer Architecture and Design)

7.3章はがっつりディープニューラルネットワークの内容だ。今回は多層パーセプトロン、畳み込みニューラルネットワーク、再帰型ニューラルネットワーク。

目次

• 7.1章 イントロダクション
• 7.2章 DSAのガイドライン
• 7.3章 DSAの用語集
• 7.3章 DNNのニューロン・トレーニングと推論
• 7.3章 多層パーセプトロン、畳み込みニューラルネットワーク、再帰型ニューラルネットワーク
• 7.4章 Google's Tensor Processing Unit
• 7.5章 Microsoft Catapult, フレキシブルなデータセンタアクセラレータ
• 7.6章 Intel Crest, 学習のためのデータセンタアクセラレータ
• 7.7章 Pixel Visual Core, パーソナルモバイル画像処理ユニット)

これは著者が読んだ内容をまとめているだけなので、誤訳、理解不足により誤っている可能性があります！鵜呑みにしないようにお願いします。 長々と書いたけれども、Tensor Processing Unitの論文よりも詳細なことは書いていない。そりゃそうか。

---

GoogleのTensor Processing Unit, データセンタの推論アクセラレータ

Tensor Processing Unit(TPU)はGoogleの最初のWSC向けDSAのASICである。 個のチップのターゲット領域はDNNの推論フェーズであり、またDNN向けに設計されたTensorFlowフレームワークを使ってプログラミングを行うことが出来る。 TPUの第1弾は2015年にGoogleのデータセンタにデプロイされた。

TPUは65,536(256×256)個の8-bit 行列積ユニットALUを備え、ソフトウェア管理のオンチップメモリが搭載されている。 TPUはシングルスレッドで、決定的な実行モデルは、99パーセンタイルのレスポンスタイムが必要とされる典型的なDNNの推論アプリケーションにうまく適合している。

TPUの由来

2006年に、GoogleのエンジニアはデータセンタにおけるGPU、FPGA、カスタムASICのデプロイについて議論を行った。 彼らは、特殊なハードウェアを必要とするアプリケーションで、大容量データセンタの過剰な計算能力を使っても仮想的に実行することが出来るものは少なく、これらを自由に改善することは難しいとの結論に達した。 2013年にはこの議論の結論は変わり、ユーザが1日に3分ほどDNNの音声検索を使用するならば、その計算能力を間に合わせるためにはGoogleのデータセンタを倍増しなければならないという結論に達した。 このデータセンタの増強は伝統的なCPUを使って達成するには非常にコストのかかるものであった。 Googleは推論処理を実行するためのカスタムASICチップを開発するためのプロジェクトを急ピッチに立ち上げた(と、トレーニングに必要なGPUの在庫を購入した)。 目標はGPUの10倍以上のコスタ性能を達成することであった。 この目標を元に、TPUは設計、検証、組み立て、デプロイがたった15ヶ月の間に行われた(Steinberg, 2015)。

TPUのアーキテクチャ

デプロイが遅延することを避けるために、TPUはPCIe I/Oバスに接続されるコプロセッサとして設計され、既存のサーバに接続されるように設計された。 さらに、シンプルなハードウェア設計とデバッグ機能をもっており、ホストサーバがPCIeバスを通じてTPUに指令を転送する。 これはTPU自身が命令をフェッチする機構を取っているわけではない。 したがって、命令をメモリから自分でフェッチするGPUとは異なり、TPUはFPU(浮動小数点ユニット: Floating Point Unit)コプロセッサのような考え方を元に設計されている。

図7.12はTPUのブロックダイアグラムを示している。 ホストのCPUはPCIeバスを通じてTPUの命令バッファに命令を転送している。 内部ブロックは256バイト(2048ビット)幅で接続されている。 右上ブロックから始まって、「行列積ユニット(Matrix Multiply Unit)」がTPUの心臓部である。 行列積ユニットは256×256個の符号あり・無しの8-bit整数積和演算ALUが内蔵されている。 32-bit幅の、4MBのアキュムレータによって、16ビットの結果が格納される。 8bitの重みと16bitの値を使って、ビット幅を混在させた場合は、行列積ユニットは半分の速度で実行する。 また、両方とも16bitの値を使った場合は行列積ユニットは4分の1の速度で計算を実行する。 1クロック当たりに256個の値を読み書きし、行列積と畳み込みのどちらかを計算することが出来る。 「活性化(Activation)」ユニットを使って、非線形関数の計算が行われる。

図7.12 TPUのブロックダイアグラム (https://cloud.google.com/blog/big-data/2017/05/an-in-depth-look-at-googles-first-tensor-processing-unit-tpu より抜粋)

行列ユニットの重みデータは、オンチップメモリの「重みFIFO(Weight FIFO)」を通じて格納される。 この値は重みメモリ(Weight Memory)と呼ばれるオフチップの8GBのDRAMから読み込まれる(推論用で、重みは読み込み専用、8GBはは多くのアクティブモデルをサポートしている)。 中間結果は24MBのオンチップの「ユニファイドバッファ(Unified Buffer)」に格納され、行列積ユニットの入力値として使用される。 プログラマブルDMAコントローラはデータをCPUホストメモリとユニファイドバッファの間を転送する役割を持っている。

TPUの命令セットアーキテクチャ

TPUの命令は比較的低速なPCIeバスを通じて転送されるため、TPU命令は伝統的なCISC命令形式であり、リピートフィールドを持っている。 TPUはプログラムカウンタを持っておらず、分岐命令も持っていない; 命令はホストCPUから転送され、これらのCISC命令の命令辺りに必要なクロックサイクル数(CPI)は典型的に10～20サイクルである。 命令種類は1ダース程度であるが、5つのカギとなる命令がある。

1. Read_Host_Memory 命令はCPUのホストメモリからユニファイドバッファにデータを転送する。
2. Read_Weights 命令は重みメモリから重みデータを重みFIFOに転送し、行列ユニットの入力値として使えるようにする。
3. MatrixMultiply/Convolve命令は行列積ユニットを実行し、ユニファイドバッファのデータを使用して行列積、ベクトル行列積、要素毎の行列積、要素毎のベクトル積、もしくはの畳み込み演算を実行し、結果をアキュムレータに格納する。行列の演算は可変サイズを取ることが出来、B*256の入力と、256×256の行列積を実行し、B*256個の結果を出力する。これにはBサイクル必要である。例えば、入力が256要素の4ベクタであるならば、Bは4であり、計算の完了までに4サイクルが必要である。
4. Activate命令は人工ニューロンのための非線形関数を実行し、ReLU、シグモイド、tanhなどを実行することが出来る。このユニットの入力はアキュムレータから提供され、計算結果はユニファイドバッファに格納される。
5. Write_Host_Memory命令はデータをユニファイドバッファからCPUのホストメモリに書き込む。

TPUのマイクロアーキテクチャ

TPUのマイクロアーキテクチャの哲学は、なるべく行列積ユニットを動作させ続けるということである。 他の命令の実行を、行列積の実行であるMatrixMultiply命令とオーバラップさせて隠すという考え方である。 したがって、実行命令以外の、4つの一般的なカテゴリの命令は個別の命令ハードウェアを持っている(ホストメモリとの読み書きは、同一のユニットを使って実行される)。 命令並列性を向上させるためには、Read_Weights命令はDecoupledアクセス/実行方式(Smith, 1982b)に則って設計されている。 これは計算を完了するのは、アドレスを転送した後であるが、計算完了前に重み情報は重みメモリからフェッチされるという方式である。 行列ユニットはNot-Ready信号をユニファイドバッファからもらっており、重みFIFOは行列ユニットにストール信号を転送し、重みデータの到着を待たせることが出来る。

ここでTPUの命令は複数クロックサイクルで実行され、伝統的なRISCパイプラインのように1サイクルで1ステージを通過していくわけではないということに注意である。

大きなSRAMから読み書きを行うことは、演算よりもコストが高いため、行列積ユニットはシストリック実行方式を取っており、ユニファイドバッファからの読み書きの回数を削減している(Kung and Leiserson, 1980; Remacher et al., 1991; Ovtcharov et al., 2015b)。 「シストリックアレイ(Systoric Array)」は2次元状に演算器を並べており、各演算器は独立して計算を実行し、関数の部分計算結果を次の演算ユニットに転送する。 データは異なる方向からセルに対して一定の間隔で転送で転送される。 アレイ内のデータフローはウェーブフロントのように転送されるため、人間の血管構造に似ており、この構造がシストリックと呼ばれる由来になっている。

図7.13はシストリックアレイが動作する様子を示している。 下側の6つの丸が積和演算ユニットであり、重みで初期化されている。 ステージングされたデータが、このアレイの上側から転送されてきている。 この図では10ステップ毎のデータ転送の様子を示している。 シストリックアレイは入力値を下側に転送し、積と和を右側に転送している。 シストリックアレイは入力データはメモリからのみ読み込まれ、出力データは1度だけメモリに書き込まれることに注意である。

TPU内では、シストリックアレイはローテートされる。 図7.14では、重みデータは上部からロードされ、入力データは左側から転送されアレイ中を通過する。 256要素の積和演算器は行列中を通過し、斜め方向にウェーブフロントとして通過する。 重みデータはあらかじめロードされ、ウェーブ中でデータブロックの最初のデータとして進んでいく。 制御情報とデータも一緒に転送され、256個のアキュムレータのメモリのデータをアップデートしてく。 正確性という面からはソフトウェアは行列ユニットのシストリックの構成については意識しないが、ユニットのレイテンシについて考慮する必要がある。

図7.13 シストリックアレイの動作

図7.14 行列積ユニットのシストリックデータフロー (https://cloud.google.com/blog/big-data/2017/05/an-in-depth-look-at-googles-first-tensor-processing-unit-tpu より抜粋)

TPUの実装

TPUチップは28-nmプロセスを使って製造されている。 動作周波数は700MHzである。 図7.15はTPUのフロアプランを示している。 ダイサイズは公表されていないが、Intel Haswellサーバマイクロプロセッサの662mm²の半分以下の大きさである。

図7.15 TPUダイのフロアプラン (https://cloud.google.com/blog/big-data/2017/05/an-in-depth-look-at-googles-first-tensor-processing-unit-tpu より抜粋)

24MBのユニファイドバッファが1/3を占めており、行列乗算ユニットが1/4、データパスが台の2/3を占めている。 24MBの大きさは行列ユニットの大きさと適合するようにできており、コンパイラを単純化することで開発期間を単純化している。 制御ユニットは全体の2%である。 図7.16はTPUのプリント基板であり、既存のSATAディスクスロットに挿入して使用する。

図7.16 TPUのプリント基板 (http://hexus.net/tech/news/industry/104299-google-benchmarks-tensor-processing-unit-tpu-chips/ より抜粋)

TPUのソフトウェア

TPUのソフトウェアスタックはCPUでの開発やGPUでの開発と互換性があり、アプリケーションを短期間で移植することが出来る。 典型的に、アプリケーションをTPU上で実行するときはTensorFlowとそのコンパイルされたAPIを使ってGPUもしくはTPU上で実行する(Larabel, 2016)。 図7.17はTesorFlowを使ってMLPを実行するコードである。

GPUのように、TPUスタックはユーザスペースドライバとカーネルドライバに分かれている。 カーネルドライバは軽量で、メモリ管理と割り込みを管理するだけである。 これは長期間において安定するように設計されている。 ユーザスペースドライバは頻繁に変更される。 ユーザスペースドライバはTPU実行のセットアップと制御を担当し、TPUの形に合うようにデータを変換し、APIコールをTPUの命令に変換してアプリケーションバイナリに変更する。 ユーザスペースドライバは、コードが評価されると、まずモデルをコンパイルし、プログラムイメージをキャッシュして重みイメージをTPUの重みメモリに書き込む; これにより2番目以降の評価ではフルスピードで動作することが出来る。 TPUは入力から出力までほとんどのモデルを完全に実行することができ、TPUの実行時間とI/O時間の比率を最大化する。 計算はしばしば一度に1レイヤで終了し、クリティカルパスではない処理によって行列演算の処理をオーバラップしてレイテンシを隠蔽する。

TPUの向上

TPUアーキテクトはTPUをより向上させるために多くのマイクロアーキテクチャを検討した。

FPUと同様にTPUのコプロセッサは比較的簡単なマイクロアーキテクチャである。したがってTPUアーキテクトは性能モデルを作成し、メモリのバンド幅、行列ユニットのサイズ、クロックレートとアキュムレータの数を利用して性能を見積もった。 TPUハードウェアカウンタを使った測定では、作成したモデルとハードウェアとは約8%の誤差が存在した。

図7.18は、これらのパラメータによりTPUの性能にどの程度影響があるかを0.25倍から4倍の間で比較したものである(7.9章のベンチマークを使用した)。 クロックレートを向上させたもの(図.18のClock)に加えて、図7.18ではクロックレートを向上させたものとアキュムレータの数を増やしたものを追加しており、したがってコンパイラは実行中のメモリ参照を増やすことが出来る。 同様に、図7.18はアキュムレータの数が2乗に増加した場合に行列ユニットを拡張した場合の性能向上を示している。 一方で"matrix"は行列ユニットを向上させただけである。

まず、メモリバンド幅を向上させた場合(memory)は最も大きなインパクトがあった: メモリバンド幅を4倍にすると、性能は4倍になり、これは重みメモリの待ち時間が削減されたからである。 次に、アキュムレータの数を変更せずにクロックレートのボトルネックはわずかであったであるということである。 3番目に、行列ユニットを256×256から512×512にしても、アキュムレータの数を増やしてもそうでなくても、図7.18の平均的な性能はわずかに低下するだけであった。 この問題は大きなページを扱った場合の内部フラグメンテーションの問題に似ている。 TPUでは2次元のデータしか扱わないためである。

LSTM1において600×600の行列を使った場合は、256×256の行列ユニットでは600×600の行列を全て計算するために9ステップ必要で、合計で18usとなる。 より大きな512×512の行列ユニットを使った場合は4ステップで完了させることが出来るが、各ステップの計算時間が4倍になり、32usの時間がかかる。 TPUのCISC命令は長いため、デコードはDRAMからのロードのオーバヘッドを隠蔽することはできない。

性能モデルから得られる結果を見て、TPUアーキテクトはもしさらに15ヶ月を使って設計した場合の仮説を立てた。 よりアグレッシブな論理合成と、ブロックデザインのクロック周波数を約50%向上することが出来たと考えている。 アーキテクトは、K80で使われているようなGDDR5メモリのインタフェースを設計することで、重みメモリのバンド幅を5倍近く向上させることが出来ると考える。図7.18に示すように、クロックレートを1050MHzにしても、メモリの支援が無ければ性能にほとんど影響を与えない。 もしクロックレートが700MHzのままでも、GDDR5を重みメモリに使用すると、性能は約3.2倍に向上する。 これらを両方実施しても、さらに性能を向上させることはない。