Scattered Mixture-of-Experts Implementation

发表时间: 2024-03 · arXiv:2403.08245

作者: Shawn Tan, Yikang Shen, Rameswar Pandarpanda, Aaron Courville
机构: mila.quebec, http://ibm.com, http://iro.umontreal.ca

速读

一句话结论 本文提出了一种名为ScatterMoE的稀疏专家混合模型高效实现方法，通过引入全新的底层原语融合分组与线性变换步骤，在保持数值精度的同时显著降低了显存占用并提升了模型的训练与推理吞吐量。

要解决什么问题 现有的稀疏专家混合（SMoE）模型在GPU上的高效实现面临严重的显存开销与架构扩展性卡点。以广泛应用于Megatron-LM模型的Megablocks为代表的现有实现，在计算前需要对输入执行分散到分组的复制操作，这在训练阶段为了保存反向传播所需的张量，会产生巨大的显存分配开销。同时，为了迎合GPU的底层计算特性，现有方法通常会将路由到各个专家的张量块强制填充到相同大小，这种不必要的填充进一步加剧了高带宽内存的浪费。此外，将SMoE转化为稀疏矩阵格式的做法虽然在多层感知机上高效，但其固化的中间表示形式极难扩展到注意力机制等其他类型的专家模块上。早期的TPU实现甚至要求在编译时定死张量容量，一旦专家负载不均，要么丢弃超出容量的令牌，要么对空闲专家进行无效填充，造成严重的算力浪费。

怎么做的 核心思路是通过融合分组与线性变换步骤，彻底避免对输入进行显式填充和冗余复制，从而绕开显存分配的卡点。标准SMoE的计算需要先根据令牌嵌入分配权重并选择头部专家、将分配给同专家的令牌分组并产生扇出、执行专家线性变换、将令牌分散回原时间步，最后进行加权求和，其输出定义为 $$\mathbf{Y}_t = \sum_{e \in \text{topk}(g(\mathbf{X}_t))} g_e(\mathbf{X}_t) \cdot f_e(\mathbf{X}_t)$$ 传统方法会在高带宽内存中实例化整个填充数组并按排序复制数据，而ScatterMoE仅根据专家对令牌排序并对索引进行填充，在加载到静态RAM时直接根据填充后的索引读取，避免了全局显存占用。关键设计是由Triton编写的新原语ParallelLinear，其底层依赖一个名为scatter2scatter的核函数，能在分散的向量上直接执行分组矩阵运算。该原语支持输入和输出在分组与分散状态之间的四种任意组合，其产生的中间表示是标准的PyTorch张量，极大地提升了扩展性。在反向传播计算梯度 $$\nabla \mathbf{X} = \nabla \mathbf{Y} \mathbf{W}^{\top}, \quad \nabla \mathbf{W} = \mathbf{X}^{\top} \nabla \mathbf{Y}$$ 时，ParallelLinear在嵌入未分组时对其进行分组以计算权重梯度，并通过复用数组来最小化显存使用，例如将输出梯度的数组直接用作后续计算的输出。在具体实现多层感知机时，ScatterMoE将两次线性变换配置为先从分散到分组、再从分组到分散，使得反向传播只需一次分组操作。而在扩展到多头注意力混合模块时，由于应用位置编码需要保持令牌的时间顺序，ParallelLinear可以直接执行从分散到分散的变换，无需像现有方法那样分配额外的数组来回倒腾数据。该注意力实现类似于分组查询注意力，每个专家拥有特定数量的键头，并从中选择对应比例的查询头，从而实现高效的注意力混合计算。

