ウェハースケールエンジン（WSE）は、2015年にアンドリュー・フェルドマン、ゲイリー・ラウターバッハらによって設立されたCerebras Systems社が開発した一連の大規模プロセッサである。シリコンウェハから切り出される従来のチップとは異なり、WSEはウェハ全体を単一のダイとして使用し、前例のない数のコアとオンチップメモリを実現する。この設計は、マルチチップシステムを悩ませる通信ボトルネックを低減することを目的としており、大規模言語モデルやその他の深層学習ワークロードのトレーニングに特に適している。
第一世代のWSE-1は2019年8月に発表された。1.2兆個のトランジスタと40万個のAI最適化コアを搭載し、TSMCの16ナノメートルプロセスで製造された。チップのサイズは46,225平方ミリメートルで、当時の最大のグラフィックス処理ユニットの50倍以上である。第二世代のWSE-2は2021年4月に続き、7ナノメートルプロセスを使用し、コア数を85万個に倍増させ、オンチップメモリを40ギガバイトに増やした。2024年3月、Cerebrasは5ナノメートルプロセスで構築されたWSE-3を発表し、90万個のコアと44ギガバイトのオンチップSRAMを備え、最大24兆パラメータのモデルのトレーニングを対象としている。
アーキテクチャと設計
WSEの核となる革新は、ウェハスケール集積である。ウェハ上に多数の小さなダイを製造して個別にパッケージングする代わりに、Cerebrasはウェハ全体を無傷のまま維持し、製造欠陥を迂回するための冗長回路を備えている。このアプローチにより、分散システムでしばしば性能ボトルネックとなるプリント回路基板上のチップ間通信の必要性が排除される。チップはコアのメッシュ接続ネットワークを使用し、各コアは独自のメモリを持ち、高帯域幅で低遅延のデータ移動を可能にする。
各コアは、ニューラルネットワークで一般的なスパース線形代数演算に最適化された簡素化されたプロセッサである。コアはFP16、BF16、FP32を含むさまざまなデータ形式をサポートし、密な演算とスパースな演算の両方を効率的に実行するように設計されている。オンチップメモリはコア全体に分散されており、従来のGPUメモリシステムをはるかに超える総帯域幅を提供する。このアーキテクチャは、AIトレーニングの制限要因となることが多い外部メモリへの依存を低減する。
WSEは単体の製品ではなく、電源供給システムと冷却ソリューションを含むCerebrasのCS-2およびCS-3システムに統合されている。2020年にリリースされたCS-2はWSE-2を搭載し、2024年に発表されたCS-3はWSE-3を搭載している。これらのシステムは既存のデータセンターに接続できるように設計されており、大規模トレーニングのためにクラスターで使用されることが多い。
性能とベンチマーク
CerebrasはWSEのいくつかの性能ベンチマークを公開している。2021年、同社は単一のCS-2システムがGPTアーキテクチャを使用して18億パラメータのモデルをトレーニングでき、複数のシステムを接続した場合にほぼ線形のスケーリングを達成することを実証した。2022年、Cerebrasとサンディア国立研究所はCS-2を使用して2億ニューロンのスパイキングニューラルネットワークをリアルタイムでシミュレートした。これは数千のGPUを備えたスーパーコンピュータを必要とするタスクである。
2023年、Cerebrasはそのシステムが通信オーバーヘッドの低減により、16台のNVIDIA A100 GPUのクラスターよりも短時間で70億パラメータのモデルをトレーニングできると報告した。同社によると、WSE-3は最大125ペタフロップスのAI計算を提供でき、AI向けの最も強力なシングルチッププロセッサの1つとなっている。ただし、独立したベンチマークは限られており、ほとんどの性能主張はCerebras自身からのものである。
ソフトウェアエコシステム
WSEをプログラムするために、CerebrasはCerebras Software Platform（CSoft）を開発した。これには、機械学習モデルをチップのアーキテクチャにマッピングするコンパイラが含まれる。CSoftはTensorFlowやPyTorchなどの人気フレームワークをサポートしており、研究者はコードを大幅に変更することなく既存のモデルを実行できる。コンパイラはデータフロースケジューリング、メモリ割り当て、ウェハスケール設計の最適化を処理する。
Cerebrasはまた、トランスフォーマー、畳み込みネットワーク、リカレントネットワークを含む、事前最適化されたモデル実装のライブラリを提供している。ソフトウェアスタックには、ウェイトストリーミングと呼ばれる技術を使用して複数のCSシステムにわたる分散トレーニングのためのツールが含まれており、モデルのウェイトをチップのメモリにパイプライン処理する。このアプローチにより、外部ストレージからウェイトをストリーミングすることで、オンチップメモリよりも大きなモデルのトレーニングが可能になる。
アプリケーションとユースケース
WSEは、商業目的と研究目的の両方でさまざまな組織に採用されている。2021年、アラブ首長国連邦のG42グループはCerebrasと提携してAI研究用のスーパーコンピュータを構築し、後にJaisおよびFalconモデルのトレーニングに使用された。2023年、CerebrasはQualcommとオンデバイスAIモデルの開発で協力することを発表したが、この提携の詳細は乏しいままである。
医療分野では、Cerebrasは製薬会社と協力して創薬シミュレーションを加速している。大規模なグラフニューラルネットワークを処理するチップの能力は、分子特性予測に使用されている。さらに、WSEは天気予報にも採用されており、GPUクラスターと比較して高解像度気候モデルのトレーニングが高速化される実験が示されている。
他のAIハードウェアとの比較
WSEは、Groqの言語処理ユニットやSambaNovaの再構成可能なデータフローユニットなど、他の専門AIプロセッサと競合している。汎用並列プロセッサであるGPUとは異なり、WSEはAIワークロードに最適化されたドメイン固有アーキテクチャである。その主な利点は、大規模トレーニングで大きなオーバーヘッドとなる可能性があるチップ間通信の排除である。
ただし、WSEのサイズと消費電力（CS-2で最大15キロワット）は、適切な冷却を備えたデータセンターへの展開を制限している。対照的に、AWS TrainiumやGoogleのTPUはクラウドサービスとして提供されており、多くの組織にとってよりアクセスしやすい可能性がある。WSEはまた、AMDのInstinctアクセラレータやIntelのGaudiプロセッサとの競争にも直面している。これらはより従来的な設計であるが、成熟したソフトウェアエコシステムの恩恵を受けている。
課題と制限
WSEの主な課題の1つは歩留まりである。欠陥のないウェハの製造は困難であるが、Cerebrasは「冗長性」と呼ばれる技術を使用して、故障したコアを無効化し通信を再ルーティングすることで欠陥を許容している。このアプローチにより、同社は廃棄されるはずのウェハを使用できるが、チップごとに有効なコア数が異なることも意味する。
もう1つの制限はメモリ容量である。オンチップメモリは大きいが、GPUで利用可能な高帯域幅メモリよりも依然として小さい。大きな埋め込みテーブルやアテンションマトリックスを必要とするモデルの場合、WSEは外部メモリに依存する必要があり、性能上の利点が低下する。さらに、ソフトウェアエコシステムはCUDAほど成熟しておらず、カスタムカーネルの移植に困難を報告する研究者もいる。
CS-2システムのコストは公表されていないが、数百万ドルと推定されており、多くの小規模研究グループにとっては手の届かないものとなっている。Cerebrasは、自社のデータセンターやOracle Cloud、Azureとの提携を通じてクラウドアクセスを提供することで、この問題に対処している。
将来の方向性
CerebrasはWSEアーキテクチャの改良を続けている。2024年に発表されたWSE-3は、最大24兆パラメータのモデルのトレーニングをサポートするように設計されており、複数のCS-3システムのクラスターが必要となる。同社はまた、チップレットなどの新しいパッケージング技術を探求して、性能をさらに向上させコストを削減することを目指している。
2023年、CerebrasはIPOを申請したが、市場状況により計画は延期された。同社はOpenAIの共同創設者であるサム・アルトマン（公式な立場ではないが）やBenchmark Capitalを含む投資家から7億ドル以上を調達している。2024年現在、WSEはAIハードウェアの分野でニッチながら影響力のある製品であり、シングルチップコンピューティングの可能性の限界を押し広げている。
評価と影響
WSEはAIコミュニティからさまざまな反応を受けている。支持者はメモリ帯域幅の問題を解決する革新的なアプローチを強調する一方、懐疑論者はGPUクラスターと比較した実用性とコスト効率に疑問を投げかけている。独立した研究では、スパースアテンションやグラフ処理などの特定のワークロードでは、WSEがGPUを大幅に上回る性能を発揮できるが、他のワークロードではその利点はあまり明確ではないことが示されている。
議論にもかかわらず、WSEは特にオンチップメモリとインターコネクトの分野で他のAIチップの設計に影響を与えている。その成功は、以前は非現実的と考えられていたウェハスケール集積の概念を検証した。生成AIモデルが成長し続けるにつれて、WSEのような専門ハードウェアへの需要は高まる可能性が高いが、競争環境は依然として動的である。
関連項目
参考文献
- Cerebras Systems. "Wafer-Scale Engine: An Introduction." 2019.
- Feldman, A. et al. "The Cerebras Wafer-Scale Engine." IEEE Micro, 2021.
- Cerebras Systems. "CS-3 System Overview." 2024.
- Sandia National Laboratories. "Real-Time Spiking Neural Network Simulation." 2022.
- G42. "Building the World's Largest AI Supercomputer." 2021.