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

大規模なAIモデルの学習は、現代のIT技術の最先端をいく分野の一つだ。ChatGPTのような巨大なAIが世の中に登場し、その能力に多くの人が驚いたことだろう。しかし、そのような高性能なAIモデルを開発し、トレーニングするには膨大な計算資源と時間がかかる。従来のコンピュータシステム、特にCPUやGPUだけでは、その要求を満たすのが難しくなってきている。ここで登場するのが、Cerebras社が開発する「Wafer Scale Engine（WSE）」という革新的なAIチップだ。

WSEは、一般的なコンピュータチップとは根本的に発想が異なる。通常のチップは、シリコンウェハーと呼ばれる丸い基板から小さな四角いチップを多数切り出して製造される。例えば、私たちが普段使うスマートフォンのプロセッサや、AIの学習によく使われるGPUなども、この小さな四角いチップの一つだ。しかし、WSEは、このウェハー全体を一つの巨大なチップとして使う。直径300mmにもなるウェハーのほぼ全面に、AI計算に特化した何十万、何百万という数の「AIコア」と呼ばれる計算ユニットをびっしりと敷き詰めているのだ。これは、これまでのチップの常識を覆す大胆な設計と言える。

なぜこのような巨大なチップが必要なのか。それは、大規模AIモデルのトレーニングにおける「データ転送」の課題を解決するためだ。AIモデルのトレーニングでは、大量のデータと計算結果がチップ間を頻繁に行き来する。従来のシステムでは、複数の小さなGPUチップを連携させて計算を行うが、このチップ間のデータ転送は、ボード上の配線を介したり、時にはネットワークを介して複数のサーバー間で行われたりする。この「チップ間の通信」が非常に遅く、計算全体のボトルネックになっていた。

WSEはこの問題を根本的に解決する。ウェハー全体が単一の巨大なチップであるため、チップ間のデータ転送はウェハー内部、つまり「オンウェハーネットワーク」と呼ばれる高速な配線上で行われる。この内部通信は、従来のチップ間通信と比較して圧倒的に高速で、データ転送の遅延が極めて少ない。これにより、何十万ものAIコアが、あたかも一つの巨大な計算ユニットであるかのように、協力して効率的に計算を進めることができる。

さらに、WSEはAIコアだけでなく、コアのすぐ隣に「オンチップメモリ」と呼ばれる非常に高速な記憶領域を大量に搭載している。従来のシステムでは、計算に必要なデータを外部のメモリ（DRAMなど）からチップに読み込む必要があり、ここでもデータ転送の遅延が発生していた。しかし、WSEでは計算に必要なデータをチップ内部の高速なメモリに保持できるため、外部メモリへのアクセスを最小限に抑え、計算の速度をさらに向上させている。この膨大なAIコアと高速なオンチップメモリ、そして超高速なオンウェハーネットワークの組み合わせが、WSEの最大の強みだ。これらにより、AI計算のレイテンシ（遅延）を極限まで低く抑え、スループット（処理量）を最大化できる。

このWSEを搭載したシステムは「CSシステム」として提供されている。例えば最新のCS-3システムは、WSE-3という第三世代のウェハーチップを搭載している。このシステムは、単一のAIチップでありながら、数兆パラメータを持つ大規模AIモデル全体をメモリ内に収め、非常に少ないデータ転送の遅延で、驚異的な速度で学習させることができる。これにより、これまで数ヶ月かかっていたAIモデルの学習時間が、数日、あるいは数時間へと大幅に短縮される可能性を秘めている。モデルの並列化も最小限で済むか、不要になる場合もあり、開発者の負担を軽減する。

システムエンジニアを目指す上で理解しておきたいのは、ハードウェアの進化とソフトウェアの進化は密接に関係しているということだ。Cerebras社は、WSEという革新的なハードウェアを提供するだけでなく、その性能を最大限に引き出すためのソフトウェアスタック（ツールやライブラリ群）の開発にも力を入れている。例えば、多くのAI開発者が利用するPyTorchなどのフレームワークから、WSEの性能を簡単に引き出せるようにすることで、開発者は複雑な並列計算の最適化に頭を悩ませることなく、モデル開発に集中できる。これは、ハードウェアの革新が、より多くの人がAI開発に参入しやすくなる環境を整えることにも繋がる。

CerebrasのWSEのようなAIチップは、現在のAI研究開発における最も重要な課題の一つ、すなわち「計算能力の限界」に挑戦している。AIモデルがさらに大規模化し、より複雑なタスクをこなせるようになるためには、WSEのような革新的なアーキテクチャを持つチップが不可欠になるだろう。このような技術の登場は、AIの可能性をさらに広げ、新たなイノベーションを生み出す原動力となる。システムエンジニアとしてAI分野に関わるならば、このような最先端のハードウェア技術が、どのような課題を解決し、どのような未来を切り開くのかを理解しておくことは、非常に重要だ。AIの進化はハードウェアとソフトウェアの両輪で進んでおり、それぞれの進歩が相互に影響し合いながら、私たちを次の時代へと導いている。