Speculative MoE: Communication Efficient Parallel MoE Inference with Speculative Token and Expert Pre-scheduling
作者: Yan Li, Pengfei Zheng, Shuang Chen, Zewei Xu, Yunfei Du, Zhengang Wang
发布时间: 2025-03-07
来源: arxiv
研究方向: 大规模语言模型（LLM）推理中的并行化和通信优化
主要内容
本文研究了大规模语言模型（LLM）推理中的并行化和通信优化问题，特别是针对MoE（Mixture of Experts）模型的并行推理。作者提出了一种名为Speculative MoE（s-MoE）的框架，通过预测并预先调度令牌和专家，有效减少了MoE推理中的通信开销，从而提高了推理效率和可扩展性。
主要贡献
1. 提出了Speculative MoE（s-MoE）框架，通过预测并预先调度令牌和专家，显著提高了MoE推理的效率。
2. 设计了两种投机并行化方案：投机令牌重排和投机专家分组，以减少EP（Expert Parallelism）的通信量。
3. 在DeepSpeed-MoE和SGLang等主流MoE推理引擎中实现了s-MoE，并验证了其有效性。
4. 通过实验表明，s-MoE在多个数据集和硬件平台上均能显著提高MoE推理的吞吐量。
5. 提出了基于概率模型的令牌-专家路由路径预测方法，并实现了高效的令牌-专家协同调度算法。
研究方法
1. 投机并行化：通过预测令牌的专家路由路径，预先调度令牌和专家，以减少通信开销。
2. 令牌-专家协同调度：通过平衡令牌和专家的分布，最大化本地激活率，减少远程激活和通信。
3. 概率模型：使用概率模型预测令牌-专家路由路径，以指导令牌和专家的调度。
4. 整数线性规划：使用整数线性规划求解令牌-专家协同调度问题，以实现负载平衡和通信优化。
实验结果
实验结果表明，s-MoE在多个数据集和硬件平台上均能显著提高MoE推理的吞吐量。在DeepSeek-V2和Mixtral-8x7B等MoE模型上，s-MoE的吞吐量比DeepSpeed-MoE提高了1.58x-2.37x倍。在SGLang上，s-MoE的吞吐量提高了1.68x-1.97x倍。
未来工作
未来研究将重点关注以下方向：1）进一步优化投机并行化方案，提高预测精度和调度效率；2）将s-MoE扩展到其他类型的MoE模型和LLM推理引擎；3）探索更有效的令牌-专家协同调度算法，以实现更好的负载平衡和通信优化。