MoonshotAI/MoonEP

解決する課題

MoonEPは、Expert Parallelism (EP) におけるMixture-of-Experts (MoE) モデルの負荷不均衡問題に対処します。従来のMoEでは、一部のエキスパートが他のエキスパートよりもはるかに多くのトークンを受け取るため、"ホットランク" が発生し、全体の学習または推論プロセスを遅くするだけでなく、GPUメモリのフラグメンテーションやメモリ不足（OOM）エラーを引き起こす可能性があります。

動作方法

MoonEPは、ルーティング分布にかかわらず、すべてのGPUランクが同じ数のトークンを受け取ることを保証します。これにより、以下の仕組みが実現されます：

• 動的冗長エキスパート：オンラインで少数の冗長エキスパートを計画し、必要とするランクに事前にプリフェッチすることで、すべてのランク間で負荷が完全にバランスされるようにします。
• オンライン計画：低遅延のGPUカーネルが、これらの冗長エキスパートの計画を最小限のオーバーヘッドで処理します。
• ゼロコピー通信：トークンをリモートランクの最終位置に直接送信するために、融合された permute/unpermute 操作を使用し、通信バッファのビューを計算に直接提供することで、高コストなデータコピーを排除します。
• 静的メモリ形状：トークン用のバッファサイズを固定することで、メモリフラグメンテーションを防ぎ、予測可能なメモリ使用量を確保します。

対象ユーザー

NVLinkを用いた複数GPUで安定した学習パフォーマンスを維持しながら、大規模なMoEモデルを構築する開発者や研究者向けです。

主な特徴

• 不均衡に強い：ルーティングの不均衡が増加しても、イテレーション時間は一定に保たれます。
• ゼロコピー効率：通信バッファとユーザーバッファ間のコピーを排除することで、通信速度が向上します。
• 静的メモリ：高不均衡状態でもGPUメモリのフラグメンテーションやOOMを防止します。
• 統合された重みプリフェッチ：冗長エキスパートのプリフェッチおよび勾配の削減を専用カーネルで実行します。