论文阅读：DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation

DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale
DeepSpeed-MoE:推进混合专家推理和训练，为下一代人工智能规模提供动力

前置信息：对于大模型相关工作在推理阶段和训练阶段的硬件制约条件是不相同的，在训练阶段，大模型受计算资源的限制更大，moe的稀疏结构可以在成倍扩充模型参数的同时，不过分的提高计算量
在推理阶段，样本量较小，大模型的制约因素变成了模型的体积（模型读取速度，模型参数对显存的占用），过大的模型会制约推理性能，在这个角度moe是有劣势的，以一个token为例，虽然参数很多，但是有相当一部分没激活（在推理计算时没用到）,参数效率低；以一组（多个token）为例，又可能会激活几乎全部参数，在卡、机之间的通信时延又会影响性能

随着大型密集模型的训练达到硬件资源可用性和能力的极限，与质量等效的密集模型相比，混合专家(MoE)模型显著降低了训练成本，成为最有前途的模型体系结构之一。

然而，由于模型尺寸更大，结构独特，如何提供快速的MoE模型推理仍然是一个挑战和未解决的问题，限制了其实际应用。

为了在不牺牲模型质量的情况下减少训练大型模型的计算需求，已经进行了大量的努力。为此，基于混合专家(mix -of- experts, MoE)的体系结构开辟了一条很有前途的道路，实现了与模型参数相关的亚线性计算需求，并允许在不增加训练成本的情况下提高模型质量。

MOE的局限性：
有限的应用范围：

基于MoE的模型在NLP领域的范围主要局限于编码器-解码器模型和序列到序列任务，在探索其在其他领域的应用方面所做的工作有限。

海量内存需求：
虽然MoE模型需要更少的计算来实现与密集模型相同的模型质量，但它们需要更多的参数。换句话说，与质量等效的密集模型相比，基于MoE的模型具有更低的“参数效率”。
更大的模型尺寸和更低的参数效率给训练和推理带来了挑战。
对于训练来说，模型尺寸的巨大增加需要设备内存的比例增加。
这种巨大的模型尺寸使得基于MoE的大规模模型训练具有挑战性，因为需要设备内存以及有效利用数千个gpu的设备内存所需的系统支持。
有限推理性能：
上述模型尺寸大，参数效率低，基于MoE的模型的快速推理更具挑战性。

一方面，更大的参数尺寸需要更多的gpu来适应，而多gpu推理技术并不是为基于MoE的模型而设计的。另一方面，由于推理通常受内存带宽限制，基于MoE的模型可能比密集模型大10倍，因此可能需要高10倍的可实现内存带宽才能实现与密集模型相似的推理延迟

对应方案：
针对应用范围有限的问题，将基于MoE的模型扩展到自回归NLG任务，训练成本降低。
针对参数效率低的问题，开发一种新的MoE架构来提高基于MoE模型的参数效率，我们称之为金字塔残差MoE (PR-MoE)。PR-MoE是使用剩余连接创建的密集和MoE混合模型，同时只在专家最有效的地方应用专家。PR-MoE可以将MoE模型参数大小**减少3倍，**而不改变模型质量，对计算要求的改变最小，分阶段的知识蒸馏，创建了PR-MoE的蒸馏版本，我们称之为混合学生(MoS)
（减少了模型参数）
针对多卡推理：我们开发了DeepSpeed-MoE推理系统，这是一种高度优化的MoE推理系统，可在数百个gpu上有效扩展推理工作负载，与现有的MoE推理解决方案相比，可将推理延迟和成本降低7.3倍。
2 相关工作
2.1大规模密集NLP模型
GPT3 1750亿
然而，由于在2000多个A100 gpu上的训练需要3个月，由于无法克服的计算需求，仅仅通过增加模型尺寸来获得更好的模型质量已不再可行。

2.2通过MoE架构降低训练成本
MoE的稀疏性在不增加计算成本的情况下显著改善了模型尺寸的缩放。
Switch Transformer

MoE模型也可以应用于自回归的自然语言生成任务，包括在多个下游评估任务上的初步结果。然而，与这些平行探索相比，我们的工作存在重大差异:

(1)我们的工作研究了MoE模型的训练、模型设计和推理机会，而主要关注MoE训练;

