HotChips 2019の様子がTwitterで流れてきた。やはり一番大きなインパクトがあったのはCerebrasのWafer Scale Engineではないだろうか。 性能、機能性、意味はとりあえずおいておいて、そのインパクトはかなり強力だった。
発表資料や論文など手に入らないかと探していると、いろいろと出てきた。
- HotChipsの発表に関してCerebrasのブログ :
- HotChipsの発表資料：
- ホワイトペーパ Cerebras Wafer Scale Engine: An Introduction
今回は、内容的に一番情報量の多そうなホワイトペーパを読んでみた。
結論からすると、Cerebras WSEの重要なポイントは以下の3つだ。
- 1チップにより多くのコアを詰め込む。
- 1チップにより多くのオンチップメモリを詰め込む。
- オフチップ通信をなるべく除去し、オンチップ通信に収める。
このアプローチは上手くいけば一応効果はありそうに思える。なぜならば小さなチップを組み合わせてオフチップ通信を行うならば、どうしても高速I/Oは必要で、そのためのIPを組み込むと面積だって馬鹿にならない。 実際発表資料を見ているとウェハレベルといいつつスクライブラインは見えるし、スクライブラインを超える通信は近距離ワイヤでつなぐらしい。 これだけのサイズのマスクなど製造できないだろうし、ワンショットで転写できる最大サイズでユニットを構成し、その間にスクライブラインを引いたのかな。
凄い技術だが、強力なブレークスルーにはならない気がする。上記のメリットに対して、組み立て、熱、ソフトウェア、制御、歩留まり、などの乗り越えなければならない問題が多すぎる。
Cerebras Wafer Scale Engine: An Introduction
1. 速度の必要性
ディープラーニングにおいて、計算速度は非常に重要であり、高速な計算機が必要となってくる。そのなかで、Google, Facebook, Baiduなどは学習時間がAIの長い学習時間がAIの進歩の足かせになっていると指摘している。 ここを解決するために、CerebrasはWafer Scale Engine(WSE)と呼ばれる超大規模なウェハレベルのチップを製造した。 このチップは46,225平方ミリメートルで、1.2兆個以上のトランジスタを内蔵している。 比較として、WSEは最大のグラフィック処理ユニットより56倍以上大きく、オンチップメモリは3,000倍、メモリ帯域幅は10,000倍以上を誇っている。
WSEは、最も複雑なAI作業にかかる時間を数か月から数分に短縮できる。
2. ディープラーニングのワークロード
まずはWSEを設計するにあたり、ディープラーニングのワークロードを定義する。ディープラーニングのベースとなるニューラルネットワークは、入力データからの出力データへ、パラメータを行列に掛け合わせるなどの、単純かつ高度に並列な操作から構成されている。 学習中は、データは入力と出力の間で双方向に流れる。学習中は出力ノードからエラーの情報が逆方向に伝搬され、パラメータの値が調整される。 ニューラルネットワークの学習に必要な時間は、このフィードバックループを介して入力を処理できる速度が重要となる。 入力がネットワーク内を高速で移動できるほど、単位時間あたりより多くの入力がループを介して処理することができ、ネットワークの学習が高速になる。
3. 性能の理解
ディープラーニングの学習時間を短縮する方法は、入力のフィードバックにかかる時間を短縮すること。つまり、計算速度を向上させる、そして通信速度を高速化させる、ということになる。
まず、計算時間の高速化は演算コアを増加させることにより達成できる。特に学習の場合は浮動小数点演算ユニットを増強することで計算を乞うおs区化させる。 次に、通信速度は帯域幅の増加、遅延の削減により達成できる。これにより、ループ内のステージの結果が後続のステージに迅速に伝達できる。通信速度の問題は非常に重要で、すでにAIの計算に必要なコアの数は、ダイに収めることのできる演算器の数を超えてしまっている。そこでNVIDIAはNVLinkと呼ばれるプロトコルを開発したり、を買収するなどして通信機能の強化を図ってきた。
今回のCerebrasのアプローチは、小さなチップを高速な通信方式を使って並べるのではなく、大きなウェハ上にチップを構築する。これにより以下のメリットがある。
- 計算に使用できるコアの数が劇的に増加する。
- チップ面積に余裕があるため、ローカルメモリなどの独自の記憶装置を配置する。
- 通信はすべてオンダイ上で行う。オフチップ技術を使う必要がない。
(個人的な感想 : ただしオンダイとはいえ1ショットでは製造できないしスクライブラインがあるだろうから、結局ワイヤで接続する必要がありオンダイほどの通信性能は出ないのでは？)
このチップを設計するにあたり、「設計・製造・電力・冷却・通信・調整」などの問題が発生する。
4. Cerebrasアーキテクチャ
Cerebrasアーキテクチャの3つのキーを説明する。この3つのキーにより、WSEは高い性能を達成する。
- コア
- コアに近いメモリ
- コア間帯域幅
4.1 シリコン領域が増えると、計算コアの面積が増える
CerebrasのWSEのコアはSLAと呼ばれている。コアはキャッシュを持たず、汎用命令(制御系命令)、線形代数命令・疎密度な線形代数命令などを備える。また、テンソル命令などを備える。
SLAの特徴としては、疎密度な線形代数に対する処理の最適化を実施している点である。疎密度な行列に対する最適化パスや命令を持っていることで性能を向上させ、消費電力を削減する。
4.2 オンチップメモリにより多くのシリコン面積を割く
ディープラーニングにおいてなぜオンチップメモリが重要なのかというと、モデルのパラメータ・アクティベーション・モデルの構成などを保持するため。このためのメモリをオンチップに持ち、高速に読み書きすることが必要になる。
これまでのチップでは大容量高速メモリは大部分がオフチップであり、HBMなどの技術を使っても非常に低速なままである。
そこで、CerebrasのWSEには18GBにもおよぶオンチップメモリを搭載し、9.6バイト(注:これはTypo?)のバンド幅を持っている。
4.3 シリコン上での高速ネットワーク : Swarm Fabric
Cerebras WSEの400,000のコアは、1秒あたり100Peta Bitの帯域幅を持つ2DメッシュのSwarm通信Fabricを介して接続されている。 Swarmは、各計算コアにハードウェアルーティングエンジンを内蔵し、レイテンシと帯域幅に最適化された短いワイヤでそれらを接続する。
また、このFabricはコンフィギャラブルで、ソフトウェアで制御することによりAIのトレーニング中に正確な通信を実現できる。
- Cerebras Wafer Scale Engineのソフトウェアスタック
WSEのソフトウェアスタックは、WSEの設計時からハードウェアと同時に開発された。 Cerebrasのグラフコンパイラは、TensorFlowやPyTorchなどのMLフレームワークとのシームレスなインタフェースを持っており、ニューラルネットワークのデータフローを抽出し、これらをマイクロコードプログラムに変換し、WSEのリソースにマッピングする。そしてSwarm Fabricを構成し、データパスを形成して性能を最大化させる。
と、良いことはいっぱい書いてあるが、すべてに同意できるわけではないし、16nmで歩留まりはどれくらいになるのだろう。
アーキテクチャとしては少し古臭く感じる。疎密度の線形代数とか、最先端のディープラーニングのアルゴリズムではもう時代遅れなのでは？
どちらにしろ、彼らがどのようにしてこの困難なチップを実用化していくのか、とても楽しみだ。
