开源 MoonEP：MoE 分布式训练通信库

Moonshot AI 开源 MoonEP：为 MoE 训练打造的完美均衡专家并行库

Moonshot AI 开源了 MoonEP，这是一个为分布式混合专家（Mixture-of-Experts, MoE）工作负载设计的专家并行（Expert Parallelism, EP）通信库。该团队宣布发布这一库是为了在规模化时实现更高效的专家并行通信。它采用 MIT 许可证发布。

MoonEP 作为 Kimi K3 开放日的一部分推出。除了 K3 模型权重和技术报告外，Moonshot 还发布了三个基础设施代码库：MoonEP、FlashKDA 和 AgentEnv。FlashKDA 已经开源；MoonEP 和 AgentEnv 在此次发布中推出。MoonEP 是 Kimi K3 宣称实现的 2.5 倍扩展效率改进背后的创新之一，Kimi K3 是一个拥有 2.8 万亿参数的 MoE 模型，具备原生视觉能力和 100 万 token 的上下文窗口。

MoonEP 针对的问题

在专家并行中，路由器将每个 token 发送给其 top-K 个专家，这些专家分布在不同的进程（rank）上。路由器很少能实现均衡。一些专家收到的 token 远多于其他专家。

该库用 maxvio 来量化不均衡程度，定义为 max_e (T_e / T̄) − 1，其中 T_e 是路由到专家 e 的 token 数，T̄ 是完美均衡下的预期数量。maxvio 为 0 表示完美均衡。

不均衡的成本是结构性的，而非偶发的。集合的延迟由其最慢的参与者决定，因此最繁忙的进程决定了迭代时间。更糟的是，每个进程的 token 数在每一步都会变化。这些动态激活形状会使 GPU 内存碎片化，并强制每层进行主机同步。

核心思想：动态冗余专家

MoonEP 的主要特性是一个硬性不变量。每个进程恰好接收 S × K 个 token，无论路由多么不均衡 —— 其中 S 是每个进程的输入 token 数，K 是每个 token 的 top-k 路由值。

它通过在线规划少量冗余专家来实现这一点，直接从当前路由器输出生成。这些重复的专家在专家计算前被预取。在反向传播中，它们的梯度被约化回各自的主进程。

该设计分为三个特性：

完美均衡：通过上述 S × K 保证，利用在线规划的冗余专家实现。

在线规划：一个近似最优的 GPU 规划内核，开销可忽略。它用 CUTLASS CuTe DSL 实现；setup.py 固定 nvidia-cutlass-dsl==4.4.2。

零拷贝和静态形状：融合排列/逆排列。Token 直接写入远程进程上专家分组的位置，缓冲区视图返回给计算使用。只需要一个固定的 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) 由单独的 reduce 缓冲区支持，而不是参数梯度。冗余专家的梯度是临时的，必须对框架自身的梯度约化不可见。每个进程将所有 R 个 reduce 缓冲区映射为一个 [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。两个库从共享种子接收相同的路由矩阵。

他们的 GitHub 页面报告了三项发现。零拷贝通过消除 comm-buffer → user-buffer 副本（主导尾声）使原始通信更快，所以 MoonEP 的通信时间在每个不均衡级别都始终低于 DeepEP v2。完美均衡使 MoonEP 几乎不受倾斜影响：随着 maxvio 增长，其通信时间保持几乎平坦，而 DeepEP v2 —— 其延迟由最繁忙的进程决定 —— 稳定下降。

• MoonEP 声称每个 EP 进程恰好接收 S × K 个 token，无论路由倾斜如何。
• 均衡来自在线规划的冗余专家，在 GPU 上规划，然后在专家计算前预取。
• 静态形状消除了每层 MoE 主机同步，并阻止了导致 DeepEP OOM 的内存碎片化。
• 训练需要 B = E/R 个预取槽；推理可以降低到 B = 3–4，无正确性成本。
• 在 Kimi K3 开放日以 MIT 许可证发布，同 FlashKDA 和 AgentEnv。

来源：GitHub 上的 MoonshotAI/MoonEP 和 @Kimi_Moonshot 公告

Michal Sutter 是一位数据科学专业人士，拥有帕多瓦大学数据科学理学硕士学位。凭借统计分析、机器学习和数据工程的扎实基础，Michal 擅长将复杂数据集转化为可行的洞察。