What Is COMET
COMET is an optimization system launched by ByteDance for Mixture-of-Experts (MoE) models, capable of solving the problem of excessive communication overhead in distributed training. Based on fine-grained computation-communication overlap technology, COMET deeply integrates communication and computation operations, eliminating resource waste and latency caused by granularity mismatch in traditional methods. COMET introduces two major mechanisms, shared tensor dependency resolution and adaptive load allocation, dynamically adjusting resource allocation to achieve efficient communication hiding and computation optimization. In large-scale production environments, COMET significantly improves the training efficiency of MoE models, achieving up to 1.96x single-layer speedup and 1.71x end-to-end speedup, and has cumulatively saved millions of GPU hours. COMET's core code has been open-sourced, supports mainstream large models and multiple parallel strategies, and can be seamlessly integrated into existing training frameworks.
COMET's Main Features
- Fine-Grained Computation-Communication Overlap: Decomposes shared tensors and reschedules computation order to achieve fine-grained alignment between computation and communication.
- Adaptive Load Allocation: Dynamically adjusts GPU thread block resources, balancing communication and computation loads according to input scale and parallel strategy, eliminating pipeline bubbles and improving overall efficiency.
- Efficient Resource Management: Encapsulates communication and computation tasks in independent thread blocks, preventing remote I/O from blocking computation cores and improving resource utilization.
- Strong Robustness and Generalization Ability: Maintains low latency under imbalanced expert loads or different hardware environments, supporting multiple parallel strategies and large-scale cluster deployment.
- Easy Integration: Directly plugs into existing MoE training frameworks as a plugin, requiring no invasive changes, and supports mainstream compilation ecosystems (such as Triton).
COMET's Technical Principles
- Shared Tensor Dependency Resolution：
- Tensor Decomposition: Splits the shared tensors passed between MoE layers along the Token dimension (M) or hidden layer dimension (N), aligning the smallest units of communication and computation. For example, decomposition along the M dimension in the first MoE layer (Layer 0) and along the N dimension in the second layer (Layer 1) achieves efficient overlap.
- Computation Rescheduling: Dynamically adjusts the computation order of data blocks, prioritizing local data blocks and asynchronously fetching remote Tokens, eliminating computation latency caused by waiting for remote data.
- Adaptive Load Allocation：
- Thread Block Isolation: Encapsulates communication and computation tasks separately in independent thread blocks, preventing remote I/O from blocking computation cores. Computation thread blocks are dedicated to GEMM operations executing asynchronous TMA instructions, while communication thread blocks implement single-Token-level data transfer based on NVSHMEM.
- Dynamic Load Balancing: Adjusts the thread block allocation ratio in real time according to input scale (such as Token length) and parallel strategy (EP/TP ratio), based on precompiling multiple versions of computation-communication fused operators to achieve "zero-overhead" dynamic switching at runtime.
COMET's Performance
- Single-Layer Speedup: Achieves 1.96x single-layer speedup in large-scale MoE models.
- End-to-End Speedup: Achieves an average end-to-end speedup of 1.71x in complete MoE models.
- Stability Across Different Models and Input Scales：
- In multiple large-scale MoE models (such as Mixtral-8x7B, Qwen2-MoE, etc.), COMET reduces forward latency by 31.8%-44.4% compared to other baseline systems.
- Under different input token lengths, COMET's execution time is significantly shorter than that of the baseline schemes, with an average speedup of 1.28x to 2.37x.
- Robustness: In scenarios with imbalanced expert loads, COMET can maintain lower latency than other baseline systems, demonstrating good robustness.
- Generalization ability: COMET can provide stable speedups under different network environments such as NVLink and PCIe, supports multiple parallelism strategies (such as EP, TP, EP+TP), and is suitable for large-scale training frameworks.
COMET project address
- GitHub repository：https://github.com/bytedance/flux
- arXiv technical paper：https://arxiv.org/pdf/2502.19811
COMET application scenarios
- Large-scale MoE model training acceleration: Optimize the distributed training of large models such as Mixtral-8x7B, save GPU hours, and improve training efficiency.
- Adaptation to heterogeneous hardware environments: Supports different GPUs such as H800 and L20 and network environments (such as NVLink and PCIe), ensuring stable acceleration.
- Support for multiple parallelism strategies: Compatible with EP, TP, and hybrid parallelism strategies, adapting to different models and hardware configurations.
- Dynamic load balancing: When expert loads are imbalanced, dynamically adjust resource allocation to maintain low-latency operation.
- Seamless integration with existing frameworks: As a plugin, it can be directly integrated into MoE training frameworks (such as Megatron-LM) without invasive changes, facilitating rapid deployment.