Verilog
DeepLearning
TensorFlow
TPU

素人がGoogle TPUを作ってみる~アーキテクチャ考察編~

目標

Google TPUの思想を理解するためにTPU likeシストリックアレイを構成することでなぜTPUがDNNを回す効率が良いかを理解するのが目標。
読む方もニューラルネットワーク(全結合、CNN)がわかっていればなんとなく理解できれば良いな。。と思ってます。

記事はアーキテクチャ編と設計編に分け、本アーキテクチャ編ではTPUの動作などを考察する。
設計編ではVerilogでシストリックアレイ部のみ記述してみる予定。ただPE部は資料がないので筆者の妄想。
ほぼ構成要素としては前記事のなんちゃってCPUと同じだと思う。

素人が作る超基本CPU@Verilog w/Vivado
https://qiita.com/arutema47/items/40dc748e349c30d7194a

今回考えるのはTPUv1

ISCA2017で論文発表され、詳細なアーキテクチャが公開されている。
TPUの特徴として:
・演算精度の緩和(8bit INT)
・シストリックアレイからなる行列演算器の採用

が挙げられる。この2つを解説するのがこの記事になる。

そもそも何故Googleは専用チップ(ASIC)を設計?

またTPUのような専用チップをソフトウェア会社であるGoogleが開発しているのはなぜか。
GPUをたくさん買ってくればいいのではないか?

まずひとつにムーアの法則の終焉がある。ムーアの法則は18ヶ月で半導体の性能は倍になるというもの。
つまりソフトウェアの実行速度も一年半待てば倍になるという事が期待できた。。昔は。
最近はもう半導体性能は向上せず、コスパも悪くなってしまっている。(これをポストムーア時代とか)

そのためDNNのような演算が重い(そしてこれからもどんどん規模が大きくなる)サービスをクラウドで提供するにはCPU/GPUじゃモノ足りない!専用チップが必要だ!と思い立つ。
GPUは科学演算をサポートするためゴツい32bit/64bit演算器があるが、それはニューラルネットにはオーバースペック。その無駄を徹底的に省いたニューラルネットワーク専用ハードウェアがTPU(TensorProcessingUnit)と言える。

一方で去年に発表されたTPUv2は32FPに対応しHBM搭載と完全に別モンになった。。

行列演算回路(Matrix Multiply Unit)に注目

TPUv1のアーキテクチャ(N.Juppi)
tpu.png

TPUv1の論文。
N. Juppi et al, In-Datacenter Performance Analysis of a Tensor Processing Unit
https://dl.acm.org/citation.cfm?id=3080246

上記論文はDNNのデータセンター利用例と興味深い事がたくさん書かれているが、今回注目するのはMatrix Multiply回路部
Google BlogでもここがDNN演算の効率化のキーと書かれている。

An in-depth look at Google’s first Tensor Processing Unit (TPU)
https://cloud.google.com/blog/big-data/2017/05/an-in-depth-look-at-googles-first-tensor-processing-unit-tpu
Calculations_in_Neural_Network.gif

行列演算回路がなぜ重要なのか?

ぶっちゃけて言ってしまうと。。
ニューラルネットワークの演算はほとんど行列演算(MatMul)(95%くらい)

じゃあ行列演算を効率化すれば全体の効率もめっちゃ良くなる!

という考え。コンピュータ・アーキテクチャの金言である共通のものを速くせよに則ってる。(Patterson)
ざっくばらんに言うと、行列演算回路をDNN専用化したからCPUやnVidia GPUに対して速い+低電力になる!という主張。

先端のCNNでは行列演算は一枚の画像に付き数10億のオーダで実行される。(VGGでは画像一枚につき1兆!)
一方で他の演算(活性化関数とか)はx万回というオーダなのでCPUで実行しても全体の効率からいうとペナルティは大したことはない。
image.png