(2)我们提出了PR-MoE架构和MoS知识蒸馏，以实现更好的MoE参数效率和更好的零射击评估质量，如第3.3节所述;

(3)开发了DeepSpeed-MoE推理系统，以高效地服务于大规模MoE模型，具有高吞吐量和低延迟。虽然最近的研究如[30]和[31]讨论了FLOPs的减少，但与训练不同的是，推理延迟和成本并不仅仅依赖于计算。
2.3 MoE训练和推理系统
DeepSpeed MoE训练系统主要针对MoE模型的大规模优化训练。这项工作的主要目标是建立一个灵活的面向用户的API，用于扩展现有的训练代码，以实现高效和可扩展的MoE模型训练。它通过利用不同类型并行性的灵活组合，包括张量切片、数据并行、ZeRO[23]驱动的数据并行和专家并行，支持高达8倍的大模型大小。它被用于训练Z-code MoE，这是一个100亿个参数的多任务多语言MoE模型(transformer encoder-decoder architecture)
现有工作：
FastMoE，Fairseq-MoE，tutel

3 用于NLG的DeepSpeed-MoE:将语言模型的训练成本降低5倍

**模型设计：**350M(24层，1024个隐藏尺寸，16个注意头)，1.3B(24层，2048个隐藏尺寸，16个注意头)，6.7B(32层，4096个隐藏尺寸，32个注意头)。我们用“350M+MoE-128”表示一个以350M密集模型为基础模型，每隔一个前馈层增加128位专家的MoE模型。也就是说，350M+MoE-128和13 b +MoE-128总共有12个MoE层。

为啥都是隔一个放一个？？
top1
对于MoE模型，我们发现使用较低的学习率和较长的学习率衰减持续时间与密集的对应部分(例如，密集的1.3B与1.3B+MoE-128)相比，可以提供更好的收敛性。
两个额外的超参数:每个MoE层的专家数量，以及将MoE层损失添加到总训练损失时的系数

3.3 MoE提高NLG模型的质量
4 PR-MoE和MoS:减小模型尺寸，提高参数效率

训练成本降低5倍的情况下达到了同样的质量，而得到的模型的参数大约是相应密集模型的8倍.
如此庞大的MoE模型在训练期间需要更多的内存，并且在推理期间满足此类模型的延迟要求是具有挑战性的，因为读取模型权重所消耗的内存带宽是推理中的主要性能瓶颈。减少参数个数，提高MoE的参数效率。
基于模型，我们提出了MoE模型架构(称为PR-MoE)的创新，它在不影响模型质量的情况下将整体模型大小减少了3倍。此外，我们设计了一种新的MoE-to-MoE知识蒸馏技术，以创建一个蒸馏版本的PR-MoE，我们称之为混合学生(MoS)，进一步减小了MoE模型的大小，优化了推理时间和成本。

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

观察与直觉:
标准的MoE架构在所有MoE层中具有相同的专家数量和结构。这提醒了我们机器学习社区的一个基本问题:深度神经网络中的所有层都学习相同的表示吗?这个问题在计算机视觉(CV)中得到了很好的研究:浅层(接近输入)学习一般表示，深层(接近输出)学习更客观的特定表示。这也启发了CV中的迁移学习，冻结浅层以微调[38]。
建了两个模型，moe集中在前半部分/后半部分，结果是
Deeper layers benefit more from large number of experts.

为了提高MoE模型的泛化性能，有两种常见的方法:

(1)在增加专家数量的同时保持专家容量(即每个令牌经过的专家数量)相同;

(2)在保持专家数量不变的情况下，以略微增加的计算量(33%)为代价，将专家的能力提高一倍。但对于(1)，由于专家数量较多，对培训资源的内存需求需要增加;对于(2)，更高的容量也会使通信量增加一倍，从而显著降低训练和推理速度。
方法一消耗内存，方法二消耗通信量

