【ITニュース解説】進化し続けるWSE（Wafer Scale Engine）、大規模AIモデルのトレーニング性能とは――Cerebras社訪問2024【前編】

Cerebras Systems社は、WSE（Wafer Scale Engine）という画期的なAIプロセッサを開発し、その技術を基盤としたシステムCS-2（最新版はWSE-3を搭載したCS-3）を提供している。WSEとは、一般的なプロセッサが小さなチップとして製造されるのに対し、シリコンウェハー全体を一枚の巨大なプロセッサとして利用するという、非常にユニークなアプローチを取っている。この技術は、大規模なAIモデルのトレーニング性能を劇的に向上させることを目指している。

AI技術の進化に伴い、開発されるAIモデルはますます巨大化している。ChatGPTに代表されるような大規模言語モデル（LLM）などは、何兆ものパラメータ（学習によって調整される数値）を持ち、そのトレーニングには膨大な計算資源と時間が必要となる。従来のCPUやGPUといったプロセッサは、通常、複数の小さなチップを並列に接続して大規模な計算に対応するが、この方式ではチップ間のデータ転送速度がボトルネックとなり、計算効率が低下するという課題があった。データがチップからチップへと移動するたびに時間がかかると、せっかく高性能なチップがたくさんあっても、その能力を最大限に引き出せなくなるのである。

Cerebras社のWSEは、この課題を根本的に解決するべく考案された。シリコンウェハーとは、半導体チップを作る際の材料となる薄い円盤状の板のことだが、WSEはこのウェハーの全面に計算ユニットやメモリ、そしてそれらをつなぐ高速なネットワークを一体として構築する。これにより、複数のチップを繋ぐ際の通信遅延をほぼゼロにし、全ての計算リソースが極めて高速に連携できるようになる。これは、まるで一つの巨大な頭脳が、その内部で一切の滞りなく思考するようなものだ。

WSE-3は、このWSEの最新世代であり、これまでの設計で培われたノウハウが凝縮されている。WSE-3の設計方針は、さらなる計算能力の向上と、効率的なデータ処理に重点を置いていると考えられる。例えば、コアと呼ばれる計算を担当する部分の数を増やし、それぞれのコアがより高速に動作できるよう改良したり、AIの計算で頻繁に使われる特殊な処理を効率良く実行するための回路を組み込んだりすることが考えられる。また、大規模AIモデルは、計算の途中で多くのデータが必要となるため、WSE-3ではプロセッサ内部に大容量かつ高速なメモリを統合することで、外部メモリへのアクセスを最小限に抑え、データ転送のボトルネックを解消している。さらに、電力効率も重要な設計要素だ。巨大なプロセッサであるWSEが消費する電力をいかに抑え、発熱をコントロールするかは、システム全体の安定稼働に直結するため、非常に高度な技術が投入されているはずである。

WSE-3を核とするシステムがCS-3である。CS-3は、単一のWSE-3プロセッサを搭載し、それ単体で非常に強力なAIトレーニング能力を持つ。しかし、AIモデルがさらに大規模化していく現代において、一つのプロセッサだけでは対応しきれない場合も出てくる。そこでCerebras社は、複数のCS-3システムを連携させる「クラスタシステム」を構築している。これは、複数のCS-3を高速なネットワークで接続し、全体として一つの巨大なスーパーコンピュータのように機能させる仕組みである。このクラスタシステムにより、何兆ものパラメータを持つような、現在の最高峰のAIモデルでも効率的にトレーニングすることが可能になる。各CS-3が連携してデータや計算結果をやり取りすることで、まるで一台の超巨大なWSEが稼働しているかのような感覚で、AIモデルのトレーニングを実行できる。このようなシステムは、AI研究者や開発者が、計算資源の制約に縛られることなく、より大規模で複雑なAIモデルの開発に挑戦できる環境を提供する。

WSE技術は、これまでの半導体設計の常識を覆し、AIモデルのトレーニングにおける性能と効率の限界を押し広げている。単一チップとしては世界最大級のプロセッサを開発し、それを複数連結させることで、今日のAIが直面する計算需要に柔軟に対応する。このアプローチは、今後のAI研究開発の加速に大きく貢献し、AIがさらに社会の様々な分野で活用される未来を切り拓くだろう。