余談:TPUを受けてか知らないけどnVidiaのVoltaアーキテクチャは4x4の行列演算回路を専用に持つことで大幅にDNN演算速度を改善すると同時に価格も大幅にアップした。8(16?)bit対応もしていたりTPUに対して敵意剥きまくり。。

行列演算回路の効率改善の手法

Google TPUは2つの効率化手法を取っている。
・演算精度の緩和
・演算器のシストリックアレイ化

まず一つ目の演算精度に関して。

従来の科学演算では32FPや64FPと高い演算精度が求められていたのに対し、DNNではあまり精度がなくても性能が出る。
ぶっちゃけるとモデルが汎化しているので演算がちょっとズレてもちゃんとモデルはちゃんと動く(認識精度は落ちない)。
今では当たり前のこの知見が3,4年前のGoogleには既にあって8bit化を決断したのが凄い。

ImageNetでは8bit精度が必要だが、MNISTなどでは3bitくらいでも結構動いてしまう。
更にBinary(1bit)まで落としても動かなくはないという研究もあるけどちょっと怪しい。。

そのためTPU内部では全てのデータは8bitを扱っている。
CPUやGPUでは基本的にデータは32bitなのでそれに比べるとデータ量は1/4になり、電力も1/4になる。
厳密には8bit化により演算電力はもっと下がるけど、メモリ電力が支配的なのだと思う。

ただ8bitは推論でしか使えないのに注意。低精度学習の研究もあるが、まだまだな印象がある。
ちなみに学習に対応するためTPUv2は32FPになった。ほぼシストリックアレイ以外はGPUとなってしまった気もするけど詳細が気になる。。

ちなみにPytorchに有名な画像認識モデルについて量子化を行った結果がある。

aaron-xichen/pytorch-playground
https://github.com/aaron-xichen/pytorch-playground

キャプチャ.JPG

これは単純なmin-maxで再学習もかけていないのでちゃんとdyanmic量子化をやればImageNetでも6bitで動く気がする(誰か試して記事書いて下さい)

シストリックアレイ化に関して

シストリックアレイの採用が個人的にはGoogleの一番思い切った一手だと思う(もちろん8bit化も凄い決断だけど)。
シストリックアレイのメリットを出来るだけわかりやすく記述する。

シストリックアレイの長所:データの局所性の最大化

シストリックアレイの長所は行列演算におけるデータ局所性を高められることである。

教科書にあった例としては図書館でレポートを書く時の例がわかりやすい。
よく使う本は一度本棚から取ったら手元の机においておきたい。本棚に戻してしまうと次使う時にまた本棚まで行く手間が生じてしまう。

コンピュータも同じでよく使うデータは本棚(HDDやDRAM)などにしまわず、手元の机(プロセッサ内のメモリやSRAM)においておこう!そうすることでデータを取ってくる電力や時間を節約することができる。
端的にはよく使うデータは近くにおこうというのがデータ局所性の概念である。

image.png
大体プロセッサから見るとSRAM(プロセッサの小さいメモリ)、DRAM(プロセッサ外のでかいメモリ)、HDD(プロセッサ外の超でかいメモリ)と順々に遠く見えている。近くにおけばおくほどデータを取ってくる手間は小さく、DRAMは先程の例では本棚でHDDは隣町の図書館くらいに遠い。

このような行列演算(Matmul)を行う時に入力行列の[1 2 3]は重みの1列目[11 13 15]と2列目[12 14 16]を演算する2回使用する。
極端な例だが、もし重みの一列目を演算した後に入力行列[1 2 3]をDRAMやHDDなどに格納してしまうと直後に再度データを取ってくる無駄な手間が生じてしまう!
普通のプロセッサ(CPUなど)は行列演算をしていると知らない。そのためこのような無駄なことをしてしまう(極端な例だけど)。
一方でシストリックアレイは行列演算をするという前提で設計されているのでデータ移動が最適化されている。

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

Googleのシストリックアレイの解説わかりにくいんだよな。。(カッコいいんだけど)

Systolic_Array_for_Neural_Network_2.gif
from google research blogs

