COMET: Fine-grained Computation-communication Overlapping for Mixture-of-Experts

COMET: Fine-grained Computation-communication Overlapping for Mixture-of-Experts

Speakers: Shulai Zhang, Ningxin Zheng, Haibin Lin, Ziheng Jiang, Wenlei Bao, Xin Liu

Conference: MLSys 2025

YouTube: https://www.youtube.com/watch?v=None

Overview

This talk introduces Comet, a novel framework designed to achieve fine-grained computation-communication overlapping in Mixture-of-Experts (MoE) models. Presented by Ningxin Zheng from Bytedance, Comet addresses a critical bottleneck in the efficient distributed training and inference of large-scale MoE architectures: the significant latency introduced by communication operations. As MoE models become increasingly prevalent in cutting-edge machine learning applications, optimizing their distributed execution is paramount for maximizing hardware utilization and reducing operational costs.

The core challenge tackled by Comet stems from the inherent nature of MoE layers, where each input token is dynamically routed to a subset of specialized "expert" neural networks. In a distributed setting, this often necessitates extensive data movement between devices to gather input tokens for experts and then combine their outputs. Historically, communication can consume up to 47% of the total execution time in MoE layers, representing a major impediment to scalability and performance. Comet's innovative approach, centered around shared tensor-based dependency resolution and adaptive workload assignment, provides a robust solution to this problem, significantly improving the efficiency of MoE model deployment.

The significance of Comet lies in its ability to unlock greater efficiency for MoE models, which are critical for scaling model capacity without a proportional increase in computational cost. By intelligently overlapping communication with computation, Comet not only reduces overall latency but also enhances the utilization of powerful accelerators like GPUs. Its adoption in production environments with thousands of GPUs and demonstrated savings of millions of GPU hours underscore its practical impact and importance for the future of large-scale AI systems.

Background

▶ Watch: MOE communication overhead and overlapping challenges (0:00)

Mixture-of-Experts (MoE) models have emerged as a powerful paradigm for scaling neural networks, offering a way to increase model capacity and achieve superior performance without incurring the prohibitive computational costs associated with densely activated large models. In an MoE layer, a router or gate network determines which subset of "experts" (typically feed-forward networks) will process each input token. This selective activation allows for models with billions or even trillions of parameters to be trained and inferred efficiently, as only a small fraction of the total parameters are activated for any given input.

However, the distributed nature of large-scale MoE models introduces significant systems challenges, particularly regarding communication overhead. In a typical distributed MoE setup, experts are sharded across multiple devices (e.g., GPUs). The execution flow for an MoE layer generally involves three phases:

1. Data Receiving: Tokens are sent to the devices hosting their assigned experts. This often involves an all-to-all communication primitive, where each device sends a portion of its tokens to all other devices and receives tokens from all other devices.
2. Expert Computation: The assigned experts process their respective tokens. This involves standard matrix multiplication operations (e.g., GEMM).
3. Data Sending: The computed results from experts are sent back to the original devices, where they are combined to form the final output. This also typically involves an all-to-all or similar collective communication operation.

The talk highlights that communication, particularly the all-to-all operations, can account for a substantial portion—up to 47%—of the total execution time in MoE layers. To mitigate this latency and maximize hardware utilization, the concept of computation-communication overlapping is crucial. The goal is to perform communication operations concurrently with computation, effectively hiding the communication latency.

Previous attempts at overlapping often employ a "coarse-grained" approach, splitting the entire computation and communication into larger, independent units, for example, along the expert dimension. However, this strategy faces several limitations:

