MegaBlocks: Efficient Sparse Training with Mixture-of-Experts
49:13 Share

从周末早午餐到晚宴，在这个假日季，凭借Finish Ultimate的力量，成为终极主人。Finish Ultimate应对最严苛条件下的顽固烧焦污渍，让您始终准备好给您的客人留下深刻印象。今天在Target以金色包装寻找它。与Rob一起阅读关于人工智能的论文，让您了解最新的研究成果。本篇阅读由Mars Race提供，在红色星球上获得赞誉。可在Android和iOS上使用。

MegaBlocks。高效的混合专家稀疏训练。2022年由Trevor Gale、D. Narayanan、C. Young和M. Zaharia撰写。Trevor Gale 1、Deepak Narayanan 2、Cliff Young 3、Matej Zaharia 1。A. B. S. T. R. A. C. T. 我们提出了MegaBlocks，这是一个在GPU上进行高效混合专家（MoE）训练的系统。

我们的系统源于当前框架的局限性，这些框架限制了MoE层中的动态路由，以满足现有软件和硬件的约束。这些公式在模型质量和硬件效率之间形成了权衡，因为用户必须在丢弃计算中的标记或在填充上浪费计算和内存之间做出选择。

为了解决这些限制，我们根据块稀疏运算重新制定了MoE计算，并开发了新的块稀疏GPU内核，可以有效地处理MoE中存在的动态性。我们的方法从不丢弃标记，并能有效地映射到现代硬件，与使用最先进的Tutel库训练的MoE相比，端到端训练速度提高了高达40%，与使用高度优化的Megatron-LM框架训练的DNN相比，速度提高了2.4倍。1.

I N T R O D U C T I O N。利用深度神经网络（DNN）的权重、激活和输入数据中的稀疏性是减少实现给定模型质量所需计算量的有效技术。Khan等人，2015年。Gale等人，2019年。

过去十年来，在算法和高性能软件方面取得了重大进展，使稀疏性在实践中变得有用。Gray等人，2017年。Naring等人，2017年。Kouchbrenner等人，2018年。Elson等人，2020年。Gale等人，2020年。稀疏性仍然面临的一个挑战是在加速器上进行模型训练。

DNN最常在像GPU（NVIDIA，2020年）和TPU（Jupy等人，2017年）这样的硬件加速器上进行训练，这些加速器利用密集计算的规律性来提供高性能。因此，细粒度的稀疏计算在这些处理器上效率较低。

为了在加速器上实现高效计算，可以在稀疏矩阵上强制执行结构，Naring等人，2017年。Gray等人，2017年。Yao等人，2019年。具有底层结构化稀疏性的一类新兴模型是混合专家（MoE），Shazir等人，2017年。NMO中的每一层都是专家集合，它们本身就是小型DNN。

当数据通过MoE层时，每个标记都会动态路由到一部分专家进行计算。通过利用这种稀疏计算，MoE在自然语言处理和计算机视觉应用中的训练时间减少了多达四倍。Artetics等人，2021年。Raquel May等人，2021年。

这些1斯坦福大学，加利福尼亚州斯坦福市，美国2微软研究院，华盛顿州雷德蒙德市，美国3谷歌研究院，加利福尼亚州山景城，美国。通信给Trevor Gale小于T Gale at cs.stanford.edu大于

收益转化为模型训练的新规模水平。将模型规模推到1万亿参数以上，Artetics等人，2021年。Do等人，2021年。FADIS等人，2022年。高效计算MoE的挑战在于处理这些架构的基本动态路由和负载不平衡计算。然而，现有的深度学习硬件和软件使得难以应对这一挑战。

例如，TPU及其XLA编译器要求所有张量形状都是静态已知的，并且经常难以处理像散射和收集这样的细粒度操作。这些约束使得难以直接在TPU上实现MoE。虽然GPU更灵活，但MoE中的稀疏计算不能很好地映射到主要框架和库中支持的软件原语。

用于MoE训练的最先进框架通过对MoE路由施加严格的约束来规避这些挑战。为了消除计算的动态性，映射到每个专家的标记集被修剪或填充到用户指定的尺寸。Lepikin等人，2020年。Fedus等人，

2022年。Wong等人，2022年。这种强制性公式在模型质量和硬件效率之间引入了权衡，因为用户必须决定是丢弃标记还是在填充上浪费计算和内存。这个决定通常是通过超参数调整来做出的，这增加了使用MoE的复杂性。为了应对这些挑战，我们开发了一种基于稀疏原语的MoE路由和计算方法。

我们的方法从不丢弃标记，并能有效地映射到现代GPU，与MoE和DNN训练的最先进框架相比，端到端训练速度分别提高了高达40%和2.4倍。我们做出了以下具体贡献。我们展示了如何将MoE层中的计算表示为块稀疏运算，以适应标记到专家的不平衡分配。我们使用这种公式来训练无丢弃MoE（DMoE）。

