Cerebras が 1 台のマシンで 200 億のパラメータ モデルをトレーニングするという新記録を樹立

今週、チップスタートアップのCerebrasは、100億を超えるパラメータを持つNLP（自然言語処理）人工知能モデルを単一のコンピューティングデバイスでトレーニングするという新たなマイルストーンを発表しました。

Cerebras によってトレーニングされた AI モデルは、複数のアクセラレータにわたってワークロードをスケーリングすることなく、前例のない 200 億のパラメータに到達しました。この作業は、現在インターネット上で最も人気のあるテキストから画像への AI 生成モデルである OpenAI の 120 億パラメータ モデル DALL-E のニーズを満たすのに十分です。

Cerebras の新しい取り組みの最も重要な側面は、インフラストラクチャとソフトウェアの複雑さに対する要件が軽減されたことです。同社が提供するチップ「Wafer Scale Engine-2（WSE2）」は、その名の通り、TSMC 7nmプロセスウエハ1枚にエッチングされており、通常であれば数百個の主流チップを収容できるほどの広さの領域で、驚異的な2.6兆個のトランジスタ、85万個のAIコンピューティングコア、40GBの統合キャッシュを備え、パッケージング後の消費電力は15kWにも達します。

Wafer Scale Engine-2 は、ほぼウエハーと同じ大きさですが、その面積は iPad よりも大きいです。

Cerebras のマシン単体のサイズはすでにスーパーコンピュータと同等ですが、最大 200 億のパラメータを持つ NLP モデルを 1 つのチップに保持することで、数千の GPU のトレーニング コストと関連するハードウェアおよびスケーリング要件が大幅に削減され、モデルを GPU 間で分割する際の技術的な困難も解消されます。セレブラス氏は、後者は「NLP ワークロードの最も面倒な側面の 1 つ」であり、「完了するまでに数か月かかる」こともあると述べています。

これは、処理される各ニューラル ネットワークだけでなく、各 GPU の仕様やそれらを結び付けるネットワークにも固有のカスタマイズの問題です。これらの要素は、最初のトレーニングが開始される前に事前に設定しておく必要があり、システム間で移植できません。

Cerebras の CS-2 は、Wafer Scale Engine-2 チップと、関連するすべての電源、メモリ、ストレージ サブシステムを含むスタンドアロンのスーパーコンピューティング クラスターです。

200億のパラメータはおおよそどのくらいのレベルですか?人工知能の分野では、大規模な事前学習済みモデルは、最近さまざまなテクノロジー企業や機関が開発に力を入れている方向です。OpenAI の GPT-3 は、人間の読者を騙すのに十分な記事全体を書いたり、数学的な計算や翻訳を行ったりできる NLP モデルです。そのパラメータは驚異の 1,750 億個に上ります。昨年末に発売されたディープマインドのGopherは、パラメータ数の記録を2800億に引き上げた。

最近、Google Brain は、1兆を超えるパラメータを使用して Switch Transformer と呼ばれるモデルをトレーニングしたことを発表しました。

「NLP 分野では、モデルが大きいほどパフォーマンスが向上することがわかっています。しかし、従来、これらの大規模なモデルを分解し、数百または数千のグラフィックス処理ユニットに分散するという大変な作業を行うためのリソースと専門知識を備えた企業はほとんどありませんでした」と、Cerebras の CEO 兼共同創設者である Andrew Feldman 氏は述べています。 「その結果、大規模な NLP モデルをトレーニングできるのはごく少数の企業に限られています。コストが高すぎるし、時間がかかりすぎるため、業界の他の企業には手が届きません。」

現在、Cerebras のアプローチにより、GPT-3XL 1.3B、GPT-J 6B、GPT-3 13B、GPT-NeoX 20B モデルの参入障壁が下がり、AI エコシステム全体が数分で大規模なモデルを構築し、単一の CS-2 システムでトレーニングできるようになりました。

ただし、フラッグシップ CPU のクロック速度と同様に、パラメータの数は、大規模モデルのパフォーマンスを示す 1 つの指標にすぎません。最近では、パラメータを減らしながらより良い結果を達成した研究もいくつかあります。例えば、今年4月にDeepMindが提案したChinchillaは、わずか700億のパラメータで、通常の状況下でGPT-3やGopherを上回りました。

この種の研究の目標は、もちろん、より一生懸命働くことではなく、より賢く働くことです。そのため、セレブラスの成果は、一見した以上に重要な意味を持つ。この研究により、既存のチップ製造プロセスがますます複雑化するモデルに対応できるという確信が得られ、同社によれば、特殊なチップをベースにしたシステムは「数千億、あるいは数兆ものパラメータ」を持つモデルをサポートする能力があるという。

単一チップ上でトレーニング可能なパラメータの数が爆発的に増加するには、Cerebras の Weight Streaming テクノロジが必要です。このテクノロジーはコンピューティングとメモリの使用を切り離し、AI ワークロードの急速に増加するパラメータ数に基づいてメモリを任意の順序で拡張できるようにします。これにより、セットアップ時間が数か月から数分に短縮され、GPT-J や GPT-Neo などのモデル間の切り替えが可能になります。研究者らはこう述べている。「数回のキー操作だけで済みます。」

「Cerebras は、大規模言語モデルを低コストかつ簡単に実行できるようにすることで、AI のエキサイティングな新時代を切り開きます。大規模モデルの競争に参加するために数千万ドルを費やすことができない組織に、簡単で安価な方法を提供します」と Intersect360 Research の最高研究責任者である Dan Olds 氏は述べています。 「CS-2 のお客様が、大規模なデータセットで GPT-3 および GPT-J レベルのモデルをトレーニングする際に、新たなアプリケーションや発見が生まれることを非常に楽しみにしています。」