DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale

作者/机构: Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He (Microsoft)

A1 主要贡献

本文旨在解决混合专家（Mixture-of-Experts, MoE）模型在实际应用中面临的挑战，使其成为训练和部署下一代大规模AI模型的可行且高效的方案。

核心问题:
随着密集模型（dense model）的训练逐渐达到硬件资源的极限，MoE模型因其相较于同等质量密集模型显著降低的训练成本而备受关注。然而，MoE模型也面临着自身的挑战，限制了其广泛应用：
1. 应用范围有限: MoE在自然语言处理领域的应用主要局限于编码器-解码器模型，在自回归自然语言生成（NLG）等计算成本更高的领域探索不足。
2. 巨大的内存需求: MoE模型虽然计算成本较低，但参数量远超同等质量的密集模型，导致“参数效率”低下。这不仅给训练带来了巨大的设备内存压力，也使模型难以在数千个GPU上有效扩展。
3. 推理性能受限: 巨大的模型尺寸和低参数效率使得MoE模型的快速推理极具挑战性。大参数量需要更多GPU来装载模型，而推理性能往往受限于内存带宽，MoE模型可能需要比密集模型高出数倍的内存带宽才能达到相似的延迟。

研究目标与主要贡献:
为了克服上述挑战，使MoE模型变得实用、易于获取和应用，本文提出了DeepSpeed-MoE，一个端到端的MoE训练和推理解决方案，其核心贡献包括：
1. 将MoE扩展至自回归NLG任务: 本文将MoE模型应用于类GPT-3和MT-NLG的自回归语言生成任务，并证明在达到相同模型质量的情况下，可将训练成本降低5倍。这为在现有硬件资源下训练更大规模、更高质量的模型开辟了道路。
2. 通过模型创新提高参数效率:
* Pyramid-Residual MoE (PR-MoE): 提出了一种新颖的MoE架构，通过残差连接将密集层和MoE层混合，并仅在最有效的网络层应用专家网络。PR-MoE能够在不影响模型质量和计算量的情况下，将MoE模型的参数量减少最多3倍。
* Mixture-of-Students (MoS): 基于PR-MoE，通过分阶段知识蒸馏技术创建了一个蒸馏版本MoS，可在保持相近模型质量的同时，将MoE模型尺寸进一步压缩，最多减少3.7倍。
3. 开发高性能MoE推理系统:
* DeepSpeed-MoE推理系统: 这是一个高度优化的MoE推理系统，支持在数百个GPU上高效扩展推理工作负载。与现有MoE推理解决方案相比，它能将推理延迟和成本降低多达7.3倍。
* 系统与模型协同优化: 结合系统和模型层面的优化，DeepSpeed-MoE为MoE模型提供了比同等质量密集模型快4.5倍、便宜9倍的推理性能，并能为万亿参数级别的MoE模型提供25毫秒以下的超快推理延迟。

总而言之，本文的创新和系统使MoE模型在训练和推理成本上都显著优于同等质量的密集模型，为AI领域从密集模型向稀疏MoE模型的转变提供了坚实的基础，使得用更少的资源训练和部署更高质量的模型成为可能。

A3 背景知识与相关工作

2.1 大规模密集NLP模型

模型规模的指数级增长。为了验证模型容量与参数数量之间的缩放定律【索引6，Scaling laws for neural language models，2020，arXiv】，预训练自然语言处理模型的规模在过去几年里以每年10倍的速度增长。早期的工作，如BERT【索引7，BERT: Pre-training of deep bidirectional transformers for language understanding，2019，NAACL-HLT (1)】、XLNet【索引8，Xlnet: Generalized autoregressive pretraining for language understanding，2019，NeurIPS】、RoBERTa【索引9，Roberta: A robustly optimized bert pretraining approach，2019，arXiv】、ALBERT【索引10，ALBERT: A lite bert for self-supervised learning of language representations，2019，ICLR】和GPT【索引11，Improving language understanding by generative pre-training，2018，OpenAI Blog】等，参数规模通常在数亿级别。随后，数十亿到数百亿参数的模型，如GPT-2【索引12，Language models are unsupervised multitask learners，2019，OpenAI blog】、TuringNLG【索引13，Turing-nlg: A 17-billion-parameter language model by microsoft，2020，Microsoft Blog】、Megatron-LM【索引14，Megatron-LM: Training multi-billion parameter language models using gpu model parallelism，2019，arXiv】、T5【索引15，Exploring the limits of transfer learning with a unified text-to-text transformer，2019，arXiv】等被相继推出，并在各种自然语言理解和生成任务上展现了更强的泛化性能【索引16，The lambada dataset: Word prediction requiring a broad discourse context，2016，arXiv】【索引17，Glue: A multi-task benchmark and analysis platform for natural language understanding，2018，arXiv】【索引18，Superglue: A stickier benchmark for generalpurpose language understanding systems，2019，arXiv】【索引19，A corpus and evaluation framework for deeper understanding of commonsense stories，2016，arXiv】【索引20，Semantic parsing on freebase from question-answer pairs，2013，EMNLP】【索引21，Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension，2017，arXiv】。GPT-3【索引22，Language models are few-shot learners，2020，arXiv】进一步将参数上限推至1750亿，并展示了通过零样本/少样本学习，其性能可以与经过微调的小规模模型相媲美甚至更优。最近，借助DeepSpeed【索引23，Zero: Memory optimizations toward training trillion parameter models，2020，SC20】和Megatron-LM【索引14，Megatron-LM: Training multi-billion parameter language models using gpu model parallelism，2019，arXiv】的软件支持，超大规模的Megatron-Turing NLG 530B模型【索引2，Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, the World’s Largest and Most Powerful Generative Language Model，2021，NVIDIA Developer Blog】达到了5300亿参数，并在零样本/少样本学习方面取得了新的SOTA。然而，该模型的训练在2000多块A100 GPU上耗时3个月，这表明仅仅通过增加模型规模来提升质量，因其巨大的计算需求而变得不再可行。

