万卡集群真实部署,已节省数百万 GPU 小时!MoE 通信优化技术 COMET 开源

万卡集群真实部署,已节省数百万 GPU 小时!MoE 通信优化技术 COMET 开源

日期

2025-03-10

分类

技术发布

当前,MoE 架构是业界拓展模型规模的重要方向,然而,其在分布式训练中存在的大量通信开销,仍严重制约了训练效率和成本。为攻克这一瓶颈,豆包大模型团队提出了一个全新的​通信优化系统 COMET​,通过更精准、细粒度的计算-通信重叠技术,在大规模 MoE 模型上可达到单层 1.96 倍加速,端到端平均 1.71 倍效率提升,且在不同并行策略、输入规模及硬件环境下均表现稳定。


目前,COMET 已实际应用于万卡级生产集群​,助力 MoE 模型高效训练,并已累计节省了数百万 GPU 小时资源。此外,COMET 还可与豆包大模型团队此前发布的新一代稀疏模型架构 UltraMem 结合,实现协同优化。


该工作获 MLSys 2025 会议 5/5/5/4 高分评审,核心代码​已开源​。


image


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


论文链接​:https://arxiv.org/pdf/2502.19811


开源地址​:https://github.com/bytedance/flux


混合专家模型(MoE)通过稀疏激活机制突破了传统稠密模型(Dense Model)的计算瓶颈,然而,MoE 的分布式训练仍面临一项严峻挑战:​跨设备通信开销巨大​。例如,Mixtral-8x7B 模型在 Megatron-LM 框架中的通信时间占比可高达 40%,严重制约了训练效率和成本。


核心问题在于,MoE 的专家网络分布在多个 GPU 上,每次计算需频繁执行 Token 分发与结果聚合,导致 GPU 计算资源大量闲置。因此,如何将通信隐藏到计算的过程中,提升模型训练效率、节省计算资源,成为了 MoE 系统优化的关键。


1. 难点:「复杂的数据依赖」与「流水线气泡」


为了掩盖巨大的通信开销,现有方案主要集中在​如何对「计算-通信」进行高效重叠​​。


一种方案是将流水线调度与通信算子结合起来,即通过定制训练中流水线并行的调度方式,将不同 microbatch 的计算和通信进行重叠,如 DeepSeek 的 DualPipe。但是,这一方式会导致较大的显存开销,并需要对现有训练框架进行复杂的​侵入性改动​。


其它 MoE 系统方案则是在 microbatch 内部采用了粗粒度的计算-通信流水线,将输入数据分割成「数据块」进行通信与计算的重叠。然而,这种粗粒度的重叠方式难以高效利用计算资源,且无法实现无缝的通信延迟隐藏,尤其在动态路由、异构硬件环境下,​性能损失显著​。


因此,团队认为现有的系统级 MoE 解决方案仍面临两大困境:


1)难以解决复杂的数据依赖

MoE 架构的稀疏特性导致计算和通信间的依赖动态且复杂。MoE 会动态地将 Token 分配给不同专家,而传统的粗粒度矩阵分块方式,会导致 GPU 频繁等待远程数据,从而造成计算资源闲置。‍


如图 1 所示,当专家 0 需要在紫色「数据块」中进行 Tile-level 的计算时,必须先通过 Token-level 的通信接收远程数据(Token B),这种由于复杂数据依赖导致的计算-通信粒度上的错配,使得效率严重下滑。


image


图 1:单层 MoE 模型示意图(专家分布在 GPU0 和 GPU1 两张卡上)


2)难以消除计算-通信流水线气泡

另一个问题是,现有方法无法精确控制计算任务和通信任务对硬件资源的使用,因而,也无法根据不同的模型结构和动态输入,来自适应地调整资源分配。这导致计算和通信无法实现无缝重叠,进而产生大量流水线气泡,增加了系统的延迟。


因此,团队认为:解决 MoE 模型中计算与通信的粒度不匹配问题是实现两者高效重叠的关键,同时,还需要根据负载情况自适应调整通信和计算的资源分配,以进一步实现无缝重叠。


2. COMET 核心方案


COMET 是一个针对 MoE 模型的通信优化系统,通过细粒度计算-通信重叠技术,助力大模型训练优化。


团队分析发现,MoE 架构包含两条不同的生产-消费流水线:「计算-通信流水线」和「通信-计算流水线」。如图 2 所示,数据在流水线中流动时,各流水线内的操作会通过一个共享缓冲区链接,该缓冲区被称作「共享张量」。


image

图 2:COMET 的设计结构