我们为块稀疏矩阵乘积开发了高性能GPU内核，可以有效地处理动态MoE计算。我们的内核使用两种技术，即块式CSR-COO编码和转置索引，以实现具有转置或非转置顺序的稀疏输入和输出的有效矩阵乘积。我们已经在一个名为MegaBlocks的系统中实现了这些技术，该系统建立在用于训练Transformer模型的最先进的Megatron-LM库之上，Shuabi等人，2019年。

我们通过微基准测试和Transformer语言模型的端到端训练来评估我们的系统。两个背景MoE层。MoE层由许多专家组成，它们本身就是小型神经网络。每个标记都会根据路由器计算的分数动态路由到一部分专家进行计算。专家通常被定义为小型多层感知器（MLP）。

标记通常被发送到少量专家，通常在1到4个之间。FATUS等人，2022年。MoE层通常与其他DNN层交错，最常用于替换自然语言中的前馈网络层。数据通常称为标记。对于视觉，数据通常是像素或补丁。Dostoevsky等人，2021年。

为简单起见，我们在本文中使用术语“标记”。工作，FFN，Transformer中的层，Shazir等人，2017年。Fedus等人，2022年。这种混合架构在自然语言和视觉任务上都取得了良好的结果，Du等人，2021年。Raquel May等人，2021年。据推测，这些改进是专家专门针对数据分布的不同部分的结果，Shazir等人，2017年。

我们在图1中说明了一个MoE层，并在本节的其余部分详细描述它。2.1路由MoE层的第一个阶段是路由器，它负责确定标记到专家的分配。除了专家分配之外，MoE路由器还会为每个分配生成概率，以反映映射的置信度。

这些权重被编码为每个标记-专家对的分数矩阵，用于线性组合每个标记的前k个专家输出。最常见的MoE路由样式是Shazir等人提出的学习路由器。在这个路由器中，标记通过与其他模型参数一起学习的权重矩阵相乘，从隐藏大小元素投影到num_expert分数。

分数用softmax进行归一化，路由决策是通过贪婪地为每个标记选择前K个得分最高的专家来做出的。2.2排列。最先进的MoE实现旨在并行计算所有专家层，以便有效利用GPU和TPU上可用的并行性。Lepikin等人，2020年。Fedus等人，2022年。Wong等人，2022年。

实现使用的标准原语是批量矩阵乘法，它计算一组相同形状的矩阵乘积。然而，将MoE计算映射到这个原语是非平凡的。为了遵守批量矩阵乘法的形状约束，专家必须被约束为具有相同形状的权重矩阵，并且分配给每个专家的标记数量必须相等。

后一个约束尤其成问题，因为上面描述的学习路由算法不保证标记到专家的负载平衡分配。为了满足这个约束，先前的工作定义了一个固定的专家容量，即每个专家可以分配的标记数量。Lepikin等人，2020年，Fedus等人，2022年。如果分配给专家的标记数量超过其容量，则会丢弃多余的标记。

也就是说，它们不会传递给任何专家进行计算，并且模型依赖于残差连接，在MoE层之后重新引入丢弃的标记的表示。如果专家层没有分配足够的标记来填充其容量，则其标记集将被填充以填充剩余的空间。

专家容量通常以容量因子超参数来指定，它是根据完美均匀分布下分配给每个专家的预期标记数量的倍数。专家容量等于num_tokens / num_experts * 容量因子。容量因子可以被认为是一个参数，它减少了丢弃标记的机会。

这个超参数代表了额外计算和模型质量之间的权衡。因此，希望最大限度地减少标记到专家的负载不平衡。这样做的典型机制是辅助负载。除了启用专家层的批量计算之外，这些约束还允许所有张量形状都是静态已知的，这是TPU和XLA所必需的。2.3计算。一旦数据被排列，专家就可以并行计算。

对于专家是MLP的模型，这需要使用批量矩阵乘法为所有专家计算每一层。对于卷积专家，可以使用分组卷积计算各层。2.4反排列。计算专家后，生成的特征向量将被反排列，使其顺序与层的输入顺序匹配。

MoE计算的最后一步是用它们被分配到各自专家的分数来缩放输出标记。当标记被路由到多个专家时，这些加权结果将被求和以生成每个标记的最终层输出。3动机。MoE中的标记丢弃。尽管使用了负载平衡损失，但先前的工作表明标记路由仍然高度不平衡。Wong等人，2022年。

为了量化标记丢弃对模型质量的影响，我们在堆栈上训练了MoE语言模型：GAO等人，2020年，具有各种容量因子。我们训练了类似于Fedus等人使用的Transformer MoE，2022年，其中每个模型都是一个Transformer，其FFN层被64个专家MoE层替换，每个专家都是一个与原始FFN维度匹配的两层MLP。

我们使用了top-1路由，并将我们的MoE模型基于表1中描述的Transformer小型模型。所有模型都使用来自GPT-2的标记化进行训练，Radford等人，2019年，100亿个标记，序列长度为1024，原子优化器，以及来自ShueiB等人，2019年的学习率和梯度裁剪设置。

