Moonshot AI 开源 MoonEP:用于 MoE 训练的完美均衡专家并行库
Moonshot AI
Moonshot AI 开源了 MoonEP,这是一个用于分布式混合专家工作负载的专家并行通信库。MoonEP 通过在线规划冗余专家,确保每个计算节点无论路由偏斜如何都能精确接收 S×K 个 token。它作为 Kimi K3 开放日的一部分,以 MIT 许可证发布。
Moonshot AI 已开源 MoonEP,这是一个用于分布式混合专家(MoE)工作负载的专家并行(EP)通信库,基于 MIT 许可证发布。MoonEP 作为 Kimi K3 开放日的一部分,与 K3 模型权重、技术报告以及其他基础设施代码库 FlashKDA 和 AgentEnv 一同发布。它通过强制执行每个秩(rank)恰好接收 S×K 个 token 这一硬性约束,解决了专家并行中的路由器不平衡问题,这一效果通过在计算前预取冗余专家的在线规划实现。该设计确保了完美平衡,通过 GPU 内核进行在线规划,开销可忽略不计,零拷贝和静态形状消除了每层的主机同步。与 DeepEP v2 的基准测试显示,随着不平衡程度增加,MoonEP 的通信时间保持稳定,而 DeepEP 则性能下降。训练时需要 B = E/R 个预取槽位;推理时可以使用 B = 3-4,而不会影响正确性。
来源: MarkTechPost —
原文