2.2 通过MoE架构降低训练成本

MoE作为降低训练成本的有效途径。一个有前景的降低训练成本的方法是使用混合专家（MoE）【索引24，Mixture of experts: a literature survey，2014，Artificial Intelligence Review】。在【索引3，Outrageously large neural networks: The sparsely-gated mixtureof-experts layer，2017，arXiv】中，作者通过在堆叠的LSTM层【索引25，Long short-term memory，1997，Neural computation】之间应用卷积MoE层，将LSTM模型扩展到1270亿参数，用于语言建模和机器翻译。MoE的稀疏性显著提升了模型规模的扩展能力，而无需增加计算成本。GShard【索引4，Gshard: Scaling giant models with conditional computation and automatic sharding，2020，arXiv】利用MoE训练了一个6000亿参数的Transformer模型【索引26，Attention is all you need，2017，NeurIPS】用于多语言翻译，并表明这个6000亿参数MoE模型的训练成本甚至低于一个1000亿参数的密集模型。Switch Transformer【索引5，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2021，arXiv】基于T5模型延续了这一思路，将模型扩展到1.6万亿参数。为了达到相同的准确率，【索引5，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2021，arXiv】展示了MoE模型相比大型密集模型有2.5倍的训练速度提升。此后，更多近期的工作【索引27，M6-10t: A sharing-delinking paradigm for efficient multitrillion parameter pretraining，2021，arXiv】【索引28，Scalable and efficient moe training for multitask multilingual models，2021，arXiv】【索引29，Taming sparsely activated transformer with stochastic experts，2021，arXiv】也展示了MoE的稀疏性优势。

与同期工作的对比。一些近期及并行的工作【索引30，Glam: Efficient scaling of language models with mixture-of-experts，2021，arXiv】【索引31，Efficient large scale language modeling with mixtures of experts，2021，arXiv】表明，MoE模型也可以应用于自回归自然语言生成任务，并给出了多个下游评估任务的初步结果。然而，我们的工作与这些并行探索有主要区别：(1) 我们的工作研究了MoE模型的训练、模型设计和推理机会，而【索引30，Glam: Efficient scaling of language models with mixture-of-experts，2021，arXiv】【索引31，Efficient large scale language modeling with mixtures of experts，2021，arXiv】主要关注MoE训练；(2) 我们提出了PR-MoE架构和MoS知识蒸馏，以实现更好的MoE参数效率和相当/更优的零样本评估质量，如第3.3节所述；(3) 我们开发了DeepSpeed-MoE推理系统，以高吞吐量和低延迟高效地服务大规模MoE模型。虽然近期研究如【索引30，Glam: Efficient scaling of language models with mixture-of-experts，2021，arXiv】和【索引31，Efficient large scale language modeling with mixtures of experts，2021，arXiv】讨论了FLOPs的减少，但需要指出的是，与训练不同，推理延迟和成本不仅取决于计算量。高效的推理还依赖于模型大小、内存带宽以及系统从内存中高效读取数据的能力。DeepSpeed-MoE推理系统是一个端到端的系统，它结合了模型架构创新（第4节）和大量的优化与技术（第5节），为吞吐量提供了超低延迟和超线性加速。

2.3 MoE训练与推理系统

现有MoE系统的概览。鉴于MoE模型近来被世界各地的许多研究人员所采纳，支持MoE模型训练的开源系统和框架正在被开发或扩展。据我们所知，目前还没有专门为推理设计和优化的MoE系统。最近已经出现了一些MoE训练系统，下面我们讨论其中的几个。

DeepSpeed MoE训练系统。【索引32，Scalable and efficient moe training for multitask multilingual models，2021，CoRR】主要致力于大规模MoE模型的优化训练。这项工作的主要目标是建立一个灵活的、面向用户的API，用于扩展现有训练代码，以实现高效和可扩展的MoE模型训练。它通过灵活组合不同类型的并行方式，包括张量切片、数据并行、由ZeRO【索引23，Zero: Memory optimizations toward training trillion parameter models，2020，SC20】驱动的数据并行以及专家并行，支持高达8倍的模型规模。该系统被用于训练Z-code MoE，一个100亿参数的多任务多语言MoE模型（Transformer编码器-解码器架构），在50种语言上，其训练时间比同等质量的密集模型减少了5倍。