我们在单个A100 GPU上训练所有模型，批量大小为512个序列。我们训练了容量因子为1、1.5和2的MoE，以及Wong等人提出的动态容量因子技术，2022年，其中容量因子动态设置为避免标记丢弃的最小值。作为基线，我们在各种规模上训练了标准Transformer模型。

所有Transformer和MoE模型的词汇量大小为51200，序列长度为1024，注意力头大小为64。我们的模型配置总结在表1中，实验结果如图2所示。对于这些模型，我们观察到标记丢弃的影响是显著的。

虽然容量因子为1的MoE实现了0.15的验证损失降低，但避免丢弃标记的MoE实现了0.26的降低，是前者增益的1.73倍，足以超过Transformer中型的质量。虽然丢弃标记会降低模型质量，但增加容量因子会增加额外的计算和内存。在这个例子中，为了避免丢弃标记，MoE层数学运算增加了2倍以上。

Wong等人，2022年表明，一些MoE需要高达11的容量因子才能避免丢弃标记，而其他模型在训练过程中避免丢弃标记所需的容量因子会不可预测地激增。除了增加容量因子的计算开销之外，必须调整额外的超参数会大大增加需要为目标任务训练的模型数量。

对于大型神经网络来说，这尤其麻烦，因为训练单个模型的成本可能高达数十万美元。Mosaic ML，2022年。可能是由于这个原因，一些关于MoE的大型研究根本没有探索不同的容量因子。Artetics等人，2021年。Clark等人，2022年。4无标记遗漏：使用块稀疏性。

本节描述了我们如何根据块稀疏计算来制定MoE层计算，以避免丢弃标记。使用块稀疏原语来表达MoE计算的动机是多方面的。首先，如下所示，块稀疏矩阵是描述MoE中动态和负载不平衡计算的自然而灵活的方式。其次，块稀疏性可以有效地映射到围绕收缩阵列矩阵乘法器（如GPU和TPU）构建的硬件加速器。

由于MoE专家的粗粒度，我们可以为我们的实现选择一个足够大的块大小，以使计算能够实现峰值设备吞吐量的很高比例。最后，块稀疏内核（如矩阵乘法和卷积）是通用的原语，可用于各种应用。Marring等人，2017年。Gray等人，2017年。Child等人，2019年。Elson等人，2020年。

这使得对高性能内核的投资更实际，因为工作可以在目标任务之间进行摊销。我们也可以类似地投资于可变大小的批量矩阵乘法内核，但这在当今设计的MoE架构中的效用将受到限制。除了这些考虑之外，MoE的块稀疏公式还提供了一个新的视角，将这些算法视为一种动态的、结构化的、激活稀疏性。

名称“无标记遗漏”指的是Fedus等人，2022年简要讨论的技术，这是试图恢复因丢弃标记而损失的质量的失败尝试。这种观点与大量关于稀疏训练算法的文献相呼应，并为利用该邻近领域的见解进一步改进MoE提供了机会。

预备知识稀疏矩阵乘积表示在本论文的其余部分中，我们经常提到矩阵乘法，其中三个矩阵（两个输入和一个输出）中的一个矩阵是稀疏的，而其他矩阵是密集的。我们借用Triton（Tillett等人，2019年）的表示法来描述这些不同的运算，每个运算都用三个字符的字符串来描述，其中每个字符都是s（稀疏）或d（密集）。字符的顺序是输出、左输入、右输入。

例如，两个密集矩阵与稀疏输出的乘积是SDD，也称为采样密集密集矩阵乘法（SDDMM）。这种表示法对于区分DSD和DDS等运算很有用，它们是稀疏矩阵密集矩阵乘法（SPMM）的不同形式。上标T表示输入参数的转置。例如，SDDT表示右输入矩阵被转置的SDD。

4.1使用块稀疏性的专家计算我们方法背后的关键见解如图3所示。与其使用批量矩阵乘法计算MoE层内的专家，不如等效地计算专家为SDD，其中输出稀疏矩阵具有块对角结构，如图3b所示。允许标记到专家的负载不平衡分配类似于允许块对角矩阵中的块具有可变数量的行。

为此，我们建议使用块稀疏矩阵乘法将每个块计算为许多较小的固定大小的块，如图3c所示。为了构建多层专家，我们可以迭代SDD和DSD操作，参见图6。

在这个公式中，我们还可以放宽对每个块中列数的约束，以构建具有可变大小专家的MoE层，如图3c所示。虽然这是未来工作的一个有趣方向，但我们没有探索这些配置，因为需要更多研究来确定如何利用这种能力来提高效率。

对于足够大的块，块稀疏矩阵乘法能够在现代GPU上达到峰值吞吐量的高比例，Gray等人，2017年，NVIDIA，2021年。

MoE中的粗粒度稀疏性适合于使用MoE-FFN层的Transformer模型中的这一要求。图3B中所示的块中的列数对应于FFN隐藏大小，通常在1024到8192之间。Vaswani等人，2017年。Radford等人，2019年。Brown等人，2020年。

