[论文解读] Scalable and Efficient MoE Training for Multitask Multilingual Models

Young Jin Kim, Ammar Ahmad Awan+7 more|arXiv (Cornell University)|Sep 22, 2021

Topic Modeling参考文献 26被引用 33

一句话总结

本文提出 DeepSpeed MoE，一套可扩展的系统，用于在多任务多语言设置下训练大规模 Mixture-of-Experts (MoE) 模型，并结合训练技术（RTS、AoE、剪枝）并在 Z-code M3 up to 10B parameters 上展示强劲的 MT 与多语言生成结果。

ABSTRACT

The Mixture of Experts (MoE) models are an emerging class of sparsely activated deep learning models that have sublinear compute costs with respect to their parameters. In contrast with dense models, the sparse architecture of MoE offers opportunities for drastically growing model size with significant accuracy gain while consuming much lower compute budget. However, supporting large scale MoE training also has its own set of system and modeling challenges. To overcome the challenges and embrace the opportunities of MoE, we first develop a system capable of scaling MoE models efficiently to trillions of parameters. It combines multi-dimensional parallelism and heterogeneous memory technologies harmoniously with MoE to empower 8x larger models on the same hardware compared with existing work. Besides boosting system efficiency, we also present new training methods to improve MoE sample efficiency and leverage expert pruning strategy to improve inference time efficiency. By combining the efficient system and training methods, we are able to significantly scale up large multitask multilingual models for language generation which results in a great improvement in model accuracy. A model trained with 10 billion parameters on 50 languages can achieve state-of-the-art performance in Machine Translation (MT) and multilingual natural language generation tasks. The system support of efficient MoE training has been implemented and open-sourced with the DeepSpeed library.

研究动机与目标

• 实现可扩展的多任务多语言 Mixture-of-Experts 模型训练，达到十亿到万亿级参数。
• 开发提高样本效率和运行时效率的训练与系统技术。
• 展示在50种语言中实现强MT和多语言生成的实用模型（Z-code M3）。

提出的方法

• 开发 DeepSpeed MoE 系统，支持五种并行形式（数据、专家、模型、张量切片、ZeRO）以及 ZeRO-Offload，以突破 GPU 内存极限。
• 引入随机令牌选择以减轻 MoE 路由中的令牌偏置。
• 提出专家聚合（AoE），通过将检查点合并来创建更大的专家池，用于初始化和训练。
• 探索专家剪枝策略（随机和基于利用率）以实现更快的推理。
• 在单一目标下训练多任务多语言模型（MT、DAE、ELECTRA、MLM），通过对任务损失求和。
• 使用带有 MoE 层的 Transformer 编码器-解码器架构，MoE 层每隔一层放置，采用 250k 的 SentencePiece 词汇表。

实验结果

研究问题

• RQ1 如何将 MoE 架构扩展到万亿级参数以实现多任务多语言训练？
• RQ2 DeepSpeed MoE 是否能突破 GPU 内存极限，以实现更大基础模型和更多专家？
• RQ3 基于 MoE 的多任务多语言训练方案是否优于密集基线，在下游的 MT 与 NLG 任务上表现更好？
• RQ4 哪些训练技术可最大化大型 MoE 模型的样本效率和推理效率？
• RQ5 多任务目标对多语言翻译与生成质量的影响如何？

主要发现

模型	X→英文	英文→X	平均
Non-MoE multilingual model	35.16	30.36	32.76
Individual bilingual models	35.70	30.40	33.05
Z-code M3 [10B]	42.09	32.22	37.15

• DeepSpeed MoE 在跨 GPU 的吞吐量几乎线性扩展，并通过 ZeRO-Offload 支持超出 GPU 内存的模型规模。
• RTS 减少令牌位置偏差，提升 MoE 训练的收敛速度与正则化。
• AoE 通过从检查点聚合参数来初始化更大模型，从而创建更大的有效专家池。
• 专家剪枝在保持较小推理成本的同时不显著降低性能，支持随机或基于利用率的选择。
• Z-code M3 (10B 参数，64 专家) 在 MT 与多语言任务上优于密集基线和较小的 MoE 配置，展示出强大的多任务多语言能力。
• 微调的 Z-code M3 模型在下游任务如 Wikilingua 与跨语言生成上取得显著改进。