FastMoE。【索引33，Fastmoe: A fast mixture-of-expert training system，2021，CoRR】是一个研究软件，旨在展示MoE模型如何在数据并行和专家（模型）并行下进行训练。它并未完全支持各种并行维度的组合。FastMoE的示例包括Transformer-XL和Megatron-LM，但尚未提供关于大规模端到端训练的结果。

Fairseq-MoE。【索引31，Efficient large scale language modeling with mixtures of experts，2021，arXiv】提供了一个MoE API以及一个用于通用语言模型的训练流程。Fairseq系统后来由Tutel【索引34，Tutel: An efficient mixture-of-experts implementation for large dnn model training，2021，Microsoft】进一步优化，性能比Fairseq提升了高达40%。

A2 方法细节

4. PR-MoE与MoS：缩小模型尺寸并提升参数效率

虽然MoE模型在NLG任务上实现了5倍的训练成本降低，但其参数量大约是同等质量密集模型的8倍（例如，6.7B密集模型 vs. 52B的1.3B+MoE-128模型）。这种巨大的模型尺寸在训练时需要更多内存，在推理时也因内存带宽瓶颈而难以满足延迟要求。为了解决这个问题，我们提出了模型架构创新（PR-MoE）和知识蒸馏技术（MoS），旨在减少参数数量并提高MoE模型的参数效率。

4.1 PR-MoE：用于更小模型尺寸和更快推理的金字塔-残差-MoE

4.1.1 两个观察与直觉

现象一：不同网络层对MoE的需求不同。首先，标准的MoE架构在所有MoE层中都使用相同数量和结构的专家。这引发了一个根本性问题：深度神经网络中的所有层是否学习相同的表示？在计算机视觉领域，这个问题的答案是：浅层网络学习通用表示，而深层网络学习更具任务特异性的表示【索引37，Visualizing and understanding convolutional networks，2014，ECCV】。这启发了计算机视觉中的迁移学习，即在微调时冻结浅层网络【索引38，How transferable are features in deep neural networks?，2014，arXiv】。然而，这一现象在自然语言处理，特别是MoE架构中，尚未得到充分探讨。为了研究这个问题，我们基于350M+MoE模型比较了两种不同的“半MoE”架构：a) 将MoE层置于模型的前半部分，后半部分为密集层（称为First-Half-MoE）；b) 将MoE层置于后半部分，前半部分为密集层（称为Second-Half-MoE）。结果如图2（左）所示，Second-Half-MoE的性能显著优于其对应方案。这证实了并非所有MoE层都学习同等级别的表示，更深的网络层从大量的专家中获益更多。我们称此为现象一。

图2：基于350M+MoE模型的First-Half-MoE/Second-Half-MoE（左）和Top2-MoE/Residual-MoE（右）的验证损失。
图2：基于350M+MoE模型的First-Half-MoE/Second-Half-MoE（左）和Top2-MoE/Residual-MoE（右）的验证损失。

现象二：残差结构可高效替代增加专家容量。其次，提升MoE模型泛化性能通常有两种方法：（1）增加专家数量，同时保持专家容量（每个token通过的专家数）不变；（2）将专家容量加倍，这会带来约33%的额外计算量，同时保持专家数量不变。方法（1）增加了内存需求；方法（2）使通信量加倍，显著拖慢训练和推理。我们的问题是：是否存在一种方法，既能提升泛化性能，又能保持训练/推理效率？我们直觉地认为，增加专家容量之所以有效，是因为额外的专家可以帮助纠正第一个专家的“表示”。但是，第一个专家是否每次都需要改变？或者我们能否固定第一个专家，只为不同的token分配不同的额外专家？为了验证这一点，我们进行了比较：(1) 将容量加倍（称为Top2-MoE）；(2) 固定一个专家（一个密集的MLP模块），并从MoE模块中选择第二个变化的专家（称为Residual-MoE）。在Residual-MoE中，一个token总是通过一个固定的密集MLP模块和一个从MoE模块中选择的专家，这可以看作是残差网络的一个特例。我们将这两个分支的输出相加，将MoE专家的输出视为对密集MLP模块的误差修正项。这样，我们就能以与Top-1门控相同的通信量，获得使用2个专家的好处。我们在350M+MoE模型（32个专家）上进行了比较，验证曲线如图2（右）所示。我们发现Top2-MoE和Residual-MoE的泛化性能相当，但我们新设计的Residual-MoE由于通信量减少，训练速度比Top2-MoE快了10%以上。我们将此现象称为现象二。

图3：标准MoE（左）和PR-MoE（右）的图示。
图3：标准MoE（左）和PR-MoE（右）的图示。

4.1.2 金字塔残差MoE架构