容量更大准确度更高的直觉：这些额外的专家可以帮助纠正第一个专家的“表现”
(所以此处的容量应该指topk)
第一个专家需要每次都更换吗?或者我们可以固定第一个专家，只将不同的额外专家分配给不同的令牌?两种方案，(1)将容量加倍(称为Top2-MoE)，以及(2)固定一个专家并在不同专家之间改变第二个专家(称为残差- moe。特别是，对于(2)，令牌总是会传递一个密集的MLP模块和一个来自MoE模块的专家，这可以看作是残差网络的一种特殊情况。之后，我们将这两个分支的输出相加，得到最终的输出。主要的直观方法是将MoE模块中的专家视为密集MLP模块的误差校正项。这样，我们就可以实现与Top-1门控功能相同的通信量，每层使用2个专家的好处，效果相当。由于减少了通信量，我们新设计的Residual-MoE的训练速度比Top2-MoE快10%以上。
4.1.2 金字塔残余MoE架构
在后面的层利用MoE会带来更多的好处，我们的新体系结构在最后几层使用了比前几层更多的专家。这就给出了Pyramid-MoE设计，我们在图3(右)中展示了一个示例——最后两层的专家数量是前两层的2倍。同时，考虑到现象II，我们提出了残差- moe架构，其中每个令牌分别通过一个固定的MLP模块和一个选定的专家。
4.1.3系统设计
如何使用专家并行性有效地训练MoE模型。
该方法在应用于PR-MoE模型时的局限性
如何扩展现有的基于专家并行的训练系统来有效地训练PR-MoE模型。
有效训练MOE模型
batch_size要足够大，因为MoE的输入令牌数量在所有专家之间进行了分区，与没有进行此类分区的模型的其余部分相比，这减少了每个专家的令牌数量与专家数量成比例。避免每个专家令牌减少的最简单方法是使用数据并行和专家并行。
（可否理解为为了保证每个专家的token数量，batch_size=num_expertdense batch）
由于在PRMoE模型的不同阶段存在不同数量的专家，因此设计一个能够有效训练PRMoE模型的训练基础设施是非常重要的。训练基于MoE的模型的最有效方法是使专家并行性等于专家的数量，以避免减少每个专家的输入令牌。然而，由于PR-MoE中专家数量的变化，没有一个单一的专家并行度对所有MoE层都是最优的。此外，如果将专家并行性设置为模型中专家数量最少，那么对于专家数量较多的MoE层，每个GPU需要多个专家，由于每个专家的批处理大小减少，导致效率低下，并且每个GPU所需的内存增加。另一方面，如果我们将专家并行度设置为模型中专家的最大数量，那么这将导致负载平衡问题，其中一些gpu要处理的专家比其他gpu多，最终限制了训练吞吐量效率
？？？？专家并行度：expert parallelism degree啥意思，有点抽象 啊？是不是等价于在每个GPU上部署的expert数
具有多专家和多数据并行支持的DeepSpeed-MoE
灵活的多专家和多数据并行设计，允许以不同的专家和数据并行度训练模型的不同部分。
对于非专家并行性，可以使用128路数据并行性进行训练，对于MoE参数，可以使用{32、64、128}专家并行性加{4、2、1}数据并行性进行训练。请注意，每个GPU现在可以在每个MoE层训练1个专家，而不管其中的专家数量如何，这不会减少每个专家的输入令牌，也不会导致负载不平衡，也不会增加每个GPU的内存需求。
用灵活的数据并行PR结构的MOE 保持通信量的稳定。
4.1.4 PR MOE评估
对于350M和13 b的情况，我们的PR-MoE使用更少的参数，但达到了与标准moe模型相当的精度。

提出的四种MOE架构：
Standard-MoE
residalmoe
Pyramid-MoE
PR-MoE
4.2混合学生:用于更小模型尺寸和更快推理的蒸馏
有很多方法进行模型压缩，如量化和修剪，但我们目前的工作重点是通过知识蒸馏进行减层–减少模型大小和模型计算，并在学生模型中保留MoE结构。
知识蒸馏是将大模型压缩成小模型的一种成功方式。
他们只考虑小型变压器（几百个参数）和基于编码器的密集型LM模型（如BERT）。相比之下，我们重点研究的是基于稀疏MoE的自动生成LMs模型在数十亿参数规模上的KD。我们所知道的唯一其他关于MoE蒸馏的分析是由[49, 31]，他们研究将MoE蒸馏成密集模型。然而，通过这样做，提炼出来的学生模型失去了MoE提供的稀疏微调和推理优势。

fig5：KD最初有帮助，但在训练结束时开始损害准确性。
4.2.1 分阶段的KD mos
架构选择与优化目标：
减少教师模型中每个专家分支的深度，以获得相应的学生。
使用交叉熵损失，学生和教师模型都是PRMOE
通过分阶段的知识蒸馏提高学生的准确率
使用教师的PR-MoE导致学生的准确率低于从头开始训练的PR-MoE学生。由于KD通常能提高学生的泛化能力，这就提出了一个问题：为什么KD不能提高生成性语言模型上预训练MoE的准确性。
虽然KD损失最初提高了验证的准确性，但在训练结束时，它开始损害准确性。
作者的假设是：PR-MoE已经通过利用体系结构的变化降低了与标准MoE相比的容量，进一步降低模型的深度会导致学生的能力不足，使其落入欠拟合状态。因此，学生的PR-MoE可能没有足够的能力来最小化训练损失和知识蒸馏损失，最终可能以牺牲另一种损失（交叉熵损失）为代价来最小化一种损失（KD损失）
解决方案：希望逐渐减弱KD的影响，或者在训练过程的早期停止KD，在剩余的训练过程中只针对标准语言建模损失进行优化
4.2.2 MOS评估
通过减少expert的深度获得student

在减少相同的深度但没有基于MoS的KD的情况下（第2行），PR-MoE模型在一些任务上遇到了明显的准确性下降，
通过分阶段的KD（第4行），我们能够提高学生PR-MoE的性能，
在整个训练过程中执行KD（完全KD，第3行）会损害LAMBADA（低0.8分）和PIQA（低0.7分）的下游任务准确性。
总结：与原始的标准MoE模型相比，在训练期间，PR-MoE导致内存消耗下降了3倍。在推理过程中，PR-MoE和MoS一起将MoE模型的大小减少了3.7倍，同时保持了模型的准确性。

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

在推理过程中，批处理量一般较小，因此MoE模型的推理延迟主要取决于从主存储器加载模型参数所需的时间，这与传统观念相反，即较少的计算量应导致更快的推理。因此，MoE推理性能取决于两个主要因素：整体模型大小和整体可实现的内存带宽。

在推理阶段缩小模型比减少计算量的优势更大。
本节主要是系统优化：
通过创建一个多GPU MoE推理系统来最大化可实现的内存带宽，利用几十个分布式GPU的聚合内存带宽来加快推理速度。
5.1 DeepSpeed-MoE推理系统设计

MoE 推理性能是一个有趣的悖论：
最好的情况下，对于一个Token,top1 的计算量和同结构的密集模型相同，但moe结构的参数远远多于密集模型
对于一组token，所需的总参数可能和整个模型的大小一样大（因为不同的标记可能激活不同的专家），这使得实现短时延和高吞吐量成为挑战。
系统优化的三个点：

• 仔细划分模型并接受不同类型的并行性；将具有相同关键数据路径的所有令牌分组和路由在一起，以减少每个设备的数据访问并实现最大的总带宽；
• 利用并行性协调优化通信调度，以有效地分组和路由令牌；
• 优化变压器和MoE相关内核，以提高每个设备的性能
  5.2 灵活组合 Tensor-Slicing, Expert-Slicing, Data Parallelism, and Expert Parallelism
  通过使用（1 ) 专家并行性和专家参数切片，以及 (2) 非专家参数的数据并行性和张量切片to最大化可实现的聚合内存带宽，并同时提供充足的聚合内存

