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

字节跳动技术团队 · 公众号 · · 2025-03-25 12:00

正文

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

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

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

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），这种由于复杂数据依赖导致的计算-通信粒度上的错配，使得效率严重下滑。

图 1：单层 MoE 模型示意图

（专家分布在 GPU0 和 GPU1 两张卡上）

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

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

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

2. COMET 核心方案

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

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

图 2：COMET 的设计结构

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

1）共享张量依赖解析

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

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

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

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

图 4：COMET 在 MoE layer0 中

分解并重新调度共享张量