PR-MoE架构的设计。基于上述观察，我们提出了新颖的MoE架构。现象一表明在模型的后期层使用MoE能带来更多好处，因此我们的新架构在最后几层使用了比前面层更多的专家，这构成了Pyramid-MoE设计，如图3（右）所示，最后两层拥有2倍于前面层的专家数量。同时，考虑到现象二，我们提出了Residual-MoE架构，其中每个token分别通过一个固定的MLP模块和一个选定的专家，如图3（右）所示（橙色块代表固定的MLP）。通过将Pyramid-MoE和Residual-MoE结合，我们得到了Pyramid-Residual-MoE模型（简称PR-MoE），其中所有标准的MoE层都被新的PR-MoE层所取代。图3展示了标准MoE和PR-MoE架构的对比。

4.1.3 系统设计

高效训练标准MoE模型。高效训练MoE模型需要在MoE模块中为每个专家提供足够大的批次大小以实现良好的计算效率。这是一个挑战，因为MoE层的输入token被分配到所有专家中，导致每个专家的token数量与专家数量成反比地减少。最简单的解决方法是结合使用数据并行和专家并行【索引32，Scalable and efficient moe training for multitask multilingual models，2021，CoRR】，且专家并行度等于专家数量。这增加了每个MoE副本批次中的总token数，这些token将在专家间分配，从而使得每个专家的token数与模型其余部分相比没有减少。

PR-MoE面临的挑战。为PR-MoE模型设计高效的训练基础设施并非易事，因为模型不同阶段的专家数量不同。如上所述，训练MoE模型最有效的方法是使专家并行度等于专家数量，以避免减少每个专家的输入token。然而，由于PR-MoE中专家数量的变化，没有一个单一的专家并行度对所有MoE层都是最优的。如果将专家并行度设置为模型中最小的专家数，那么对于拥有更多专家的MoE层，每个GPU将需要处理多个专家，导致每个专家的批次大小减小，效率降低，同时每个GPU所需的内存也会增加。反之，如果将专家并行度设置为最大的专家数，则会导致负载不均衡问题，一些GPU需要处理比其他GPU更多的专家，最终限制训练吞吐效率。

DeepSpeed-MoE支持多专家和多数据并行。为应对这些挑战，我们在DeepSpeed-MoE之上开发并实现了一个灵活的多专家和多数据并行设计，允许模型不同部分使用不同的专家和数据并行度进行训练。例如，一个在128个GPU上运行的PR-MoE模型，在不同MoE层分别有32、64和128个专家，可以对非专家部分使用128路数据并行，对MoE参数使用{32, 64, 128}路专家并行加上{4, 2, 1}路数据并行。请注意，现在每个GPU在每个MoE层都可以精确地训练1个专家，无论该层有多少专家，从而避免了每个专家输入token的减少、负载不均衡或每个GPU内存需求的增加。通过这种灵活的扩展，DeepSpeed-MoE【索引32，Scalable and efficient moe training for multitask multilingual models，2021，CoRR】可以高效地训练PR-MoE模型，以及未来任何可能在模型不同阶段需要不同专家数量的MoE变体，而不会牺牲训练效率或内存需求。

4.2 混合专家-学生模型：通过蒸馏获得更小模型和更快推理

通过知识蒸馏进一步优化。模型压缩和知识蒸馏为进一步提升推理性能提供了机会。尽管模型压缩方法多样，如量化【索引39，Q-BERT: hessian based ultra low precision quantization of BERT，2020，AAAI】【索引40，HAWQ: hessian aware quantization of neural networks with mixed-precision，2019，ICCV】【索引41，HAWQ-V2: hessian aware trace-weighted quantization of neural networks，2020，NeurIPS】和剪枝【索引42，Generating long sequences with sparse transformers，2019，CoRR】【索引43，Block pruning for faster transformers，2021，EMNLP】，我们目前的工作重点是通过知识蒸馏（KD）【索引44，Distilling the knowledge in a neural network，2015，CoRR】进行层数削减，从而减少模型大小和计算量，同时保留学生模型中的MoE结构。

MoE模型蒸馏的独特性。知识蒸馏已被证明是成功的模型压缩方法，能将大模型压缩成参数和计算量都少得多但性能依然有竞争力的小模型。已有工作将KD应用于特定任务，将预训练的大型语言模型蒸馏为小模型【索引45，Distilbert, a distilled version of BERT: smaller, faster, cheaper and lighter，2019，CoRR】【索引46，Patient knowledge distillation for BERT model compression，2019，EMNLP-IJCNLP】【索引47，Minilm: Deep self-attention distillation for task-agnostic compression of pre-trained transformers，2020，NeurIPS】【索引48，Mobilebert: a compact task-agnostic BERT for resource-limited devices，2020，ACL】。然而，这些工作仅考虑了小型Transformer（几百个参数）和基于密集编码器的语言模型（如BERT）。相比之下，我们专注于研究数十亿参数规模的、基于稀疏MoE的自回归语言模型的知识蒸馏。我们所知的唯一其他关于MoE蒸馏的分析来自【索引49，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2021，CoRR】和【索引31，Efficient large scale language modeling with mixtures of experts，2021，arXiv】，他们研究了将MoE蒸馏为密集模型。然而，这样做会使蒸馏出的学生模型失去MoE带来的稀疏微调和推理优势。我们的研究表明，通过知识蒸馏预训练的较小MoE模型，可以在许多下游任务的零样本评估上达到相似性能，从而得到在推理时更轻、更快的模型。

