Exploiting Inter-Layer Expert Affinity for Accelerating Mixture-of-Experts Model Inference

TL;DR

ExFlow introduces context-coherent expert parallelism and exploits inter-layer expert affinity to accelerate GPT-style Mixture-of-Experts inference on distributed hardware. By modeling the cross-layer routing likelihood $P(E_{p,j+1}|E_{i,j})$ and solving an offline integer-programming placement, ExFlow reduces cross-GPU and inter-node communication, achieving up to 67% Alltoall reduction and up to 2.2x end-to-end throughput across GPT MoE models with $E$ experts per layer. The approach does not require fine-tuning or retraining and shows affinity stability across training and out-of-distribution datasets, enabling practical deployment on varied hardware topologies. Overall, ExFlow delivers a topology-agnostic, deployment-friendly method to scale MoE inference efficiency in large language models, with clear gains for practical, large-scale systems.

Abstract

In large language models like the Generative Pre-trained Transformer, the Mixture of Experts paradigm has emerged as a powerful technique for enhancing model expressiveness and accuracy. However, deploying GPT MoE models for parallel inference on distributed systems presents significant challenges, primarily due to the extensive Alltoall communication required for expert routing and aggregation. This communication bottleneck exacerbates the already complex computational landscape, hindering the efficient utilization of high-performance computing resources. In this paper, we propose a lightweight optimization technique called ExFlow, to largely accelerate the inference of these MoE models. We take a new perspective on alleviating the communication overhead by exploiting the inter-layer expert affinity. Unlike previous methods, our solution can be directly applied to pre-trained MoE models without any fine-tuning or accuracy degradation. By proposing a context-coherent expert parallelism on distributed systems, our design only uses one Alltoall communication to deliver the same functionality while previous methods all require two Alltoalls. By carefully examining the conditional probability in tokens' routing across multiple layers, we proved that pre-trained GPT MoE models implicitly exhibit a strong inter-layer expert affinity. We then design an efficient integer programming model to capture such features and show that by properly placing the experts on corresponding GPUs, we can reduce up to 67% cross-GPU routing latency. Our solution beats the cutting-edge MoE implementations with experts from 8 to 64, with up to 2.2x improvement in inference throughput. We further provide a detailed study of how the model implicitly acquires this expert affinity at the very early training stage and how this affinity evolves and stabilizes during training.

Figures (16)

• Figure 1: Given a pre-trained MoE model, (a) vanilla placement strategy causes intensive cross-GPU communication. (b) leveraging inter-layer expert affinity can avoid unnecessary Alltoall communication.
• Figure 2: Heatmaps illustrating the distribution of inter-layer expert routing preference. Color intensity represents the magnitude of the likelihood, with white signifying low values and red indicating high values. We measure the conditional probability of expert routing in different parts of a pre-trained GPT model with 12 MoE layers, and 32 experts per layer. For each row, we can observe only a few columns are red, indicating a strong affinity.
• Figure 3: Due to current expert parallelism consisting of data parallel, different GPUs do not share the context of tokens. Therefore, $T_2$ needs to go back to GPU 2 for performing attention in the next layer.
• Figure 4: Before inference, we use Allgather to ensure every GPU has all contexts. After each iteration, another Allgather is performed on the newly generated tokens, we then append them to the current contexts for the next iteration.
• Figure 5: Each GPU has a capacity of 2 experts per layer. (a). Vanilla expert parallelism. Each token needs to come back to its original GPU to perform attention computation with its context. (b). Enabling token context coherence across all GPUs. Tokens do not need to go back to the original GPU to perform attention computation, because they can attend to their context on the local GPU. (c). Exploiting expert affinity to further reduce token communication. The placement of experts on each layer is now following an optimal pattern such that tokens will remain on local GPUs with the maximum probability.