Moonshot AI 開源 MoonEP：面向 MoE 訓練的完美平衡專家並行庫

Moonshot AI 在 Kimi K3 開放日上正式開源了 MoonEP，這是一個專為分佈式混合專家（MoE）工作負載設計的專家並行（EP）通信庫。該庫旨在使大規模專家並行通信更加高效，並以 MIT 許可證發佈。

MoonEP 是 Kimi K3 開放日的一部分。除了 K3 模型權重和技術報告外，Moonshot 還發布了三個基礎設施代碼庫：MoonEP、FlashKDA 和 AgentEnv。其中 FlashKDA 此前已開源，而 MoonEP 和 AgentEnv 則是本次新發布的。MoonEP 是聲稱將 Kimi K3 的擴展效率提升 2.5 倍的關鍵創新之一。Kimi K3 是一個擁有 2.8 萬億參數、原生視覺能力和 1M 令牌上下文窗口的 MoE 模型。

MoonEP 針對的問題：在專家並行中，路由器將每個令牌發送給其 top-K 個專家，這些專家位於不同的排名上。路由器很少是平衡的，一些專家接收的令牌遠多於其他專家。該代碼庫使用 maxvio 來量化偏斜，定義為 max_e (T_e / T̄) − 1，其中 T_e 是路由到專家 e 的令牌數，T̄ 是完美平衡下的預期數量。maxvio 為 0 表示完全平衡。不平衡的成本是結構性的，而非偶然。集合的延遲由其最慢的參與者決定，因此最熱的排名決定迭代時間。更糟的是，每個步驟的令牌計數都會變化，這些動態激活形狀會碎片化 GPU 內存，並強制每層主機同步。

核心思想：動態冗餘專家。MoonEP 的主要創新是一個硬性不變量：無論路由多麼偏斜，每個排名接收恰好 S × K 個令牌——其中 S 是每個排名的輸入令牌數，K 是每個令牌的 routed top-k。它通過直接從當前路由器的輸出在線規劃少量冗餘專家來實現這一點。這些重複的專家在專家計算之前被預取。在反向傳播中，它們的梯度被歸約回其主排名。

該設計具有三個屬性：完美平衡（通過在線規劃的冗餘專家實現 S × K 保證）、在線規劃（一個近最優的 GPU 規劃內核，開銷可忽略，使用 CUTLASS CuTe DSL 實現，setup.py 鎖定 nvidia-cutlass-dsl==4.4.2）、零複製和靜態形狀（融合的排列/反排列，令牌直接寫入遠程排名上的專家分組位置，緩衝區視圖返回給計算，只需固定的 S × K 緩衝區，靜態已知形狀消除每層 MoE 主機同步）。

內存契約：MoonEP 與訓練或推理框架的契約是具體的：每個專家投影一個連續的對稱內存權重張量，加上一個規劃器產生的 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，因為規劃器每個排名最多從一個遠程主組複製專家，該邊界保證組 GEMM 接觸的每個專家都是本地的。推理允許 B < E/R，README 推薦 B = 3–4。如果一個排名需要多於 B 的不同遠程專家，組 GEMM 會通過對稱映射直接從主排名讀取溢出權重——稍慢，但不影響正確性。

訓練中，權重佈局在 fp32 中鏡像，每個投影有一個 [E+B, H, H'] 梯度緩衝區。關鍵的是，行 [E, E+B) 由一個單獨的歸約緩衝區支持，而非參數梯度。重複專家的梯度是臨時的，必須對框架自身的梯度歸約不可見。每個排名將所有 R 個歸約緩衝區映射為一個 [R, B, H, H'] 視圖，然後 reduce_grad 從每個排名通過 NVLink 讀取其自己專家的槽，累積到本地參數梯度，並清零已消耗的槽。

與 DeepEP v2 的基準測試：兩個已發佈的基準測試在 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。兩個庫接收來自共享種子的相同路由矩陣。報告了三個發現：零複製使原始通信更快，消除了困擾 DeepEP v2 後段的複製開銷，因此 MoonEP 的通信時間在所有不平衡水平上都低於 DeepEP v2；完美平衡使 MoonEP 幾乎不受偏斜影響，通信時間隨 maxvio 增長几乎不變，而 DeepEP v2 的延遲由最熱排名決定，穩定下降。

關鍵要點：MoonEP 聲稱每個 EP 排名接收恰好 S × K 個令牌，無論路由器偏斜如何；平衡來自在線在 GPU 上規劃的冗餘專家，然後在專家計算前預取；靜態形狀消除每層 MoE 主機同步，並阻止導致 DeepEP 內存溢出的碎片；訓練需要 B = E/R 預取槽，推理可降至 B = 3–4 而無正確性成本；以 MIT 許可證在 Kimi K3 開放日發佈，與 FlashKDA 和 AgentEnv 一同。