Cerebrasが3兆パラメータモデルを59枚のウェハーに配置する方法

兆単位のパラメータを持つモデルでは、計算性能だけでなく、膨大な重みとリクエスト状態をどこに保持するかが大きな課題になる。Cerebrasは技術記事で、約3兆パラメータのオープンモデルKimi K3を例に、WSE-3Tウェハースケールプロセッサーへモデルを分散する考え方を説明した。記事に記載された性能やベンチマーク比較は、Cerebrasによる説明であり、独立した検証結果ではない。

層単位でモデルを分割

1枚のWSE-3Tには44 GBのオンチップメモリがある。一方、兆単位のモデルの重みはテラバイト規模になるため、1枚に収めることはできない。Cerebrasは4ビット量子化を前提に、層ごとの保存容量を見積もる。Kimi K3では、3兆パラメータを93層に分けると、平均で1層あたり約16 GBになる。

ただし、すべての層を同じように配置できるわけではない。Kimi K3は69個の線形アテンション層と24個の通常のアテンション層で構成される。線形アテンションは、リクエストごとに比較的小さく固定的な状態を保持するため、例の構成では1枚に2層を置ける。通常のアテンション層は、コンテキストに応じて増えるKVキャッシュを必要とするため、1枚に1層を置き、残りのメモリをキャッシュに割り当てる。

この計算では、線形層に35枚、通常の層に24枚、合計59枚のウェハーが必要になる。

ウェイトではなく活性値を転送

モデルを配置した後、各ウェハーに割り当てた重みは基本的にその場に留まる。リクエストは組立ラインのように段階を順番に通過し、ある段階で計算された活性値が次の段階へ送られる。Cerebrasは、ウェハー間の通信にRDMA over Ethernetを使い、1ホップあたり約2マイクロ秒という遅延を示している。

重要なのは、完全な層を1枚のウェハーに収められる場合が多い点だ。重みをチップ間で何度も移動させる必要がなく、通信の中心は前段の出力である活性値になる。これは、重み転送の負担を抑えながら、大規模モデルをパイプライン化する設計である。

パイプラインをリクエストで満たす

パイプライン方式でも、同時に1件しか処理できないわけではない。あるリクエストが後段にいる間、別のリクエストは中段、さらに別のリクエストは前段で処理できる。実運用では、記事の単純な層単位の例よりも細かい段階に分けるという。各段階は、単一リクエストではなく小さなマイクロバッチを処理することもできる。

小さなマイクロバッチは個々の応答遅延に有利で、大きなマイクロバッチは全体のスループットを高めやすい。一方、通常のアテンション層に残されたKVキャッシュ容量は、同時にパイプライン上に置けるリクエスト数を制限する。重み容量だけでなく、キャッシュ、同時実行数、通信、遅延を一体で設計する必要がある。

意義と注意点

この例は、ウェハーを追加することで単一アクセラレーターを超えるメモリと計算資源を確保し、パイプラインで処理を重ね合わせる方法を示している。ただし、59枚という数は普遍的な要件ではない。コンテキスト長、量子化方式、モデル構造、キャッシュ形式、リクエスト分布によって必要な構成は変わる。

出典：Hacker News