基于此,COMET 引入​两项关键机制​,以最小化整体延迟并提升流水线性能。


1)共享张量依赖解析

通过分解和重调度共享张量,解决通信与计算之间的粒度错配问题,实现细至单 Token 级的重叠。


张量分解:​将 MoE 层间传递的共享张量沿 Token 维度(M)或隐层维度(N)进行切割,使通信与计算的最小单元对齐。例如,在 MoE 第一层(Layer 0,图 3 左)沿 M 维度分解,使通信和计算在 M 维度进行对齐;在 MoE 第二层(Layer 1,图 3 右)沿 N 维度分解,细粒度传输 Token 结果,保证计算和通信的高效重叠。


image

图 3:COMET 对共享张量进行依赖解析和分解


计算重调度​:为了更好地隐藏计算与通信的延迟,COMET 会动态调整数据块的计算顺序。例如,优先计算本地数据块,同时异步拉取远程 Token。当某个专家需处理 Token A(本地)和 Token B(远程)时,系统会优先启动 Token A 的计算线程,并与 Token B 的通信线程并行执行,从而消除等待延迟。


image

图 4:COMET 在 MoE layer0 中分解并重新调度共享张量


2)自适应负载分配

动态分配 GPU 线程块资源,精准平衡通信与计算负载,消除流水线气泡。


线程块隔离​:将通信与计算任务分别封装在独立线程块中,避免远程 I/O 阻塞计算核心。在 Nvidia Hopper 架构中,计算线程块专用于执行异步 TMA 指令的 GEMM 运算,通信线程块通过 NVSHMEM 实现单 Token 级数据传输,这种设计赋予了系统在算子级别进行资源管理的能力。


image

图 5:COMET 的计算/通信线程块隔离设计


动态负载平衡​:根据输入规模(如 Token 长度 M)、并行策略(EP/TP 比例)实时调整线程块分配。如图 6 所示,当 TP=8、EP=1 时,通信线程块占所有线程块的比例为 19.7%,而当 TP=4、EP=2,该比例需提升至 34.8%,系统通过预编译多个版本的计算-通信融合算子实现在运行时的「零开销」算子动态切换,并始终提供低延迟的算子。


image

图 6:单个 MoE 层使用不同数量的通信线程块的时延结果


3. 大规模落地验证


团队在多个大规模 MoE 模型中评估了 COMET 的端到端性能。结果表明,COMET 在 8 卡 H800 的实验集群中,​端到端 MoE 模型​(Mixtral-8x7B、Qwen2-MoE 等)的前向时延较其他基线系统可降低 31.8%-44.4%​,且在不同并行策略、输入规模及硬件环境下均表现稳定。


image

图 7:COMET 在多个 MoE 模型中的测评结果


在单个 MoE 层上,当输入 Token 数量不同的情况下,COMET 的执行时间均显著短于基线方案,平均实现了 1.28 倍到 2.37 倍的速度提升。


image

图 8:COMET 在单个 MoE 层不同输入 Token 长度下的延迟情况


目前,COMET 已实际应用于万卡级生产集群,助力 MoE 模型高效训练,并已累计节省数百万 ​GPU​​ 小时。​该工作在 MLSys 2025 会议获得 5/5/5/4 的评审高分,并被认为在大规模生产环境中极具应用潜力。


具体表现为——


强鲁棒性​:COMET 采用的细粒度计算-通信重叠方案,即使在专家负载不均衡的场景下,也能保持低于其它基线系统的延迟,具有较好的鲁棒性;


强泛化能力​:COMET 在 NVLink 与 PCIe 等不同网络环境下均能提供稳定的加速比;在使用不同并行策略时均能生成低时延算子,以供大规模训练框架使用。


4. 核心代码开源


COMET 包含约 1.2 万行 C++ 和 CUDA 代码,以及 2 千行 Python 代码,并向开发者提供了一套友好的 Python API。


image

图9:COMET 开源页面


此外,COMET 建立了面向 MoE 的细粒度流水线编程范式,通过深度融合 NVSHMEM 通信库与 CUTLASS 高效计算算子,实现了通信操作与 GEMM 计算的算子内融合。例如,MoE Layer 1 的 GEMM 计算与 Token 聚合通信可在单一 GPU 算子内完成。这与此前提到的 Deepseek DualPipe 方案并不冲突,两者结合或将带来更好的优化空间。


此外,COMET 可直接接入已有的 MoE 训练框架,支持 TP/EP/EP+TP 多种并行模式,并提供了灵活的插拔式部署方案。


核心代码​现已开源​,并计划兼容 Triton 等编译生态,欢迎大家一起讨论交流!