效果如何 实验在同一计算节点上的8张Nvidia A100 GPU上搭建了1.5B参数的Mixtral配置模型，训练数据设置为有效批量大小256、序列长度2048个令牌，并在单卡上进行了多层感知机和注意力模块的单元测试。对比基线包括代表基础实现的HuggingFace朴素版本、代表当前高效稀疏路线的Megablocks稀疏模式及其内存高效模式，以及一种高效的PyTorch实现。在整体训练吞吐量上，ScatterMoE比稀疏Megablocks基线高出38.1%。在多层感知机单元测试中，ScatterMoE的训练显存占用仅为Megablocks的66.2%，推理显存占用更是低至53.6%。在测试细粒度专家设置时，随着粒度的增加，ScatterMoE展现出比Megablocks更好的吞吐量扩展性，特别是在批处理推理中优势明显。在降低稀疏性的实验中，即使将选中的专家数量增加到30个，ScatterMoE的吞吐量依然优于Megablocks，且两者均高于等效参数的密集模型。在注意力混合模块的测试中，当每次激活8个专家时，其推理吞吐量比Megablocks高出24.0%，且性能差距随粒度增加而拉大。此外，在winogrande、hellaswag等十余个语言模型评估基准上，ScatterMoE与HuggingFace标准实现的数值精度完全一致，绝对误差可忽略不计。作者也承认了当前方法的局限性，目前尚未实现用于加速解码的专用内核，且在多节点环境下的并行化支持仍需进一步开发。

A1 主要贡献

核心问题: 尽管稀疏专家混合（SMoE）模型在扩展Transformer方面越来越受欢迎，但其在GPU上的高效实现仍面临挑战。现有的实现方式存在一些局限性：

内存开销大: 现有的SMoE实现（如Megablocks）通常需要在计算前对输入进行scatter-to-group的复制操作，这在训练中会因为存储用于反向传播的张量而产生显著的内存分配开销。
不必要的填充（Padding）: 为了便于GPU计算，一些实现会将路由到各个专家的张量块填充到相同大小，这进一步增加了内存开销。
扩展性受限: 将SMoE问题转化为稀疏矩阵格式（如Megablocks和PIT）虽然高效，但其稀疏矩阵的中间表示形式使其难以扩展到其他类型的专家模块（如注意力机制）。
TPU实现的局限: 早期的TPU实现要求在编译时指定所有张量的大小（容量），导致专家负载不均衡时，要么丢弃超过容量的令牌，要么对未充分利用的专家进行填充，造成资源浪费。

研究目标与创新点: 本文提出了ScatterMoE，一种旨在克服上述局限性的SMoE实现。其核心目标是改善批处理推理、训练速度和内存占用。

避免内存开销: ScatterMoE通过避免对输入进行填充和过多的复制操作来最小化内存开销。其关键创新在于融合了分组和线性变换步骤。
引入ParallelLinear原语: 本文引入了一个名为ParallelLinear的新原语，它可以在分散的向量上执行分组矩阵运算。这个模块是实现ScatterMoE的核心，它将专家线性变换和重排序操作融合在一起。
提升扩展性: ParallelLinear产生的中间表示是标准的PyTorch张量，这使得将现有的SMoE方法轻松扩展到其他类型的专家模块成为可能。
功能验证: 本文通过使用ParallelLinear实现一个注意力混合（Mixture of Attention）模型，展示了其扩展性。同时，通过与Megablocks进行基准测试，证明了ScatterMoE在吞吐量和内存占用方面的优势。

图 1: 当前的SMoE多层感知机（MLP）实现需要在分组时复制嵌入（左图），而ScatterMoE则融合了分组和线性变换步骤（右图），从而减少了我们方法的内存占用。不同的颜色代表不同的专家，垂直的矩形框代表嵌入，其上方或下方标注了相关的时间步。

A3 相关工作

其他实现与依赖关系。ScatterMoE的核心部分是使用Triton【【索引18，Triton: an intermediate language and compiler for tiled neural network computations，2019，Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages】】实现的，这是一种基于图块（tile-based）的GPU编程语言，在Python中编写，使其易于修改和扩展。本文的主要比较对象是Megablocks，它使用同样基于Triton的STK框架实现。Megablocks被广泛应用于Megatron-LM模型【【索引16，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，arXiv】；【索引12，Efficient large-scale language model training on gpu clusters using megatron-lm，2021，Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis】；【索引9，Reducing activation recomputation in large transformer models，2023，Proceedings of Machine Learning and Systems】】，是训练SMoE的一种高效方法。另一个流行的SMoE实现库是CUTLASS【【索引8，Who says elephants can’t run: Bringing large scale moe models into cloud scale production，2022，arXiv】】，Megablocks也使用它作为其分组计算的选项。

