我要投稿

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

发布日期：2025-03-25 12:25:23 浏览次数： 2358

作者：字节跳动技术团队

微信搜一搜，关注“字节跳动技术团队”

当前，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 中

分解并重新调度共享张量

2）自适应负载分配

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

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

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

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

图 6：单个 MoE 层使用不同数量的

通信线程块的时延结果

3. 大规模落地验证

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

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

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

图 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。

图9：COMET 开源页面

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

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