• Kernel Efficiency Degradation: Splitting large computations into smaller chunks can lead to a decrease in the efficiency of individual computation kernels. The sum of the execution times of smaller kernels (T1 + T2) might exceed the original single, larger kernel's time (T), due to overheads or suboptimal hardware utilization for smaller workloads.
• Dynamic Nature of MoE: The dynamic routing in MoE models means that the workload distribution across experts and devices is unpredictable. This variability in communication and computation patterns can lead to "bubbles" or idle periods in overlapping pipelines, where one stage (e.g., computation) finishes early while waiting for the other (e.g., communication) to complete, or vice-versa.
• Granularity Mismatch: Computation within MoE experts is typically performed at a "tail" level (referring to a batch of tokens processed by an expert), while communication often occurs at a "token" level. This mismatch makes fine-grained overlapping difficult.
• Complex Data Dependencies: As illustrated in the talk, a single computation "tail" might require data elements that are both local to the GPU and remote (requiring communication). This interdependency prevents straightforward concurrent execution.
• Harm to Computation Efficiency: Naively fusing fine-grained, token-level communication directly into computation kernels can negatively impact the performance of the highly optimized computation kernels themselves, as they might lose their regular access patterns or benefit from specialized hardware units.
• Variability in Configuration and Input Loads: Different MoE model configurations (number of experts, hidden dimensions) and varying input batch sizes lead to diverse communication and computation loads, demanding an adaptive approach to resource allocation for effective overlapping.

These challenges underscore the need for a more sophisticated, fine-grained, and adaptive mechanism for overlapping communication and computation in MoE models, which Comet aims to provide.

Key Findings

▶ Watch: Challenges in MOE communication-computation overlapping (4:00)

Comet delivers significant advancements in optimizing MoE model execution through its novel approach to computation-communication overlapping. The core findings and contributions are:

• Fine-grained Overlapping: Comet successfully enables fine-grained overlapping of communication and computation within MoE layers, addressing the traditional granularity mismatch and complex data dependencies that hinder efficiency.
• Shared Tensor-Based Dependency Resolving: A key innovation is the use of a shared tensor concept to model and decompose data dependencies between communication and computation stages. By analyzing these dependencies and rescheduling expert computations along independent dimensions (token dimension for the first MoE layer, hidden dimension for the second), Comet effectively breaks down barriers to concurrent execution.
• Adaptive Workload Assignment: Comet introduces an adaptive mechanism for resource allocation, specifically using warp and thread block specialization. This allows for dynamic assignment of separate thread blocks to communication and computation tasks, optimizing resource utilization based on varying model configurations, input loads, and parallelism strategies.
• Significant Performance Speedups:
• End-to-End Latency: Comet achieves an average speedup of 1.7 times for the end-to-end latency of MoE models integrated into the Megatron framework.
• Single MoE Layer Latency: For a single MoE layer, Comet demonstrates an even more impressive average speedup of 1.9 times.
• Communication Hiding: The framework is highly effective at hiding communication latency, achieving up to 90% concealment of communication time within a single MoE layer.
• Robustness Across Parallelism Strategies: Comet exhibits low latency and high efficiency across diverse parallelism strategies, including tensor parallelism, expert parallelism, and combined tensor and expert parallelism.
• Production Readiness and Scalability: Comet has been successfully adopted in production environments, operating on clusters with 10,000-scale GPUs, and has yielded substantial savings of millions of GPU hours.

These findings collectively demonstrate Comet's efficacy in tackling one of the most pressing performance bottlenecks in large-scale MoE deployments, offering a practical and highly effective solution for improving efficiency and reducing operational costs.

Technical Deep Dive

▶ Watch: Shared tensor-based dependency resolving for MOE pipelines (8:00)

Comet's technical innovation revolves around two primary designs: shared tensor-based dependency resolving and adaptive workload assignment. These mechanisms work in concert to achieve fine-grained computation-communication overlapping.

Shared Tensor-Based Dependency Resolving

At its core, Comet leverages the concept of a shared tensor to bridge the gap between communication and computation operations within an MoE layer. The talk models these operations as a producer-consumer pattern, where one operation generates data that the other consumes. By carefully analyzing and decomposing the shared tensor, Comet identifies dimensions along which data elements remain independent for the consumer operator, enabling effective rescheduling and overlapping.

The MoE layer typically consists of two main stages:

1. First Layer (or "Forward Pass" within the expert): Tokens are routed to experts, undergo an all-to-all communication operation to gather on the correct devices, and then proceed to the first GEMM operation (matrix multiplication) of the expert. Here, the all-to-all operation acts as the producer, and the GEMM operation is the consumer.
2. Second Layer (or "Backward Pass" / "Reduction" within the expert): After the first GEMM, tokens undergo another GEMM, and then their results are combined (e.g., via Top-K reduction and another all-to-all to return results to original token locations). In this stage, the GEMM operator is the producer, and the Top-K reduction followed by all-to-all communication functions as the consumer.