图5：从零开始训练（无蒸馏）与在整个预训练过程中进行知识蒸馏的验证损失曲线对比。图中，学生PR-MoE模型为21层。知识蒸馏在初期有帮助，但在训练后期开始损害准确率。
图5：从零开始训练（无蒸馏）与在整个预训练过程中进行知识蒸馏的验证损失曲线对比。图中，学生PR-MoE模型为21层。知识蒸馏在初期有帮助，但在训练后期开始损害准确率。

4.2.1 通过分阶段知识蒸馏实现“混合专家-学生”模型

架构选择与优化目标。为了将知识蒸馏应用于MoE，我们首先训练一个教师MoE模型。我们通过减少教师模型中每个专家分支的深度来获得相应的学生模型。这样，最终的学生模型与教师MoE具有相同的稀疏门控架构，只是每个专家分支的深度较小。因此，我们称之为混合专家-学生（Mixture-of-Students, MoS）模型。由于MoE结构通过稀疏训练和推理带来了显著优势，我们经过任务无关蒸馏得到的MoS模型继承了这些好处，同时保留了相对于同等质量密集模型的推理优势。创建MoS后，我们强制MoS在训练数据集上模仿教师MoE的输出。我们采用KD损失的一般形式【索引50，Kl-divergence regularized deep neural network adaptation for improved large vocabulary speech recognition，2013，ICASSP】为：

公式1：知识蒸馏损失函数。
公式1：知识蒸馏损失函数。

该损失是模型预测与给定硬标签之间的交叉熵损失，以及模型预测与教师模型软标签之间的KL散度损失的加权和。此外，鉴于PR-MoE的出色性能，我们将PR-MoE与KD结合（PR-MoS），以进一步减小MoE模型的大小。换言之，我们选择教师模型和学生模型均为PR-MoE。

通过分阶段知识蒸馏提升学生模型准确率。在蒸馏MoE模型时，一个有趣的发现是，使用教师PR-MoE进行蒸馏得到的学生模型准确率，低于从零开始训练的PR-MoE学生模型（见表5第2行和第4行）。由于KD通常能提高学生模型的泛化能力，这就引出了一个问题：为什么在预训练生成式语言模型的MoE上，KD没有提高准确率？由于之前没有关于蒸馏MoE的实验结果报道，我们深入研究了结果。图5比较了从零开始训练的PR-MoE和使用其教师模型进行知识蒸馏的PR-MoE的验证损失。我们发现，虽然KD损失在初期能改善验证准确率，但在训练后期（例如，400K步之后）开始损害准确率。我们推测，这是因为PR-MoE本身已经通过架构改变（例如，减少底层专家的数量）降低了模型容量，进一步减小模型深度导致学生模型容量不足，陷入欠拟合状态。因此，学生PR-MoE可能没有足够的能力同时最小化训练损失和知识蒸馏损失，最终可能以牺牲交叉熵损失为代价来最小化KD损失。这一假设表明，我们应该在训练过程中逐渐减弱KD的影响，或者提前停止KD，在剩余的训练中仅针对标准语言模型损失进行优化。

5. DeepSpeed-MoE推理：以前所未有的规模和速度服务MoE模型

MoE推理的挑战与目标。优化推理延迟和成本对于MoE模型的实际应用至关重要。在推理过程中，批次大小通常较小，因此MoE模型的推理延迟主要取决于从主内存加载模型参数所需的时间，这与传统观念中计算量越少推理越快的看法相反。因此，MoE推理性能取决于两个主要因素：整体模型大小和可实现的内存总带宽。在前一节中，我们提出了PR-MoE和MoS来减小MoE模型的大小，同时保持模型准确性。本节将介绍我们的系统优化解决方案，通过创建一个多GPU MoE推理系统，利用数十个分布式GPU的聚合内存带宽来加速推理。综合来看，DeepSpeed-MoE（简称DS-MoE）提供了前所未有的规模和效率来服务大规模MoE模型，与基线MoE系统相比，延迟降低了7.3倍，成本更低；与同等质量的密集模型相比，MoE推理速度提高了4.5倍，成本降低了9倍。DS-MoE是【索引51，Deepspeed inference: Enabling efficient inference of transformer models at unprecedented scale，2022，arXiv】中介绍的更广泛的DeepSpeed-inference工作的一部分。

5.1 DeepSpeed-MoE推理系统设计

MoE推理性能的悖论。MoE推理性能存在一个有趣的悖论：
- 从最佳情况看：MoE模型（采用top-1门控）的每个输入token在每个MoE层只激活一个专家，这导致其关键数据路径等同于基础密集模型的大小，比实际模型大小小数个数量级。例如，在对一个1.3B+MoE-128模型进行推理时，每个输入token只需要13亿参数，尽管总参数量是520亿。
- 从最坏情况看：处理一批token（例如一个句子或一段文本）所需的总参数量可能与完整模型大小相当（因为不同的token可能激活不同的专家），即前述例子中的全部520亿参数，这使得实现短延迟和高吞吐量变得极具挑战性。

