Hot Chips 31はマシンラーニングが花盛り(13) NVIDIAの次世代AI用研究チップはスケーラブルな推論アクセラレータ

次の図のように、右側に描かれた分散型の推論アクセラレータ16個がチップに搭載され、パッケージには36チップ搭載できる。各アクセラレータには8並列8レーンのベクタ演算器が入っているので、全部で64×16×36=36864個のMACがあることになる。

左上の図が16PEを集積するチップ。左下は36個のチップを搭載するパッケージ。右の図は、PEのブロック図である。重みのバッファと8個のMAC×8、Poolingなどを行うPPUなどが入っている

そして分散型のアクセラレータであるので、まず、重みをPEに分散し、入力をグローバルPEに入れRISC-Vでコントロールレジスタをセットしてストリーム入力を各PEに配る。そして、PEで演算を行い、結果をグローバルPEに格納する。

左に書いたフローチャートのようにして全部のPEで並列に計算を行わせる

左上の計算を4個チップに分割して計算させる例である。チップ内でも右の図のように多数のPEに処理を分散させて計算している。

左は4チップに分割して推論計算を行う例。右はチップ内の16PEに並列計算を行わせる例である

最終的にできたのは次の図のようなものである。中央の写真がMCMで36個のチップが載っている。右側の写真は1個のチップを拡大したものである。16個のPEの左にグローバルバッファとRISC-Vが置かれている。

そして、4つの角に2つずつ置かれているのがGRS信号の送信と受信部で、メッシュネットワークを構成する東西南北の送受信回路がある。上辺の中央はGPIO、可変の中央はJTAGとなっている。

そして、GRSの送受信は予定通り100Gbpsを達成し、コンピュートタイルは128TOPS、9.5TOPS/Wの効率を達成した。また、チップのコア部の面積は111.6mm²で、電源電圧は0.52～1.1Vで動作し、動作周波数は0.48～1.8GHzである。

そして、ターゲット仕様が決まってからテープアウトまでの設計期間は、10人以下のエンジニアで6カ月で行うことができた。これは驚異的に効率の高い設計である。

左の図がパッケージに36チップが搭載された写真。右はチップのダイ写真である。チップ間通信のGRSリンクは100Gbps/laneで動作した。128TOPSのシステムを作った時、9.5TOPS/Wのエネルギー効率。高生産性設計で、10人以下のエンジニアで6カ月で設計できた

高効率設計を行うため、ハードウェア設計の抽象度を高めた。論理設計はVerilogではなくC++で記述した。そして、C++の設計からハイレベル合成ツールで論理合成を行った。また、既存のライブラリの利用やジェネレータを使用して設計作業を低減した。

そして、小さな設計チームでアーキテクチャから実装、VLSI設計までを行った。並行して、継続的に自動化したツールフローの組み込みを行った。さらに、アジャイルなプロジェクトマネジメント技法を使うことにより、C++の設計からレイアウトまで24時間で行えるようになった。

VerilogではなくC++言語で設計を記述し、ハイレベル合成ツールを使ってLSIを設計した。また、できる限り既存のライブラリなどを利用した。そして、小さなチームでアーキテクチャからVLSIまで一貫して設計を行った。自動化されたツールフローの組み込みを継続して行い、C++からGDSまでを24時間で行えるようになった

次の図は、自動運転のDRIVENETが弱スケール(ハードウェアの性能向上に比例して仕事も増やした場合は実行時間は一定になる)するかを調べたもので、左の図は、チップ数を1、4、32と変え、バッチ数はチップ数に比例して増やした場合のスループットを示している。オレンジ色の折れ線グラフはレーテンシを示しており、チップ数が増えるとわずかにレーテンシが増えるものの、ほぼ一定であり、弱スケールしているといえる。

右の図は、各コアの消費エネルギーをプロットしたもので、青で示したコアエネルギーは若干の増加で、弱スケールしているといえる。一方、オレンジ色で示したGRS部分のエネルギーはチップ数が増えると増加しているが、これはクロックゲーティングを改良することで改善できるとのことである。

左はチップ数を1、4、32とした場合のDRIVENETのスループットの変化を示す棒グラフで、オレンジ色の折れ線はレーテンシを示す。右はDRIVENETを画像認識の消費エネルギーを示す棒グラフである。レーテンシはほぼ一定でありWeak Scalingになっている。また、画像1枚当たりのコアエネルギーもほぼ一定であるが、GRSの通信エネルギーが増えている。これはクロックゲートを改善すれば低減できるとのことである

ResNet-50でのイメージ認識の場合は、32チップを使うとレーテンシは約1/12となっており、1/32とはかなり差があるが、まあ、強スケール(ハードウェア性能をn倍にすると、同じ仕事をする時間が1/nになる)に近い性能改善が得られている。完全な強スケールにならない主因はチップ数が増えると通信オーバヘッドが増えるためであるとのことである。

左は1、4、32チップでのResNet-50での画像認識のスループット、右は消費エネルギーを示す棒グラフである。チップ数を増やすとレーテンシもそれに逆比例して短くなる。通信オーバヘッドなどの影響で完全な強スケールにはなっていないが、まあ、それに近い結果である

まとめであるが、NVResearchのテストチップは、スケーラブルな推論アクセラレータで、MCMを使うことにより、異なる市場向けの推論エンジンが作れることを示した。1パッケージに36チップを搭載し、ResNet-50で2.5Kイメージ/秒のスループットとバッチ=1で 0.4msのレーテンシを達成した。高効率設計では、ASICの設計と検証の工数を1/10に削減し、製品の市場投入期間を短縮し、あるいはより多くの機能をチップに盛り込むことを可能にした。