专家参数的专家并行和专家切片
虽然每个令牌仅在每个 MoE 层激活一个专家，但对于具有多个令牌的批量推理，所有令牌所需的聚合参数可以大到整个参数集，使得同时实现低延迟和高吞吐量具有挑战性。为了解决这个问题，我们跨设备划分专家，将分配给相同专家的所有输入令牌分组在相同的关键数据路径下，并使用专家并行性在不同设备之间并行处理具有不同关键路径的令牌组。即每个GPU只需要处理与该设备上的专家对应的一个数据组,
critical data path：关键数据路径
专家并行性具体是如何实现的？？
提出“专家切片”以利用专家参数的张量切片概念，水平/垂直划分专家参数。
跨GPU??
具体而言,专家并行可以最大限度减小每个设备需要处理的关键路径数据;专家切片进一步减小每个设备上的专家参数量;数据并行用于降低非专家参数的规模;张量切片也用于进一步减小各个维度上非专家参数的规模。

非专家参数的数据并行和张量切片 虽然专家并行减少了每个设备关键路径中专家参数的数量，但它并没有减少关键路径中的非专家参数。这导致了两个限制：（1）MoE 模型中可推断的非专家参数的最大大小受限于单个设备内存，以及（2）模型的非专家组件的执行延迟是受限于单个设备内存带宽。

