今日の人工知能(AI)システムを、継続的な進歩に必要な規模で機能させることが、ますます困難になっている。これらのシステムは、すべての処理チップが生成するデータを迅速に共有し、一つのユニットとして機能させるために、膨大な量のメモリを必要とする。
過去10年間のディープラーニングブームを主に支えてきたチップは、グラフィックス・プロセッシング・ユニット(GPU)と呼ばれるものである。これらは元々ゲーム用に設計されたもので、思考プロセスの各ステップが1ミリ秒未満で行われなければならないAIモデル用ではなかった。
各チップには限られた量のメモリしか搭載されていないため、AIシステムを支える大規模言語モデル(LLM)は、高速ネットワークで接続された多数のGPUに分割されなければならない。LLMは膨大な量のテキストでAIを訓練することで機能し、そのあらゆる部分がチップ間のデータ移動を伴う。このプロセスは低速でエネルギー集約的であるだけでなく、モデルが大きくなるにつれて、ますます多くのチップを必要とする。
例えば、OpenAIは約20万個のGPUを使用して最新のモデルGPT-5を作成したが、これは3年前のChat-GPTの最初のバージョンを動かしたGPT-3モデルで使用された数の約20倍である。
GPUの限界に対処するため、カリフォルニアに拠点を置くCerebrasなどの企業は、ウェハースケール・プロセッサと呼ばれる異なる種類のチップを構築し始めている。これらはディナープレートほどの大きさで、GPUの約5倍大きく、最近になってようやく商業的に実用可能になった。それぞれが膨大なオンチップメモリと数十万個の個別プロセッサ(コアと呼ばれる)を搭載している。
その背後にある考え方はシンプルである。数十個の小さなチップを調整する代わりに、すべてを一つのシリコン上に保つことで、データがハードウェアのネットワークを横断する必要がなくなる。これは重要である。なぜなら、AIモデルが答えを生成する際(推論として知られるステップ)、あらゆる遅延が積み重なるからである。
モデルが応答するのにかかる時間はレイテンシと呼ばれ、このレイテンシを削減することは、チャットボット、科学分析エンジン、不正検出システムなど、リアルタイムで動作するアプリケーションにとって極めて重要である。
しかし、ウェハースケール・チップだけでは十分ではない。そのアーキテクチャ専用に設計されたソフトウェアシステムがなければ、理論的な性能向上の多くは決して現れない。
より深い課題
ウェハースケール・プロセッサは、特殊な特性の組み合わせを持っている。各コアのメモリは非常に限られているため、チップ内でデータを共有する必要性が極めて高い。コアは自身のデータにナノ秒でアクセスできるが、各チップ上に非常に広い領域にわたって多数のコアがあるため、ウェハーの反対側のメモリを読み取るのは1000倍遅くなる可能性がある。
各チップのルーティングネットワークの制限により、コア間のすべての通信を一度に処理することもできない。要するに、コアは十分な速度でメモリにアクセスできず、自由に通信できず、最終的にはほとんどの時間を待機に費やしているのである。
ウェハースケール・チップは通信遅延によって速度が低下する。Brovko Serhii
我々は最近、WaferLLMと呼ばれる解決策に取り組んでいる。これはエジンバラ大学とMicrosoft Researchの共同事業で、ウェハースケール・チップ上で最大規模のLLMを効率的に実行するように設計されている。そのビジョンは、チップ上の各コアが主にローカルに保存されたデータを処理するように、LLMの実行方法を再編成することである。
ソフトウェアの観点からこの問題を探求した最初の論文において、我々はモデルの大規模な数学演算をはるかに小さな断片に分割する3つの新しいアルゴリズムを設計した。
これらの断片は、隣接するコアが一緒に処理できるように配置され、次のコアには小さなデータの断片のみを渡す。これにより、情報はウェハー全体でローカルに移動し続け、チップ全体を遅くする長距離通信を回避できる。
また、ウェハーの大部分を遊休状態にすることなく、LLMの異なる部分(または層)を数十万個のコアに分散させる新しい戦略も導入した。これには、あるコアグループが計算しているときに、別のグループがデータを移動し、3番目のグループが次のタスクを準備するように、処理と通信を調整することが含まれる。
これらの調整は、Edinburgh International Data Facilityにあるヨーロッパ最大のウェハースケールAI施設で、MetaのLlamaやAlibabaのQwenなどのLLMを使用してテストされた。WaferLLMは、ウェハースケール・チップによるテキスト生成を以前の約100倍高速化した。
16個のGPUクラスターと比較すると、これはレイテンシが10分の1に削減されたことになり、エネルギー効率も2倍向上した。したがって、AIパフォーマンスの次なる飛躍はLLM専用に設計されたチップから生まれる可能性があると主張する者もいるが、我々の結果は、代わりに既存のハードウェアの構造に適合するソフトウェアを設計できることを示唆している。
短期的には、低コストでの高速推論により、毎秒より多くの仮説を評価できる応答性の高いAIツールの可能性が高まる。これは、推論アシスタントから科学分析エンジンまで、あらゆるものを改善するだろう。不正検出やシミュレーションを通じたアイデアのテストなど、さらにデータ量の多いアプリケーションでさえ、大規模なGPUクラスターを必要とせずに、劇的に大きなワークロードを処理できるようになる。
今後の展望
GPUは柔軟性があり、広く入手可能で、成熟したソフトウェアエコシステムによってサポートされているため、ウェハースケール・チップがそれらに取って代わることはない。代わりに、医薬品発見や金融取引など、超低レイテンシ、極めて大規模なモデル、または高エネルギー効率に依存するワークロードに対応する可能性が高い。
一方、GPUも停滞していない。より優れたソフトウェアとチップ設計の継続的な改善により、より効率的に動作し、より高速を実現できるようになっている。時間の経過とともに、さらなる効率性が必要とされる場合、一部のGPUアーキテクチャもウェハースケールのアイデアを採用する可能性がある。
より広範な教訓は、AIインフラストラクチャが共同設計の問題になりつつあるということである。ハードウェアとソフトウェアは共に進化しなければならない。モデルが成長するにつれて、単により多くのGPUでスケールアウトするだけでは十分ではなくなる。WaferLLMのようなシステムは、次世代のAIパフォーマンスを引き出すために、ソフトウェアスタックの再考が不可欠であることを示している。
一般市民にとって、その恩恵は棚に並ぶ新しいチップとして現れるのではなく、以前は遅すぎるか高価すぎて実行できなかったアプリケーションをサポートするAIシステムとして現れるだろう。科学的発見、公共サービス、大量分析のいずれにおいても、ウェハースケール・コンピューティングへのシフトは、AIシステムの構築方法と、それらが達成できることの新しい段階を示している。