Comet's strategy for decomposing the shared tensor and rescheduling computation differs for these two layers due to distinct data dependencies:

First MoE Layer (All-to-All to GEMM)

• Dependency Analysis: In the first layer, the output of the all-to-all operation (the shared tensor) provides input tokens for the subsequent GEMM. Crucially, the communication and computation of different data elements along the token dimension are independent of each other. Each token can be processed by its assigned expert without direct dependency on other tokens being processed by different experts at that immediate stage.
• Rescheduling Strategy: To exploit this independence, Comet reorders the tokens. Instead of processing a mixed batch of local and remote tokens, which would require waiting for all remote tokens to arrive, Comet sorts tokens by their source ranks. The GEMM operation then starts computing "tails" (batches of tokens) with local tokens first. Simultaneously, the communication hardware can be busy fetching and receiving remote tokens. This allows computation to begin immediately on available data while communication for other data proceeds in parallel.

Second MoE Layer (GEMM to Top-K Reduction + All-to-All)

• Dependency Analysis: The second layer is more complex. The GEMM operator produces intermediate results (the shared tensor), which are then consumed by the Top-K reduction and subsequent all-to-all operations. The critical difference here is that the Top-K reduction typically performs reduction along the token dimension (the M dimension in the GEMM operation). This means that the shared tensor is not independent along the token dimension anymore, as the reduction operation requires results across tokens to determine the top-K.
• Rescheduling Strategy: Given the token-dimension dependency, Comet decomposes and reschedules the computation along the hidden dimension (the N dimension in the GEMM output, representing features). Instead of computing expert results entirely, expert by expert, Comet reorganizes the GEMM to follow a column-major order. This means computing the "yellow columns" (representing a subset of the hidden dimension for all tokens processed by an expert) first. Once these partial results are available, the Top-K reduction and all-to-all communication can immediately begin processing them, even while the "green columns" (the remaining hidden dimensions) are still being computed. This enables overlapping computation of different parts of the hidden dimension with the communication and reduction of already computed parts.

Adaptive Workload Assignment with Warp and Thread Block Specialization

After resolving data dependencies, Comet employs a hardware-aware technique called warp and thread block specialization to enable efficient overlapping. This approach is conceptually similar to horizontal fusion but is applied dynamically.

• Thread Block Specialization: The fused kernel within Comet contains two distinct types of thread blocks:

1. Computation Thread Blocks: These are dedicated to performing the expert computations (GEMM operations). They leverage highly optimized warp specialization implementations, which are efficient for high-throughput arithmetic.
2. Communication Thread Blocks: These are dedicated to handling the communication operations (all-to-all, data movement).

• Pipelining via Global Memory: These specialized computation and communication thread blocks are pipelined together, exchanging data through global memory. This separation ensures that each type of task can utilize its allocated hardware resources optimally without interfering with the other, while still cooperating to process the data flow.
• Adaptive Resource Allocation: A crucial aspect of Comet is its adaptive workload assignment. The number of thread blocks allocated to computation versus communication tasks is not fixed but is dynamically determined. This adaptability is vital because communication and computation workloads vary significantly with:
• Model configurations: Different numbers of experts, hidden dimensions, and model sizes.
• Input sizes: Varying batch sizes and sequence lengths.
• Parallelism strategies: Different distributions of experts and data across devices (tensor parallelism, expert parallelism, etc.).

The talk illustrates that the optimal division between computation and communication resources changes across these variables. To achieve this, Comet employs an offline profiling step. This profiling explores various resource splits (e.g., 70/30, 60/40, 50/50 for compute/comm thread blocks) to identify the optimal configuration for a given workload and hardware setup. Once the optimal split is determined, it is deployed to the online production system. This dynamic tuning ensures that Comet can consistently deliver high performance across a wide range of operational scenarios.

Together, the shared tensor-based dependency resolution and adaptive thread block specialization form the backbone of Comet, allowing it to intelligently manage and overlap the complex interplay of data movement and computation inherent in distributed MoE models.

