摘要
Mixture-of-Experts (MoE) 架构在不成比例增加计算成本的情况下提升模型容量，已成为扩展大规模语言模型（LLMs）至万亿参数规模的关键组成部分。MoE 模型的高效部署依赖于在多 GPU 上的分布式执行，每个 MoE 层涉及两个全到全的通信：将令牌分发到专家排名和将专家输出返回其源排名。传统的 MoE 实现会在专家计算完成后启动返回的全到全通信，这会在关键路径上暴露通信延迟并降低 GPU 利用率。
我们提出了一种细粒度方法，通过 tile-level 信号和调度重叠专家计算与第二次全到全通信。我们的生产者-消费者共同设计结合了：
- 持久的每排名计算内核（生产者），覆盖该排名的所有本地专家，以消除重复内核启动开销并优先处理远程关键 tiles；
- 在少量专用流式多处理器（SM）上的持久通信内核（消费者），随着 tiles 准备好而发起分段粒度的传输。
我们的共同设计避免对底层计算操作符或通信原语进行侵入性更改，使其在多 GPU 系统中提高分布式 MoE 执行效率变得实用。在一个 4-A100 GPU 平台上，我们对三种 MoE 模型进行了评估，与四个最先进的 MoE 系统相比，我们的方法实现了最高 2.64 倍的端到端加速和 2.74 倍的 MoE 层加速。与传统的非重叠基线相比，我们的方法在不同的 GEMM 形状、路由模式以及广泛的生产者/消费者 SM 分区下，始终提高了操作符和 MoE 层级别的性能，同时保持了正确性。
博主点评： 本文通过引入细粒度的调度策略，成功提升了 Mixture-of-Experts 模型在多 GPU 环境下的执行效率，展现了在大规模计算场景中优化通信与计算重叠的重要性。该方法不仅提高了性能，还保持了系统的灵活性与可扩展性，为未来的 MoE 模型设计提供了新的思路。