DeepSpeed-MoE的设计目标与策略。DeepSpeed-MoE推理系统的设计目标是引导性能趋向于最佳情况。这是通过三组精心协调的优化实现的：
- 模型分区与并行：仔细划分模型，并采用不同类型的并行策略；将具有相同关键数据路径的所有token进行分组和路由，以减少每个设备的数据访问量，并实现最大的聚合带宽。
- 通信调度优化：通过并行协调优化通信调度，以有效地对token进行分组和路由。
- 核函数优化：优化Transformer和MoE相关的核函数，以提高单设备性能。

我们将在接下来的三节中深入探讨这些优化。

5.2 张量切片、专家切片、数据并行和专家并行的灵活组合

多维并行策略。为了在MoE模型上实现前所未有的低延迟和高吞吐量，我们的推理系统设计旨在最小化每个设备上的关键数据路径，最大化可实现的聚合内存带宽，并同时提供充足的聚合内存以支持海量模型。这是通过以下方式实现的：(1) 对专家参数使用专家并行【索引32，Scalable and efficient moe training for multitask multilingual models，2021，CoRR】和切片；(2) 对非专家参数使用数据并行和张量切片。图7展示了一个单一的MoE Transformer层，其中包含专家（如MLP）和非专家参数（如Attention），以及我们如何使用并行策略组合来处理每个组件。

图7：DS-MoE设计，它为模型的不同分区（专家和非专家）采用了复杂的多维并行。
图7：DS-MoE设计，它为模型的不同分区（专家和非专家）采用了复杂的多维并行。

专家参数的专家并行与专家切片。正如MoE性能悖论所述，虽然每个token只激活一个专家，但在多token的批处理推理中，所需总参数可能与整个专家集合相当。为了解决这个问题，我们将专家划分到不同设备上，将分配给相同专家的所有输入token分组到相同的关键数据路径下，并利用专家并行在不同设备上并行处理具有不同关键路径的token组。例如，对于1.3B+MoE-128模型，当专家并行度为128时，每个GPU只处理对应于其上专家的一个token组。这使得每个设备上的序列化路径仅为13亿参数，比其同等质量的6.7B参数密集模型小5倍。因此，在不考虑通信开销的情况下，理论上基于MoE的模型使用专家并行可以比其同等质量的密集模型快5倍。此外，我们提出了“专家切片”，将张量切片的概念应用于专家内部的参数，即将专家参数在多个GPU上进行水平/垂直切分。这个额外的并行维度对于需要扩展到比专家数量更多的设备的延迟敏感场景非常有帮助。

非专家参数的数据并行与张量切片。专家并行减少了关键路径中专家参数的数量，但并未减少非专家参数。这导致两个限制：(1) MoE模型中可推理的最大非专家参数受限于单个设备的内存；(2) 模型中非专家组件的执行延迟受限于单个设备的内存带宽。我们使用节点内的张量切片来解决这些瓶颈，通过利用聚合的GPU内存，支持数百亿的非专家参数，同时也利用了节点内所有GPU的聚合内存带宽。虽然可以跨节点进行张量切片，但其通信开销和计算粒度的减小通常使其不可行。为了跨多个节点扩展非专家参数，我们使用数据并行，通过创建非专家参数的副本在不同节点上处理不同的批次，这种方式不会产生通信开销或计算粒度的减小。

多维并行的协同作用。通过将专家并行和专家切片与张量切片和数据并行相结合，DS-MoE推理系统能够将一个数万亿参数的MoE模型（拥有万亿级专家参数和千亿级非专家参数）扩展到数十甚至数百个跨节点的设备上。这些设备的聚合带宽和每个设备上最小化的关键数据路径为实现前所未有的低延迟和高吞吐量推理创造了机会。然而，要实现这一点，还需要高性能的通信集合操作和单设备核函数，我们将在下文讨论。

5.3 优化的通信子系统：更高效地分组和路由Token

通信开销的挑战与优化。专家并行需要在所有专家并行设备之间进行all-to-all通信。DS-MoE默认通过“torch.distributed”接口使用NCCL进行通信，但我们观察到在大规模使用时存在显著开销。为了优化这一点，我们开发了一个自定义通信接口来使用Microsoft SCCL【索引52，SCCL: Synthesizing Optimal Collective Algorithms，2020，CoRR】，并取得了比NCCL更好的性能。尽管有插件式优化，但将专家并行扩展到许多设备仍然很困难，因为延迟会随着设备数量的增加而线性增长。为了解决这个关键的扩展挑战，我们设计了两种新的通信优化策略，它们利用底层的点对点NCCL操作和自定义CUDA核函数来执行必要的数据布局转换。

分层All-to-all。分层树状算法常用于all-reduce、broadcast等通信集合操作，以减少通信跳数。我们实现了一个分层的all-to-all，它分为两步：首先是数据布局转换，然后是节点内all-to-all，接着是第二次数据布局转换，最后是节点间all-to-all。这将通信跳数从 $O(p)$ 减少到 $O(G + p/G)$，其中 $G$ 是一个节点中的GPU数量，$p$ 是GPU设备总数。图8展示了该实现的设计概览。尽管通信量增加了2倍，但这种分层实现对于小批量大小的扩展性更好，因为在这种消息大小下，通信更多地受延迟而非带宽的限制。

