MoonshotAI/MoonEP

解決的問題

MoonEP 解決了混合專家（MoE）模型中專家並行（EP）的效率瓶頸，特別是負載不均衡問題。在標準 EP 中，某些 GPU 排名可能接收遠多於其他排名的令牌（"最熱排名"問題），導致延遲增加、記憶體碎片化，以及訓練過程中可能出現的記憶體溢出（OOM）錯誤。

工作原理

MoonEP 確保無論路由器如何分配，每個 GPU 排名處理的令牌數量完全相同。它透過以下方式實現：

• 動態冗餘專家：該庫在線規劃少量冗餘專家並預取它們，以在各排名間實現完美的負載均衡。
• 零複製通訊：使用融合的 permute/unpermute 操作，將令牌直接發送到遠端排名的最終專家分組位置，並返回緩衝區視圖給計算，從而消除不必要的資料複製。
• 線上規劃：高性能 GPU 內核以可忽略的開銷確定專家和令牌的最佳分配。
• 靜態記憶體形狀：透過保持令牌的固定緩衝區大小，消除了每層主機同步的需要，並防止記憶體碎片化。

適用對象

該庫專為需要優化通訊開銷並確保在多個 NVIDIA GPU（以及未來支援 Zhenwu PPUs）上實現穩定、均衡訓練和推論的大規模 MoE 模型開發者和研究人員設計。

主要亮點

• 完美負載均衡：每個排名接收完全相同的令牌數量，使迭代時間不受路由不均衡的影響。
• 零複製分發：消除了通訊緩衝區到使用者緩衝區的複製，加快了原始通訊速度。
• 防止 OOM：靜態記憶體形狀防止了在負載不均衡的 MoE 訓練中通常因激活形狀變化而產生的記憶體碎片化。
• 整合權重預取：包含專用機制，將冗餘專家權重預取到本地槽位，以實現高效計算。