Moonshot AI、MoEトレーニングのための完全にバランスの取れたエキスパート並列ライブラリ「MoonEP」をオープンソース化

Moonshot AIは、Kimi K3オープンデイにおいて、分散混合エキスパート（MoE）ワークロード向けのエキスパート並列（EP）通信ライブラリ「MoonEP」を正式にオープンソース化しました。このライブラリは、大規模なエキスパート並列通信をより効率化するために構築され、MITライセンスで公開されています。

MoonEPは、Kimi K3オープンデイの一環として登場しました。K3モデルの重みと技術レポートに加え、MoonshotはMoonEP、FlashKDA、AgentEnvの3つのインフラストラクチャコードベースをリリースしました。FlashKDAは既にオープンソース化されており、MoonEPとAgentEnvは今回新たに公開されました。MoonEPは、2.8兆パラメータ、ネイティブビジョン、1Mトークンコンテキストウィンドウを持つMoEモデル「Kimi K3」のスケーリング効率を2.5倍向上させたとされる革新的技術の1つです。

MoonEPが対象とする問題：エキスパート並列では、ルーターが各トークンをそのtop-Kエキスパートに送信しますが、これらは異なるランクに存在します。ルーターのバランスが取れることは稀で、一部のエキスパートが他のエキスパートよりもはるかに多くのトークンを受け取ります。リポジトリでは、maxvio（max_e (T_e / T̄) − 1、T̄は完全バランス時の期待値）で偏りを定量化しています。maxvioが0は完全バランスを示します。不均衡のコストは構造的であり、偶発的なものではありません。コレクティブのレイテンシは最も遅い参加者によって決まるため、最もホットなランクがイテレーション時間を決定します。さらに、トークン数はステップごとに変化し、動的なアクティベーション形状がGPUメモリを断片化し、レイヤーごとのホスト同期を強制します。

核となるアイデア：動的冗長エキスパート。MoonEPの主要な革新はハードな不変条件です。ルーティングがどれほど偏っていても、各ランクは正確にS × K個のトークンを受け取ります（Sはランクあたりの入力トークン数、Kはトークンあたりのルーティングtop-K）。これは、現在のルーター出力から直接オンラインで少数の冗長エキスパートを計画することで実現されます。これらの複製されたエキスパートは、エキスパート計算の前にプリフェッチされます。逆伝播では、それらの勾配はホームランクに還元されます。

設計は次の3つの特性で構成されます：完全バランス（上記のS × K保証をオンライン計画の冗長エキスパートで実現）、オンライン計画（オーバーヘッドが無視できるほどのほぼ最適なGPU計画カーネル。CUTLASS CuTe DSLで実装され、setup.pyはnvidia-cutlass-dsl==4.4.2に固定）、ゼロコピーと静的シェイプ（融合されたpermute/unpermute。トークンはリモートランク上のエキスパートグループ位置に直接書き込まれ、バッファビューが計算に返される。固定のS × Kバッファのみが必要で、静的に既知のシェイプによりレイヤーごとのMoEホスト同期が排除されます）。

メモリ契約：MoonEPとトレーニングまたは推論フレームワークとの契約は具体的です。各エキスパートプロジェクションにつき1つの連続した対称メモリ重みテンソルと、プランナーが生成するcu_seqlensです。VMグループGEMMは、単一の[E+B, H, H']重みテンソルを消費します。ここで、Eは総ルーティングエキスパート数、Bはランクあたりのプリフェッチスロット数、Hは隠れサイズ、H'はエキスパートFFN中間サイズです。dispatchによって返されるcu_seqlens[E+B]が、どのエキスパート行がアクティブかを選択します。連続性はハード要件です。グループGEMMは行インデックスでエキスパートをアドレス指定するためです。レイアウトは明確に分割されます：行[0, E)は全ランクのローカルエキスパートを保持し、ランクあたりE/R行。各チャンクは物理的にホームランクのパラメータメモリであり、対称メモリを介してすべての場所にマッピングされます。行[E, E+B)はローカルプリフェッチスロットであり、buffer.prefetch_weightで満たされます。プリフェッチスロットはすべてのレイヤーで共有されるプロセスグローバルプールから描画されます。この詳細は重要です。追加のメモリコストは、プロジェクション全体でB個のエキスパート重みであり、レイヤーごとではありません。

Bの設定方法はワークロードに依存します。トレーニングではB = E/Rを使用する必要があります。これは、プランナーがランクあたり最大1つのリモートホームグループからエキスパートを複製するためであり、この境界によりグループGEMMが触れるすべてのエキスパートがローカルであることが保証されます。推論ではB < E/Rが可能で、READMEはB = 3–4を推奨しています。ランクがBよりも多くの異なるリモートエキスパートを必要とする場合、グループGEMMは対称マッピングを介してホームランクから直接オーバーフロー重みを読み取ります。これは少し遅くなりますが、正確性には影響しません。

トレーニングでは、重みレイアウトはfp32でミラーリングされ、プロジェクションごとに[E+B, H, H']勾配バッファがあります。重要なのは、行[E, E+B)は別の還元バッファによってバックアップされており、パラメータ勾配ではありません。複製されたエキスパートの勾配は一時的であり、フレームワーク自身の勾配還元からは見えないようにする必要があります。各ランクはすべてのR個の還元バッファを1つの[R, B, H, H']ビューとしてマッピングし、reduce_gradは各ランクからNVLinkを介して自身のエキスパートのスロットを読み取り、ローカルパラメータ勾配に累積し、消費されたスロットをゼロにします。

DeepEP v2とのベンチマーク：公開された2つのベンチマークはH20上でEP=8、ルーター不均衡をスイープして実行されます。比較スクリプトbenchmarks/bench_vs_deepep.pyはデフォルトでS=8192, E=384, H=7168, K=8, H'=2048, 32 SM、maxvioターゲット0.2, 1, 10, 20を使用します。両ライブラリは共有シードからの同一のルーティングマトリックスを受け取ります。3つの結果がGitHubページで報告されました：ゼロコピーにより生の通信が高速化され、DeepEP v2のエピローグを支配するcomm-buffer→user-bufferコピーが排除されるため、MoonEPの通信時間はすべての不均衡レベルで一貫してDeepEP v2を下回ります；完全バランスによりMoonEPは偏りの影響をほとんど受けず、maxvioが増加しても通信時間はほぼ一定ですが、DeepEP v2のレイテンシは最もホットなランクによって決まるため、着実に低下します。

主なポイント：MoonEPは、ルーターの偏りに関係なく各EPランクが正確にS × K個のトークンを受け取ると主張；バランスはオンラインでGPU上で計画され、エキスパート計算前にプリフェッチされる冗長エキスパートから得られる；静的シェイプによりレイヤーごとのMoEホスト同期が排除され、DeepEPをOOMさせるメモリ断片化が防止される；トレーニングではB = E/Rのプリフェッチスロットが必要、推論ではB = 3–4で正確性に影響なし；Kimi K3オープンデイでMITライセンスでリリースされ、FlashKDAおよびAgentEnvと共に公開。