COMET：字节跳动开源MoE训练加速神器，单层1.96倍性能提升，节省百万GPU小时

❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发感兴趣，我会每日分享大模型与 AI 领域的开源项目和应用，提供运行实例和实用教程，帮助你快速上手AI技术！

🥦 AI 在线答疑 -> 智能检索历史文章和开源项目 -> 尽在微信公众号 -> 搜一搜：蚝油菜花 🥦

🚀 「MoE训练卡在通信瓶颈？字节黑科技让GPU效率飙升196%！」

大家好，我是蚝油菜花。当你还在为这些场景抓狂时：

👉 70%训练时间浪费在All-to-All通信
👉 专家并行时GPU利用率不足40%
👉 万亿参数模型迭代动辄消耗百万算力

COMET 带着字节跳动最新研究杀到了！这个专治MoE训练低效的「手术刀」，通过三大颠覆性创新：

✅ 原子级计算-通信缝合术：把张量切割到Token级实现精准对齐
✅ 自适应负载动态平衡：智能调配8.7万个线程块资源
✅ 零侵入式部署：像装插件一样接入Megatron框架

实测单层加速达1.96倍，万卡集群训练成本直降42%。某头部AI公司用它把70亿MoE模型训练周期砍半------你的GPU准备好解除性能封印了吗？

🚀 快速阅读

COMET是字节跳动推出的针对Mixture-of-Experts（MoE）模型的优化系统。

核心功能：通过细粒度的计算-通信重叠技术，显著提升分布式训练效率。
技术原理：引入共享张量依赖解析和自适应负载分配两大机制，动态调整资源分配，实现高效的通信隐藏和计算优化。

COMET 是什么

COMET是字节跳动推出的针对Mixture-of-Experts（MoE）模型的优化系统，能解决分布式训练中通信开销过大的问题。基于细粒度的计算-通信重叠技术，COMET将通信与计算操作深度融合，消除传统方法中因粒度不匹配导致的资源浪费和延迟。

COMET引入共享张量依赖解析和自适应负载分配两大机制，动态调整资源分配，实现高效的通信隐藏和计算优化。在大规模生产环境中，COMET显著提升了MoE模型的训练效率，单层加速可达1.96倍，端到端加速1.71倍，已累计节省数百万GPU小时。COMET的核心代码已开源，支持主流大模型和多种并行策略，能无缝集成到现有训练框架中。

COMET 的主要功能

细粒度计算-通信重叠：分解共享张量和重新调度计算顺序，实现计算与通信的细粒度对齐。
自适应负载分配：动态调整GPU线程块资源，根据输入规模和并行策略平衡通信与计算负载，消除流水线气泡，提升整体效率。
高效资源管理：将通信和计算任务封装在独立线程块中，避免远程I/O阻塞计算核心，提升资源利用率。
强鲁棒性和泛化能力：在专家负载不均衡或不同硬件环境下，保持低延迟，支持多种并行策略和大规模集群部署。
易于集成：作为插件直接接入现有的MoE训练框架，无需侵入性改动，支持主流编译生态（如Triton）。

COMET 的技术原理

共享张量依赖解析 ：
- 张量分解：将MoE层间传递的共享张量沿Token维度（M）或隐层维度（N）进行切割，让通信与计算的最小单元对齐。
- 计算重调度：动态调整数据块的计算顺序，优先计算本地数据块，异步拉取远程Token，消除因等待远程数据导致的计算延迟。
自适应负载分配 ：
- 线程块隔离：将通信与计算任务分别封装在独立线程块中，避免远程I/O阻塞计算核心。
- 动态负载平衡：根据输入规模（如Token长度）和并行策略（EP/TP比例），实时调整线程块分配比例，基于预编译多个版本的计算-通信融合算子，实现运行时的"零开销"动态切换。

如何运行 COMET

1. 从源码安装

bash 复制代码

git clone --recursive https://github.com/bytedance/flux.git && cd flux

# 安装依赖
bash ./install_deps.sh

# 对于Ampere(sm80) GPU
./build.sh --arch 80 --nvshmem
# 对于Ada Lovelace(sm89) GPU
./build.sh --arch 89 --nvshmem
# 对于Hopper(sm90) GPU
./build.sh --arch 90 --nvshmem

2. 在虚拟环境中安装

bash 复制代码

conda create -n flux python=3.11
conda activate flux
pip3 install packaging
pip3 install ninja
pip3 install torch==2.6.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

./build.sh --clean-all
./build.sh --arch "80;89;90" --nvshmem --package

3. 从PyPI安装

bash 复制代码

# 确保已安装PyTorch
pip install byte-flux

4. 自定义安装

构建选项

添加 --nvshmem 以支持NVSHMEM。
设置环境变量 FLUX_BUILD_SKIP_CMAKE 为1以跳过cmake。
添加 --package 以构建wheel包。

依赖

核心依赖包括NCCL、CUTLASS和NVSHMEM，位于3rdparty文件夹中。

5. 快速开始

bash 复制代码

# gemm only
python3 test/python/gemm_only/test_gemm_only.py 4096 12288 6144 --dtype=float16

# all-gather fused with gemm (dense MLP layer0)
./launch.sh test/python/ag_gemm/test_ag_kernel.py 4096 49152 12288 --dtype=float16 --iters=10

# gemm fused with reduce-scatter (dense MLP layer1)
./launch.sh test/python/gemm_rs/test_gemm_rs.py 4096 12288 49152 --dtype=float16 --iters=10

# all-gather fused with grouped gemm (MoE MLP layer0)
./launch.sh test/python/moe_ag_scatter/test_moe_ag.py

# grouped gemm fused with reduce-scatter (MoE MLP layer1)
./launch.sh test/python/moe_gather_rs/test_moe_gather_rs.py

资源

GitHub 仓库 ：github.com/bytedance/f...
技术论文 ：arxiv.org/pdf/2502.19...

🥦 AI 在线答疑 -> 智能检索历史文章和开源项目 -> 尽在微信公众号 -> 搜一搜：蚝油菜花 🥦