Source-linked AI summary
Comet: Fine-grained Computation-communication Overlapping for Mixture-of-Experts
Shulai Zhang, Ningxin Zheng, Haibin Lin, Ziheng Jiang, Wenlei Bao, Chengquan Jiang, Qi Hou, Weihao Cui, Size Zheng, Li-Wen Chang, Quan Chen, Xin Liu
TL;DR
Distributed MoE models face substantial communication overhead, and coarse-grained overlap can impair computational efficiency and conceal latency imperfectly. COMET uses shared-tensor dependency resolution, task rescheduling, and adaptive workload assignment for fine-grained communication-computation overlap. It achieves 1.96× speedup for a single MoE layer and 1.71× for end-to-end MoE execution on average.
Problem
MoE communication can occupy 47% of execution time, while coarse-grained overlap suffers from computational inefficiency, non-overlapping phases, and dynamic dependency and workload challenges.
Method
COMET decomposes and reschedules shared tensors and adaptively assigns GPU thread blocks within fused kernels to balance communication and computation.
Results
1.96× speedup is achieved for a single MoE layer and 1.71× speedup for end-to-end MoE execution on average compared with existing systems.
Takeaways & Limitations
COMET has been deployed in production clusters with over ten thousand GPUs, achieving savings of millions of GPU hours.
Takeaways & Limitations
Thread-block specialization must contend with hardware resource limits: communication warps may underuse bandwidth and interfere with computation warps.
Abstract
from arXiv · showhide
Mixture-of-experts (MoE) has been extensively employed to scale large language models to trillion-plus parameters while maintaining a fixed computational cost. The development of large MoE models in the distributed scenario encounters the problem of large communication overhead. The inter-device communication of a MoE layer can occupy 47% time of the entire model execution with popular models and frameworks. Therefore, existing methods suggest the communication in a MoE layer to be pipelined with the computation for overlapping. However, these coarse grained overlapping schemes introduce a notable impairment of computational efficiency and the latency concealing is sub-optimal. To this end, we present COMET, an optimized MoE system with fine-grained communication-computation overlapping. Leveraging data dependency analysis and task rescheduling, COMET achieves precise fine-grained overlapping of communication and computation. Through adaptive workload assignment, COMET effectively eliminates fine-grained communication bottlenecks and enhances its adaptability across various scenarios. Our evaluation shows that COMET accelerates the execution of a single MoE layer by $1.96\times$ and for end-to-end execution, COMET delivers a $1.71\times$ speedup on average. COMET has been adopted in the production environment of clusters with ten-thousand-scale of GPUs, achieving savings of millions of GPU hours.
1 Introduction
Distributed MoE execution incurs substantial communication overhead, while coarse-grained overlap can reduce computational efficiency and leave latency insufficiently concealed. COMET addresses these issues with fine-grained dependency-aware overlapping and adaptive workload assignment, improving both layer and end-to-end execution.
- 47% of total execution time is consumed by inter-device communication in the forward pass of several popular MoE models.
- Coarse-grained overlap partitions input data into chunks, but smaller expert chunks can reduce GPU utilization and leave non-overlapping communication phases.Partitioned expert computation can take longer than the original computation, while initial and final communication phases create GPU idle time.
- COMET resolves communication-computation dependencies and adaptively assigns GPU thread blocks to balance workloads within kernels.Its dependency method structures the pipeline, while adaptive assignment improves latency concealment under dynamic workloads.
- 1.96× speedup is achieved for typical MoE layers, while end-to-end MoE model execution gains 1.71× on average versus SOTA MoE systems.
- COMET has been deployed in production clusters with over ten thousand GPUs, saving millions of GPU hours.
2 Background and Motivation
MoE routes tokens to selected experts distributed across GPUs, making communication a major part of execution. Fine-grained overlap is difficult because computation tiles depend on irregularly distributed tokens, remote communication is costly, and workload demands vary dynamically; COMET targets these constraints with shared-tensor dependency resolution and adaptive in-kernel resource allocation.
- 2.1 MoE Structure: MoE activates only selected experts for each token, using routing and feed-forward computations before gathering and reducing the topk outputs.
- 2.1 MoE Structure: Expert parallelism distributes intact expert weights across GPUs, while tensor parallelism partitions all expert weights along the hidden dimension; deployments often combine both.
- 2.2 Background and Motivation: Fine-grained overlap is motivated by the increasing communication share of larger, sparser MoE models and the limited optimization potential of coarse-grained schemes.
- 2.2.1 Granularity mismatch between computation and communication: A computation tile cannot start until all of its required tokens arrive, because runtime routing distributes those tokens irregularly across devices.Fine-grained communication and data reorganization are proposed to expose only the data each tile needs and hide preparation with computation.
- 2.2.1 Granularity mismatch between computation and communication: Remote token I/O can block computation thread blocks, so integrating token-wise communication with tile-wise computation risks reducing kernel efficiency.
- 2.2.2 Diverse loads of computation and communication: COMET resolves granularity mismatch through shared-tensor decomposition and reorganization, then adapts in-kernel resource allocation for communication and computation.
- 2.2.2 Diverse loads of computation and communication: Dynamic routing produces varying expert input shapes, compute demands, and communication demands across GPUs and hardware environments.
3 Design of Comet
Comet resolves MoE communication–computation dependencies by decomposing and rescheduling shared tensors, then uses adaptive resource assignment and specialized kernels for fine-grained overlap.
- 3.1 Shared Tensor Based Dependency Resolving: Comet analyzes shared-tensor dependencies to decompose data along consumer-independent dimensions and reschedule computation for overlap.Shared tensors bridge producer and consumer operators; layer0 decomposes along M, while layer1 decomposes along N because top-k reduction creates M-dimension interdependencies.
- 3.1.2 How to reschedule the decomposed shared tensor?: Rescheduled sub-tensors align with computation tiles while prioritizing producer outputs that the consumer can use immediately.Fine-grained row- or column-level decomposition alone can reduce GEMM efficiency, so Comet reorganizes sub-tensors into tiles.
- 3.1.2 How to reschedule the decomposed shared tensor?: Layer0 sorts tokens by source rank and prioritizes tiles containing local tokens while remote-token transfers proceed concurrently.GroupGEMM computes experts on the current rank, reducing dependence on remote data during the communication-computation pipeline.
- 3.1.2 How to reschedule the decomposed shared tensor?: Layer1 executes GroupGEMM operations column-wise so reduction and communication can begin after the first TN columns are produced.This rescheduling lets the consumer start before all expert computations finish, unlike sequential expert execution.
- 3.2 Adaptive Workload Assignment: Comet combines fused communication and computation with thread-block specialization to isolate communication effects and balance workload latencies.Adjusting thread-block allocation reduces overlap bubbles, while specialized kernels preserve computation performance by separating communication and computation resources.
- 3.2.1 Thread block specialization: Thread-block specialization remains constrained by hardware resources because communication warps can limit bandwidth utilization and interfere with computation warps.Integrating communication and computation warps could reduce redundant global-memory accesses, but warp-count restrictions create competing resource demands.
- 3.2 Adaptive Workload Assignment: The optimal communication–computation resource division changes with input length and parallel configuration.For TP = 8, the optimal nc changes from 18 to 26 as M increases from 4096 to 16384.
4 Implementation
Comet provides Python APIs and CUDA/C++ implementations for MoE integration, using optimized GEMM kernels and NVSHMEM for fine-grained GPU communication.
- 4 Implementation: Comet optimizes MoE GEMM kernels with CUTLASS templates and register caching of row indices to reduce global-memory access costs.The row indices are otherwise accessed from global memory at each K iteration in MoE layer0.
- 4 Implementation: Comet uses NVSHMEM to provide fine-grained GPU-initiated communication through a global address space spanning multiple GPUs.Compared with NCCL’s higher-level operations, NVSHMEM offers a more composable low-level API for fine-grained data access within kernels.
5 Evaluation
Comet is evaluated against established MoE systems across end-to-end models, single-layer workloads, parallelism strategies, parameter settings, token distributions, and cluster environments. It consistently improves latency and communication-computation overlap, including bandwidth-limited settings.
- Overall Performance: 34.1%, 42.6%, 44.4% and 31.8% lower end-to-end latency than Megatron-Cutlass, Megatron-TE, FasterMoE and Tutel, respectively.Comet outperforms the baselines in all tested configurations.
- Detailed Evaluation on a Single MoE Layer: 1.28× to 2.37× speedup over baselines across varying input token lengths.The advantage is especially prominent for small M, where host-side scheduling dominates and Comet moves scheduling into fused kernels.
- Detailed Evaluation on a Single MoE Layer: 86.5% of communication latency hidden by Comet, compared with 29.2% for FasterMoE and 68.6% for Tutel.Comet maintains expert computational efficiency while overlapping communication.
- Detailed Evaluation on a Single MoE Layer: Comet maintains low latency across diverse parallelism strategies, while tensor parallelism increases latency for other supported baselines.Comet reschedules shared tensors to preserve computational efficiency and eliminates weight-switching overhead.
- Scaling to Distinct Clusters: 1.19× to 1.46× average speedup over baselines on an L20 cluster with PCIe-connected GPUs and roughly 25 GB/s GPU-to-GPU bandwidth.The result covers a bandwidth-limited environment distinct from the H800 cluster.
6 Related Work
Related work reduces MoE communication cost through faster or hierarchical communication, compression, and communication-computation pipelining. Existing overlap methods use coarse or manually tuned pipeline strategies, motivating more fine-grained approaches.
- Communication optimization: Communication optimization methods accelerate data transmission, use 2D-hierarchical all-to-all algorithms, or reduce communication volume through compression.These approaches target communication cost directly rather than primarily changing computation scheduling.
- Computation-communication overlapping: FasterMoE and Tutel pipeline all-to-all communication with expert computation using fixed, manually set, or heuristically searched pipeline degrees.FasterMoE uses a pipeline degree of 2, while Tutel supports manual settings or heuristic search under a limited search space.
7 Conclusion
Comet targets fine-grained communication-computation overlap in MoE through shared-tensor dependency resolution and adaptive workload assignment. It reports speedups for both single-layer and end-to-end execution relative to existing systems.
- Conclusion: Comet uses shared-tensor dependency resolution to enable fine-grained overlap while eliminating fine-grained communication I/O bottlenecks.The design is paired with workload assignment for adaptive operator overlap and latency concealing.
- Conclusion: 1.96× speedup for a single MoE layer and 1.71× speedup for end-to-end MoE execution compared with existing literature.These are the paper’s reported aggregate speedups.