Cerebras、兆パラメータAIモデルを展開

Cerebras wafer-scale チップ上で兆パラメータの最先端モデルを展開する方法

私たちがよく受ける質問の一つに、Cerebrasの wafer-scale チップ上で大規模モデルをどのように展開しているかというものがあります。単一の Cerebras WSE-3T wafer は 44 GB のオンチップメモリを持っていますが、兆パラメータの最先端モデルは、その重みだけでテラバイトを必要とします。

この記事では、複数の WSE-3T wafer に大規模な最先端モデルを分割し、毎秒 1000〜2000 トークン以上でマルチ兆パラメータモデルを実行する方法を説明します。

例として、約3兆パラメータを持つオープンソースモデル Kimi K3 を使用します。

最終的には、以下の点を理解できるようになります。

• モデル全体をオンチップメモリにどのように収めるか
• モデルを複数のチップにどのようにマッピングするか
• 大量の同時リクエストをどのように処理するか

Kimi K3 で分解する

Kimi K3 は、69 個の線形アテンション層と 24 個の密層を持つ、オープンソースの3兆パラメータモデルです。

［画像: Kimi K3 モデルのアーキテクチャを示す図］

［画像: Kimi K3 モデルのレイヤー構造を示す図］

この例では、以下を使用します。

• 256k という想定平均リクエストサイズ
• 4ビットの重み量子化（Hugging Face の実装と同様）
• 8ビットの KV キャッシュ

#1: 各レイヤーのサイズ

Kimi K3 をサポートするために必要な Cerebras WSE-3T wafer の総数を計算するために、まず各モデルレイヤーのサイズを計算しましょう。

Kimi K3 は 3兆パラメータを持っています。もし 3兆パラメータすべてが各 4ビットで保存される場合、生のパラメータペイロードは以下のようになります。

［画像: 3兆パラメータを4ビットで保存した場合のペイロードサイズを計算する式］

このペイロードを 93 層で割ると、以下のようになります。

［画像: 各レイヤーのサイズを計算する式］

#2: wafer の数

各レイヤーが約 16 GB を必要とすることがわかったので、44 GB の wafer にいくつのレイヤーが収まるかを計算できます。答えはレイヤーの種類によって異なります。

• 密層は KV キャッシュのために追加の領域を必要とします。なぜなら、各密層（フルアテンション）は、アクティブな各リクエストの各トークンのキーと値を保存するからです。そのため、wafer ごとに 1 つの密層を配置し、残りの SRAM を KV キャッシュ用に確保します。
• 線形アテンション層は KV キャッシュのための追加領域を必要としません。各トークンのキーと値を保存する代わりに、線形アテンション層は、コンテキスト長に関係なく、リクエストごとに小さく固定サイズの状態を保持します。これにより、wafer ごとに 2 つの線形レイヤーを配置できます（2 × 16 GB = 32 GB）。

したがって、Kimi K3 の 69 個の線形レイヤーと 24 個の密層 を使用すると、以下のようになります。

［画像: wafer の数を計算する式］

これにより、wafer の総数は 35 + 24 = 59 wafer となります。

#3: wafer を接続する

Cerebras はパイプラインでモデルを実行します。Cerebras コンパイラがモデルレイヤーを wafer の領域にマッピングした後、ランタイムがトークンが wafer 間を流れるパスをマッピングします。超高速レイテンシのために、アクティベーションのみが wafer 間を通過する必要があり、これは Ethernet 上の RDMA を介して行われます。これはホップあたり約 2 マイクロ秒です。

［画像: wafer 間の接続とデータフローを示す図］

各 Cerebras wafer は 44 GB のオンチップ SRAM を持ち、各レイヤーはわずか 16 GB であるため、モデルのレイヤー全体を直接チップに収めることができる場合が多いことを強調しておくことが重要です。

これは、ほとんどの場合、モデルの重みが移動する必要がなく、レイヤーが wafer をまたいで分割される必要がないことを意味します。アクティベーションのみが次の wafer に渡されます。

#4: リクエストを 1 つ処理する方法

wafer が接続され、モデルが配置されたので、各 wafer を組み立てラインの 1 ステージと考えることができます。各ステージのモデルの重みは所定の位置に留まり、リクエストの中間結果がレイヤーのステージを通過します。

