字节砍MoE训练成本，节省数百万GPU小时

admin管理员组

文章数量:1446759

字节砍MoE训练成本，节省数百万GPU小时

字节跳动豆包大模型团队开源新一代MoE训练优化技术COMET，成功实现模型训练成本再降40%。该技术已在其万卡GPU集群中完成生产验证，累计节省数百万GPU算力资源。此前该团队曾推出推理优化架构UltraMem（降低83%推理成本），结合此次训练端突破，形成从训练到推理的全链路降本组合方案。

技术突破方面，COMET聚焦解决MoE模型分布式训练中的通信瓶颈问题，通过独创的细粒度计算-通信折叠算法，配合动态GPU资源调度机制，有效利用专家网络的闲置算力资源。实测显示，该技术使MoE单层执行速度提升1.96倍，端到端训练效率平均提高71%。值得关注的是，该方案可与DeepSeek早前开源的DualPipe+DeepEP形成互补，共同突破MoE训练效率天花板。

相较于现有方案，COMET具备三大核心优势：

1. 非侵入式架构设计，支持插件化部署，适配主流大模型框架
2. 计算-通信算子深度融合，避免传统方案中的流水线排布限制
3. 动态资源管理系统实时优化硬件利用率，突破静态资源分配瓶颈

该技术凭借其工业级实用价值，已获机器学习系统顶会MLSys 2025评审团高度认可（均分5/5/5/4），成为首个在该会议获得满分的国产MoE优化方案。开源地址现已在GitHub平台发布，为行业提供高效易用的MoE训练加速解决方案。

背景和动机

MoE 结构：MoE 由多个专家组成，输入 token 仅路由到部分专家计算，通过门网络选择 topk 个专家。其包含两层，layer0 有 token 通信和第一层专家计算，layer1 有第二层专家计算、token 合并等操作。同时，MoE 采用专家并行和张量并行两种策略，实际常采用混合并行。

计算和通信重叠：随着 MoE 模型规模增大，通信时间占比增加。粗粒度重叠优化潜力有限，细粒度重叠虽有机会但面临挑战。一方面，计算和通信的粒度不匹配，GEMM 计算以 tile 为单位，而数据传输以 token 为单位，复杂数据依赖影响重叠效率，且整合细粒度通信与计算困难；另一方面，MoE 中 token 动态路由导致计算和通信负载多样，需动态调整 GPU 资源分配，但现有方法难以实现。

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

Comet 设计

基于共享张量的依赖解决：共享张量是实现重叠的关键，通过分析其访问模式，在消费者操作数据独立的维度分解。如 layer0 沿 M 维度分解，layer1 沿 N 维度分解。分解后，重新组织和调度子张量成 tile 计算，遵循与原计算 tile 粒度对齐和优先处理可立即使用数据的原则，减少远程数据依赖，提前计算。

自适应工作负载分配：为平衡细粒度通信和计算延迟，Comet 采用线程块专业化和自适应线程块分配。线程块专业化将通信和计算工作负载隔离，避免相互干扰，GEMM 线程块采用默认实现，通信线程块负责数据读写。自适应线程块分配根据输入形状和模型配置确定最优线程块分配比例，Comet 预编译多个内核，运行时选择最优内核执行。

Comet 系统中基于共享张量的依赖解决方法，体现了生产者（Producer）和消费者（Consumer）之间的数据交互与计算关系，以及共享张量在其中的作用，具体如下：

生产者和消费者组件：图的上方 “Producer” 代表生产者部分，包含 All2all/AllGather 操作和 GEMM（通用矩阵乘法）操作 ；图下方 “Consumer” 代表消费者部分，包含 GEMM 操作以及 TopK-reduce + All2all/ReduceScatter 操作。这些操作是 MoE（混合专家）模型计算和通信过程中的关键步骤。

共享张量（Shared tensor）：图中间两个方框内标注的 “Shared tensor” 表示共享张量，是连接生产者和消费者的关键数据结构。图中箭头和文字表明共享张量在 M 维度和 N 维度上的独立性。在生产者到消费者的数据传递过程中，共享张量沿 M 维度在生产者端是相互独立的，沿 N 维度在消费者端是相互独立的。这为基于共享张量进行依赖分解和计算调度提供了基础，有助于解决计算和通信中的数据依赖问题，实现细粒度的计算 - 通信重叠 。

数据流向和依赖关系：从生产者到消费者的箭头指示了数据的流向，通过共享张量建立起两者之间的联系。生产者先通过 All2all/AllGather 操作收集数据，再进行 GEMM 计算；消费者接收数据后，先进行 GEMM 计算，然后进行 TopK-reduce + All2all/ReduceScatter 操作。共享张量在其中起到了数据共享和传递的作用，并且利用其在不同维度上的独立性，打破了原有的粗粒度数据依赖，使得计算和通信可以更细粒度地重叠进行 。

MoE（混合专家）模型 Layer0 在 Rank0 设备上的执行过程，重点体现了共享张量的分解与计算重调度，具体如下：

数据来源：左侧标注 “Local data (Rank0)” 和 “Remote data (Rank1)”，表示数据一部分来自本地设备 Rank0，另一部分来自远程设备 Rank1。

共享张量分解：中间部分 “original shared tensor” 是原始的共享张量，包含多个专家（Expert #0、Expert #1、Expert #2）的数据。这些数据沿 M 维度进行分解（Decomposed along M），将原始共享张量拆分成更小的子张量。

计算重调度：右侧展示了重新调度（Reschedule）的过程，根据源 token（by source tokens）对分解后的子张量进行排序（sort），然后分别送入不同的 GEMM 计算单元（GEMM #0、GEMM #1、GEMM #2）进行分组 GEMM 计算（GroupGEMM），最终得到计算序列（compute sequence） 。这一过程通过对共享张量的分解和重调度，优化了计算流程，有助于实现细粒度的计算 - 通信重叠，提升 MoE 模型的执行效率。

大规模落地验证

不同模型在不同并行策略下的延迟（Latency）对比，评估 Comet 等系统在 MoE（混合专家）模型执行中的性能

对比模型：图表对比了 Megatron-TE、Megatron-Cutlass、FasterMoE、Tutel、Comet 以及 Attention 六种方法在不同模型上的延迟表现。涉及的模型有 Mixtral、Qwen、Phi 。

实验设置

输入长度（M）：分为 M=4096 和 M=8192 两种情况，每一组柱状图上方都有相应标注。

并行策略（TP,EP） ：TP 代表张量并行度，EP 代表专家并行度，分别设置了 (1,8)、(2,4)、(4,2)、(8,1) 四种不同的组合，不同并行策略组合也在每一组柱状图上方标注 。

结果分析：从图表中可以看出，在不同模型、不同输入长度和不同并行策略下，Comet（浅黄色柱状）相比其他对比方法，大部分情况下延迟更低，体现出 Comet 在优化 MoE 模型执行延迟方面的优势。而 Attention（带斜线柱状）的延迟相对稳定，作为对比基准存在 。

核心代码开源

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

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

Flux 作为字节跳动推出的用于 GPU 上密集型 / MoE 模型的通信重叠库，优势体现在多方面，能为模型训练和推理带来显著性能提升。

GitHub地址

论文地址

本文标签： 字节砍MoE训练成本，节省数百万GPU小时