高速ディープラーニングが可能なWSEをCerebrasはどうやって実現したのか

ネットワークのコア利用率を高める工夫

オンチップのメッシュネットワークで送受するメッセージのサイズは非常に小さいという。このため、コア間の接続を頻繁に切り替え、通信の自由度を高めてコアの利用率を高めているという。ネットワークはパケットスイッチとサーキットスイッチの混合というようなアーキテクチャになっているという。

CerebrasのWSEは処理単位が小さく、コアを遊ばせずに使うことができるので、性能が高いという。通常のGPUベースのシステムでは学習時のコアの利用効率は30%程度であるが、WSEのコア利用効率はその2倍以上になるという。

また、GPUなどを使うAIアクセラレータでは、1つのネットワークで多数の入力を纏めて計算するサイズの大きなバッチ処理を行うことが多い。大きなサイズのバッチ処理を行うと、バッチの中では、ネットワークの接続や重みの値を変えずに次々と入力だけを変えて計算すれば良いので、効率が良い。しかし、それぞれの入力に対する中間結果を記憶しておく必要があり、そのために大量のメモリが必要となる。このためにメモリを必要とするので、大きなネットワークをチップ内に記憶して処理することができないという問題が出てくる。

これに対してWSEは大量のコアとオンチップのメモリを持っており、大規模なニューラルネットをオンチップで扱うことができ、他社の製品のように、大きなニューラルネットを分割して処理するオーバヘッドが発生しないという。

WSEの隣接コア間の通信は1クロックサイクルである。そして、WSE全体ではおおよそ600×600のコアがあるので、対角線の位置にあるコア間の通信では、おおよそ1200クロックかかることになる。

1つの計算に使うすべての入力が揃わないと計算が行なえないので、入力が揃うのを待ち合わせる必要がある。そのためのデータバッファとして、1GBを超えるメモリを持っているとのことである。ただし、このメモリはデータフロー処理を実現するためのものであり、18GBのデータメモリとは別に存在するという。

高効率を実現したデータフロー処理

GPUやGoogleのTPUでも重みや入力がゼロの場合は計算を省いて、消費電力を抑えることはできる。しかし、WSEでは演算を必要としないコアに別の入力を与えて、有効な計算をやらせることができる。

したがって、重みや入力がゼロの部分は計算を省くようにネットワークを作れば、全要素が非ゼロであると想定して固定の順序の行列演算を行うGoogleのTPUなどと比べて実行できる計算量を大幅に増やすことができる。そして、チップ全体で見ればコアのビジー率が上がり、演算効率を上げることができる。

このため、WSEはチップに搭載されているコア数に比べて高いAI演算性能を持っているはずであるが、Cerebrasは　MLPerfなどのベンチマークの性能を公表していない。Cerebrasは、現在の初期の顧客の抱えている問題を処理する実行性能を改善するのに全力を注いでいて、現在はベンチマークなどの性能を上げることには力を割く余裕は無いとのことである。