SMoE的其他应用。除了多层感知机（MLP），研究人员也提出了SMoE版本的注意力模块【【索引19，Mixture of attention heads: Selecting attention heads per token，2022，arXiv】；【索引2，Switchhead: Accelerating transformers with ´ mixture-of-experts attention，2023，arXiv】】。这些注意力混合（MoA）实现已被用于扩展通用Transformer（Universal Transformers）【【索引3，Universal transformers，2018，arXiv】；【索引17，Sparse universal transformer，2023，arXiv】】，并应用于全模块化Transformer的持续学习【【索引15，Moduleformer: Learning modular large language models from uncurated data，2023，arXiv】】。高效地计算SMoE将为这些模型的训练和推理带来巨大好处。

A2 方法细节

在本文中，我们将遵循矩阵使用粗体（例如X）的符号约定。除非另有说明，为便于理解，第一个维度是批处理-时间维度（batch-time，即将批处理和时间维度展平）。此外，$X_i$表示X的第i行。

3.1 稀疏专家混合（Sparse Mixture-of-Experts）

SMoE模块的构成与计算流程。SMoE模块由E个专家（experts）组成，这些专家通常是具有相似架构的子模块。输入中的T个令牌（tokens）中的每一个都通过一个路由模块进行路由，然后根据路由器的输出权重分配给k个专家，其中k ≤ E。然而，朴素的SMoE计算方法（即遍历所有令牌并评估相应专家的输出）速度太慢，无法充分利用GPU的并行计算能力。在实践中，SMoE的实现通常执行以下主要步骤：

路由（Routing） – 路由器根据每个令牌嵌入$X_t$为每个专家分配权重$g(X_t)$，并仅选择top-k个专家。
分组（Grouping） – 此步骤将分配给同一专家的所有令牌组合在一起。如果k > 1（通常是这种情况），这也会导致令牌的“扇出”（fan out），总共产生kN个嵌入。
专家变换（Expert transform） – 既然令牌已按专家分组，每个专家（一个线性变换）可以通过批处理向量变换（即矩阵-矩阵乘法）来高效计算。
分散（Scattering） – 此步骤将每个令牌返回，使其按其原始时间步进行分组。这仍然总共产生kN个嵌入。
加权求和（Weighted sum） – 此步骤通过原始的路由权重将每个令牌的k个输出组合起来，

$$\mathbf{Y}_t = \sum_{e \in \text{topk}(g(\mathbf{X}_t))} g_e(\mathbf{X}_t) \cdot f_e(\mathbf{X}_t)$$

最终产生N个嵌入。

专家模块的典型结构。通常，每个专家是一个多层感知机（MLP），带有一个维度为$d_{expert}$的隐藏层。

Megablocks与ScatterMoE的对比。在Megablocks中，步骤（1）和（4）会导致对原始输入的复制。它们进一步对每个专家的令牌块进行填充（padding），使其数量相等，以便于GPU计算（见图1）。这个过程会在高带宽内存（HBM）中为填充后的嵌入分配数组，并将原始嵌入按排序顺序复制进去。然后，对这个按专家排序的数组执行分组通用矩阵乘法（Grouped GeMM）。相比之下，ScatterMoE避免了在HBM中实例化整个填充数组。我们不是将所有嵌入复制到一个填充数组中，而是根据专家对令牌进行排序，并转而对索引进行填充。当将一个图块（tile）加载到静态RAM（SRAM）时，我们根据填充后的索引进行加载，从而得到一个填充后的图块。

3.2 ParallelLinear操作

ParallelLinear的核心功能。我们的SMoE实现依赖于ParallelLinear，它允许不同组合的分组通用矩阵乘法（Grouped GeMMs）。为了实现这一点，我们编写了一个名为scatter2scatter的Triton核函数，它支持图2中所示的所有操作组合。该操作融合了分组GeMM和分散的读写操作，使我们能够跳过一个中间的分组和复制步骤。ParallelLinear为输入和输出都提供了分组（grouped）和分散（scattered）的选项，从而产生了图2中看到的四种可能组合。通过这些操作的组合，我们可以实现ParallelLinear的前向和反向传播。

图 2: ParallelLinear允许执行不同组合的SMoE变换，使得输入和输出既可以是分组的也可以是分散的。这个基本功能构成了ScatterMoE前向和反向传播的基础。与现有实现不同，这些操作在执行时无需对输入和输出张量进行额外的复制（或填充）。

