Cerebras Systemsは2016年にAndrew Feldman氏とGary Lauterback氏、Michael James氏、Sean Lie氏、Jean-Philippe Fricker氏の5人によって創立された。

この5名の創立者はSeaMicroという多数の小規模サーバをネットワークで接続して性能の高いサーバを作るという開発を行っていたが、SeaMicroは2012年に3億3400万ドルという高値でAMDに買収された。当初はFeldman氏等はAMDに移ったが、2016年に昔のSeaMicroの5人の仲間で独立しCerebrasという会社を創立した。

なお、AMDはSeaMicroの技術をもとにInfinity Fabricを開発し、現在では同社のチップ間ネットワークとして広く活用している。

一方、Lauterbach氏は、サーバやワークステーションの大手のSun MicrosystemsのCPUエンジニアでUltraSparc IIIチップのチーフアーキテクトを務めた。Sun Microsystemsは大手であったが、現在ではOracleに吸収されて無くなってしまった。

新しく創立された会社の名前は「Cerebras Systems」という。この名前はCerebral Cortex(大脳皮質)から取られており、マシンラーニングのアクセラレータの開発を狙った新会社の旗揚げで有る。

国内では、東京エレクトロンデバイスが代理店となってCerebrasのAIアクセラレータ「CS-2」の取り扱いを行っている。CerebrasのCS-2はシリコンウェハ1枚をまるごとチップにしたという他社に例を見ないものであり、巨大チップの使用により、性能や消費電力低減などの点で大きなメリットが出ている。特に、他社のシステムでは扱いが難しい巨大モデルの学習に強みを持っている。

しかし、価格/性能ではCS-2は優れているのであるが、絶対値の価格ではかなりの金額であり、ユーザが手を出しづらいのも事実であろう。そこで、東京エレクトロンデバイスは横浜市都筑区にAIラボを開設し、CS-2を設置して顧客がCS-2を使ってみることができる環境を用意した。

CerebrasのAIエンジンの特徴は、1枚の300mmウェハに載せられる最大のサイズのチップ(約215mm角)を作っている点である。巨大なチップを作ればチップ間の配線を短くでき、寄生容量が減り、動作クロックを速くできたり、消費電力を減らしたりできるというメリットがある。ということで、大きなチップを作るというアプローチは以前から例があり、Trilogyという会社は63.5mm角のECL LSIを開発しようとした。しかし、当時の半導体製造技術では十分な歩留まりが得られず、Trilogy社は1984年にこのLSIの製造に見切りをつけてしまった。

Cerebrasのような巨大なチップを作るとき、一番難しいのは如何にして熱膨張を吸収するかという点である。Cerebrasのウェハ上の2つの点の最大距離は約300mmである。ガラスエポキシ基板の熱膨張率は2.1×10-5/℃程度であるので、2点の最大温度差が40℃とすると、約0.252mm伸び縮みすることになる。これは高密度のコネクタでは隣のピンにつながってしまう程度の大きなズレである(シリコン基板の熱膨張率は小さいので無視している)。これだけ距離がずれると、ガラスエポキシ基板側のコネクタ端子のはんだ付け部にストレスがかかり壊れてしまう。

Cerebrasのホワイトペーパーには解決法は書かれていないのであるが、Cerebrasの米国特許(US 10,468,369 B2特許)を見ると、金属粒子をシリコンゴムのようなエラストマに分散して混ぜ込んだ材料に触れられている。このような材料は通常は絶縁体であるが、圧力をかけて金属粒子の密度を高めると、金属粒子が接触して導通するという材料を作る事ができる。Cerebrasは、このような材料を使って熱膨張を吸収するフレキシブルなコネクタを開発したようである。

マシンラーニングの計算の大部分は積和演算で、計算自体は比較的単純である。LSIの集積度が低い場合は、積和演算の入力の前段からの接続や積和演算の出力の次段への接続にも接続の自由度が高く遅延の大きいスイッチが必要で、信号を接続するスイッチ部分の遅延時間が馬鹿にならず、ML(機械学習)アクセラレータのクロック周波数が下がり、計算性能が低下してしまう。

これに対して、WSI(Wafer Scale Integration)のように演算器の個数が十分にある場合は、積和演算を何段も積み重ねるスイッチのロスの少ない大きな計算ネットワークを作る事ができる。

Cerebrasはこの巨大LSIというアプローチに挑戦して、Wafer Scale Engine(WSE)を作り上げた。同社のホワイトペーパーに書かれているように、通常のLSIの製造方法では900mm2程度のLSIが露光できる最大のチップである。そして、ウェハに数十個のこのチップを露光する。さらに、ウェハからチップを切り離すためのスクライブ領域をチップの周りに設ける。

しかし、 Cerebrasは隣接したチップ間を接続する配線をスクライブ領域に作ることにした。

LSIを製造する台湾のTSMCにとってはやったことのない製造法であったが、Feldman氏のチームはInfinity Fabricのチップ製造で付き合いがあり、信頼できるチームであることから、スクライブ領域に隣接チップ領域間を接続するパターンを作ることを承諾してくれたという。もちろん、ウェハスケールのチップはTSMCにとっても誇示できる製品で、現在では同社の製品展示室に展示されているという。

