MoonshotAI/MoonEP

解决的问题

MoonEP 解决了混合专家（MoE）模型中专家并行（EP）的效率瓶颈，特别是负载不均衡问题。在标准 EP 中，某些 GPU 排名可能接收远多于其他排名的令牌（"最热排名"问题），导致延迟增加、内存碎片化，以及训练过程中可能出现的内存溢出（OOM）错误。

工作原理

MoonEP 确保无论路由器如何分配，每个 GPU 排名处理的令牌数量完全相同。它通过以下方式实现：

• 动态冗余专家：该库在线规划少量冗余专家并预取它们，以在各排名间实现完美的负载均衡。
• 零拷贝通信：使用融合的 permute/unpermute 操作，将令牌直接发送到远程排名的最终专家分组位置，并返回缓冲区视图给计算，从而消除不必要的数据复制。
• 在线规划：高性能 GPU 内核以可忽略的开销确定专家和令牌的最佳分配。
• 静态内存形状：通过保持令牌的固定缓冲区大小，消除了每层主机同步的需要，并防止内存碎片化。

适用对象

该库专为需要优化通信开销并确保在多个 NVIDIA GPU（以及未来支持 Zhenwu PPUs）上实现稳定、均衡训练和推理的大规模 MoE 模型开发者和研究人员设计。

主要亮点

• 完美负载均衡：每个排名接收完全相同的令牌数量，使迭代时间不受路由不均衡的影响。
• 零拷贝分发：消除了通信缓冲区到用户缓冲区的复制，加快了原始通信速度。
• 防止 OOM：静态内存形状防止了在负载不均衡的 MoE 训练中通常因激活形状变化而产生的内存碎片化。
• 集成权重预取：包含专用机制，将冗余专家权重预取到本地槽位，以实现高效计算。