提议的分层AlltoAll设计

图8：提议的分层all-to-all设计示意图。
图8：提议的分层all-to-all设计示意图。

并行协调的通信优化。在单个模型中结合专家并行、张量切片和数据并行，在有效处理通信方面并非易事。张量切片将单个算子分割到多个GPU上，并需要在它们之间进行all-reduce；而专家并行将专家算子放置在不同GPU上而不分割它们，并需要在它们之间进行all-to-all。一种朴素的方法是将每种并行视为一个黑盒，独立地执行所需的通信，但这会导致次优性能。张量切片中的all-reduce操作在相关设备间复制了数据。当执行张量并行算子后跟专家并行算子时，这种复制允许我们为all-to-all算子创建一个优化的通信调度，该调度不需要在所有专家并行进程之间进行通信。相反，all-to-all只需在共享相同张量切片秩的设备子集内进行，因为跨张量并行秩的数据是复制的（图9）。结果，all-to-all的延迟被限制在 $O(p/L)$ 而不是 $O(p)$，其中 $L$ 是张量切片并行度，$p$ 是GPU设备总数。类似地，当执行专家并行算子后跟张量切片算子时，最终的all-to-all可以以同样的方式完成，但这次之后需要一个tensor并行秩之间的all-gather操作来复制张量切片所需的数据（图9）。这将延迟开销从 $O(p)$ 减少到 $O(p/L) + O(L)$。这种减少的延迟开销允许更好地扩展到大量设备。例如，当扩展到128个GPU，使用8路张量切片和128路专家并行时，由于8路张量切片，这种方法将all-to-all的延迟开销从 $(128C_1 + C_2)$ 减少到 $(16C_1 + C_2)$，其中 $C_1$ 和 $C_2$ 是由点对点延迟、消息大小和带宽决定的常数。

图9：并行协调通信示意图。
图9：并行协调通信示意图。

5.4 高度优化的Transformer和MoE相关核函数

核函数优化。DS-MoE推理系统包含高度优化的多GPU Transformer核函数以及高度优化的MoE相关核函数。我们使用DeepSpeed推理核函数来最大化非专家Transformer层的带宽利用率【索引53，DeepSpeed: Accelerating large-scale model inference and training via system optimizations and compression，2021，Microsoft Research Blog】。本文重点关注用于执行门控和不同数据布局转换的MoE相关算子，这些在文献中通常使用稀疏-密集einsum实现。在高层次上，我们通过将这些算子实现为显式的数据布局转换，而不是高度稀疏-密集的einsum，来优化它们，从而将计算复杂度从立方级降低到平方级。我们还将大部分这些算子融合成一个单一的核函数。

MoE计算的优化细节。具体来说，MoE相关的计算包括三个主要部分：
1. 门控函数：确定token到专家的分配，结果表示为一个稀疏张量（一个one-hot向量）。
2. 稀疏einsum算子：在one-hot张量和所有token之间进行，根据分配的专家ID对token进行排序。
3. 最终einsum：在MoE计算结束时，对token进行缩放并重新排序回其原始顺序。

从稀疏到稠密的优化。门控函数中的稀疏张量表示和稀疏einsum算子引入了显著的延迟开销。首先，门控函数包含创建token掩码、选择top-k专家和执行累积和（cumsum）等众多操作，这些操作不仅因稀疏张量表示而浪费，而且由于多次核函数调用而非常缓慢。其次，稀疏einsum的复杂度为 $S \times E \times M \times c_e$，其中S是token总数，E是专家数，M是模型隐藏维度，$c_e$是专家容量。在这个方程中，每个token的E个操作中有 $(E-1)$ 个是与零的乘法和加法。这是因为门控操作的泛化导致了对多个掩码矩阵或one-hot向量的einsum，产生了大量不必要的零计算。我们使用密集表示和核函数融合来优化这些算子。

具体实现。首先，我们将门控函数融合成一个单一的核函数，并使用一个密集的token-专家映射表来表示从token到专家的分配，大大减少了核函数启动开销以及稀疏表示带来的内存和计算开销。具体来说，门控核函数包括top-k、cumsum和scatter操作。我们使用Blelloch扫描算法在GPU上并行化cumsum。最后，我们使用映射表和token ID来将正确的token路由到MoE专家。其次，为了优化剩下的两个稀疏einsum，我们使用上述映射表将它们实现为数据布局转换，首先根据专家ID排序，然后再恢复原始顺序，而无需任何稀疏einsum。这将这些操作的复杂度从 $S \times E \times M \times c_e$ 降低到 $S \times M \times c_e$。结合数据转换，我们使用相应的门控logits（在概率域）来更新专家输出。总的来说，这些优化使得MoE核函数相关的延迟减少了6倍以上。

A4 实验

实验环境

表1：不同密集和MoE NLG模型的超参数。
表1：不同密集和MoE NLG模型的超参数。

实验结果

MoE在NLG任务中的有效性