Experimental Setup & Results

▶ Watch: Rescheduling computation to improve MOE layer overlapping (9:00)

The evaluation of Comet was conducted on a cluster equipped with eight A100 GPUs. The Mixture-of-Experts (MoE) models used for testing were integrated into the Megatron framework, a popular toolkit for training large language models developed by NVIDIA. This choice of framework and hardware reflects a realistic high-performance computing environment for large-scale AI.

Comet's performance was benchmarked against several existing MoE frameworks and optimization techniques, including:

• Megatron: The baseline implementation within the Megatron framework.
• FastMoE: A well-known open-source library for MoE training.
• To Tail: Another technique or framework for MoE optimization, though less detailed in the talk.

The primary metrics used to assess Comet's effectiveness were:

• End-to-End Latency: The total time taken for a complete MoE model inference or training step.
• Single MoE Layer Latency: The latency specifically for executing one MoE layer, which is the core bottleneck targeted by Comet.
• Communication Latency Hiding: The percentage of communication time that is successfully overlapped and hidden by concurrent computation.

Headline Results

The experimental results demonstrated significant performance gains achieved by Comet:

• End-to-End Latency Speedup: On average, Comet achieved a 1.7 times speedup in end-to-end latency compared to baseline implementations. This indicates a substantial improvement in the overall efficiency of MoE model execution.
• Single MoE Layer Latency Speedup: Focusing on the critical MoE layer, Comet delivered an even greater average speedup of 1.9 times. This highlights its effectiveness directly at the source of the communication bottleneck.
• Communication Hiding: A deep dive into the execution profile of a single MoE layer revealed that Comet can hide up to 90% of the communication latency. This is a testament to the success of its fine-grained overlapping strategies.

Robustness Across Parallelism Strategies

The evaluation also confirmed that Comet's benefits are not confined to a specific distributed setup. It consistently achieved low latency across various parallelism strategies commonly employed for MoE models:

• Tensor Parallelism: Where parts of individual tensors (e.g., weight matrices) are sharded across devices.
• Expert Parallelism: Where different experts are assigned to different devices.
• Combined Tensor and Expert Parallelism: A hybrid approach leveraging both strategies.

This versatility underscores Comet's general applicability and robustness in diverse distributed training and inference environments.

Scalability and Production Deployment

Beyond controlled benchmarks, Comet has proven its mettle in real-world production environments. It has been adopted within clusters comprising 10,000-scale GPUs, leading to documented savings of millions of GPU hours. This real-world validation is a strong indicator of its practical utility and efficiency at massive scales.

During the Q&A, the speaker addressed the generalizability of the adaptive profiling approach for arbitrary sequence lengths and parallelism degrees. It was clarified that while the approach is general, offline profiling is indeed required again when parallelism strategies or input sizes vary to find the optimal division of thread blocks for communication and computation. The paper, however, includes details on this profiling and potential equations for computing thread block numbers.

Regarding multi-node deployments, the speaker confirmed that Comet can be applied to multi-node scales and has been implemented in such environments. However, a nuanced observation was made: when communication scales to several nodes, the communication overhead can grow significantly faster than computation. In such extreme cases, while fusing smaller computations into larger communication windows can still work, the beneficial impact might be less pronounced compared to single-node or smaller-scale multi-node scenarios where the communication-to-computation ratio is more balanced. This suggests that while Comet provides a strong foundation, the relative gains might diminish in scenarios dominated by extremely high inter-node communication costs.

Practical Implications

▶ Watch: Adaptive resource allocation for communication and computation (12:00)

Comet's advancements in fine-grained computation-communication overlapping for MoE models carry significant practical implications for practitioners, infrastructure teams, model builders, and those responsible for deploying large-scale AI systems.

First and foremost, Comet directly addresses a critical bottleneck in the efficiency of MoE models: the substantial communication overhead. By achieving a 1.7x end-to-end speedup and hiding up to 90% of communication latency in a single MoE layer, Comet dramatically reduces the overall time required for training and inference. This translates into faster model development cycles, quicker iteration on experiments, and reduced latency for real-time inference applications. For model builders, this means they can explore larger and more complex MoE architectures with less concern about the system-level performance implications.