这些块中的行数对应于分配给每个专家的标记数量，在均匀分布下，这预计等于标记数量除以专家数量。这可以从每个专家几千个到几万个标记不等。Lepikin等人，2020年。Artetics等人，2021年。Fedus等人，2022年。

这些粗粒度块比用于密集矩阵乘法内核的最大平铺尺寸大许多倍，这使我们能够灵活地选择一个可以匹配其吞吐量的块大小。5 MegaBlocks。用于高效MoE训练的框架。我们在一个名为MegaBlocks的系统中实现了我们的技术，该系统建立在Megatron-LM（Shuabi等人，2019年）和PyTorch（Poshke等人，2019年）的基础上。

除了高性能的无丢弃MoE层之外，我们的系统还支持使用数据和专家模型并行性的MoE分布式训练。本节讨论了我们的DMoE实现的设计，包括我们的块稀疏内核以及构建高效系统的其他考虑因素。5.1.1节讨论了现有块稀疏内核的局限性。

5.1.2节分析了块大小对块稀疏乘积性能的影响。5.1.3节描述了我们的混合块式CSR-COO稀疏矩阵格式，它可以实现具有稀疏输入和输出操作数的有效矩阵乘积。5.1.4节介绍了转置索引，作为有效迭代转置顺序的块稀疏矩阵的机制。

最后，5.2节讨论了DMoE的有效路由和排列。预备知识。GPU上的矩阵乘法。GPU上的矩阵乘法内核利用平铺，其中输出矩阵被分解成静态大小的二维值块，NVIDIA，2022C。这些平铺的计算可以并行化，并且可以调整单个平铺的大小以权衡算术强度和并行性。

分配给平铺的线程组称为线程块。5.1 MoE的高效块稀疏内核。为了使用块稀疏内核训练MoE，我们需要用于前向和反向传递的原语。考虑一个MoE-FFN层，其中每个专家都是一个两层MLP。对于此配置，前向传递需要SDD操作，然后是DSD，图6。

对于反向传递，我们分别为第二层数据梯度和权重梯度计算SDDT和DSTD，然后分别为第一层数据梯度和权重梯度计算DSD和DDTS。5.1.1现有块稀疏原语我们考虑了GPU上块稀疏矩阵乘法的两个现有库，NVIDIA CUSPARSE（NVIDIA，2022B）和Triton Blocksparse（Tillett等人，2019年）。

CUSPARSE支持用于DSD的块式L稀疏矩阵格式。但是，截至CUDA 11.8，此操作不支持稀疏矩阵输入的转置。CUSPARSE也没有提供具有块式L矩阵的SDD原语。除了这些限制之外，块式L格式还要求稀疏矩阵中的所有行具有相同数量的非零元素，这将违背我们支持负载不平衡矩阵的目标。

Blocksparse支持SDD、DSD和DDS以及所有转置和非转置输入的组合。但是，这些原语假设稀疏矩阵的拓扑在调用之间不会改变3。库API采用描述稀疏操作数的位掩码，然后预计算查找表和块分组以加速计算。对于我们的用例，稀疏矩阵拓扑在训练的每次迭代和模型中的每个MoE层都会发生变化。

为了使用Blocksparse，我们将不得不反复支付这些预处理步骤的成本。基于此分析，我们选择编写我们自己的块稀疏原语，以便根据MoE专家计算的动态性对其进行调整。我们实现了针对NVIDIA GPU的SDD、DSD和DDS操作。我们的内核支持所有转置和非转置输入的组合。本节的其余部分详细介绍了内核的设计和实现。

5.1.2为MoE选择块大小。为了有效地使用现代GPU，我们希望使用具有足够算术强度的稀疏块。3. 这可能是因为它们是为稀疏注意力等应用程序编写的，其中稀疏矩阵拓扑是在训练之前确定的。Child等人，2019年

保持矩阵乘法单元繁忙。大型块也希望摊销存储和操作稀疏矩阵元数据的成本，因为元数据（如列索引）只需要为每个非零块保留。为了选择我们的目标块大小，我们研究了来自NVIDIA Cutlass（NVIDIA-2022C）的不同平铺尺寸的密集矩阵乘法内核的性能。

我们在具有从512到16384的二次幂边长的方阵上对混合精度（FP16 + FP32累加）矩阵乘法进行了基准测试，这些方阵具有Cutlass支持的每一组平铺尺寸。对于矩形平铺，我们只显示第一个平铺维度较大的配置，因为我们发现对于这些问题，这些配置略微优于替代排序。

我们在具有CUDA 11.5和Cutlass 2.5的A100SXM4 80GB GPU上运行所有基准测试。这些基准测试如图4所示。在这些基准测试中，我们观察到128x128平铺始终与其他配置一样好或更好。轶事地，我们观察到NVIDIA Kubla的NVIDIA 2022A为我们研究的密集Transformer模型通常选择相同的配置。

