Cerebrasが3兆パラメータ級モデルをウェハ59枚で動かす手法を公開

何が示されたのか

Cerebrasはブログ記事「How we deploy Trillion-parameter frontier models on Cerebras」で、ウェハスケールチップ「WSE-3T」上に数兆パラメータのモデルを配備する方法を公開した。単一のWSE-3Tが持つオンチップメモリは44GBで、兆パラメータ級モデルの重みはテラバイト規模になる。そのため1枚には収まらず、複数枚への分割が前提となる。記事では、複数枚に分割して毎秒1000〜2000トークン以上で動かす方法を説明している。

具体例として使われるのが、オープンソースの「Kimi K3」だ。約3兆パラメータを持ち、線形アテンション層69層とdense層(通常の全アテンション層)24層で構成される。記事では、このモデルがTerminal-Bench 2.1で88.3%、BrowseCompで91.2%を記録したと紹介されている。あくまでモデル側のベンチマークスコアとしての言及で、Cerebras上で測定した値かどうかは明記されていない。

1層16GBという前提

まず各層のサイズを見積もる。3兆パラメータをすべて4ビットで保持すると仮定し、これを93層(69+24)で割ると、1層あたり約16GBになる。重みの量子化方式はHugging Face実装に合わせた4ビットと説明されている。

次に、44GBのオンチップメモリに何層載るかを層の種類ごとに決める。dense層は各トークンのキーと値を保持するKVキャッシュ用の領域が要るため、1枚に1層だけ置き、残りをKVキャッシュに充てる。一方の線形アテンション層は文脈長によらず固定サイズの状態しか持たないため、1枚に2層(16GB×2=32GB)を置ける。結果として69層を2で割って切り上げた35枚と、dense層24枚の合計59枚になる。

重みを動かさないパイプライン

Cerebrasのコンパイラが層をウェハ群へ割り当て、ランタイムがトークンの通り道を設定する。ウェハ間でやり取りするのはアクティベーションだけで、RDMA over Ethernetを使い1ホップあたり約2マイクロ秒と説明されている。各層16GBに対して1枚44GBあるため、層を分割せず丸ごと1枚に載せられる点が効いている。重みは動かさず、層の間をアクティベーションだけが流れる。

推論は工場のラインに似ている。ウェハ1が層1〜2を処理して結果をウェハ2へ渡し、ウェハ2が層3〜4を処理する、という具合に最後のウェハが次のトークンを出す。ここで「1リクエストずつ順番に流すなら、各ウェハは待ち時間ばかりになるのではないか」という疑問が出る。記事はこれを否定し、各ステージに別々のリクエストを流してパイプラインを埋めることで、全体を遊ばせないと説明する。

同時実行数と評価軸

同時に処理できるリクエスト数を決めるのは、KVキャッシュに割り当てられるメモリ量だ。dense層を1層載せたウェハでは、44GBから層の16GBを引いた残りがKVキャッシュに使える。256kトークンのリクエスト1件が使うキャッシュ量を求め、残り容量を割ることで同時実行数を見積もる、という考え方が示されている。ただし原文では計算式の結果にあたる具体的な数値が示されていないため、何件まで処理できるかは読み取れない。

記事は、同時実行数やバッチサイズだけでは推論システムの効率は測れないと指摘する。より意味のある指標として挙げているのが、同じ出力生成速度で測った「1メガワットあたり毎秒何トークン処理できるか」だ。電力あたりのトークン処理量という軸で性能を見る姿勢が示されている。

現時点の留意点

記事で示されるのはあくまでKimi K3を例にした構成の説明であり、実際の顧客環境でのスループットや遅延の測定値は記載されていない。1リクエストあたり256kトークンという平均値を前提に置いている点も、実際の利用形態によって変わる条件だ。より大きなモデルにも同じ考え方で拡張できるとしているが、具体的な上限は示されていない。

また、料金や提供地域、既存システムからの移行方法については触れていない。無料APIと営業窓口の案内があるのみで、日本からの利用可否や価格はこの記事だけでは判断できない。