算法1：ParallelLinear前向传播。算法1提供了ParallelLinear的伪代码。它是scatter2scatter核函数的一个薄封装，分组选项作为参数提供给ParallelLinear。这是我们SMoE实现的主要工作引擎，可用于实现MLP和注意力层。我们独立于其下游用途实现了ParallelLinear的反向传播，使其可以作为构建其他专家的原语。

算法 1 ParallelLinear 前向传播

3.2.1 反向传播

梯度计算。在一个典型的批处理线性变换$Y = XW$中，我们需要计算梯度$∇X$和$∇W$。

$$\nabla \mathbf{X} = \nabla \mathbf{Y} \mathbf{W}^{\top}, \quad \nabla \mathbf{W} = \mathbf{X}^{\top} \nabla \mathbf{Y},$$

ParallelLinear中的梯度计算。在ParallelLinear中，我们将需要为E个专家中的每一个计算这些梯度。虽然这可以在$X$和$∇Y$都是分散的情况下计算，但当$X$和$∇Y$都是分组的时，该操作的实现速度最快。

算法 2 ParallelLinear 反向传播

内存优化策略。算法2在嵌入未分组时对其进行分组，以便计算$∇W$。groupXTY是实现这种分组矩阵乘法的核函数。虽然这种分组会为可能非常大的矩阵带来额外的内存分配，但这些数组分配可以被重用。一旦梯度$∇p$被计算出来，$Ŷ$的数组就可以用作$∇Ȳ$的输出。$X̄$可以被重用于$∇X̄$，因为它们的维度相同。然后，我们可以通过重用用于分组操作的数组来进一步最小化反向传播期间的内存使用。我们在算法2中分别用蓝色和橙色标记了被重用的数组。

3.2.2 SMoE 多层感知机 (SMoE MLP)

内存占用优化。在SMoE MLP的背景下，我们可以进一步减少内存占用。MLP需要两次线性变换，可以天真地用两个设置为执行scatter-to-scatter变换的ParallelLinear操作来实现。然而，我们可以将这两个线性变换配置为scattered-to-grouped，然后是grouped-to-scattered。这将使得SMoE MLP中的每个ParallelLinear变换在反向传播期间只需要一次分组操作。

算法 3 SMOE 多层感知机

3.3 可扩展性：注意力混合（Mixture-of-Attention, MoA）

MoA的背景与挑战。已经有几种将SMoE应用于注意力层的提议【【索引19，Mixture of attention heads: Selecting attention heads per token，2022，arXiv】；【索引2，Switchhead: Accelerating transformers with ´ mixture-of-experts attention，2023，arXiv】；【索引17，Sparse universal transformer，2023，arXiv】】。无论采用何种形式，在应用注意力层之前，嵌入都应按时间顺序排列（即分散状态），以便应用位置嵌入并计算注意力权重和值嵌入的结果。使用现有的SMoE实现，这将需要额外的一对分组-分散操作，从而产生额外的内存成本。

算法 4 SMOE 多头注意力

ScatterMoE的优势。ScatterMoE提供了一个优势。由于我们可以通过ParallelLinear变换保持分散的顺序，因此我们可以实现MoA，而无需为分组和分散分配额外的数组。图3展示了用于SMoE注意力的操作。在本报告中，我们具体实现并基准测试了在【【索引17，Sparse universal transformer，2023，arXiv】】中发现的多头注意力混合（Mixture of Multi-head Attention）变体。

MoMHA的实现细节。这种实现类似于分组查询注意力（Grouped-query Attention, GQA）【【索引1，Gqa: Training generalized multi-query transformer models from multi-head checkpoints，2023，arXiv】】，其中每个键（key）头有多个可能的查询（query）头。而在SMoE设置中，将有$h_{expert}$个键头，并从可能的$E \cdot h_{expert}$个头中选择$k \cdot h_{expert}$个查询头，其中$h_{expert}$是每个专家的头数。

标准注意力实现约定。遵循标准的注意力实现约定，我们设置$d_{out} = d_{expert} \cdot d_{head}$，其中$d_{head}$是每个头的维度。然后在执行注意力操作时，我们相应地重塑张量，以使独立的头能够独立交互。