基于此分析，我们选择使用128x128块稀疏性。虽然块稀疏矩阵乘法的平铺尺寸和稀疏矩阵中的块大小不需要相等，但我们发现对于128x128块，我们的工作负载中性能最高的平铺尺寸也是128x128。

为了实现我们的内核，我们扩展了Cutlass（NVIDIA，2022C），以支持块稀疏矩阵，并重用了它们用于具有不同数据类型和GPU架构的高性能矩阵乘法的机制。MegaBlocks。高效的混合专家稀疏训练。5.1.3使用混合块式CSR-COO计算稀疏输出。我们使用块压缩稀疏行（BCSR）作为我们的主要稀疏矩阵格式。

BCSR使迭代行中的非零元素变得简单，这对于DSD和DDST等操作是必要的。使用BCSR迭代块的开销也很小，因为识别块在矩阵中的位置只需要加载其列索引一次。我们在5.1.4节中讨论了使用此格式有效迭代列中非零元素的方法。BCSR稀疏矩阵的一个挑战是有效地并行计算SDD操作。

在内核启动时，每个线程块都需要识别其输出块的行和列，以便知道需要哪些输入矩阵的行和列来计算它。因为BCSR只为每个块编码列索引，所以识别非零块的行索引需要搜索行偏移量。解决这个问题的一个方法是启动可能需要计算输出每一行的最大数量的线程块，如果它是完全密集的。

启动时，每个线程块都可以检查其列偏移量是否超出其行中非零元素数量的范围，如果没有任何工作要做，则返回。Gale等人，2020年表明，对于中等稀疏矩阵（50-90%的零），启动额外线程块引入的开销可以忽略不计。我们尝试了这种方法，但观察到对于MoE，启动这些未使用的线程块的成本是显著的，特别是对于专家数量较多的模型，其中块稀疏矩阵的稀疏性水平非常高。

为了有效地并行化SDD，我们额外地将每个非零块的行索引具体化，以便线程块可以轻松查找输出矩阵中稀疏块的坐标。由于我们只需要为128x128块中每16384个非零值存储一个索引，因此此额外元数据的存储需求可以忽略不计。

即使有了这些额外的元数据，我们仍然保持非零块的行顺序，以便矩阵可以作为BCSR或块坐标格式BCOO进行操作。我们在图5中说明了这种混合块式CSR-COO编码。5.1.4块稀疏转置与转置索引。模型训练的正向和反向传递需要稀疏矩阵转置。

然而，以转置顺序迭代BCSR矩阵需要搜索每一行以识别目标列中的块是否为非零。Bullock等人，2009。我们可以明确地具体化稀疏矩阵的转置版本，但这会产生运行时和存储成本，因为矩阵中的所有非零值都需要被复制。

为了能够以转置顺序有效地迭代BCSR矩阵，我们构建了转置矩阵的元数据，但没有显式转置非零值。相反，我们构建了一个索引数组，每个非零块一个，这些索引按转置顺序存储，并包含内存中每个非零块的偏移量。此附加元数据允许通过间接寻址以转置顺序有效地迭代矩阵，如图5所示。

这个想法类似于数据库中的二级索引，它允许以与主索引不同的顺序有效地访问条目。与我们的混合块式CSR-COO编码类似，这项技术依赖于这样一个事实：由于我们的块大小很大，元数据的存储和计算成本比非零值低很多倍。

5.2高效路由和置换按照目前的实现，我们的块稀疏矩阵乘法内核要求分配给每个专家的令牌数量必须是块大小的倍数。为了满足这个约束，我们将每组令牌用零填充到最接近的128的倍数，并将此操作融合到自定义置换内核中。

我们可以通过支持问题边缘处的部分块来消除这个约束，这类似于矩阵乘法处理不能被平铺维度整除的矩阵的方式。然而，鉴于我们预计分配给每个专家的令牌数量将达到数千甚至数万，因此此功能的性能影响将微乎其微。

MegaBlocks。高效的混合专家稀疏训练。表3。用于模型训练的微批量大小。在所有实验中，我们使用了适合内存的最大微批量大小。模型微批量大小。Megatron LM Transformer XS64 Transformer Small 32 Transformer Medium 16 Transformer Large 16 Transformer XL8。 Megablocks DMOE XS64 DMOE Small 32 DMOE Medium 8。

TUTL DMOE XS32 DMOE SMALL 8 DMOE MEDIUM 1。一旦路由器计算出专家分配，我们就使用自定义CUDA内核创建块稀疏矩阵的元数据。我们还在此时构建转置元数据，以在使用它的多个块稀疏矩阵乘法中分摊成本，这些乘法跨越正向和反向计算。6。

本节分析了我们的系统与最先进的库（Microsoft Tuttle，Wong等人，2022年，以及NVIDIA Megatron LM，Shueibe等人，2019年）在训练Transformer MoE和标准Transformer方面的性能比较。为了确保公平比较，我们扩展了Megatron LM以使用Tuttle的MoE层额外支持MoE训练。

