Moonshot AI Open-Sources MoonEP: A Perfectly Balanced Expert Parallelism Library for MoE Training

Moonshot AI 正式將其專家平行（Expert Parallelism，EP）通訊函式庫 MoonEP 開源，以 MIT 授權釋出。這套函式庫專為分散式混合專家（Mixture-of-Experts，MoE）模型的訓練與推論設計，目標是讓大規模專家平行運算的通訊效率大幅提升。MoonEP 在 Kimi K3 開放日（Kimi K3 Open Day）中同步亮相，與 K3 模型權重、技術報告，以及另外兩套基礎架構程式碼 FlashKDA 與 AgentEnv 一同公開。其中 FlashKDA 先前已開源，而 MoonEP 與 AgentEnv 則是本次新增的開源專案。MoonEP 正是 Moonshot AI 宣稱能讓 Kimi K3 模型擴展效率提升 2.5 倍的關鍵技術之一。Kimi K3 是一個擁有 2.8 兆參數的 MoE 模型，具備原生視覺能力與 100 萬 token 的上下文長度。 專家平行架構在 MoE 模型中扮演核心角色，但長期面臨一個結構性問題：路由器的分配極不平均。在專家平行機制中，路由器會將每個 token 送到其 Top-K 個專家，而這些專家分散在不同的運算節點（rank）上。然而路由器的決策很少能達到完美平衡，某些專家收到的 token 數量遠多於其他專家。Moonshot AI 在 GitHub 儲存庫中以 maxvio 指標量化這種偏差，其定義為 max_e (T_e / T̄) − 1，其中 T_e 是路由到專家 e 的 token 數，T̄ 是完美平衡下的期望值。maxvio 為 0 代表完全平衡，數值越大表示偏差越嚴重。這種不平衡的代價並非偶然，而是結構性的——集體通訊的延遲由最慢的參與者決定，因此最熱門的 rank 會直接拖慢整個迭代時間。更糟的是，每個 rank 的 token 數量每步都在變化，動態的激活形狀會導致 GPU 記憶體碎片化，並迫使每層進行主機端同步。 MoonEP 的核心創新在於引進動態冗餘專家（dynamic redundant experts）。其設計建立一個硬性不變量：無論路由偏差多大，每個 rank 都會收到恰好 S × K 個 token（S 是每個 rank 的輸入 token 數，K 是每個 token 路由的 Top-K 專家數）。這項保證是透過直接從當前路由器輸出中，即時規劃少量冗餘專家來達成。這些被複製的專家會在專家計算之前被預取（prefetch）。在反向傳播階段，這些冗餘專家的梯度會被歸約回它們原本所在的 rank。 MoonEP 的設計可歸納為三個特性：第一，完美平衡，即上述 S × K 保證，透過線上規劃的冗餘專家實現。第二，線上規劃，採用一個近乎最優的 GPU 規劃核心，運算開銷極低，基於 CUTLASS CuTe DSL 實作，並且在 setup.py 中固定使用 nvidia-cutlass-dsl==4.4.2 版本。第三，零拷貝與靜態形狀，合併了置換（permute）與反置換（unpermute）操作，token 會直接寫入遠端 rank 上已按專家分組的對應位置，然後將緩衝區的視圖返回給計算單元。這只需要一個固定大小的 S × K 緩衝區，且靜態已知的形狀消除了每層 MoE 所需的主機端同步。 MoonEP 與訓練或推論框架之間有明確的記憶體合約：每個專家投影必須對應一個連續的對稱記憶體權重張量，加上一個由規劃器產生的 cu_seqlens。VM 群組 GEMM 會接收一個單一的 [E+B, H, H'] 權重張量，其中 E 是總路由專家數，B 是每個 rank 的預取槽位數，H 是隱藏層大小，H' 是專家 FFN 中間層大小。由 dispatch 回傳的 cu_seqlens[E+B] 則負責選取哪些專家行是活躍的。連續性是硬性要求，因為群組 GEMM 僅透過行索引來定址專家。記憶體佈局清晰分割：行 [0, E) 存放所有 rank 的本機專家，每個 rank 擁有 E/R 行。每個 chunk 實際上是該 rank 的參數記憶體，透過對稱記憶體映射到所有 rank。行 [E, E+B) 則是本機的預取槽位，由 buffer.prefetch_weight 填入。預取槽位來自一個所有層共享的處理程序全域池，這意味著額外的記憶體成本是每個投影 B 個專家權重，而非每層 B 個。 如何設定 B 取決於工作負載。訓練階段必須使用 B = E/R，因為規劃器最多只會從每個 rank 的一個遠端主機群組複製專家，這個保證確保群組 GEMM 觸及的每個專家都是本機的。推論階段則允許 B < E/R，官方 README 建議 B = 3 至 4。如果某個 rank 需要的不同遠端專家數量超過 B，群組 GEMM 會直接透過對稱映射從主機 rank 讀取溢出的權重——速度會稍慢，但不影響正確性。訓練階段，權重佈局在 fp32 中會對應一個每個投影的 [E+B, H, H'] 梯度緩衝區。關鍵是，行 [E, E+B) 是由一個獨立的歸約緩衝區支持，而非參數梯度。被複製的專家的梯度是暫時的，必須對框架本身的梯度歸約不可見。每個 rank 會將所有 R 個歸約緩衝區映射為一個 [R, B, H, H'] 視圖，然後 reduce_grad 會從每個 rank 的槽位中讀取自己的專家數據，透過 NVLink 累加到本機參數梯度，並清空已消耗的槽位。 Moonshot AI 在 H20 硬體上，以 EP=8 的設定，對路由器偏差進行了全面掃描，並與 DeepEP v2 進行了基準測試。比較腳本 bench_vs_deepep.py 的預設參數為 S=8192、E=384、H=7168、K=8、H'=2048、32 個 SM，maxvio 目標設定為 0.2、1、10 和 20。兩個函式庫都接收來自相同種子的相同路由矩陣。測試結果顯示兩項重要發現：第一，零拷貝架構使通訊速度更快，因為消除了傳統通訊中從通訊緩衝區複製到使用者緩衝區的耗時步驟，因此在所有偏差等級下，MoonEP 的通訊時間都穩定低於 DeepEP v2。第二，完美平衡讓 MoonEP 幾乎不受偏差影響，其通訊時間隨著 maxvio 增加幾乎保持不變，而 DeepEP v2 的延遲則由最熱門 rank 決定，因此隨著偏差增加而持續惡化。 總結而言，MoonEP 的核心貢獻在於：每個 EP rank 都能確保收到恰好 S × K 個 token，無論路由偏差多大；平衡來自於在 GPU 上即時規劃的冗餘專家，並在專家計算前預取；靜態形狀消除了每層 MoE 的主機同步，並阻止了導致 DeepEP OOM 的記憶體碎片化；訓練階段需設定 B = E/R 預取槽位，推論階段可降至 B = 3 至 4 而不影響正確性。MoonEP 在 Kimi K3 開放日以 MIT 授權釋出，與 FlashKDA 和 AgentEnv 共同構成 Moonshot AI 開源基礎架構的最新陣容。