例えば、Wafer 1 はモデルレイヤー 1〜2 を実行し、そのアクティベーションを wafer 2 に送信します。Wafer 2 はレイヤー 3〜4 を実行し、その結果を Wafer 3 に渡します。これを最後の wafer が次のトークンを生成するまで繰り返します。

［画像: 単一リクエストの処理フローを示す図］

#5: 多数のリクエストを処理する方法

ここで一般的な誤解が生じます。トークンが順に各 wafer を通過する必要がある場合、各 wafer は順番待ちでアイドル状態にならないのでしょうか？実質的に、システムはバッチサイズ 1 で実行されるのではないでしょうか？

いいえ！もし一度に 1 つのリクエストだけをラインに送り、それが完了するのを待ってから次のリクエストを開始した場合、パイプラインの非効率的な使用となります。代わりに、パイプライン全体を異なるリクエストで満たしておきます。モデルはすでにパイプラインステージに分割されているため、各ステージは同時に異なるリクエストを処理できます。

ステージ 3 がリクエスト A のトークンを処理している間に、ステージ 2 はリクエスト B のトークンを、ステージ 1 はリクエスト C のトークンを処理します。この例では、各ステージはレイヤーです。本番環境では、Cerebras はより細かい粒度のステージを、さらに多く使用します。

［画像: 動画サムネイル - 同時リクエスト処理のデモ］

各ステージは、単一のリクエストではなく、マイクロバッチ（リクエストの小さなグループ）を処理することもできます。したがって、インフライトリクエストの数は、ステージ数にマイクロバッチサイズを掛けたものにほぼ等しくなります。

［画像: 動画サムネイル - マイクロバッチ処理のデモ］

トレードオフを理解するために、マイクロバッチサイズが小さいほど各リクエストはさらに高速になりますが、マイクロバッチサイズが大きいほど全体の処理能力が増加します。

#6: どれだけ のリクエストを処理できるか？

KV キャッシュに割り当てられたオンチップメモリの量は、同時リクエストの最大数を決定します。インフライト中の各リクエストは、各密層 wafer 上で独自の KV キャッシュを必要とするため、オンチップメモリのサイズが同時リクエストの上限を設定します。前述のように、密層の KV キャッシュ要件のみを計算する必要があります。

既に 1 つの密層を持つ各 44 GB wafer について、残りの KV キャッシュ領域は以下のように計算できます。

総 wafer ストレージ - 1 つの密層のサイズ。この場合：

［画像: KV キャッシュ用に使用できるメモリ容量を計算する式］

同時リクエスト数の計算

特定のモデルマッピングでサポートされる同時リクエスト数を把握するために、まず 256k リクエスト 1 つが占めるキャッシュ量を計算します。

［画像: 単一リクエストあたりの KV キャッシュサイズを計算する式］

次に、この構成がサポートする同時リクエストの総数を計算するために、以下を割ります。

［画像: 同時リクエスト数を計算する式］

同時リクエストの総数は、任意の時点でパイプラインを通過している同時リクエストの数を示します。GPU バッチとは異なり、これらのリクエストはお互いの完了を待つのではなく、パイプラインを継続的にストリーミングします。

もちろん、同時実行性もバッチサイズも、推論システムがどれだけ効率的にパフォーマンスを発揮するかを正確には示しません。推論システムの処理能力を評価する際には、同じ出力生成速度で測定した場合、単位電力あたりの総トークン処理能力（例: メガワットあたりの秒間トークン数）の方がより意味のある指標となります。

3兆パラメータを超えるモデルへのスケーリング方法

同じアプローチがより大規模なモデルにも拡張されます。重みとリクエスト状態に必要なメモリを計算し、モデルを wafer に分散させ、それらをパイプラインに接続します。追加の各 wafer は、コンピューティングとオンチップメモリを提供し、単一チップの容量を超えるモデルに対応できるようにします。

ご清聴ありがとうございました！Cerebras でモデルを展開することや、当社の推論クラウドで構築することに興味がありますか？

• 無料 API で始めましょう: cloud.cerebras.ai
• 営業チームにお問い合わせください: sales@cerebras.ai

デザイン: Halley Chang。レビューとフィードバックを提供してくれた Atharva Talpade、Zhenwei Gao、Angela Yeung、Jessica Liu、Joyce Er、Emma Call、Rohan Gupta に感謝します。