所有实验都在配备CUDA 11.5、Cutlass 2.5的NVIDIA A100 SXM 480GB GPU上进行，并使用了Megatron LM中实现的混合精度训练（Misikovicius等人，2018年）。

6.1不丢弃令牌的MoE训练。为了评估我们的避免令牌丢弃技术的效率，我们将其与Wang等人（2022年）提出的DMOE方法进行了比较，其中容量因子动态设置为避免令牌丢弃的最小值。我们在Pile（Gao等人，2020年）上训练了仅解码器的Transformer语言模型，使用了第3节中描述的相同超参数。

对于Transformer MoE，我们训练了从RxS、小型和中型模型缩放的模型，其中每个FFN层都被替换为使用top-1的64个专家MoE层。0.00.51.01.52.02.5训练时间，天，2.22.32.42.52.62.72.82.9。

验证损失。Transformer。Megatron LM。DMOE。Tootle。E64。Top 1。DMOE。Megablocks。E64。Top 1。图7。在Pile上训练的Megablocks DMOE、Tootle DMOE和Megatron LM Transformer。

MegaBLOX使用块稀疏操作来处理MoE中动态的和负载不平衡的计算，与TUTL使用的基于填充的方法相比，这使得MoE XS、MoE-small和MoE-medium的端到端训练加速分别达到1.38倍、2.0倍和4.35倍。随着模型大小的增加，我们方法的优势也随之增加，因为填充专家批次的内存需求迫使TUTL使用更小的微批量大小，从而降低了硬件效率。

与使用Megatron LM训练的密集Transformer语言模型相比，Megablocks在这些模型上实现了相同的验证损失，端到端训练加速分别为1.8倍和2.4倍。路由。我们还训练了参数从46M到13亿的参数标准Transformer模型，相当于Transformer base（Vaswani等人，2017年）到GPT-3 XL（Brown等人，2020年），作为密集基线。

我们使用8个A100 SXM4 80GB GPU在MoE层使用8路专家模型并行化，在所有其他层使用数据并行化来训练所有模型。我们在训练中对所有模型使用梯度累积，批量大小为512个序列，并且使用不会耗尽内存的最大微批量大小，norionon at all。2021A。

我们的模型配置总结在表1和表2中。对于每个模型，我们在图7中报告了端到端训练时间和在验证集上达到的最终损失。与用于避免令牌丢弃的流行的基于填充的方法相比，我们使用块稀疏性的自适应MoE计算技术使MoE XS、MoE-small和MoE-medium的端到端训练加速分别达到1.38倍、2.0倍和4.35倍。

除了计算开销外，在Toodle中实现的基于填充的方法还会显著增加存储MoE层中激活所需的内存量。这尤其成问题，因为与标准Transformer相比，MoE已经需要多倍的存储空间来存储其大型权重矩阵。

对于这些模型，我们观察到这种内存使用量的增加使Tootle能够使用的最大微批量大小比Megablocks分别减少了2倍、4倍和8倍（对于MoE XS、MoE-small和MoE-medium）。这反过来又会由于硬件效率降低而增加训练时间。因此，我们观察到Megablocks相对于Tootle的优势随着模型大小的增加而增长。表3显示了每个模型配置使用的微批量大小。

与使用Megatron LM训练的Transformer模型相比，使用Megablocks训练的DMOE将达到给定验证损失所需的训练时间减少了1.8到2.4倍。这种比较中的差异主要源于MoE模型增加的权重内存使用量，这迫使Megablocks对MoE-medium使用比类似的Transformer模型小2倍的微批量大小。这些结果突出了减少MoE中内存使用量作为未来研究方向的重要性。

对于这些Transformer模型，我们观察到Megatron LM维持了该8 GPU系统2.5 petaflop峰值吞吐量的21%到48%，效率随着模型大小的增加而提高。Megablocks相对于这个最先进框架所取得的加速证明了我们系统在MoE功效方面的效率。6.2带有令牌丢弃的MoE训练。我们还将我们的DMOE模型与使用TUTL训练的令牌丢弃MoE进行了比较。

为了找到最有效的配置，我们训练了具有1倍、1.5倍和2倍容量因子的MoE XS、MoE-small和MoE-medium模型，总共9个附加模型。对于这些配置，所有令牌丢弃MoE模型都能够使用与类似的DMOE相同的微批量大小，而不会耗尽GPU内存。我们在图8中报告了这些模型的端到端训练时间和验证损失，以及我们的DMOE和标准Transformer结果。

比较MoE和DMOE的相同6.3块稀疏矩阵乘法性能。为了评估我们块稀疏矩阵乘法内核的质量，我们对训练MoE XS、MoE-small和MoE-medium模型中使用的配置进行了基准测试，并与Kubla的批量矩阵乘法进行了比较。这包括每个FFN层中两层的正向传递、反向权重和反向数据操作。