图1：不同模型大小的密集和MoE NLG模型的token级验证损失曲线。
图1：不同模型大小的密集和MoE NLG模型的token级验证损失曲线。

表2：不同密集和MoE NLG模型的零样本评估结果（后六列）。所有零样本评估结果均使用准确率指标。PR-MoE NLG和Mixture-of-Students NLG的结果详情见第4节。最后6行为来自Google [30]和Meta (Facebook) [31]的相关工作。
表2：不同密集和MoE NLG模型的零样本评估结果（后六列）。所有零样本评估结果均使用准确率指标。PR-MoE NLG和Mixture-of-Students NLG的结果详情见第4节。最后6行为来自Google [30]和Meta (Facebook) [31]的相关工作。

表3：在128个A100 GPU上，能够达到相同模型质量的MoE模型与密集模型的训练吞吐量对比。
表3：在128个A100 GPU上，能够达到相同模型质量的MoE模型与密集模型的训练吞吐量对比。

PR-MoE和MoS的模型压缩效果

表4：标准MoE与PR-MoE的零样本评估对比（后六列）。
表4：标准MoE与PR-MoE的零样本评估对比（后六列）。

图4：基于350M+MoE的不同MoE模型的验证曲线。
图4：基于350M+MoE的不同MoE模型的验证曲线。

图6：在350M+PR-MoE和1.3B+PR-MoE上，使用分阶段知识蒸馏训练的学生PR-MoE获得了与教师PR-MoE几乎相同的验证损失。
图6：在350M+PR-MoE和1.3B+PR-MoE上，使用分阶段知识蒸馏训练的学生PR-MoE获得了与教师PR-MoE几乎相同的验证损失。

表5：PR-MoE与PR-MoE + MoS的零样本评估对比（后六列）。
表5：PR-MoE与PR-MoE + MoS的零样本评估对比（后六列）。

DeepSpeed-MoE推理系统性能

图10：DeepSpeed-MoE推理系统（优化后）相较于PyTorch（基线）在8到64个GPU上对520亿参数、128个专家的MoE模型的延迟和吞吐量改进。
图10：DeepSpeed-MoE推理系统（优化后）相较于PyTorch（基线）在8到64个GPU上对520亿参数、128个专家的MoE模型的延迟和吞吐量改进。

图11：DeepSpeed-MoE（优化后）相较于PyTorch（基线）在不同MoE模型大小（1070亿至2万亿参数）下的延迟和吞吐量改进。此处显示的吞吐量为单GPU吞吐量。
图11：DeepSpeed-MoE（优化后）相较于PyTorch（基线）在不同MoE模型大小（1070亿至2万亿参数）下的延迟和吞吐量改进。此处显示的吞吐量为单GPU吞吐量。

图12：使用PR-MoE+MoS时，MoE推理所需资源减少2倍。图标题表示标准MoE架构下的模型大小。
图12：使用PR-MoE+MoS时，MoE推理所需资源减少2倍。图标题表示标准MoE架构下的模型大小。

图13：在不同GPU数量和模型大小下，标准MoE与PR-MoE和PR-MoE+MoS压缩的推理延迟对比。
图13：在不同GPU数量和模型大小下，标准MoE与PR-MoE和PR-MoE+MoS压缩的推理延迟对比。

图14：MoE模型与同等质量的6.7B参数密集模型的推理延迟对比。
图14：MoE模型与同等质量的6.7B参数密集模型的推理延迟对比。

图15：MoE模型与同等质量的175B参数密集模型的实测推理延迟对比。
图15：MoE模型与同等质量的175B参数密集模型的实测推理延迟对比。

A5 结论

随着模型规模的指数级增长，我们正面临现代超级计算集群在训练和部署大型模型方面的能力边界。由于对硬件资源的巨大需求，仅仅通过增加模型规模来提升模型质量变得越来越困难。

本文与近期其他研究【索引30，Glam: Efficient scaling of language models with mixture-of-experts，2021，arXiv】【索引31，Efficient large scale language modeling with mixtures of experts，2021，arXiv】共同证明，基于MoE的模型能够将大型NLG模型的训练成本降低数倍，为在当前硬件上训练下一代AI模型提供了可能性。然而，在此工作之前，如何以与密集模型相当甚至更优的延迟和成本来服务参数量大得多的MoE模型，仍然是一个悬而未决的挑战，阻碍了大规模MoE模型的实际部署。

为了实现MoE模型的实用和高效推理，我们提出了新颖的PR-MoE模型架构和MoS蒸馏技术，显著减少了这些模型的内存需求。同时，我们开发了一个MoE推理框架，在空前的模型规模下实现了极低的延迟和成本。通过结合这些创新，我们不仅使MoE模型变得可以服务，而且其推理延迟和成本均低于同等质量的密集模型。

总而言之，本文的创新和基础设施为在不增加计算资源的情况下，训练和推理下一代AI规模的模型提供了一条充满希望的道路。从密集模型向稀疏MoE模型的转变，可以为大型模型领域开辟新的方向，使得用更少的资源部署更高质量的模型变得更加普及。

引用文献汇总