Cerebras Systemsは300mmウェハの中で取り得る215mm角という最大サイズのチップ(chipの語源から、これをチップと呼ぶのは違和感があるが)のAIアクセラレータ「Wafer Scale Engine(WSE)」を開発した。次の写真で、右側のAndy Hock氏が持っているのがそのチップである。

2019年末、東京エレクトロン デバイス(TED)が代理店契約を結び、CerebrasのAIアクセラレータである「CS-1」を日本でも販売するという発表を行うため、Cerebrasの幹部2人が来日した際に、詳しく話を聞く機会をいただいた。

  • WSE
  • WSE
  • Cerebrasのウェハスケールエンジン(WSE)を搭載するCS-1システムの発表のために2019年末に来日した、同社の創立者の1人でチーフテクニカルオフィサー(CTO)のGary Lauterbach氏(左)とプロダクトマネジメントディレクタのAndy Hock氏(右)

世界最大サイズのチップは誰が考えたのか?

まず、Cerebrasという社名は、どのように決められたのかと聞いたところ、Cerebral Cortex(大脳皮質)から取ったとのことで、人間の大脳に匹敵するようなAIアクセラレータを作ることを目指しているという。

次の図に、Wafer Scale Engineの諸元とNVIDIAのV100 GPUとの比較を示す。チップサイズは56.7倍、コア数は78倍で、オンチップのメモリ容量は3000倍、メモリバンド幅は1万倍、ファブリックのバンド幅は3万3000倍と桁外れの規模である。

  • WSE

    CerebrasのWafer Scale Engineの諸元とNVIDIAのV100 GPUとの比較を示したスライド (出典:2019年12月に開催された記者発表会でのHock氏の発表スライドを撮影したもの)

Wafer Scale Engineという巨大チップを誰が考えたのかと質問すると、CTOのLauterbach氏が、自分が提案したと答えてくれた。Lauterbach氏は、創立者でCEOのAndrew Feldman氏が、巨大チップはリスクが大きいと反対するだろうと思っていたのであるが、その案が通ってしまったという。

  • WSE

    WSEの実チップ。100円玉と比べると、その大きさが良くわかる

巨大チップを作る上での問題

巨大チップを作る場合、一般に認識されている問題は、歩留まりと熱膨張係数のミスマッチである。

WSEは40万コアもの回路を集積しているので、不良が全く無い完全良品が採れるとは考え難い。そのため、一般的に取られる対策は、予備的に冗長の回路を作り込んで置き、不良のコアは予備の良品コアと繋ぎ換えて、見かけ上、不良の無いウェハスケールのチップを作り出すという方法である。WSEの場合は、同じ回路が沢山繰り返されているので、比較的少ない予備回路を置くだけでよいので、作りやすい。

もう1つの問題は、シリコンチップの熱膨張率は比較的小さいが、プリント基板の熱膨張率は何倍も大きく、温度が変わると、チップの端子とプリント基板の接触が確保できないという問題である。

不良コアの置き換えで歩留まりの問題を解決

ちょっと話が変わるが、CerebrasのCEOを務めるAndrew Feldman氏は多数の小型CPUを使うサーバを作ったSeaMicro社のCEOで、Lauterbach氏もSeaMicroのCTOであり、プロセサネットワークの専門家である。この技術蓄積がWSEのチップ内ネットワークに上手く生きていると思われる。

なお、SeaMicro社はその後AMDに買収され、そのネットワーク技術は、AMDのEPYCサーバのプロセサチップ間のネットワーク接続技術として使われている。

そして、そのネットワーク技術はWSEにも使われているのではないかと思われる。WSEには40万個のAIコアを集積しており、それらを2次元メッシュのネットワークで接続している。さらに、一部のコアに不良があっても、それを避けて不良の無いコアだけを使って、2次元メッシュネットワーク接続を作れば、見かけ上不良の無い良品のWSEチップを作ることができる。

電源がショートしたコアがあったとしても、1コアだけの電源ショートで流れる電流は僅かであり、チップ全体のトランジスタのリーク電流と区別がつかず、WSEとしては致命的な不良にはならない。

WSEは約1.5%のコアを不良コアの置き換えのために作り込んでおり、これを使って大部分のWSEウェハを良品とすることができるとのことである。

WSEチップはTSMCの16nmプロセスで製造されているが、WSEのような巨大チップをテストするのは大変である。このため、CerebrasはノーテストでTSMCの製造したウェハを受け入れているという。

通常のチップの場合はウェハからチップに切り分けるためのスクライブラインという領域が取ってあるが、WSEではその領域にも配線を作り、隣のチップとの接続を行っている。これには通常とは異なるウェハ処理が必要となるが、TSMCは喜んで引き受けてくれたという。Cerebrasの幹部はSeaMicroの時代にTSMCとは良い関係を築いており、その時代の信用が役立っているという。そして、WSEチップは、TSMCのHall of Fameに飾られているという。

熱膨張でずれるチップの端子位置

そして、このような巨大チップを使う上でのもう1つの大きな問題は、熱膨張率の違いである。シリコン基板の熱膨張率は2.4×10-6/℃であり、50℃の温度差があればチップの対角線の長さである約300mmでは、36μm伸びる。しかし、WSEを搭載するFR4プリント板は熱膨張率が大きく、200μmあまり伸びてしまう。これは高密度のコネクタでは隣のピンに繋がってしまうような大きなずれが生じてしまいコネクタが壊れるという事態にもなる。

次の図はCerebrasの特許に書かれているものであるが、シリコンゴムのような弾性のある材料に微細な金属球を混ぜたシートを使ってコネクタを作っているようである。

弾性体に圧力が掛かっていない状態では、金属球同士は離れていて電流は流れない。しかし、圧力が増え、金属球の密度が高くなると、金属球同士が接触して電流が流れるようになる。

プリント板側に適当な厚みの金属パッドを作っておけば、パッドの部分には大きな圧力が掛かるので弾性体シートは導通する。一方、パッドが無い部分では圧力は小さく弾性体シートは導通しない。このようにして作られるコネクタでは、異なる熱膨張によって、シリコンチップ上の端子の位置とプリント基板上の配線パッドの位置が多少ずれても接触を保つことができる。

  • WSE

    CerebrasのUS 10,468,369 B2特許の図。金属球を含んだ弾力性のある材料を使うコネクタ

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

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

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は、現在の初期の顧客の抱えている問題を処理する実行性能を改善するのに全力を注いでいて、現在はベンチマークなどの性能を上げることには力を割く余裕は無いとのことである。

WSEの学会発表は2020年Q2~Q3の予定

これまで、Cerebrasはアーキテクチャの詳細や性能ついては発表していない。しかし、2020年のQ2~Q3ころには詳しい学術論文を出し、詳細を発表するつもりであるという。それまで、実行性能の公表はお預けになりそうである。