已节省数百万GPU小时！字节再砍MoE训练成本，核心代码全开源

白交
2025-03-11
16:34:41

来源：量子位

C++OMET团队 投稿

量子位 | 公众号 QbitAI

字节对MoE模型训练成本再砍一刀，成本可节省40%！

刚刚，豆包大模型团队在GitHub上开源了叫做COMET的MoE优化技术。

COMET已应用于字节的万卡训练集群，在真实的生产环境中，累计帮助节省了数百万GPU小时。

早前，豆包团队发布了新一代稀疏架构UltraMem，将模型推理成本砍掉 83%，此次，又开源了COMET，向模型训练成本出手。从技术理念上看，两者还可以结合使用，组成一套“砍价刀法”。

具体来看，COMET主要针对的是MoE模型在分布式训练中，仍存在大量通信开销的问题。

COMET内部通过一套细粒度计算-通信折叠技术，并结合GPU资源的动态分配，极致压榨了MoE专家“摸鱼闲置”的时间，在大规模MoE的单个执行层上可提速1.96倍，端到到平均提速1.71倍。

有趣的是，此前DeepSeek也专门针对MoE的通信瓶颈，开源了DualPipe+DeepEP方案，通过排布算子来掩盖通信开销。豆包团队则直接喊话，两种方案一起开挂，可能带来更大的提升空间。

不过，COMET这种直接将计算-通信算子融合的方法，可以在MoE训练框架中像插件一样直接插拔使用，无需侵入性改动，部署更加方便、灵活，且支持业界绝大部分主流大模型。

因简洁、高效的设计理念，该工作5/5/5/4高分入选了全球机器学习系统顶级会议 MLSys 2025，并被认为在实际的大规模生产环境中极具应用价值。

接下来，详细看下COMET的技术解读。

MoE紧迫的通信开销问题

混合专家模型（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的通信接收远程数据（TokenB），这种由于复杂数据依赖导致的计算-通信粒度上的错配，使得效率严重下滑。

△ 1：单层MoE模型示意图（专家分布在GPU0和GPU1两张卡上）

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

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

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

2、COMET：最小化整体低延迟，提升性能

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

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

△图2：COMET的设计结构

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

1）共享张量依赖解析

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

张量分解：将MoE层间传递的共享张量沿Token维度（M）或隐层维度（N）进行切割，使通信与计算的最小单元对齐。例如，在MoE第一层（Layer0，图3左）沿M维度分解，使通信和计算在M维度进行对齐；在MoE第二层（Layer1，图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+DeepEP方案并不冲突，两者结合或将带来更好的优化空间。

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

目前，COMET核心代码已开源，并计划兼容Triton等编译生态。

论文链接：
https://arxiv.org/pdf/2502.19811
开源地址：
https://github.com/bytedance/flux

Read More