Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts

专家混合（Mixture-of-Experts, MoE）模型被广泛应用于扩展大型语言模型至万亿参数以上，同时保持固定的计算成本。然而，在分布式场景中开发大型MoE模型时遇到了通信开销过大的问题。MoE层的跨设备通信可能占用整个模型执行时间的47%，即使使用流行的模型和框架也是如此。因此，现有方法建议将MoE层中的通信与计算流水线化以实现重叠，但这些粗粒度的重叠方案会导致计算效率显著下降，并且延迟隐藏效果不理想。 为此，我们提出了COMET，这是一种具有细粒度通信-计算重叠优化的MoE系统。通过数据依赖性分析和任务重新调度，COMET实现了通信与计算的精确细粒度重叠。借助自适应工作负载分配，COMET有效消除了细粒度通信瓶颈，并增强了其在各种场景下的适应性。我们的评估显示，COMET可以加速单个MoE层的执行1.96倍，而在端到端执行中，COMET平均提供了1.71倍的速度提升。COMET已经在拥有数千规模GPU的生产环境中得到应用，节省了数百万个GPU小时。

论文试图解决在大规模分布式环境下，混合专家模型（MoE）中存在的通信开销过大的问题。具体来说，MoE层的设备间通信占用了整个模型执行时间的47%，这显著影响了模型训练和推理的效率。虽然已有方法尝试通过粗粒度的通信与计算重叠来缓解这一问题，但效果并不理想。

论文提出了一种名为COMET的优化系统，旨在通过细粒度的通信与计算重叠来提高MoE模型的执行效率。COMET利用数据依赖性分析和任务重新调度技术，实现了精确的通信与计算重叠，并通过自适应工作负载分配消除了细粒度通信瓶颈。相比现有的粗粒度方案，COMET能够更有效地隐藏通信延迟并提升计算效率。

该研究展示了COMET在单个MoE层上的加速比为1.96倍，在端到端执行中平均加速比为1.71倍。此外，COMET已经在生产环境中被广泛采用，特别是在拥有数千规模GPU的集群中，节省了数百万个GPU小时。实验设计方面，作者使用了多种流行的大型语言模型和框架进行测试，确保了结果的普适性和可靠性。目前尚未提及是否有开源代码，但这项工作的创新点在于其对通信瓶颈的细粒度处理，值得进一步研究。

最近在这个领域，相关研究还包括： 1.《Scaling Mixture of Experts to Trillion Parameters》探讨了如何将MoE模型扩展到万亿参数。 2.《Efficient Training of Mixture-of-Experts with Sparse Communication》提出了稀疏通信机制以减少MoE模型中的通信成本。 3.《Optimizing Communication in Distributed Deep Learning》研究了分布式深度学习中的通信优化策略。 这些研究都旨在解决分布式训练中的通信效率问题，而COMET则提供了更加细粒度和高效的解决方案。