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

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

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

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

• SC19でCerebrasのブースに展示されたWS-1システム

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

ウェハスケールの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個の電源が故障しても運転が継続できると思われる。

• WS-1の前面ドアを開けたところ。下の大きな取っ手のついたパンチングメタルの奥に冷却水を冷やす4つの大きなファンがある。左上の16個のモジュールは電源で右上の2つの大きな箱は、水冷のためのポンプである

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

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

• SC19でのExhibitor ForumでCS-1を発表したCerebrasの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のコアやインタコネクトについて公表していないが、疎な計算を効率的に処理できると書かれており、グラフコアなどと類似のグラフ処理を行っていると思われる。

• 12インチウェハから1個採りのWafer Scale Engine。215mm角のチップに1200Bトランジスタを集積する。搭載コア数は40万コア、メモリは合計18GBのSRAMを搭載している (出典:このレポートのすべての図は、SC19におけるCerebrasのJ.P.Fricker氏とA. Hock氏の発表スライドを筆者が撮影したものである)

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

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

• SC19での発表の中でWafer Scale Engineを見せるHock氏

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

• CerebrasのソフトウェアはTensorFlowやPyTorch、Caffeなどのフレームワークで書かれたニューラルネットからカーネルグラフを抽出し、それをWSEに割り当てて実行させる

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

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

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

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

• 通常はスクライブラインの領域にダイの間を接続する配線を作っている。I/OはWSEチップの周辺に配置し、PCBとWSEを接続する専用のコネクタを自社開発した。電源は表面から供給し、排熱は裏面から行っている

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

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

• 冷却の構造を示す図。青の矢印が空気、水色の矢印が水の流れを示す。WSEチップはポンプユニットの水色の矢印の部分に水平に置かれていると考えられる。冷却水はWSEを冷やして温まり、ファンの後ろの上向きの水色矢印の部分の熱交換器で冷却されて循環している。熱交換器はファンからの風で冷却される

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

• 全体の制御とツールを動かすローカルCPUは筐体の左上の部分の中央付近に置かれている。外部との通信を行うTCPはCPUは制御だけを行い、データはCPUを通らずネットワークインタフェースとWSEの間で直接転送される