進化し続けるWSE(Wafer Scale Engine)⁠大規模AIモデルのトレーニング性能とは ――Cerebras社訪問2024後編】

Cerebras Systems, Inc.はWSEというシリコンウェハーまるごと一枚を使った巨大なAIプロセッサを独自に開発し、それを使ったAIトレーニング向けのシステムCS-2を販売しているシリコンバレーのスタートアップです。前編に続いて、2024年夏のCerebrasの取り組み、WSEの開発の状況とこれからについてお届けします。

AIトレーニングにおける「性能」

議論は再び「性能」の話に戻ります。しかしシステムのスループットとはまた異なる側面の、です。

Andy: ここは本当に(gihyo.jpの読者向けに)強調したいです。計算能力の「量」は重要です。CS-3クラスタはLlamaを1日でトレーニングすることを可能にします。GPUシステムで1ヵ月かかるのに対して。これが計算能力の「パフォーマンス」です。
その一方で、もう1つ大きな要素は、それらのモデルの「解(適切なモデルのデザインとその実装⁠⁠」に辿りつくまでの間、どのくらいプログラミングが容易か、です。我々のマシンでは、1台のCS-3を使っていても、2,048台のCS-3でも同じ、つまり1台のデスクトップマシン向けのコードが使えます。
私自身、リサーチャーとして思うのは、これが新しいモデルを迅速に作るために非常に大きい、ということです。私のコンピュータがどれだけ高速であろうとも、もし私がそれを実行するためのコードを書くために3ヵ月掛かってしまうなら、システム自体の高速さには意味がないからです。
3ヵ月掛かってコードを書いて、実行8時間。それならもっと短い時間で開発をやめて、2ヵ月掛けてトレーニングしますよね。
Yasu: そうなりますね。
Andy: これはGPUクラスタを使ったAI処理の「隠された真実[1]だと思います。マーケティングでは時折、いくつのGPUを使って、どのくらい早く1つのモデルを1つの製品としてトレーニングできたかが話されますが、しかしそのためにプログラミング、とくに分散学習の作業にどれだけ多くの時間を費やしたかを言う人はほんのわずかです。
しかし、これはとても重要なことです。AI 研究コミュニティのリーダーたちも、この「GPUをAIのためにプログラムすることの困難さ」を言い始めています。これはAndrej Kapathy、つまりOpenAIのCo-Founderですが、彼は大規模モデルをGPUクラスタでトレーニングする作業がいかに「もどかしい(frustrating⁠⁠」か言っています図5⁠。
図5 “大規模モデルのトレーニングには日々イライラさせられがちだ。注意深くモニタリングを行わないと10,000個のGPUがアイドリング状態になってしまう。”
引用 ※(元のXポストはすでに削除されています)
Andy: このGPT-1の論文を見てください。GPT modelは120Mパラメタ、今から見ると小さなものでした。著者は4人です。GPT-4 はその10倍を超えています。著者は240人、そのうち35人がただ分散学習とスパコン的な作業のために呼ばれています。
これがソフトウェアの問題です。もうGPUクラスタの上ではモデルが大きすぎて、GPUのメモリに収まり切らないため、モデルを小さく分割しなければなりません。そして小さなモデルを全てのGPUに乗せて、処理を分散し、マシン間通信のオーケストレーションをしなければなりません。これがとても複雑な作業です。
大規模AIモデルのGPUクラスタ上でのトレーニングには、標準的なソフトウェアライブラリとしてNVIDIAのMegatronがよく使われます。Megatronは2万行を超えるコード、それもC++とPython、そしてとても低レベルで不透明(obscura)なCUDAコードです。とても大規模で複雑なものです。
GPUクラスタ上で実行するためのこれらの作業とは対照的に、私たちのシステムで同じモデルを実行するには僅か数百行のコードしか必要としません。しかもPythonとPyTorchだけであり、低レベルの特別なコードはありません。よほどシンプルです。我々のシステムの上では、それらのモデルは十分小さなものですからね。我々のシステムでは現時点での世界最大のモデルでも収まります。
つまり我々のシステムのユーザは1台のシステムでトレーニングしようと、1,000台だろうと、あるいは1 Billion(10億)のパラメータであろうと、Trillion(兆)のパラメータであろうと、デスクに座るユーザには関係ないのです。
つまり(単体の)デスクトップマシンのプログラマビリティで、スパコンのパフォーマンスを得られるのです。
この「パフォーマンス」「プログラマビリティ」という2つのファクターが両立しているために、モデルをより早く組み上げられるのです。

AndyはCerebrasの戦略的パートナーでもあるUAEのG42が、Condor Galaxy 1と呼ぶCS-2 64台構成のクラスタを用いて行ったトレーニングの実例を紹介してくれました。

Andy: 彼らはGalaxy 1を使って世界最高のOpen Sourceベースのアラビア語言語モデルを作りました。13 billion パラメタの、アラビア語と英語のChatGPTですね。もともと彼らはGPU 64台構成のクラスタを持っていましたが、その13 billionパラメタのモデルをトレーニングするのに(開発工程を含めて)およそ68日かかりました。まあ2ヵ月ちょっとですね。さてGalaxy 1の1/4を使って、我々は同じモデルを2日でやりました。
Yasu: おお。2日ですか。
Andy: つまり64 GPUのシステムより35倍速かったわけです。Galaxyシステムの1/4で。これが私の言う重要な「Xファクター」⁠エックス、つまり何倍早いか、という意味)です。
実際に顧客に重要なのは「このモデルのトレーニング作業が(開発工程を含めて)2日ちょっとで完了する」ということです。これは彼らがより多くの実験を繰り返せることを意味し、それによってモデルをどんどん良くしていけるのです。
我々はこうやって世界クラスの言語モデルをいくつも作りました。アラビア語、薬学、日本語、そして物理学、ファイナンス、バンキング(銀行業務⁠⁠。何に適用できるかは問題ではありません。重要なのはそれぞれの大規模言語モデルで開発サイクルをクイックに繰り返せることです。これが「パフォーマンス」です。そしてこんなことを可能にしてくれるものはこの世界には他にありません。
Yasu: 良い事例と、良い説明ですね。
Andy: 私たちはこのことをなんとかして伝えるために2年間取り組みました。というのも、GPUクラスタ上でプログラミングやデリバリー、デバッグ、最適化を行うことがいかに難しいかについては、AI研究の「隠された真実」のようなもので、誰も語りたがらないからです。誰もがあたりまえのようにやっていることですが、それが一番難しいのです。

筆者はこのことを聞いて、1960年代あたりのとても古い時代のコンピューティングと同じことが起きているように思えました。コンピュータが速くなるにつれてプログラムは規模が大きくなり、正しく動作するものを書くことがだんだん難しくなります。するとプログラミングに掛かる時間と処理時間(データを与えて結果が得られるまでの時間)の比が問題になります。

Cerebrasも同様です。とても高速に、たとえば1ヵ月掛かるトレーニング処理を1日に圧縮できた。するとプログラミングにかかる時間がより重要な要素になってきたわけですね。今後、AI応用が各所で進むにつれ、この重みはより大きくなっていくでしょう。

AIコンピューティングの進化

Yasu: あなた方はウェハーサイズのデータフローシステムを使ってAIトレーニングをしているわけですが、一般的な実装はGPUを使ったものです。ディープラーニングが始まったのは10年ほど前で、10年前に最も安価なSIMDエンジンとしてGPUがあったから使われ、今も主としてGPUが使われている状況になりました。ただ、私はこれをあまり良い解決策だとは思っていません。
Andy: このことはいつも考えています。GPUは好きなんですが(本当?嫌ってない?)ええ、ええ、そんな。だって私の初期のキャリアは画像処理の研究でした。例えばレイトレーシングのシミュレーションや、画像やビデオ処理のアルゴリズムの並列化が書けるようになった頃を覚えていますよ。でもGPUは本当に密な行列の乗算のために作られているのですよね。
Yasu: そうですね。
Andy: GPUは基本的に画像を取り込んで処理するように作られています。画像やビデオデータは密な2次元の行列です。それを他の高密度な2次元配列と掛け合わせます。ただディープラーニングには2つの問題がありました。まず、ディープラーニングは行列ではありません。
Yasu: 理論的には行列として表現できるけれど、、
Andy: 大抵はベクトル、まあスカラーベクトルでなくてもいいのですが、しかもたいていはスパース(疎、密度が低い)なんです。スパースすぎる。しかし(データとしては)スパースであることが望ましい。これがミスマッチなのです。
Yasu: そうですね。だからあなた方はデータフローを選んだ。
Andy: その通りです。データが本当にスパースで、あなたが事前にそのスパース性を知ることができないのなら、スパースにベクトル積を得られるようにするのが良い。我々のシステムはそのために設計されており、メモリの帯域幅が桁違いです。
GPUはこれを解決できません。実行はできますがゼロを掛けるばかりで、時間を無駄にしてしまいます。画像処理、つまり動画やレイトレーシングなどであればとてもうまくいくのですが、ディープ・ラーニングではそうはいきません。だからミニバッチなんかがあるわけですね。
あなたがはじめに言ったとおりだと思います。ディープラーニングが始まったころ、行列のベクトル乗算をより多く実行する必要がありました。そのため高性能のSIMDエンジンが必要だったのですが、そこにGPUがあったのですね。
写真 お互いにホワイトボードに図を描きながらの議論
Yasu: 10年前、GPUはCPUよりも優れた性能を発揮できました。それがGPUを使った理由です。でも今、あるいは10年後の未来では、GPUを使い続ける理由がなくなっていると思います。
たとえば少し前に、Preferred Networksの取材をしました(⁠高性能・高効率なAIチップ⁠MN-Core 2⁠の設計アプローチ -Preferred Networksによる新しいハードとソフトの役割分担⁠。⁠記事を示しながら)こんな構造です。GPUとはまったく別のものです。
Andy: 違いは重要です。私はよく、これを生物学的進化に似たプロセスとして考えることがあります。生物学では、物理的な環境が変化するとさまざまな種類の生物が出現し、新しい環境に適した生物が成功して増殖します。
コンピューティング・システムには、物理的な環境はありませんがマーケット(市場)があります。重要で価値のあるコンピューティング・ワークロードとは何か、ですね。これがGPUの誕生につながりました。グラフィックスが重要視されるようになったからです。今起きているのはAIです。これまでと異なるこのワークロードとGPUはまあまあ⁠OK⁠ですが(⁠⁠まあダメってわけじゃないけどね」的なニュアンスを込めた発音の「オーケイ」でした⁠⁠、この新しい環境には最適化されていませんよね。
この新しい環境の中で我々のシステムが増殖し、生き残り、成功する生命体になるかどうかはわかりません。しかし、私はこれは本当にいいチャンスだと思いますし、チップは新しいコンピューティング環境、新しい市場環境に適応し、そのために発明・設計されています。Preferredの人たちがどんなことをしているか興味がわきますね。
スパース性はこれからのAIコンピューティングで重要で、Memory Xを使うことで許容できると思います。メモリの中にほぼ0のデータを大量に保持するのは簡単ですから、次世代のモデルでパラメータを倍、またその倍、と増やしていくことができます。
Yasu: Memory Xのメモリを増やすだけですからね。
Andy: そうです。そしてもっともっとスパースになる。ウェハーはスパース計算のために設計されており、ウェハーに積和演算が来てもゼロ乗算は避ける(avoid)ことができますから、結果的にスパース性によって加速することになります。
研究者としての私にはこれは本当にパワフルなことです。というのもこれまで私たちが行ってきたAIの研究はすべて、スパースAIの研究だったからです。高密度のプロセッサーで研究しているため、非常に効率が悪いのです。ゼロをたくさん足しても、時間の節約にはならない。そのため、人々はスパースAIをあまり効率的に研究できていなかったと思います。しかし、私たちのようなプラットフォームがあれば、研究者がスパースAIをよりよく研究し、スパース性を用いてより効率的なAIモデルを構築できるようになります。これからの2、3年はとてもおもしろくなりそうですよ。
Yasu: はい。楽しみにしています。

いつもそうなのですが、今回の取材でもAndyとは製品を超えた技術のことを多く話しました。彼は「現在のAIの世界のほとんどがGPUで構築されたモデルであり、それは素晴らしいことです」とした上で「しかしもし別のツールがあったら人々は何を作るでしょうか?私は、違うツールで作ることによって広がるAIの幅のようなものを世界が発見する手助けをしたいのです」と答えてくれました。

Cerebrasは2015年の起業、そして2019年のWSE-1の発表以来、ものすごいスピードで技術開発を行い、製品を更新し続けています。Andyの言うとおり、これからの2, 3年を楽しみに見守りたいと思います。まだまだ新しい技術開発を彼らは続けていますから。

ニュース⁠

CerebrasはもともとAI学習に特化したアクセラレータとしてWSEを設計・開発しています。しかし2024年8月27日に、Cerebrasはその飛び抜けた処理能力を学習ではなく推論に使ったチャットサービスを発表しました。

"Cerebras Launches the World’s Fastest AI Inference"

Cerebrasのトップページからアクセスできます。日本語にも対応しています。ぜひその反応の早さを体験してください。

おすすめ記事

記事・ニュース一覧