自然言語処理を中心的なアプリケーションとするTransformer系AIモデルの精度を改善するため、図1に示すように、2年ごとに1000倍というペースで、AIモデルのサイズは急速に大きくなっている。

  • AIモデルは2年ごとにサイズが1000倍になり、学習に必要な計算量も1000倍に増加している

    図1 AIモデルは2年ごとにサイズが1000倍になり、学習に必要な計算量も1000倍に増加している (Cerebras Systemsによる)

GPT-3モデルには175B(1750億)個のパラメタが必要であり、モデルを記憶するだけでも大量のメモリが必要である。さらに、学習のための3.2Mトークンの1つのバッチの記憶に7.5TBのメモリを必要とする。また、GPT-3の学習計算を行うには3.15×1023回の浮動小数点数の演算と2.8TBのメモリが必要と見積もられる。

大量の計算能力や大量のメモリが必要な場合は並列処理が用いられることが多い。図2の左側に並んだ3つの写真は、下から順にNVIDIAのA100 GPU、2番目はA100 GPUを8個使用するDGX A100サーバ、一番上はDGX A100を80システム使用するSuperpodの写真である。

  • NVIDIAのA100 GPUでデータ並列の処理を行なう場合の性能のグラフ

    図2 NVIDIAのA100 GPUでデータ並列の処理を行なう場合の性能のグラフ。540GPUで並列処理を行っても性能は84倍にしか向上しない

データ並列の処理では、全てのA100 GPUは同じプログラムを実行し、アクチベーションのデータ開始位置だけをずらせたデータを使用する。

図2は、このようなデータ並列処理を行った場合の並列動作のシステムの性能の伸びを示す図である。これによると、540台のA100 GPUを使ってもシステム全体の性能は84倍にしか向上していない。

データ並列の場合は、各GPUが担当する領域のサイズは固定で、何個のGPUを使うかでシステムの規模が決まる。したがって、規模の大きなGPT-3のようなモデルでは多数のGPUを使う並列度の高いモデルを使う必要がある。

しかし、データ並列の場合はそれぞれの領域ごとの計算を行い、それが終わると全ての領域のGradientを足し合わせて全体のGradientを求めるというReduction処理が必要になる。この計算は完全に直列に行う必要はないが、完全に並列には行うことはできないので、GPUの個数に比例した性能向上は得られない。

問題はReduction処理が必要になるだけではない。A100 GPUからReduction結果を別のA100チップに送って次のA100 GPUがReduction処理を行うので、チップ外のバッファを通して信号を他のチップに送るための遅延は、ウェハ内で別のブロックに信号を送る場合と比べて何倍も大きい。

そして、WSEの場合は集積度が高いので、ウェハを跨ぐ遅延の大きい信号伝送の回数を小さく抑えることができる。この為、通常のGPUに比べてWSEでは、AIエンジンの遅延時間を小さく抑えることができる。ただし、CerebrasはWSE、WSE-2ウェハスケールLSIの遅延時間を公表していないので、具体的な性能の違いは分からない。

また、データ並列の処理ではモデルに含まれるGPUの数によってFlops性能やメモリ性能などが決まってしまい、メモリ性能を満足させる構成とするにはFlops性能は余っているというような構成になってしまうこともある。

もちろん、データ並列ではなくAIモデルを縦に切ってパイプライン型のモデル並列やテンソル型のモデル並列とすることで無駄の少ない構成を作れる場合もあるが、一般にモデル並列でモデルの規模の変更に追従したモデルを作るのは簡単ではないことが多いという問題もある。

  • モデルが複雑になると構成を変えた場合のチューニングなどが複雑になる

    図3 データ並列以外にもパイプライン型モデル並列、テンソル型モデル並列、さらには、データ並列とモデル並列を組み合わせたモデルも作れるが、モデルが複雑になると構成を変えた場合のチューニングなどが複雑になる

これに対してCerebrasはWeight Streamingという方法を編み出した。

1回の学習を行うにあたってWeightをMemoryXからCS-2に送って計算を行う。最初は前進のためのActivationの計算を行う。2回目は後退のためのActivation Gradientの計算のためである。MemoryXとCS-2をつなぐのはSwarmXである。MemoryXは1つのweightの値をSwarmXに送り、SwarmXは全部のCS-2にその値をブロードキャストする。後退側ではMemoryXはSwarmXから一組のGradientsを受け取る。そして、SwarmXから部分的なGradientsを受け取り足し込んでいく。つまり、GradientsのreductionはSwarmXの中で行われる。このように計算を行って行くと、前回、MemoryXに書き出したWeightなどは、次の回には上書きされてしまい無駄なメモリを使う事が無く、より使用メモリ量の巨大な問題を学習することができる。

  • Cerebrasの巨大モデルの学習法「Weight Streaming」

    図4 Cerebrasの巨大モデルの学習法「Weight Streaming」。SwarmXと呼ぶチップはweightを全てのCS-2にブロードキャストする。また、CS-2からのGradientsをリダクションして纏める (Cerebras Systemsによる)

CerebrasのCS-2チップは、300mmウェハを1チップとした巨大LSIを作っており、一般のLSIに比べて無駄の少ない高性能のAIエンジンを実現している。また、CerebrasはWeight Streamingという学習方式を考案し、メモリ容量や演算性能の自由度の高い学習方式を考案した。そして、CerebrasのWeight Streamingは前進処理で書き込んだメモリを後退処理で上書きするという効率の高いメモリの使い方を行っている。このため、メモリの利用効率が高く、より大きなモデルの学習を行うことができるというメリットがある。