MoonshotAI/MoonEP

解决的问题

MoonEP 解决了在专家并行（EP）过程中混合专家（MoE）模型的负载不均衡问题。在传统 MoE 中，某些专家接收的令牌数量远超其他专家，导致出现“热排名”现象，从而拖慢整个训练或推理过程，并可能引发 GPU 内存碎片化或内存溢出（OOM）错误。

工作原理

MoonEP 确保无论路由分布如何，每个 GPU 排名都接收相同数量的令牌。它通过以下方式实现这一目标：

• 动态冗余专家：在线规划少量冗余专家并预取到需要的排名，确保所有排名之间的负载完全均衡。
• 在线规划：高性能 GPU 内核以极低开销处理这些冗余专家的规划。
• 零拷贝通信：使用融合的 permute/unpermute 操作将令牌直接发送到远程排名的最终位置，将通信缓冲区的视图直接提供给计算，从而消除昂贵的数据复制。
• 静态内存形状：通过保持令牌的固定缓冲区大小，防止内存碎片化并确保可预测的内存使用。

适用对象

需要优化通信开销并在使用 NVLink 的多 GPU 上保持稳定训练性能的大规模 MoE 模型开发者和研究人员。

主要亮点

• 抗不均衡：即使路由不均衡增加，迭代时间也保持平稳。
• 零拷贝效率：通过消除通信缓冲区与用户缓冲区之间的复制，实现更快的原始通信速度。
• 静态内存：在高不均衡情况下防止 GPU 内存碎片化和 OOM 错误。
• 集成权重预取：包含专门的内核，用于预取冗余专家的权重并减少其梯度。