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.

在这篇论文中，我们提出了一种轻量级的优化技术 called ExFlow，用于大大加速 Mixture of Experts 模型的推理过程，并通过利用层间专家亲和力来减少跨 GPU 路由延迟，取得了显著的推理吞吐量提升效果。