GoogleのTPUはWeight Stationary(重みデータを動かさない)というデータフローを用いている。(有名なMIT Eyerissもweight stationary)
入力データと出力データを1コマずつ移動させ、シストリックアレイを構成している。
そもそもシストリックというのは心臓が脈打つことを言うらしいが(カッコイイ)、個人的にデカイ演算器つきFIFOっぽい動作なのかなと思っている。

ここで簡単のため3x2行列演算するためだけに設計したぼくのかんがえたTPUを例に考える。何に使うんですかね
image.png
PEは計算を行うProcessing Elementの略。

それではこのTPUに
image.png
の演算をマッピングする。

image.png

まずは入力行列の1行目[1, 2, 3]をPE1-3に入力する。
そして予め重み行列の1列目[11, 13, 15]もPE1-3に入れておく。

PE1は入力と保持している重み信号同士の掛算(1*11)を行い、↓方向に出力をパスする。
PE2も同様に入力と保持している重み信号同士の掛算(2*13)を行い、PE1から受け取った出力と足し合わせ(1*11+2*13)↓方向に受け渡す。
PE3も同様の動作を行う。
するとPE3の出力からは所望の結果(1*11+2*13+3*15)が得られる!!
紺色の枠の演算ができました。ぱちぱち。

image.png
ただシストリックの真髄はここからである。

シストリックアレイでは出力だけではなく、入力信号を→方向に受け渡す。
この時PE1-3では 入力2行目[5, 6, 7]x重み1列目[11, 13, 15] の演算が実施され、(紺の枠)
PE4-6では 入力1行目[1, 2, 3]x重み2列目 の演算が実施される。(緑の枠)

重要なのは入力1行目[1, 2, 3]はPE4-6に最短経路で受け渡される
つまり受け渡し時にメモリとのやり取りはなく、演算回路内で受け渡されるので速くてコストが低い。

考察

シストリックアレイでは入力行列のデータが最大限に使い回され、重みデータはPE内から動かないのが特徴!

通常のプロセッサ(PE)では演算したデータはレジスタやメモリなどに保存されにいき、次の演算するデータをメモリから読み出す。
一方でシストリックアレイは前のPEからデータを受け渡されるため、無駄が最小限。工場のラインに似ている。
これが可能なのはデータが行列しか来ないと絞り込んでいるから。
言い換えると汎用性を犠牲にして行列演算に特化する事で性能向上を達成している。犠牲なくして性能向上がないのが設計です。

#これを256x256のPEでやっているのが本家TPUとなる。
#実際はx方向のシストリックだけではなく、速度を上げるためにy方向(出力信号)もシストリックになっているが簡単のため省略した。

次の記事ではこれをVerilogにて設計する。

まとめ

最後にTPU(シストリックアレイ)のメリットデメリットをまとめる。

メリット:行列演算にとっては究極のデータ局所性のため高速かつ低電力。
デメリット:行列演算しか実行できず、汎用性はCPU/GPUに負ける。(から今まで使われていなかった)
DNNの市場が圧倒的だから最近復権していると思われる。

 感想

行列演算器も噛み砕いてしまえばそんなに複雑な事をしていないが、如何せんTPUは規模が大きいので理解が難しい。(演算器なんて掛算と足し算しているだけ)
将来DNNハードウェアを設計する人の役になにかたてればと思いかなり噛み砕いた図を作ってみた。

・実際にTensorFlowからTPUまで演算をマッピングする仮定(コンパイラ、リンカ)が物凄く大変だと思う。さすがgoogle。
・実際のアレイは65536個の演算器を並べる。するとクロック分配一つとってもかなり大変そうだけど一体どうしてるんだ。
・命令セットについて考えなきゃ

lastly but not least

関係ないけどJeff Deanが書いたGoogle Brainの2017年の研究活動を振り返った記事もめっちゃ面白い。\(^o^)/
https://research.googleblog.com/2018/01/the-google-brain-team-looking-back-on.html