Introducing MegaScale-Infer: Optimizing Large Language Model Performance
Large language models (LLMs) have become essential in various applications, including chatbots, code generation, and search engines. However, as these models grow to billions of parameters, the challenge of efficient computation intensifies. Maintaining low latency and high throughput while scaling these systems requires innovative solutions in algorithm design and system optimization.
The Challenge of Sparsity and Resource Utilization
A significant issue in LLMs is the concept of sparsity, particularly in Mixture-of-Experts (MoE) models. These models activate only a subset of their components for each input, which reduces the overall computational load. However, this selective activation can lead to underutilization of hardware resources. During inference, memory access to key-value caches can create bottlenecks, while feed-forward networks (FFNs) may remain idle due to their limited token allocation. This inefficiency can lead to substantial drops in GPU utilization and increased operational costs.
Current Solutions and Their Limitations
Existing methods, such as vLLM and TensorRT-LLM, have attempted to enhance inference scaling through parallelism and optimized kernels. However, these solutions often treat the model as a single entity, making it difficult to scale individual components effectively. As MoE models continue to grow, this limitation results in smaller active batches per expert, diminishing the advantages of batching for FFNs. Furthermore, tensor and pipeline parallelism approaches introduce additional communication overhead, particularly in multi-GPU environments, which can hinder performance.
Introducing MegaScale-Infer
Researchers from ByteDance and Peking University have developed MegaScale-Infer, a groundbreaking system that redefines MoE serving architecture. Instead of treating the model as a single block, MegaScale-Infer disaggregates the attention and FFN modules and assigns them to separate GPUs. This approach allows for tailored scaling and parallelism strategies based on the specific needs of each module. Memory-intensive attention modules can be replicated to handle multiple requests, while FFNs can leverage expert parallelism for enhanced performance. Additionally, the system supports heterogeneous GPU deployments, optimizing resource allocation based on the nature of the tasks.
Performance Optimization Strategies
To further enhance performance, MegaScale-Infer utilizes a ping-pong pipeline parallelism strategy. This involves dividing request batches into smaller micro-batches that alternate between attention and FFN modules, ensuring that both components remain active. The system intelligently determines the optimal number of micro-batches required to maintain high utilization based on compute time and communication latency. For instance, if communication time is significantly lower than compute time, the system will use multiple micro-batches to maximize efficiency.
Furthermore, MegaScale-Infer incorporates a high-performance M2N communication library that minimizes unnecessary data transfers between GPUs and CPUs. This innovation reduces latency and improves stability by replacing traditional communication methods with a more efficient sender-receiver model designed for MoE and token dispatch patterns.
Case Study and Results
In practical tests with various large-scale MoE models, such as Mixtral 8×22B and a custom model with 317 billion parameters, MegaScale-Infer demonstrated remarkable improvements. In homogeneous setups using NVIDIA Ampere GPUs, the system achieved up to 2.56 times higher decoding throughput compared to vLLM and 1.28 times higher than TensorRT-LLM. In heterogeneous clusters, MegaScale-Infer provided up to 3.24 times higher throughput per dollar than baseline models, showcasing its cost-effectiveness. The M2N communication library also yielded up to 4.2 times higher throughput and 68.2% lower latency than traditional methods.
Conclusion
The research presented in this paper highlights a critical issue of underutilized GPU resources during MoE inference and offers a practical solution through architectural modularization. By implementing a disaggregation strategy, micro-batch pipelining, and a custom communication protocol, MegaScale-Infer significantly enhances serving efficiency and reduces operational costs. Businesses looking to leverage AI can learn from this approach to optimize their own systems and maximize resource utilization.
