TPUその後 ～ シストリックアレイがすごい ～
以前、TPUについて書きましたが、論文が難しくてモタモタしているうちに、Google自身が論文の解説をするブログを公開しました。なんとそれが翌日には日本語訳されて公開されています。Google、偉すぎ!
結構長いので、ブログの最後に書いてある結論を先にバラしてしまいますと、高性能のカギは「NNの推論処理への徹底した特化」ことだということです。要らない機能を削り、NNの計算にひたすら特化したハードウェアとしたのですね。しかし、特定目的へのハードウェア化が行き過ぎると、他のことができなく (できにくく) なりますから、今のNNの計算方法 (ニューロン間の接続に重みを付け、ニューロンへの入力を加算し、活性化関数を適用する、というパーセプトロンの処理) で十分実用になる (少なくともしばらくは)、という判断の下で巨額の開発費を投じたということなんでしょう。やはりGoogleはお金ありますねえ。
加えて個人的に興味深かったのは、「RISC」「MIPS」「積和演算」「ベクタープロセッサ」など、懐かしい言葉が沢山出てきたことです。私はその昔、ベクタープロセッサをマルチで搭載したArdent社の3Dグラフィックスワークステーション「TITAN」に関わっていたので、この辺の用語には思い入れがあります。
原文は長いので、キーワードだけご紹介します。
開発期間15ヶ月
設計開始からシリコン作ってデータセンターへ導入するまでに15ヶ月というのは確かに速いですね。ただ、ASIC化の検討は2006年から始まっていたということですので、仕様などについてはかなりできていたのではないかと思います。(しかし、2006年ってヒントン教授がディープラーニングの論文を書いた直後くらいですよね。早い。。) 開発はMIPSの主任アーキテクトが主導したそうです。さすがに良い人材を揃えています。
行列演算!
NNの演算で最も時間がかかるのは、行列積 (積和演算) ということです。これは汎用のスカラプロセッサは苦手 (というか真面目にやってると時間がかかりすぎる) な分野で、スパコンやGPUなども、この計算を高速に行うためのハードウェアとして産み出されたと言っても良いでしょう。スパコンの演算速度を測る指標としてLINPACKが使われますが、これは元々行列計算用のライブラリです。
8ビット
前回の記事でも触れましたが、NNの推論部分にはそれほどの精度は要らない、ということのようです。(量子化は行っているようですが) この辺はやはり長年の研究で、8ビットでいける、ということになったのでしょう。精度を割り切ることで、データ量が1/4になり、回路も小さくすみ、全体のチップサイズを抑え、消費電力を抑えることができています。
RISCではなくCISC
これも懐かしいキーワードですね。RISCというのは当時主流だったCISCへのアンチテーゼ (アーキテクチャが複雑化しすぎて、1命令で様々な処理ができる反面、1命令の処理が何クロックもかかる) で、1クロックで1命令を実行することを目指したアーキテクチャです。しかしTPUでは、1命令で数千回の演算を行う事ができ、それがCISCと呼ぶ所以でしょう。ベクタープロセッサもGPUも、CISCということになるのでしょうね。
＃しかしTPUは、1命令で数千回の演算を「1クロックで」できるとも書いてあります。この辺、どう解釈するのかなあ。まあ、RISCが目指したシンプルさとは対局にあるから、CISCで良いのかもしれません。
シストリックアレイ
ここらへんから、かなり難しくなってきます。ベクタープロセッサという、これまた懐かしい言葉が出てきます。TPUのMXUはベクタープロセッサとも違うアーキテクチャで、ALU (Alithmetic Logic Unit: 算術論理演算器) をマトリックス状に配置しているということです。最初ベクタープロセッサの2次元版かと思ったのですが、これが「シストリック」な構成になっており、ここにデータを流し込むと、次から次へと連続的に演算が行われる構造になっているということのようです。ブログに貼ってあるアニメーションを見ていると、パラメータが流れ作業で処理されていくのがわかります。「データストリーム」と呼ぶようですね。まだ理解が及んでいませんが、演算の間、レジスタへの読み書きを発生させず、ひたすらNNの計算をよどみなく実行するというのが、TPUのキモなのでしょう。
しかし、本当にNNに最適化されている感じですね。他の計算ができるのか心配になります。(する必要も無いのでしょうが) ブログにも、「最も高い電力効率と実装密度で行列積を計算するという目的に特化して設計されています。」と書いてありますが、他の行列積にもうまく使えるのでしょうか。
ミニマル設計
目的を絞り込んだことにより、普通のプロセッサで必要な様々な機能 (キャッシュ、分岐予測、アウトオブオーダー実行、マルチスレッド、投機プリフェッチ、アドレス計算、コンテキスト スイッチ等々) を排除することができ、制御ロジックを最少化できました。その結果もあってチップ面積が縮小し、製造コストを下げることができたということです。
シストリックアレイは面白いです。調べてみると、1980年代からあるんですね。でも、商用製品に使われたことはあるんでしょうか。1993年の研究で
シストリックアレイは、ハードウェアアルゴリズムに基づくアーキテクチャで、アプリケーションごとに構成が異なったものとなり汎用性に乏しい。
なんて言われてたりします。今回、NNの計算でやっと実用化、ということなのかも知れません。