图 3: ParallelLinear允许从分散到分散的转换，从而保持了时间顺序。

算法4的说明。在算法4中，我们还考虑了时间维度，因此我们将输入和中间数组表示为具有批次、时间和嵌入维度（$B \times T \times d$）的三维张量。在实践中，我们假设输入是连续的并且是按批次-时间排序的，这使我们能够展平张量并像在MLP的情况下一样继续进行。请注意，对于SMoE注意力，一个关键的区别是，我们要求ParallelLinear在第一次变换后给出未分组的输出，并且它在输出变换时接受未分组的输入，这意味着两个ParallelLinear变换都使用scattered-to-scattered配置（图2c）。

A4 实验环境

模型架构

整体测试模型: 使用一个约15亿参数的Mixtral【【索引7，Mixtral of experts，2024，arXiv】】配置进行整体集成测试。其关键参数如下：

$$d_{\text {model }}=1024, \quad d_{\text {expert }}=3584, \quad k=2, \quad E=8, \quad L=16,$$

SMoE MLP单元测试模型:

$$d_{\text {model }}=4096, \quad d_{\mathrm{ff}}=2 d_{\text {model }}, \quad d_{\text {expert }}=d_{\mathrm{ff}} / k, \quad E=8 k$$

* MoMHA单元测试模型:

$$ d_{\text{model}} = 4096, \quad d_{\text{head}} = 128, \quad T = 16 \cdot 2048, \quad h = 32, \quad h_{\text{expert}} = h/k, \quad E = 8k $$

硬件配置

多GPU: 在同一计算节点上的8个Nvidia A100 GPU。
单GPU: 单元测试在单个Nvidia A100 GPU上进行。

软件配置

实现与依赖: ScatterMoE的核心部分使用Triton【【索引18，Triton: an intermediate language and compiler for tiled neural network computations，2019，Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages】】实现，基于PyTorch【【索引13，Pytorch: An imperative style, high-performance deep learning library，2019，Advances in neural information processing systems】】。
对比基线:
1. HuggingFace的朴素实现 (Naive HF impl.)
2. Megablocks【【索引5，MegaBlocks: Efficient Sparse Training with Mixture-of-Experts，2023，Proceedings of Machine Learning and Systems】】的稀疏实现 (MB (Sparse))
3. Megablocks的内存高效分组实现 (MB (Mem. eff.))
4. 一个更高效的PyTorch实现，来自【【索引17，Sparse universal transformer，2023，arXiv】】（用于单元测试）。

数据集与评估

训练设置: 对于整体测试，训练进行了100个更新步骤，有效批量大小为256，每个实例2048个令牌。
评估基准: 使用LM Evaluation Harness【【索引6，A framework for few-shot language model evaluation，2023，https://zenodo.org/records/10256836】】在多个基准上进行评估，包括winogrande, sciq, race, piqa, openbookqa, hellaswag, copa, boolq, arc_easy, arc_challenge（评估准确率）和wikitext（评估困惑度）。

A5 实验结果

图 4: (a) 1.5B模型训练吞吐量 (b) SMoE MLP单元吞吐量 (c) SMoE MLP单元内存使用

整体训练性能

实验内容: 在一个约1.5B参数的Mixtral配置下，比较ScatterMoE与HuggingFace朴素实现、Megablocks（稀疏和内存高效模式）的训练吞吐量。
实验结果: 如图4a所示，ScatterMoE的吞吐量优于所有基线，特别是在此设置下比稀疏Megablocks实现高出38.1%。
分析结论: 结果表明，较小的内存占用对于提升整体吞吐量至关重要。但在更大的模型维度和等效批量大小下，这种优势可能不那么显著。

SMoE MLP 单元基准测试

实验内容: 在单个A100 GPU上对SMoE MLP模块进行单元测试，比较ScatterMoE、Megablocks和高效PyTorch实现的吞吐量和内存使用。
实验结果: 如图4b所示，ScatterMoE在训练（“Both”）和推理（“Forward”）中的吞吐量略高于Megablocks。在内存消耗方面（图4c），ScatterMoE的优势巨大：训练中内存使用量为Megablocks的66.2%，推理中仅为53.6%。
分析结论: ScatterMoE在吞吐量上略有优势，在内存效率上优势显著，尤其是在推理场景下。

