arXiv:2503.08467cs.LGcs.AI2025-03被引 10

通过专家张量分片实现多卡推理负载均衡,提升MoE模型效率。

Accelerating MoE Model Inference with Expert Sharding

  • 将专家矩阵按行和列分片,实现计算均匀分配。
  • 在不平衡路由下仍保持全令牌保留,TTFT提速最高6.4倍。
  • 适合大规模编码器类MoE模型的多卡推理场景。

混合专家(MoE)模型在语言建模中表现优异,但因令牌路由不均和通信开销导致硬件利用率低。现有工作多关注训练优化与解码器架构,而基于编码器的MoE模型在多卡专家并行下的推理仍缺乏研究。本文提出MoEShard,通过专家张量分片实现完美负载均衡。不同于依赖启发式容量因子或丢弃令牌的方法,该系统均匀分配计算任务,确保所有令牌被保留,显著提升硬件利用率,不受路由偏差影响。其核心是专家矩阵的行列分解策略,减少空闲时间,避免由专家分配不均引发的瓶颈。同时,通过融合分解后的专家计算,大幅降低内核启动次数,提升吞吐量。在深度学习框架DeepSpeed对比测试中,针对编码器结构的MoE模型,最高实现6.4倍的首次生成时间(TTFT)加速,证明了合理应用张量分片是高效推理的有效路径。

原文摘要 · Abstract (English)

Mixture of experts (MoE) models achieve state-of-the-art results in language modeling but suffer from inefficient hardware utilization due to imbalanced token routing and communication overhead. While prior work has focused on optimizing MoE training and decoder architectures, inference for encoder-based MoE models in a multi-GPU with expert parallelism setting remains underexplored. We introduce MoEShard, an inference system that achieves perfect load balancing through tensor sharding of MoE experts. Unlike existing approaches that rely on heuristic capacity factors or drop tokens, MoEShard evenly distributes computation across GPUs and ensures full token retention, maximizing utilization regardless of routing skewness. We achieve this through a strategic row- and column-wise decomposition of expert matrices. This reduces idle time and avoids bottlenecks caused by imbalanced expert assignments. Furthermore, MoEShard minimizes kernel launches by fusing decomposed expert computations, significantly improving throughput. We evaluate MoEShard against DeepSpeed on encoder-based architectures, demonstrating speedups of up to 6.4$\times$ in time to first token (TTFT). Our results show that tensor sharding, when properly applied to experts, is a viable and effective strategy for efficient MoE inference.

MoE推理加速张量分片多卡并行

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。