Capacity-Aware Inference: Mitigating the Straggler Effect in Mixture of Experts
作者: Shwai He, Weilin Cai, Jiayi Huang, Ang Li
发布时间: 2025-03-10
来源: arxiv
研究方向: 大规模语言模型（LLM）的推理效率优化
主要内容
该研究针对Mixture of Experts (MoE)架构在推理过程中存在的专家负载不均衡问题，即部分专家过载而其他专家未被充分利用，导致资源利用效率低下和延迟增加。研究提出了容量感知推理方法，通过丢弃过载专家的冗余令牌和重新路由令牌到未充分利用的专家，来优化MoE推理流程。
主要贡献
1. 识别了MoE推理中由令牌分配不均导致的Straggler Effect，突出了减少延迟的优化潜力。
2. 针对令牌分配不均问题，提出了Token Drop和Token Reroute技术，以增强专家的均衡利用。
3. 实验结果验证了Token Drop和Token Reroute的有效性，实现了专家利用率和推理效率的显著提升。
研究方法
1. 容量感知令牌丢弃：通过丢弃过载专家的冗余令牌来减轻负载不均衡，同时尽量减少对模型性能的影响。
2. 容量感知令牌重路由：将丢弃的令牌重新分配给未充分利用的专家，以平衡令牌分配。
3. 基于令牌重要性的评分函数：用于评估每个令牌的重要性，并据此丢弃或重新路由令牌。
实验结果
实验结果表明，与基线模型相比，所提出的方法在推理效率上取得了显著的提升。例如，在Mixtral-8 7B-Instruct模型上，平均性能提高了0.2%，推理速度提高了1.94倍。
未来工作
未来工作将扩展研究范围，将所提出的方法应用于其他MoE模型架构，并探索将容量感知令牌丢弃和令牌重路由技术集成到训练过程中，以进一步提高LLM的推理效率和性能。