COMET

COMET – 字节开源的通信优化系统

COMET是字节跳动推出的一款专为Mixture-of-Experts（MoE）模型设计的优化系统，旨在解决分布式训练中的通信开销过大的问题。通过细粒度的计算-通信重叠技术，COMET将通信与计算操作深度融合，有效消除传统方法因粒度不匹配所造成的资源浪费与延迟。该系统的引入使得MoE模型的训练效率显著提升，单层加速可达1.96倍，端到端加速可达1.71倍，累计节省了数百万GPU小时。

COMET是什么

COMET是字节跳动推出的一种针对Mixture-of-Experts（MoE）模型的优化工具，旨在解决分布式训练过程中的高通信开销问题。它依托于细粒度的计算-通信重叠技术，将通信和计算操作深度融合，消除了传统方法所引发的资源浪费和延迟。通过引入共享张量依赖解析和自适应负载分配机制，COMET能够动态调整资源分配，从而实现高效的通信隐藏和计算优化。在大规模的生产环境中，COMET展现出显著的训练效率提升，单层加速达到1.96倍，端到端加速为1.71倍，积累节省的GPU小时数达到数百万个。COMET的核心代码已经开源，支持主流大模型和多种并行策略，能够无缝集成到现有的训练框架中。

COMET

COMET的主要功能

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

COMET的技术原理

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

COMET的性能表现

单层加速：在大规模MoE模型中，实现单层加速1.96倍。
端到端加速：在完整的MoE模型中，端到端平均实现加速1.71倍。
不同模型和输入规模下的稳定性：
- 在多个大规模MoE模型（如Mixtral-8x7B、Qwen2-MoE等）中，COMET的前向时延相比其他基线系统降低31.8%-44.4%。
- 在不同输入Token长度下，COMET的执行时间显著短于基线方案，平均速度提升1.28倍到2.37倍。
鲁棒性：在专家负载不均衡的场景下，COMET能保持低于其他基线系统的延迟，表现出良好的鲁棒性。
泛化能力：COMET在NVLink和PCIe等不同网络环境下均能提供稳定的加速比，支持多种并行策略（如EP、TP、EP+TP），适用于大规模训练框架。