COMET - ByteDance's open-source communication optimization system

What is COMET?

COMET, developed by ByteDance, is an optimization system for Mixture-of-Experts (MoE) models, addressing the issue 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 mismatches in traditional methods. COMET introduces two mechanisms: shared tensor dependency resolution and adaptive load balancing, dynamically adjusting resource allocation to achieve efficient communication hiding and computational optimization. In large-scale production environments, COMET significantly improves the training efficiency of MoE models, achieving up to 1.96x speedup for single layers and 1.71x speedup for end-to-end training, saving millions of GPU hours cumulatively. COMET's core code is open-source, supporting mainstream large models and various parallel strategies, and can be seamlessly integrated into existing training frameworks.

COMET's main functions

• Fine-grained computation - communication overlapDecompose shared tensors and reschedule computation order to achieve fine-grained alignment between computation and communication.
• Adaptive load balancingDynamically adjust GPU thread block resources to balance communication and computing loads based on input size and parallelism strategy, eliminate pipeline bubbles, and improve overall efficiency.
• High-efficiency resource managementEncapsulate communication and computing tasks in independent thread blocks to avoid remote I/O blocking the computing core and improve resource utilization.
• Strong robustness and generalization abilityMaintain low latency and support multiple parallel strategies and large-scale cluster deployments under uneven expert loads or different hardware environments.
• Easy to integrateIt can be directly integrated into the existing MoE training framework as a plugin without intrusive modifications and supports mainstream compilation ecosystems (such as Triton).

COMET's technical principles

• Shared tensor dependency resolution:
  • Tensor decompositionThe shared tensors passed between MoE layers are cut along the Token dimension (M) or the hidden layer dimension (N) to align the smallest unit of communication and computation. For example, the MoE layer 1 (Layer 0) is decomposed along the M dimension, and the layer 2 (Layer 1) is decomposed along the N dimension to achieve efficient overlap.
  • Computational rescheduling: Dynamically adjust the calculation order of data blocks, prioritize the calculation of local data blocks, asynchronously pull remote tokens, and eliminate the calculation delay caused by waiting for remote data.
• Adaptive load balancing:
  • Thread block isolationThe communication and computation tasks are encapsulated in separate thread blocks to avoid remote I/O blocking the computation core. The computation thread block is dedicated to executing GEMM operations using asynchronous TMA instructions, while the communication thread block implements single-token-level data transmission based on NVSHMEM.
  • Dynamic load balancingBased on the input size (such as token length) and parallel strategy (EP/TP ratio), the thread block allocation ratio is adjusted in real time, and dynamic switching with "zero overhead" is achieved at runtime based on multiple pre-compiled versions of computation-communication fusion operators.

COMET's performance

• Single-layer accelerationIn large-scale MoE models, a single-layer speedup of 1.96x is achieved.
• End-to-end accelerationIn the full MoE model, the end-to-end average implementation speedup is 1.71x.
• Stability under different models and input sizes:
  • In several 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.
  • With different input token lengths, COMET's execution time is significantly shorter than the baseline solution, with an average speed improvement of 1.28 to 2.37 times.
• robustnessIn scenarios with uneven expert loads, COMET maintains lower latency than other baseline systems, demonstrating good robustness.
• Generalization abilityCOMET provides stable speedup in different network environments such as NVLink and PCIe, supports multiple parallel 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

Application scenarios of COMET

• Accelerated training of large-scale MoE modelsOptimize distributed training of large models such as Mixtral-8x7B, saving GPU hours and improving training efficiency.
• Heterogeneous hardware environment adaptationSupports different GPUs and network environments (such as NVLink and PCIe) such as H800 and L20, ensuring stable acceleration.
• Multi-parallel strategy supportIt is compatible with EP, TP and hybrid parallel strategies, and adapts to different models and hardware configurations.
• Dynamic load balancingWhen the expert load is uneven, dynamically adjust resource allocation to maintain low latency operation.
• Seamless integration with existing frameworksIt can be directly integrated into the MoE training framework (such as Megatron-LM) as a plugin without any intrusive modifications, making it easy to deploy quickly.