BigMac通过新通信机制显著降低MoE模型训练与推理延迟。
BigMac: A Communication-Efficient Mixture-of-Experts Model Structure for Fast Training and Inference
- 采用降维通信的DCCA机制替代传统全连接通信
- 训练延迟降低3.09倍,推理吞吐提升3.11倍
- 适合追求高效推理与快速训练的大型语言模型部署
Mixture-of-Experts(MoE)结构能以近似线性增长的计算资源扩展Transformer类大语言模型并提升性能。近期提出的细粒度DeepSeekMoE进一步提升了计算效率,但其引入的All-to-All通信成为瓶颈,尤其在激活更多专家时加剧通信开销。本文提出新型细粒度MoE结构BigMac,摒弃传统通信模式CDAC(通信-下降-上升-通信),改用DCCA(下降-通信-通信-上升)机制。具体地,在专家入口和出口分别加入降维与升维投影,使通信在极低维度进行。同时重构小专家结构,确保其具备处理输入令牌的能力。实验表明,BigMac在相同专家数量和参数量下,模型质量与细粒度MoE相当甚至更优。更重要的是,在Megatron、Tutel和DeepSpeed-Inference等主流框架上,训练端到端延迟降低最高达3.09倍,推理吞吐量提升最高达3.11倍。
原文摘要 · Abstract (English)
The Mixture-of-Experts (MoE) structure scales the Transformer-based large language models (LLMs) and improves their performance with only the sub-linear increase in computation resources. Recently, a fine-grained DeepSeekMoE structure is proposed, which can further improve the computing efficiency of MoE without performance degradation. However, the All-to-All communication introduced by MoE has become a bottleneck, especially for the fine-grained structure, which typically involves and activates more experts, hence contributing to heavier communication overhead. In this paper, we propose a novel MoE structure named BigMac, which is also fine-grained but with high communication efficiency. The innovation of BigMac is mainly due to that we abandon the \textbf{c}ommunicate-\textbf{d}escend-\textbf{a}scend-\textbf{c}ommunicate (CDAC) manner used by fine-grained MoE, which leads to the All-to-All communication always taking place at the highest dimension. Instead, BigMac designs an efficient \textbf{d}escend-\textbf{c}ommunicate-\textbf{c}ommunicate-\textbf{a}scend (DCCA) manner. Specifically, we add a descending and ascending projection at the entrance and exit of the expert, respectively, which enables the communication to perform at a very low dimension. Furthermore, to adapt to DCCA, we re-design the structure of small experts, ensuring that the expert in BigMac has enough complexity to address tokens. Experimental results show that BigMac achieves comparable or even better model quality than fine-grained MoEs with the same number of experts and a similar number of total parameters. Equally importantly, BigMac reduces the end-to-end latency by up to 3.09$\times$ for training and increases the throughput by up to 3.11$\times$ for inference on state-of-the-art AI computing frameworks including Megatron, Tutel, and DeepSpeed-Inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。