总共，我们对18个问题进行了基准测试，每个模型6个问题。为了与批量矩阵乘法进行比较，我们使用表3中列出的相同微批量大小对具有均匀令牌到专家分布的每个问题进行基准测试。这些基准测试可以被视为消融研究，评估如果使用我们的块稀疏内核来实现标准的、令牌丢弃的MoE将会引入的开销。

对于每个问题，我们对100次执行的吞吐量取平均值。我们没有将构建稀疏矩阵元数据所需的时间包含在这些基准测试中，因为这些操作在FNN层内的所有六个问题中分摊。这些基准测试的结果如图9所示。在这些问题上，我们观察到我们的块稀疏内核能够实现Kubla的吞吐量的98.6%，标准差为4%。最大相对吞吐量为104%，最小值为91%。

总的来说，我们的内核在半数问题上略微优于Kubla，在另一半问题上略微逊色。在对Cutlass进行基准测试时，我们观察到改变计算输出矩阵平铺的顺序可能会由于L2缓存效应而使操作的吞吐量改变多达10%。我们认为这些结果中大部分性能差异可以归因于块稀疏矩阵中发生的计算重新排序，尽管还需要进一步研究。

我们注意到额外开销的一个情况是在用于计算权重梯度的DST运算中。因为我们使用二级索引以转置顺序迭代稀疏操作数，所以迭代此矩阵时的访问模式几乎没有空间局部性，这反过来又会降低整个操作的吞吐量。

虽然这是一个值得进一步研究的有趣问题，但由于改进的机会有限（不到10%），以及这两个操作所代表的端到端运行时间相对较少，因此对模型性能的整体影响最小。7。

R相关的W-ORC。MoE路由。改进MoE的路由算法是一个活跃的研究领域。基础层将MoE路由制定为线性分配问题，试图在完美平衡分配的约束下最大化聚合令牌专家亲和力。Lewis等人，2021年。这种方法通过根据需要将令牌重新路由到不同的专家来保证不丢弃令牌。

Clark等人（2022年）发现基础层可能会产生大量的运行时开销，并提出了一种使用Synchern算法的近似版本。因为他们的近似不再保证避免令牌丢弃。Clark等人（2022年）在所有实验中都使用2的容量因子。其他技术已被提出，以基于哈希函数提前静态地决定令牌到专家的映射。Roller等人，2021年。

然而，Clark等人观察到这种方法的性能不如他们研究的其他路由算法。最近，Joe等人提出反转路由问题，以便每个专家选择其得分最高的K个令牌。虽然这保证了令牌到专家的负载平衡分配，但这种方法仍然会遭受令牌丢弃，因为相同的令牌可以被多个专家选择。

我们预计改进的路由算法可以补充我们对高效灵活的专家计算方法。探索如何结合这些方法是未来研究的一个有趣方向。高性能MoE。为了扩展MoE训练。Tutel实现了针对MoE的优化的分布式通信原语和隐藏专家模型并行化通信成本的技术。Wong等人，2022年。

He等人（2022年）提出了FasterMoE，这是一个基于高效通信策略和MoE路由算法的更改以避免网络拥塞的MoE分布式训练系统。我们的实现还可以从这些技术中受益，特别是对于大规模分布式训练。稀疏内核。允许高效转置访问的稀疏矩阵格式已被充分研究。Bullock等人，2009年；Smith和Kharapis，2015年；Lee等人，2018年。

探索如何将这些格式适应现代GPU上的大块稀疏性是一个有趣的研究方向。8。结论。我们介绍了Megablocks，一个在GPU上进行高效MoE训练的系统。我们的系统基于MoE根据块稀疏操作和新的块稀疏GPU内核的重新制定，这些内核有效地处理MoE中存在的动态性。

我们的方法从不丢弃令牌，并且可以有效地映射到现代硬件加速器，与使用最先进的Tuttle库训练的MoE相比，端到端训练加速高达40%，与使用高度优化的Megatron LM框架训练的DNN相比，加速2.4倍。

Megablocks。高效的混合专家稀疏训练。参考文献。Artetics。M. Beausoleil。S. Goyle。N. Mihailov。T. Ott。M. Schleifer。S. Lin。X.V. Du。J. Iyer。S. Passanuru。R. Anantharaman。G. Lee。X. Chen。S. Akin。H. Baines。M. Martin。

L. Joe。X. Cora。P.S. Ohoro。B. Wong。J. Zetlimoyer。L. Diab。M. T. Kozareva。Z.和Stoyanov。V. 使用混合专家进行高效的大规模语言建模。核心。ABS。2112.10684。2021。

Brown。T.B. Mann。B. Ryder。N. Sabia。M. Kaplan。J. Dariwal。P. Neelakanton。A. Shyam。P. Sastry。G. Askell。A. Agarwal。S. Herbert Voss。A. Kruger。G. Hennigan。T. Child。R. Ramesh。A. Ziegler。D. M. Wu。J. Winter。C. Hess。C.

