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

📄 arXiv: 2502.19811v3 📥 PDF

作者: Shulai Zhang, Ningxin Zheng, Haibin Lin, Ziheng Jiang, Wenlei Bao, Chengquan Jiang, Qi Hou, Weihao Cui, Size Zheng, Li-Wen Chang, Quan Chen, Xin Liu

分类: cs.DC, cs.AI, cs.LG

发布日期: 2025-02-27 (更新: 2025-03-04)

💡 一句话要点

COMET：面向混合专家模型，实现细粒度计算-通信重叠优化。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 分布式训练 通信-计算重叠 细粒度优化 任务调度 数据依赖分析 自适应负载均衡

📋 核心要点

  1. 现有MoE模型在分布式训练中面临严重的通信开销，粗粒度的通信-计算重叠方案效率低下。
  2. COMET通过数据依赖分析和任务重调度，实现了通信与计算的细粒度精确重叠。
  3. 实验表明，COMET显著加速了MoE模型的训练，单层加速1.96倍，端到端加速1.71倍。

📝 摘要（中文）

混合专家模型(MoE)已被广泛应用于扩展大型语言模型至万亿级参数，同时保持固定的计算成本。然而，分布式场景下大型MoE模型的开发面临着巨大的通信开销问题。在流行的模型和框架中，MoE层内的设备间通信可能占据整个模型执行时间的47%。因此，现有方法建议将MoE层中的通信与计算进行流水线处理以实现重叠。然而，这些粗粒度的重叠方案显著降低了计算效率，并且延迟隐藏效果欠佳。为此，我们提出了COMET，一个优化的MoE系统，具有细粒度的通信-计算重叠。通过数据依赖性分析和任务重新调度，COMET实现了精确的细粒度通信和计算重叠。通过自适应工作负载分配，COMET有效地消除了细粒度的通信瓶颈，并增强了其在各种场景中的适应性。评估表明，COMET将单个MoE层的执行速度提高了1.96倍，并且对于端到端执行，COMET平均提供了1.71倍的加速。COMET已被应用于拥有万级GPU集群的生产环境中，节省了数百万GPU小时。

🔬 方法详解

问题定义：论文旨在解决分布式MoE模型训练过程中，由于通信开销过大导致训练效率低下的问题。现有方法采用粗粒度的通信-计算重叠，但会降低计算效率，无法充分隐藏通信延迟。

核心思路：COMET的核心思路是实现细粒度的通信-计算重叠，通过精确地分析数据依赖关系和重新调度任务，使得通信和计算能够更紧密地交织在一起，从而最大限度地隐藏通信延迟，提高整体训练效率。

技术框架：COMET的整体框架包括数据依赖性分析、任务重调度和自适应工作负载分配三个主要模块。首先，进行数据依赖性分析，确定哪些计算任务依赖于哪些通信数据。然后，基于依赖关系，对任务进行重新调度，使得通信任务尽可能地与不依赖于该通信数据的计算任务并行执行。最后，通过自适应工作负载分配，平衡各个设备上的计算负载，避免出现通信瓶颈。

关键创新：COMET的关键创新在于实现了细粒度的通信-计算重叠。与现有方法的粗粒度重叠相比，COMET能够更精确地控制通信和计算的执行顺序，从而更有效地隐藏通信延迟。此外，自适应工作负载分配进一步优化了通信效率。

关键设计：COMET的关键设计包括：(1) 精确的数据依赖性分析算法，用于确定计算任务和通信数据之间的依赖关系；(2) 基于依赖关系的动态任务调度策略，用于最大化通信和计算的并行性；(3) 自适应工作负载分配机制，用于平衡各个设备上的计算负载，避免通信瓶颈。具体的参数设置和损失函数等细节在论文中未详细说明，属于未知信息。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

COMET在实验中表现出色，将单个MoE层的执行速度提高了1.96倍，端到端执行速度平均提高了1.71倍。该系统已在拥有数万个GPU的大规模集群的生产环境中采用，并节省了数百万GPU小时，证明了其在实际应用中的有效性和价值。

🎯 应用场景

COMET可广泛应用于大规模分布式MoE模型的训练，尤其是在需要高性能和高效率的场景下。例如，它可以用于训练更大规模的语言模型、推荐系统和图像识别模型。通过降低训练成本和时间，COMET能够加速AI技术的研发和部署，并推动相关领域的创新。

📄 摘要（原文）

Mixture-of-experts (MoE) has been extensively employed to scale large language models to trillion-plus parameters while maintaining a fixed computational cost. The development of large MoE models in the distributed scenario encounters the problem of large communication overhead. The inter-device communication of a MoE layer can occupy 47% time of the entire model execution with popular models and frameworks. Therefore, existing methods suggest the communication in a MoE layer to be pipelined with the computation for overlapping. However, these coarse grained overlapping schemes introduce a notable impairment of computational efficiency and the latency concealing is sub-optimal. To this end, we present COMET, an optimized MoE system with fine-grained communication-computation overlapping. Leveraging data dependency analysis and task rescheduling, COMET achieves precise fine-grained overlapping of communication and computation. Through adaptive workload assignment, COMET effectively eliminates fine-grained communication bottlenecks and enhances its adaptability across various scenarios. Our evaluation shows that COMET accelerates the execution of a single MoE layer by $1.96\times$ and for end-to-end execution, COMET delivers a $1.71\times$ speedup on average. COMET has been adopted in the production environment of clusters with ten-thousand-scale of GPUs, achieving savings of millions of GPU hours.