❤️ 如果你也关注 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. 从源码安装
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. 在虚拟环境中安装
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安装
# 确保已安装PyTorch
pip install byte-flux
4. 自定义安装
构建选项
- 添加
--nvshmem以支持NVSHMEM。 - 设置环境变量
FLUX_BUILD_SKIP_CMAKE为1以跳过cmake。 - 添加
--package以构建wheel包。
依赖
核心依赖包括NCCL、CUTLASS和NVSHMEM，位于3rdparty文件夹中。
5. 快速开始
# 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 仓库：https://github.com/bytedance/flux
- 技术论文：https://arxiv.org/pdf/2502.19811
❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发感兴趣，我会每日分享大模型与 AI 领域的开源项目和应用，提供运行实例和实用教程，帮助你快速上手AI技术！
🥦 AI 在线答疑 -> 智能检索历史文章和开源项目 -> 尽在微信公众号 -> 搜一搜：蚝油菜花 🥦
