Scalable and Efficient MoE Training for Multitask Multilingual Models
Young Jin Kim∗, Ammar Ahmad Awan∗ Alexandre Muzio, Andres Felipe Cruz Salinas, Liyang Lu, Amr Hendy Samyam Rajbhandari, Yuxiong He, Hany Hassan Awadalla Microsoft, One Microsoft Way, Redmond, WA 98052, USA
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.
中文速览
让大模型更大、更省钱一直是AI领域的核心难题——现有的万亿参数稠密模型光训练就要耗费超过一年的GPU算力。混合专家（Mixture of Experts, MoE）架构天生具备"参数多但计算量不增"的特性，是突破这一瓶颈的有力候选，但如何在系统层面高效支撑万亿级MoE训练、以及如何让模型训练更稳定、推理更轻量，至今仍是未解难题。本文提出DeepSpeed MoE，将五种并行策略（数据并行、张量切片、专家并行、ZeRO冗余消除、ZeRO-Offload异构内存卸载）无缝整合，使同等硬件上可训练的MoE模型规模达到现有方案的8倍，突破GPU显存上限，支持超过3.5万亿参数的模型；同时引入更高效的token路由、专家剪枝等训练与推理方法，并设计了兼顾多任务与多语言的预训练方案。基于上述系统与方法训练的Z-code M3模型（100亿参数、覆盖50种语言）在机器翻译和多语言自然语言生成任务上达到了最优水平，证明MoE路线能以远低于稠密模型的算力代价实现更强的模型能力，相关代码已通过DeepSpeed开源发布。
原文 arXiv:2109.10465;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2109.10465v1