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.

翻译：在生成式预训练Transformer等大语言模型中，混合专家范式已成为增强模型表达能力和准确性的强大技术。然而，在分布式系统上部署GPT MoE模型进行并行推理面临重大挑战，主要原因是专家路由和聚合需要大量Alltoall通信。这一通信瓶颈加剧了本已复杂的计算环境，阻碍了高性能计算资源的有效利用。本文提出一种名为ExFlow的轻量级优化技术，可大幅加速这类MoE模型的推理。我们通过利用层间专家亲和性来缓解通信开销，这一视角与以往方法截然不同。我们的解决方案可直接应用于预训练的MoE模型，无需任何微调，也不会造成精度损失。通过在分布式系统上提出一种上下文一致的专家并行机制，我们的设计仅需一次Alltoall通信即可实现同等功能，而此前所有方法均需两次Alltoall通信。通过仔细研究令牌跨层路由中的条件概率，我们证明预训练的GPT MoE模型隐式地展现出强烈的层间专家亲和性。进而设计了一种高效的整数规划模型来捕捉此类特征，并证明通过将专家合理部署在对应GPU上，最多可减少67%的跨GPU路由延迟。我们的解决方案在专家数量从8到64的范围内均优于最先进的MoE实现，推理吞吐量提升高达2.2倍。此外，我们还详细研究了模型如何在早期训练阶段隐式获得这种专家亲和性，以及该亲和性在训练过程中如何演化并趋于稳定。