图 5: 在固定活跃参数和总参数的情况下，增加k和E。我们发现我们的实现随着k的增加扩展性更好。推理粒度扩展性能。如果我们只考虑前向传播，相对吞吐量的差异会更大。

粒度与吞吐量关系

实验内容: 根据【【索引10，Scaling laws for fine-grained mixture of experts，2024，arXiv】】中定义的粒度（granularity）$G = d_{ff} / d_{expert}$，在固定总活跃参数（$d_{ff}$）的情况下，通过改变$k$和$E$来测试不同粒度对吞吐量的影响。

$$G = \frac{d_{\text{ff}}}{d_{\text{expert}}}$$

实验结果: 如图5所示，ScatterMoE随着粒度$G$的增加表现出更好的吞吐量扩展性。与Megablocks相比，性能差距随着$G$的增大而拉大，尤其是在推理场景下。
分析结论: 这可能是因为随着专家数量$E$的增加，Megablocks需要引入更多的零填充。研究表明更高的$G$值对SMoE模型有利，而我们的实现似乎非常适合这种高粒度设置，特别是在批处理推理中。

图 6: 随着我们降低稀疏性（增加k）的相对吞吐量曲线

稀疏性降低的影响

实验内容: 通过不断增加选中的专家数量$k$（最高至30，受限于Megablocks的设备内存），同时保持总专家数$E=64$不变，来测试降低稀疏性对吞吐量的影响，并与一个等效参数量的全密集模型进行比较。
实验结果: 如图6所示，ScatterMoE的吞吐量始终略优于Megablocks。两者都比等效参数的全密集MLP更高效。然而，当$k=30$时，吞吐量已接近密集模型的水平。
分析结论: SMoE在非极端稀疏的情况下仍然比密集模型有性能优势，而ScatterMoE在此场景下表现更佳。

图 8: MoMHA实现的粒度增加曲线。在这种情况下，由于专家之间共享键和值向量，活跃参数基线的吞吐量是变化的。

注意力混合（MoMHA）性能

实验内容: 对本文实现的MoMHA变体进行基准测试，并与使用Megablocks“密集”配置实现的基线进行比较。同样测试了不同粒度下的性能。
实验结果: 对于$k=8$的设置，ScatterMoE的推理吞吐量比Megablocks高出24.0%（图8a）。此外，如图8b和8c所示，随着粒度的增加（即每个专家的头数$h_{expert}$变小），ScatterMoE与Megablocks之间的性能差距也随之增大。
分析结论: ScatterMoE在MoA场景下同样具有优势，尤其适合高粒度的设置。

Mixtral 推理结果一致性验证

实验内容: 将Mixtral 8x7B模型转换为使用ScatterMoE，并在LM Evaluation Harness上运行多个基准测试，与Hugging Face的朴素实现进行比较。
实验结果: 如表1所示，ScatterMoE实现与Hugging Face实现在各项基准上的结果差异极小，绝对误差可忽略不计。
分析结论: 这证明了ScatterMoE是稀疏MoE的一个正确实现，其数值精度与标准实现一致。

表 1: 语言模型评估框架结果比较：Hugging Face 与 ScatterMoE 实现。两种实现之间的结果差异可以忽略不计。
daae0a80a2cefe4be44bd04c92664ea20d5424233e48e50696f52d90b8b6dd8d.jpg

A6 结论与局限性

结论: ScatterMoE是一种在Triton中实现的SMoE，与现有解决方案相比，它减少了内存占用并提供了稍高的GPU吞吐量。我们还设计了ScatterMoE以使用ParallelLinear原语，我们设想该模块可以被扩展，用于构建其他需要分组或分散线性变换的SMoE风格模块。目前，我们已经提供了基于ScatterMoE的MLP和Attention层的实现，希望这将有助于未来基于专家混合模型的实现，并作为将SMoE概念扩展到其他线性变换专家变体的实例。

局限性与未来工作: 目前，ScatterMoE没有实现用于加速解码的专门内核，并且在多节点环境中的并行化还需要进一步的工作。这些额外功能将在未来的迭代中添加，我们相信我们和开源社区的进一步测试将解决任何剩余的错误和未优化的性能问题。