通过优化专家协作通信,让MoE模型训练推理快1.5倍以上
Occult: Optimizing Collaborative Communication across Experts for Accelerated Parallel MoE Training and Inference
- 设计新通信策略,提升同一设备上专家协同比例
- 在保持精度前提下降低通信开销,大规模训练节省超40%时间
- 适用于大模型训练与推理,适合追求高效部署的研究者
Mixture-of-experts (MoE) 架构在专家并行下可实现优异计算效率,但严重依赖跨设备全对全通信。此类通信开销通常占总运行时间超过40%,阻碍现代MoE模型的分布式训练与推理扩展。本文首次定义协作通信以揭示此固有瓶颈,并提出系统与算法级创新。针对被同一令牌激活的一对专家,若其位于同一设备则称为“内协作”,否则为“外协作”。初步研究发现,提升内协作比例可显著加速大规模专家并行。据此提出优化框架Occult,支持精确计算(通信成本更低)或可控剪枝(协作剪枝),通过修改微调实现。在多种MoE-LLMs上的实验表明,Occult比主流先进框架更快1.5倍以上,且质量相当或更优。代码已公开。
原文摘要 · Abstract (English)
Mixture-of-experts (MoE) architectures could achieve impressive computational efficiency with expert parallelism, which relies heavily on all-to-all communication across devices. Unfortunately, such communication overhead typically constitutes a significant portion of the total runtime, hampering the scalability of distributed training and inference for modern MoE models (consuming over $40\%$ runtime in large-scale training). In this paper, we first define collaborative communication to illustrate this intrinsic limitation, and then propose system- and algorithm-level innovations to reduce communication costs. Specifically, given a pair of experts co-activated by one token, we call them "collaborated", which comprises $2$ cases as intra- and inter-collaboration, depending on whether they are kept on the same device. Our pilot investigations reveal that augmenting the proportion of intra-collaboration can accelerate expert parallelism at scale. It motivates us to strategically optimize collaborative communication for accelerated MoE training and inference, dubbed Occult. Our designs are capable of either delivering exact results with reduced communication cost or controllably minimizing the cost with collaboration pruning, materialized by modified fine-tuning. Comprehensive experiments on various MoE-LLMs demonstrate that Occult can be faster than popular state-of-the-art inference or training frameworks (more than $1.5\times$ speed up across multiple tasks and models) with comparable or superior quality compared to the standard fine-tuning. Code is available at $\href{https://github.com/UNITES-Lab/Occult}{https://github.com/UNITES-Lab/Occult}$.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。