Cerebras、数兆パラメータ級AIモデルを複数のWSE-3Tウェハーに分散配置する方法を解説

Cerebrasは、オープンソースモデルのKimi K3を複数のWSE-3Tウェハーに分割し、レイヤーをパイプライン状に配置して実行する方法を解説した。ウェハー間で受け渡されるのは活性化データのみで、オンチップのKVキャッシュの容量が同時に処理できるリクエスト数を決める。

Cerebrasは、同社のブログ記事で、大規模モデルを複数のWSE-3Tウェハーに分散して動かす仕組みを説明した。記事によると、WSE-3T1枚のオンチップメモリは44 GBであり、トリリオンパラメータ級のフロンティアモデルは重みだけで数テラバイトを必要とする。

モデルの構成と配置

Kimi K3は、3Tパラメータのオープンソースモデルで、69層のlinear-attention層と24層のdense層から成る。記事は、各レイヤーに約16 GBが必要だと算出している。

• dense層は、KVキャッシュ用の領域をウェハー上に確保する必要があるため、1ウェハーに1層を配置する。
• linear-attention層は、追加のKVキャッシュ領域を必要とせず、リクエストごとに小さな固定サイズの状態を保持する。そのため、1ウェハーに2層を収容できる。

この例では、ウェハーの総数は35 + 24 = 59枚とされている。

パイプライン実行とデータ転送

Cerebrasは、モデルをパイプラインとして実行し、ウェハー間では活性化データのみを受け渡す。活性化データは、RDMA over Ethernetを介して1ホップあたり約2マイクロ秒で転送できるとされる。

また、異なるリクエストでパイプラインを常に満たすことで、各ステージが同時に別のリクエストを処理できるようにしている。同社は、数兆パラメータ級モデルの実行について、毎秒1000〜2000+トークンという速度を目標値として示している。

KVキャッシュと同時処理数

オンチップメモリのうちKVキャッシュに割り当てる量が、同時に処理できるリクエストの最大数を決定する。記事の試算例では、重みを4ビット、KVキャッシュを8ビットとし、平均リクエストサイズを256kトークンと仮定している。

筆者の見立て

• 推論システムの評価には、同じ出力生成速度で測定した電力単位あたりの総トークンスループットの方が、より意味のある指標だと論じている。
• ほとんどのケースでは、ウェハー間を流れるのは活性化データのみであるため、モデルの重みは移動せず、レイヤーをウェハー間で分割する必要もないと解釈している。

この記事は元記事の事実のみに基づいて自動生成されました。

出典

• 媒体名: Cerebras Blog
• タイトル: How Cerebras Deploys Trillion-Parameter AI Models
• URL: https://www.cerebras.ai/blog/how-we-deploy-trillion-parameter-frontier-models-on-cerebras
• (Hugging Face) の報道による