MoonshotAI/MoonEP

何を解決するか

MoonEPは、Mixture-of-Experts (MoE) モデルにおけるExpert Parallelism (EP) の効率的なボトルネック、特に負荷の不均衡問題に対処します。標準的なEPでは、一部のGPUランクが他のランクよりも著しく多くのトークンを受け取る（"ホットランク"問題）ため、遅延の増加、メモリフラグメンテーション、およびトレーニング中のOut-of-Memory (OOM) エラーが発生する可能性があります。

動作方法

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

• 動的冗長エキスパート: ライブラリはオンラインで少数の冗長エキスパートを計画し、事前に読み込んでランク間の負荷を完全にバランスさせます。
• ゼロコピー通信: トークンをリモートランク上の最終的なエキスパートグループ位置に直接送信するために、融合されたpermute/unpermute操作を使用し、バッファビューを計算に返すことで、不要なデータコピーを排除します。
• オンライン計画: 高性能なGPUカーネルが、最小限のオーバーヘッドでエキスパートとトークンの最適な分配を決定します。
• 静的メモリ形状: トークン用のバッファサイズを固定することで、レイヤーごとのホスト同期の必要性を排除し、メモリフラグメンテーションを防止します。

対象ユーザー

このライブラリは、大規模なMoEモデルを構築し、複数のNVIDIA GPU（および将来的にはZhenwu PPUs）上で通信オーバーヘッドを最適化し、安定したバランスの取れたトレーニングと推論を確保したい開発者および研究者向けです。

主な特徴

• 完全な負荷バランス: すべてのランクが正確に同じ数のトークンを受け取り、イテレーション時間はルーティングの不均衡に影響されません。
• ゼロコピーディスパッチ: 通信バッファからユーザーバッファへのコピーを排除し、通信の速度を向上させます。
• OOM防止: 静的メモリ形状により、負荷が不均衡なMoEトレーニングで通常発生するアクティベーション形状の変動に起因するメモリフラグメンテーションを防止します。
• 統合された重みプリフェッチ: 冗長エキスパートの重みをローカルスロットに効率的にプリフェッチする専用メカニズムを備えています。