For infrastructure teams and deployers, the demonstrated savings of millions of GPU hours in production environments with 10,000-scale GPUs highlight a massive reduction in operational costs. GPU hours are a primary expense in large-scale AI, and any efficiency gain directly impacts the bottom line. Comet enables more efficient utilization of expensive GPU resources, effectively getting more "work" out of the same hardware footprint. This can defer hardware upgrades or allow for larger workloads to be run on existing infrastructure.

The robustness of Comet across various parallelism strategies (tensor, expert, and combined) is also highly practical. It means that teams are not locked into a single parallelism scheme to benefit from Comet. They can choose the most appropriate strategy for their specific model and hardware configuration, knowing that Comet will still provide significant performance improvements. This flexibility simplifies the deployment of diverse MoE models.

However, Comet's approach also introduces certain tradeoffs and considerations for practitioners:

• Offline Profiling Overhead: The need for offline profiling to adaptively assign resources (thread blocks) for communication and computation is a practical consideration. While necessary for optimal performance across varied workloads and hardware, it adds an initial setup and tuning step. Infrastructure teams will need processes in place to perform and manage this profiling, especially when model configurations, input sizes, or parallelism strategies change. The paper's inclusion of profiling details and potential equations for resource allocation would be valuable for streamlining this process.
• Diminishing Returns in Extreme Multi-Node Scenarios: While Comet scales to multi-node deployments, the speaker noted that in scenarios where communication overhead grows disproportionately fast across many nodes, the benefits of overlapping might be less pronounced. This implies that for extremely wide multi-node MoE deployments, other optimizations focusing on reducing raw communication volume (e.g., advanced routing algorithms, sparse communication patterns) might still be critical in conjunction with Comet. Practitioners should be aware of this nuance when designing ultra-large-scale systems.
• Integration Complexity: Adopting a framework like Comet into existing ML pipelines (e.g., within Megatron) requires careful integration. The open-source availability of Comet is a significant advantage, as it lowers the barrier to adoption and allows teams to inspect, customize, and integrate the solution more readily.

In summary, Comet offers a powerful, validated solution to a critical performance challenge in MoE models. Its practical implications include substantial cost savings, faster execution, and greater flexibility in deploying large-scale AI, albeit with the need for thoughtful integration and profiling in production environments.

Key Takeaways

• Communication is a Major MoE Bottleneck: Distributed Mixture-of-Experts (MoE) models suffer from significant communication overhead, accounting for up to 47% of execution time, hindering performance and scalability.
• Comet Enables Fine-Grained Overlapping: The Comet framework introduces novel techniques for fine-grained computation-communication overlapping within MoE layers, overcoming limitations of previous coarse-grained approaches.
• Shared Tensor for Dependency Resolution: Comet uses a shared tensor concept to analyze and decompose data dependencies, allowing for intelligent rescheduling of computation (e.g., sorting tokens, column-major processing) to maximize overlap.
• Adaptive Resource Allocation is Key: Warp and thread block specialization with adaptive workload assignment dynamically allocates GPU resources to communication and computation tasks based on offline profiling, optimizing performance for varying workloads.
• Significant Performance Gains: Comet achieves an average 1.7x end-to-end speedup and 1.9x speedup for single MoE layers, successfully hiding up to 90% of communication latency.
• Proven in Production: Adopted on 10,000-scale GPU clusters, Comet has demonstrated real-world impact by saving millions of GPU hours, validating its efficiency and scalability.

About the Speaker(s)

The talk was introduced by Ningxin Zheng from Bytedance, who presented the research on Comet. The work is a collaborative effort by a team including Shulai Zhang, Ningxin Zheng, Haibin Lin, Ziheng Jiang, Wenlei Bao, and Xin Liu. While specific individual bios were not provided in the transcript, their affiliation with Bytedance indicates their involvement in industrial-scale AI/ML research and development, particularly in areas related to large-scale model optimization and distributed systems. The successful deployment of Comet in Bytedance's production environment further underscores their expertise in building and optimizing high-performance machine learning infrastructure.