[论文解读] M6-T: Exploring Sparse Expert Models and Beyond

Yang An, Junyang Lin+13 more|arXiv (Cornell University)|May 31, 2021

Recommender Systems and Techniques参考文献 34被引用 14

一句话总结

本文提出专家原型化（expert prototyping）方法，将专家拆分为k个原型，并采用k个top-1路由机制，以在保持计算成本恒定的前提下提升稀疏专家模型的性能。该方法使一个1.1万亿参数的模型在仅480块V100 GPU上训练速度比同规模基线模型快5倍，性能优于此前在2048个TPU核心上训练的SOTA模型。

ABSTRACT

Mixture-of-Experts (MoE) models can achieve promising results with outrageous large amount of parameters but constant computation cost, and thus it has become a trend in model scaling. Still it is a mystery how MoE layers bring quality gains by leveraging the parameters with sparse activation. In this work, we investigate several key factors in sparse expert models. We observe that load imbalance may not be a significant problem affecting model quality, contrary to the perspectives of recent studies, while the number of sparsely activated experts $k$ and expert capacity $C$ in top-$k$ routing can significantly make a difference in this context. Furthermore, we take a step forward to propose a simple method called expert prototyping that splits experts into different prototypes and applies $k$ top-$1$ routing. This strategy improves the model quality but maintains constant computational costs, and our further exploration on extremely large-scale models reflects that it is more effective in training larger models. We push the model scale to over $1$ trillion parameters and implement it on solely $480$ NVIDIA V100-32GB GPUs, in comparison with the recent SOTAs on $2048$ TPU cores. The proposed giant model achieves substantial speedup in convergence over the same-size baseline.

研究动机与目标

• 探究影响混合专家（MoE）模型质量与训练效率的关键因素。
• 解决在超过1000亿参数规模下高效扩展稀疏专家模型的挑战。
• 在提升超大规模MoE架构性能的同时降低计算成本。
• 证明在消费级GPU集群而非TPU系统上训练万亿参数MoE模型的可行性。

提出的方法

• 专家原型化将每个专家拆分为k个独立原型，支持k个top-1路由而非标准的top-k路由。
• 每个输入标记通过门控函数被独立路由至k个不同原型，路由决策按每个原型独立进行。
• 最终输出通过softmax门控注意力机制对k个原型输出进行加权求和计算。
• 通过避免完整top-k计算，仅使用k个独立的top-1路由操作，保持FLOPs恒定。
• 采用Adafactor优化器进行训练，学习率降低至0.005，并采用缩放的权重初始化（σ=0.002）以确保稳定性。
• 在配备RDMA互连的480块NVIDIA V100-32GB GPU集群上，使用Whale分布式框架进行训练。

实验结果

研究问题

• RQ1 负载不均是否显著降低MoE模型的质量？
• RQ2 激活专家数量（k）与专家容量（C）如何影响MoE性能？
• RQ3 专家原型化是否能在不增加计算成本的前提下提升模型质量？
• RQ4 专家原型化在训练超大规模模型（如1000亿参数以上）时效果如何？
• RQ5 是否能高效地在GPU集群而非TPU集群上训练万亿参数的MoE模型？

主要发现

• 与先前假设相反，负载不均并非显著影响MoE模型质量的主要问题。
• 增加激活专家数量（k）与专家容量（C）可显著提升模型性能。
• 专家原型化使1.1万亿参数MoE模型的收敛速度相比同规模基线提升5倍。
• 1.1万亿参数的专家原型化模型在480块V100-32GB GPU上实现更优性能，超越此前在2048个TPU核心上训练的SOTA模型。
• 该方法在提升模型质量的同时保持计算成本恒定，尤其在更大规模设置下优势明显。
• 尽管FLOPs完全相同，模型收敛速度显著加快，表明训练动态效率得到提升。