Chen。M. Sigler。E. Litwin。M. Gray。S. Chess。B. Clark。J. Berner。C. McCandlish。S. Radford。A. Sutskever。I.和Amaday。D. 少样本学习的语言模型。

在神经信息处理系统进展33中。神经信息处理系统2020年年会。NeurIPS 2020。2020年12月6日至12日。虚拟。2020。Bullock。A. Feynman。J.T. Frigo。M. Gilbert。J.R.和Leiserson。CE 使用压缩稀疏块进行并行稀疏矩阵向量和矩阵转置向量乘法。

在SPAA 2009中。第21届ACM算法与体系结构并行性研讨会论文集，加拿大卡尔加里，2009年8月11日至13日，第233-244页。ACM，2009。doi。10.11451583991。

1,584,053。Child，R.，Gray，S.，Radford，A.和Sutzkever。I. 使用稀疏Transformer生成长序列。核心，摘要，1904.10509。2019。

Clark。A. De Los Casas。D. Guy。A. Mensch。A. Paganini。M. Hoffman。J. DeMock。B. Heckman。B. A. Kai。T. Bourgeois。S. Van Dendriche。G. Rutherford。E. Hennigan。T. Johnson。M. Millikan。K. Casseret。A. Jones。C. Bachatskaya。E. Budden。D. Seifer。L.

Osindaro，S. Vinyals，O. Ray，J. W. Elson，E. Cavicuoglu，K.和Simonian，K. 路由语言模型的统一缩放定律。核心，摘要，2202.01169，2022。

Dosovitsky，A. Bayer，L. Kolesnikov，A. Weissenburn，D. Jai，X. Untertheiner，T. Degani，M. Minderer，M. Heigold，G. Gelli，S. Uskarite，J.和Holzby，N. 一张图片值16x16个单词。用于大规模图像识别的Transformer。

在第9届国际学习表征会议，ICLR 2021，奥地利虚拟活动，2021年5月3日至7日。

OpenReview.net。2021。Du。N. Huang。Y. Dai。A. M. Tong。S. Lepikin。D. Xu。Y. Kirkin。M. Zhou。Y. Yu。A. W. Farat。O. Zouf。B. Fedus。L. Bosma。M. Zhou。Z. Wang。T. Wang。Y. E. Webster。K. Pellet。M. Robinson。

K. Meyer-Helstern。K. Duke。T. Dixon。L. Zhang。K. Le。Q. V. Wu。Y. Chen。Z.和Q. C. Glam。使用混合专家高效缩放语言模型，2021。Elsin。E. Duhan。M. Gale。T.和Simonian。K. 快速稀疏卷积。

在2020年IEEE，CVF计算机视觉和模式识别会议。CVPR 2020。美国华盛顿州西雅图。2020年6月13日至19日。第14617-14626页。计算机视觉基金会。IEEE。2020。DOI。10.1109.CVPR42600。

2020.01464。FADUS，W，ZOEF，B和Shazir，N-Switch Transformer。使用简单高效的稀疏性扩展到万亿参数模型。机器学习研究杂志。23，120，1-39，2022。Gale，T.，Elson，E.和Hooker，S.

深度神经网络中稀疏性的现状。核心。ABS。1902.09574。2019。Gail。T. Zaharia。M. Young。C.和Elson。e深度学习的稀疏GPU内核。在高性能计算、网络、存储和分析国际会议论文集中。南卡罗来纳州2020年。虚拟活动。美国佐治亚州亚特兰大。2020年11月9日至19日。

IEEE，ACM，2020。Gao。L. Biderman。S. Black。S. Golding。L. Hop。T. Foster。C. Feng。J. He。H. Thight。A. Nabashima。N. Presser。S.和Leahy。C. Pile。用于语言建模的800GB多样化文本数据集。

档案预印本档案。2101.00027。2020。Gray，S.，Radford，A.和Kingma，D.P. 块稀疏GPU内核。https://blog.openai.com。块稀疏GPU内核。2017。Hahn，S.，Poole，J.，Tran，J.和Dally。WJ 学习权重和连接以实现高效的神经网络。

在神经信息处理系统进展28中。神经信息处理系统2015年年会。2015年12月7日至12日。加拿大魁北克省蒙特利尔。2015。He。J. Jai。J. Antunes。T. Wong。H. Luo。F. Shi。S.和Li。Q. FastMoE。建模和优化大型动态预训练模型的训练。

在第27届ACM SIGPLAN并行编程原理与实践研讨会论文集中。PPOPP。22。第120-134页。美国纽约州纽约。2022。计算机械协会。ISBN 9781450392044。DOI。10.1145350321sts。

3,508,418。Wong。C. Kui。W. Shang。Y. Yang。Z. Lu。Z. Hu。H. Wong。Z. Solace。R. Jose。J. Ram。P. Chow。J. Cheng。P. Yang。F. Yang。M.和Shang。Y. Tuttle。大规模自适应混合专家，2022。

感谢您收听本次朗读。如需完整论文及更多内容，请访问我们的主页。