ウェハスケールAIエンジンを搭載するCerebrasのAIサーバ「CS-1」 - SC19

ウェハスケールAIエンジンを搭載したサーバが登場

ウェハスケールのAIエンジン(WSE)を作って度肝を抜いたCerebrasが、デンバーで開催されたSC19で「CS-1」というWSEを使うAI処理用のサーバの発表と展示を行った。

何しろ、20kWの消費電力のシステムであるので、それなりの大きさになり、次の写真に見られるように15U(高さ667mm)というサイズになっている。

水を噴射してチップを冷却

ウェハスケールのAIエンジンの消費電力は15kWで、CS-1全体の消費電力は20kWである。この発熱に対し、水をウェハスケールのチップの裏側に垂直に噴射して冷却しているという。

次の写真はCS-1の前面の扉を開けた写真である。右上の部分に見えるいかつい大きなハンドルが付いているのが水冷のポンプユニットである。2つのポンプユニットになっているのは、ポンプが故障して水の循環が止まると短時間でWSEが過熱してしまうので、ポンプを二重化しているためであると思われる。

WSEを冷却して温まった水を下側にある4つの大きなファンで冷却して循環させている。この部分も4個のファンで作られているので、1個のファンが故障しても問題にならない造りになっているのであろう。

そして、電源も12個に分かれている。これらの電源で15kWを供給してるとすると、1個当たり1.25kWであり、これは0.8V程度で1500Aとなり、1台の電源で供給できる最大の値である。また、12個の電源を並列に使っているので、1個、2個の電源が故障しても運転が継続できると思われる。

12インチウェハをまるまる1チップとして活用

SC19のExhibitor ForumでWSEとCS-1を発表したのは、同社の創業者の一人でチーフシステムアーキテクトのJ.P.Fricker氏とプロダクト担当ディレクタのA.Hock氏である。

Cerebrasのシステムの特徴は、12インチのウェハから1個採りの215mm角という巨大なシリコンチップを使っている点である。もちろん、この巨大サイズでは不良のコアが混じるのは避けられないが、40万＋α個のコアを搭載しているので、不良のコアを避けて配線してやれば、高い歩留まりで良品のWafer Scale Engine(WSE)が採れる。

40万個のコアを集積しているので、全体では18GBのSRAMが載っている。これは1コアあたり45KBのSRAMがあるという計算になる。そして、WSE全体では100Pb/sのインタコネクトがある。これは1コアあたり250Mb/sのインタコネクトがあるという計算になるが、インタコネクトのトポロジが分からないので、これがどの程度意味がある数字であるかは分からない。

しかし、これらの数字は他のシステムで言えばクラスタで実現するレベルの計算資源量であり、WSE 1個で複数個のマシンラーニングモデルを並列に実行することができる。したがって、学習時間を短縮することができるし、入力のバッチが1であってもコアを有効に利用できる。

CerebrasはWSEのコアやインタコネクトについて公表していないが、疎な計算を効率的に処理できると書かれており、グラフコアなどと類似のグラフ処理を行っていると思われる。

パッケージングはすべて自社開発

SC19の発表ではHock氏はWafer Scale Engineを掲げて見せた。Cerebrasの社内では、すでに何個ものWSEが動いており、顧客のジョブを実行しているという。

Cerebrasのソフトウェアは、フレームワークでの記述からLAIRというネットワークを抽出し、そこからカーネルグラフに変形する。このカーネルグラフをどのように実行するかというプランを作り、これをCS-1で実行する。

このようなウェハサイズのLSIのパッケージングはこれまでに例が無く、すべて自社で開発する必要があった。

そのため、通常は小さなチップに切り分けるために設けるスクライブラインの領域にも配線を作って、チップ間の接続を行っている。そして、オンチップのSRAMは1クロックでアクセスできるようにしたという。ただし、WSE全体に散らばる18GBのSRAMすべてを1クロックでアクセスできるとは思えず、どのような構造になっているかが分からないと何を意味しているのか分からない。

そして、外部へのI/OはWSEチップの周辺に配置し、PCBとWSEをつなぐ専用のコネクタを自社開発した。

また、WSEの表面から電源を供給し、裏面から冷却(廃熱)を行う構造となっている。

CPUを経由しないTCPネットワークのデータ

筐体内部の冷却については、次の図が示された。前面から4つの大型ファンで空気を吸い込み、その後ろにある熱交換器でWSEを冷やして温まった水を冷却している。そして、冷却された水は、筐体上部に水平に置かれたWSEを冷却して熱交換器の下部に戻る。このあたりの細かい配管は描かれておらず、この図からは詳細な水の流れは分からない。

そして、ポンプユニットの左側の電源の上にローカルのCPUが置かれている。このスライドには、TCPネットワークのデータはローカルのCPUを経由しないと書かれており、通信に関してCPUが行うのは制御だけで、大量の入出力データはネットワークインタフェースユニットからWSEにDMAで転送されていると思われる。