我们在节点内使用张量切片来解决这些瓶颈，通过利用聚合 GPU 内存允许数千亿个非专家参数，同时还利用节点内所有 GPU 的聚合 GPU 内存带宽。虽然可以跨节点执行张量切片，但张量切片的通信开销以及计算粒度的降低通常使得节点间张量切片不可行。为了跨多个节点扩展非专家参数，我们通过创建非专家参数副本来处理跨节点的不同批次，从而使用数据并行性，这不会产生通信开销或减少计算粒度。
在node内对非moe参数做模型并行，机内通信带宽大，时延不会很高
在跨机的非moe参数上使用参数复制，减少通信量

多维并行的协同作用 所有并行策略结合后能够承载很大的模型参数量，关键在于高性能的通信集体和单设备内核。
5.3 优化的通信子系统： 更有效地分组和路由令牌
现有的torch.dis在大规模使用时通信开销很大。为了优化这一点，我们开发了一个定制的通信接口，使用微软的SCCL[52]，并取得比NCCL更好的性能。但是扩展到几十个设备上仍然很难，设计了两个新的通信优化策略，利用底层的点对点NCCL操作和定制的CUDA内核来执行必要的数据布局转换。
高性能ALL 2ALL：分层ALL 2ALL实现为一个两步过程，先进行数据布局转换，然后进行节点内ALL 2ALL，再进行第二次数据布局转换，最后进行节点间ALL 2ALL。
将通信跳数从O§减少到O(G+ p/G)，其中G是一个节点中的GPU数量，p是GPU设备的总数量。
尽管通信量增加了 2 倍，但这种分层实现允许更好地扩展小批量大小，因为在这种消息大小下的通信更受延迟限制而不是带宽限制。
并行协调通信优化
5.3和5.4都是围绕AI系统通信和kernel的优化，高效的爱因斯坦和等等，5.5是系统优化方面的评估

每个 GPU 专家数量的减少有利于数据局部性，因为每个 GPU 现在从内存中读取更少的数据。
由于专家之间的交流增加，GPU 数量的增加可能会导致性能下降。

5.5.3 系统优化+模型优化
模型优化
（1）减少最小数量对这些模型执行推理所需的 GPU

(2) 进一步提高各种 MoE 模型尺寸的延迟和吞吐量

6 期待下一代AI规模
随着最近模型规模的指数增长，我们已经达到了现代超级计算集群可以用来训练和服务大型模型的边界。由于对硬件资源的不可逾越的要求，通过简单地增加模型大小来获得更好的模型质量变得越来越困难。我们的选择是等待下一代硬件，或者使用现有硬件进行创新并提高训练和推理效率。

我们与最近的文献 [30、31] 一起展示了基于 MoE 的模型如何将大型 NLG 模型的训练成本与其质量相当的密集模型相比降低数倍，从而在当前这一代的硬件上为训练下一个规模的 AI 模型提供了可能性。然而，在这项工作之前，据我们所知，目前还没有关于如何以与密集模型相当或更好的延迟和成本为 MoE 模型（具有更多参数）提供服务的现有工作。这是一个具有挑战性的问题，阻碍了大规模 MoE 模型在现实世界中的部署。

为了对 MoE 模型进行实用和高效的推理，我们提供了新颖的 PR-MoE 模型架构和 MoS 蒸馏技术，以显着降低这些模型的内存需求。我们还提供了一个 MoE 推理框架，以前所未有的模型规模实现令人难以置信的低延迟和成本。结合这些创新，我们能够使这些 MoE 模型不仅可以提供服务，而且能够以比其质量等效的密集模型更低的延迟和成本用于推理。

总的来说，新的创新和基础设施为下一代 AI 规模的训练和推理提供了一条有前途的途径，而无需增加计算资源。从密集的 MoE 模型到稀疏的 MoE 模型的转变可以在大型模型领域开辟一条通往新方向的道路，在这个领域，用更少的资源部署更高质量的模型变得更加可能。
11:36 2023/4/28