Optimizing Mixture-of-Experts Inference Time Combining Model Deployment and Communication Scheduling

TL;DR

Aurora optimizes mixture-of-experts inference by strategically deploying models and scheduling communication, significantly reducing latency and improving GPU utilization across diverse hardware setups.

Contribution

It introduces Aurora, the first method to jointly optimize model deployment and communication scheduling for MoE inference, with proven theoretical and practical speedups.

Findings

Achieves up to 3.54x inference speedup in heterogeneous environments.

Improves GPU utilization by up to 1.5x.

Provides optimal or near-optimal solutions across various GPU configurations.

Abstract

As machine learning models scale in size and complexity, their computational requirements become a significant barrier. Mixture-of-Experts (MoE) models alleviate this issue by selectively activating relevant experts. Despite this, MoE models are hindered by high communication overhead from all-to-all operations, low GPU utilization due to the synchronous communication constraint, and complications from heterogeneous GPU environments. This paper presents Aurora, which optimizes both model deployment and all-to-all communication scheduling to address these challenges in MoE inference. Aurora achieves minimal communication times by strategically ordering token transmissions in all-to-all communications. It improves GPU utilization by colocating experts from different models on the same device, avoiding the limitations of synchronous all-to-all communication. We analyze Aurora's optimization strategies theoretically across four common GPU cluster settings: exclusive vs. colocated models on GPUs, and homogeneous vs. heterogeneous GPUs. Aurora provides optimal solutions for three cases, and for the remaining NP-hard scenario, it offers a polynomial-time sub-optimal solution with only a 1.07x degradation from the optimal. Aurora is the first approach to minimize MoE inference time via optimal model deployment and communication scheduling across various scenarios. Evaluations demonstrate that Aurora significantly accelerates inference, achieving speedups of up to 2.38x in homogeneous clusters and 3.54x in heterogeneous environments. Moreover, Aurora enhances GPU utilization by up to 1.5x compared to existing methods.