论文探索：TUTEL Adaptive Mixture-of-Experts at Scale

TUTEL: Adaptive Mixture-of-Experts at Scale

TUTEL:大规模专家自适应混合

前置知识：

moe是什么 Mixture-of-Experts (MoE) 经典论文一览 - 知乎 (zhihu.com)

分布式并行的逻辑：数据并行，模型并行、流水线并行（属于模型并行）

AI系统通信原语：all2all，all reduce 分布式训练硬核技术——通信原语 - 知乎 (zhihu.com)

摘要：

混合专家(MoE)已经成为一种很有前途的深度学习技术，它可以将模型容量扩展到万亿多个参数，同时通过稀疏计算降低计算成本。虽然MoE开辟了一个超大型模型的新领域，但由于MoE的动态特性与系统的静态并行/流水线之间的不匹配，它在数千个gpu上的实现受到了限制。

TUTEL，这是一种针对 MoE 的高度可扩展堆栈设计和实现，具有动态自适应并行性和流水线。 TUTEL 在运行时提供自适应并行切换和自适应流水线。分别实现高达1.74×和2.00×的单MoE层加速。我们还提出了一种新的二维层次算法，用于MoE通信加速，其性能超过了之前最先进的2048个gpu，最高可达20.7×。

实验是在swing-T moe上做目标检测的实验。在效率方面，TUTEL加速了SwinV2-MoE，在训练和推理方面分别比Fairseq提高了1.55倍和2.11倍。在有效性方面，SwinV2-MoE模型在预训练和下游计算机视觉任务(如COCO对象检测)中都比对应的密集模型具有更高的准确性，这表明TUTEL已经为端到端现实世界模型训练和推理做好了准备

f1:在Swin Transformer V2[22,23]的MoE版本的完整端到端训练过程中，动态改变MoE层的工作负载。y轴是运行时所需的专家能力，相当于规范化的工作量.（证明工作负载在动态变化）

1 intro

为什么要用moe:注册更多的DNN模型参数是提高ML算法性能最直接但不太复杂的方法之一[15]。然而，DNN模型的容量往往受到计算资源和能量成本[39]的限制。根本原因是dnn的密集架构，其中计算成本通常与参数数量成线性比例。（扩充参数量可以提升模型性能，但消耗算力）混合专家 (MoE) 被采用到 DNN 中，它通过使用称为专家的多个并行子模型引入了稀疏架构，其中每个输入仅转发给基于智能门函数的少数专家.与密集层不同，这种方法只需要一点点额外成本就可以扩大模型容量（因此模型精度更高），因为它无需额外计算即可注册更多模型参数（即专家）。

问题：MOE的动态性质，每个MoE层由一定数量的并行专家组成，这些并行专家分布在加速器(本工作中的GPU)上，每个GPU根据智能门控函数将每个输入数据分发给几个最合适的专家，并获得相应的输出返回并组合它们。这意味着每个专家的工作量基本上是不确定的——它取决于输入数据和门控函数。在实践中，它们都在每次迭代中发生变化。在我们的实验中，单次训练的工作量变化高达4.38×**(见图1**)，不同的层也有很大的差异。

挑战：现有的DL系统，包括最新的MoE框架，大多基于静态运行时执行，不适合动态MoE特征。这个陷阱来自于专家经常不能利用最佳并行性和流水线策略，因为最佳并行性和流水线策略会随着动态工作负载的不同而不同。在运行时动态调整并行度和流水线并非易事，因为在现有系统中，这通常会导致大量运行时开销或GPU内存消耗。此外，计算端优化高度依赖于通信端优化，通信端优化需要同时进行。（AI系统通信可以参考：ZOMI酱的个人空间_哔哩哔哩_bilibili）

TUTEL:是一种 MoE 系统，可通过自适应方法针对任何规模的动态 MoE 工作负载全面优化 MoE 层性能。该机制包括两个关键技术：用于高效 MoE 调度/组合的自适应流水线和用于高效专家执行的自适应并行切换。TUTEL还引入了一种新颖的二维分层(2DH) All-to-All算法和灵活的All-to-All算法，在超大规模(4096个A100 gpu)中实现高效的MoE调度/组合。

工作：

•详细分析MoE的动态特性以及现有ML框架中的以下挑战。（intro）

•提出自适应并行切换和自适应流水线，有效处理MoE的动态工作负载，在单个MoE层上分别实现高达1.74×和2.00×的加速。

•提出了新颖的2DH All-to-All算法、灵活的All-to-All算法

•TUTEL已被用于实现和运行最先进的视觉模型SwinV2MoE的稀疏MoE版本，用于现实世界的计算机视觉问题。与之前的框架(如Fairseq)相比，它分别实现了高达1.55×和2.11×的训练和推理加速。我们还证明了稀疏模型比对应的密集模型具有更高的准确性，表明TUTEL在训练现实世界AI模型方面的准备就绪。

Mixture-ofExperts的动态特性及其在大规模训练中的低效性

通过放置一个跨 GPU 层来应用于大规模分布式 DNN 模型 [10、20、36]，该层可部分交换来自不同 GPU 的隐藏特征，通常称为 MoE 层。

2 Background & Motivation

本节介绍了moe的动态性及其在大规模训练中的低效性。

深度学习专家混合。mix -of experts (MoE) 是一个使用多个专家模型的ML概念，这些专家模型分别处理各自的专门子任务，从而共同解决整个任务。虽然MoE已被许多经典ML算法[44]广泛采用，但它的概念最近被应用于大规模分布式DNN模型 ，通过放置一个跨gpu层，部分交换来自不同gpu的隐藏特征，这通常被称为MoE层。图2给出了一个例子。首先，它运行一个门控函数[19,37,43]，在接下来的all-to-all集合通信(all-to-all)中确定每个输入token的目标GPU。在All-to-All(称为dispatch)之后，每个GPU运行自己的expert，这是一个前馈网络层(fflayer)，然后进行第二次All-to-All(称为combine)，将每个令牌的对应输出发送到令牌所在的GPU。门控函数和fflayer延迟的细节取决于模型算法。

**MoE是超大规模深度学习的关键。**MoE与现有的dnn放大方法(即增加dnn的深度或宽度)不同，其成本低。具体来说，在MoE层中注册更多的模型参数(专家)并不会增加每个令牌的计算成本。如今，MoE被认为是超大规模DL的关键技术，其最先进的结果在前人的工作中得到了证明[9,10,18,36]。

MoE 的动态工作负载。 MoE 动态工作负载的根本原因有两个。

1）令牌路由机制。 MoE 层将每个令牌动态地路由给多个专家，并且令牌在专家之间的分布通常是不均匀的。这使得每个专家的工作量在每次迭代时动态变化，

2）动态的专家能力。每个专家的工作量都受到专家能力的限制。f在训练过程中是动态调整的，这有助于动态工作负载-当token分布不均匀/均匀时，它会增加/减少

2.2 静态并行：

在大规模训练中，与 GPU 的数量相比，MoE 层通常使用相对较少的专家，并且多个 GPU 被分配给一个专家以获得更高的吞吐量。*自适应并行的必要性*

静态并行方法在动态工作负载下并不总是有效工作。我们使用优化的并行性 P1 和 P2 运行单个 MoE 层（详见第 3.2 节）。(P1和P2是什么：见3.2)

f3:两种不同并行方法的运行时首选项。该图比较了不同容量因子f(即不同工作量)和不同top-k配置下的吞吐量，其中高于1.0意味着P2优于P1，

自适应并行切换在现有系统中很困难。不同的并行策略需要在每个 GPU 中存储不同的模型参数集。因此，在运行时切换并行性会导致参数迁移的大量运行时开销。此外，不同的并行策略可能在访问输入、同步专家梯度等方面有自己的限制，这也使它们与其他策略的切换变得复杂**.难的不是确定是否切换，而是因切换而产生的参数迁移量，以及输入输出限制**

2.3 静态流水线

图2所示的MoE层通常没有充分利用gpu，因为它们运行All-to-All和fflayer，按顺序分派和组合。由于All-to-All主要由gpu之间的数据拷贝组成，这些数据拷贝不是计算密集型的，我们可以通过将其与运行数值计算的fflayer管道连接来更好地利用gpu的计算能力。

我们观察到用于调度和合并的静态流水线策略，即静态All-to-All算法和流水线度，在处理动态工作负载时效率较低。如图5所示，根据不同的MoE设置和规模，相应的最优流水线策略由各种All-to-All算法和流水线度组成。这意味着在不同的MoE设置和规模下，单一的静态策略并不总能达到最佳性能，需要在运行时采用动态流水线策略来适应不同的设置。

更糟糕的是，计算和通信之间的干扰使我们很难找到最优的流水线策略，如果我们只考虑每个单独的方面。这是因为在同一GPU上同时运行NCCL内核和计算内核所带来的减速很难估计。在我们的大量实验中，即使两种不同的All-to-All算法具有相似的吞吐量，当引入相同的并发计算内核时，它们的吞吐量通常会有很大差异，并且任何一种算法都可能在具体情况下优于另一种算法。这意味着动态调整应与计算和通信相结合。

图5 MOE策略对最佳样本数的影响 什么意思？

2.4 不可扩展的 MoE 调度与合并

All-to-All 是 MoE 中用于调度和组合的关键集体通信原语，观察到现有的 All-to-All 实现在大规模情况下表现不佳。

小规模信息传输：

大多数流行的 DL 框架利用 NCCL，最先进的 GPU 集体通信库的点对点 (P2P) API 来实现线性 All-to - all算法。它在 n 个 GPU 上运行，其中每个 GPU 将其总 S 字节的数据分成 n 个块（每个块为 S/n 字节）并与所有其他 GPU 执行 P2P 通信。当我们向外扩展（更大的 n）时，任意两个 GPU 之间传输的 P2P 块大小 S/n 将变得更小，这很难在大规模上饱和 NVLink 和 HDR InfiniBand 等高速链路。 S 是固定的，仅由模型本身决定。

与根据数据大小和网络拓扑选择不同通信算法的最先进的 all-reduce 实现不同 ，现有 DL 系统中的 All-to-All 一直只使用单一算法，该算法在大工作量和小规模，但在****小工作量和大规模时表现不佳。****这使得 MoE 的通信难以适应动态工作负载，尤其是在基于 MoE 的模型通常针对的大规模环境中。

**不适当的计算布局。**我们观察到现有系统中的可扩展性问题不仅与通信有关，还来自计算。

在图7中，我们用不同数量的gpu测量DeepSpeed MoE层上的纯fflayer计算时间。在2048个GPU上，计算时间增加到30.2 ms，与单个GPU相比，速度下降了11.3倍。经过分析，我们发现性能下降是由于All-to-All原语产生的刚性输出布局。具体地说,在图7中,当GPU的增长的数量从1到2048年,矩阵乘法由fflayer变化从一个(1、∆E 16384米)·W(∆E M V) B(2048年,∆E 8米)·W(∆E M V)(括号表示一个张量的形状,见符号描述在表2)。注意输入的三维张量的巨大差异,有大影响的效率矩阵乘法在GPU -例如,PyTorch计算这些使用bgemm_strided_batched操作,后者的计算吞吐量仅为前者的8.8%。这一发现意味着我们需要关注All-to-All引起的张量布局转换，以实现MoE层的高可扩展性。

3 TUTEL 系统设计

在 3.1 节中，我们设计了 Flexible All-to-All 以确保在不同尺度下不会出现计算下降，并保证自适应优化的鲁棒性。在第 3.2 节和第 3.3 节中，我们介绍了自适应并行切换和自适应流水线，以优化任何规模的动态工作负载的端到端 MoE 性能。最后，在第 3.4 节中，我们提出了二维分层 All-to-All 以实现 exa 级的 MoE 训练。

3.1***Flexible All-to-All****

All-to-All执行的数据交换可以表示为张量布局之间的转换–由MoE调度使用，All-to-All基元将其输入的张量布局从**（E,∆C,M）转换为（W,∆E,∆C,M）**，其形状在不同尺度上有所不同。虽然大多数现有框架直接使用这个输出进行后续计算，但这是在大规模下由于某些形状的专家计算效率低下而减慢MoE吞吐的一个潜在原因。这个w表示的是the world size,具体来说就是有多少需要交换数据的Gpu,

TUTEL通过使用不同的All-to-All接口的抽象来解决这个效率退步的问题，这个接口被称为Flexible All-to-All。

table 3 Flexible All-to-All在原有All-to-All的基础上又扩展了两个参数。除了它的第一个参数是要转换的输入数据外，第二个参数指定要连接的张量维度，第三个参数指定要分割的张量维度。

Flexible All-to-All所扮演的特殊角色是保持输出布局（∆E,C,M）不依赖W。这保证了Flexible All-to-All之后的后续计算在不同尺度下都能保持不变的布局。此外，为了避免动态容量变化引起的计算回归，Flexible All-to-All可以进一步平铺工作负载，并在必要时将计算绑定到一个经过良好调整的大小，例如(∆E,T,M)，其中T是用于将工作负载分割为的恒定平铺大小。

3.2 Adaptive Parallelism Switching 自适应并行切换

希望多种并行在运行时可以切换，以满足不同MoE训练迭代中的最优策略，这就是所谓的可切换并行。

TUTEL专门设计了一对混合并行策略，充分利用了数据和模型并行的优势，最重要的一点是它们可以零成本地动态切换。不同的并行策略被设计成在令牌供给、梯度更新和参数放置方面有相同的偏好，允许它们根据top-k和容量因子的动态变化即时切换到彼此之间。

琐碎的专家+数据并行将GPU分为多个组：a.在每个组内，作为标准的专家并行工作；（几个GPU共同训练一个专家,每个gpu只有一个专家的数据）b.不同的组复制专家副本以数据并行工作，期望每个复制的组为复制的专家产生一个子梯度（每个GPU上训练几个专家,一组多个gpu是相同布局）

f4:上角标表示分区（模型并行），下角标表示专家，M1分别在两个gpu上复制每个专家，M2分别在四个gpu上划分每个专家

****可切换的专家+数据并行化。****f4

如图4所示，普通专家+数据并行将gpu划分为多个组:a.在每个组中，它作为标准专家并行;B.不同的组复制专家副本以实现数据并行，期望每个被复制的组为被复制的专家产生一个子梯度。(M1)

每个数据复制组内的All-to-All被融合为一个单一的全局All-to-All。同时，每个GPU只维护ZeRO风格的专家

参数片，其副本格式可以通过全集通信临时访问。

*ZeRO: Memory Optimization Towards Training A Trillion Parameter Models Samyam.*

TUTEL重新设计了这种并行性，不仅为了更好的覆盖率，而且为了可切换性。根据图11，每个数据副本组中的All-to-All被融合为单个全局All-to-All。同时，每个GPU只维护一个ZeRO风格的专家参数切片，其副本格式可以被全集通信临时访问。

这样的设计保证了每个GPU上的专家参数是全局唯一的，而理论通信大小仍然与琐碎的专家+数据并行度保持等效

*可切换专家+模型并行*

可切换专家+模型并行是另一种在运行时与P1对齐的方法。它在专家参数上消耗的通信量较少，但在输入令牌上消耗的通信量较大。因此，当令牌大小相对小于专家参数大小时，P2可以比P1执行得快。

MoE调度需要在All-to-All之前进行本地重复操作，将令牌复制n-sharded次，其中n-sharded代表一个专家被分成多少个片子。在灵活的All-to-All之后，令牌立即被派发到一个合适的格式，以便进行后续的并行操作：每个GPU从All-to-All获得部分令牌的集合，可以以张量并行的方式与本地专家片进行计算。在MoE组合中，每个设备上的结果将使用全局的All-to-All来收集，然后进行局部的总和还原。这种方法需要Tmodel = O(n-sharded * ∆E*C *M)的通信规模。

内联平行性路由器 根据令牌数量和专家数量的规模差选择p1或p2，

自适应流水线：

首先，我们介绍了我们划分令牌的方法，以实现多流计算通信管道化。其次，我们展示了在线算法，以搜索最佳的流水线策略，以适应不同的MoE模型设置和动态运行时间能力。

先，我们介绍了我们的分区令牌方法，以实现多流计算通信流水线。其次，我们演示了在线算法来搜索最佳流水线策略以适应不同的 MoE 模型设置和动态运行时容量。

用于多流流水线的令牌分区：fig14

在前向传递中，在每个GPU上，形状为(E,∆C,M)的输入被沿维度C分割成两个虚拟部分.分割之后，每个虚拟分区Ci被异步地发送到通信流上，按照i的顺序执行All-to-All操作。All-to-All被定制为接受被分离的数据块作为输入，并执行在线数据洗牌，生成形状为(ΔE,C/2,M)的输出。

接下来，两个All-to-All的输出被编程为，一旦它们之前对应的All-to-All完成，就被发送到计算流上执行专家计算，而专家计算的输出被再次编程为，一旦之前对应的专家计算完成，就被发送到通信流上执行第二个All-to-All。最后，在第二个All-to-Alls之后设置一个障碍，在障碍之后，分区被合并以产生形状为（E,∆C,M）的最终输出。

这样的分布方式为什么可以达到效果？提升在哪

所有的分区和重塑操作都是通过自定义操作内联完成的。因此，与不重叠的情况相比，没有额外的数据复制开销。

在线流水线策略搜索，f是某种系数，解集空间大。基于一个直觉的算法2：两个接近的f应该有类似的工作负载模式，最佳的流水线策略应该是类似的→同一个桶里所有的fs只有一个策略：进了同一个桶的f是相似的f，相似的f有相似的开销，降低解空间的大小。

1）根据当前的f获得策略s（GETSTRATEGY）；2）用s和f运行MoE工作负载，并对其开销进行采样，即t；3）用s和t更新尝试过的策略备忘录的运行时间。在第一阶段，该算法将首先尝试为f分配一个桶，如果还没有。之后，它将在f已经尝试过的策略中挑选一个最优策略作为第一优先级，然后从f的桶中挑选第二优先级。如果两者都不是，就会让f尝试一些尚未尝试过的桶的新策略，并在以后通过OPTIMIZESTRATEGY把性能带回来。这就保证了每个桶都会以最小的重复率检查所有的策略。

3.4 2DH A2A

为了解决第2.4节中的低效率问题（小规模不可拓展的两次ALL2ALL），我们的方法是将从多个本地GPU发送到同一个远程GPU的多个数据块聚合起来。这就避免了通过网络发送多个小消息，将小块数据合并为一个大块数据，从而大大改善了链路带宽的利用率。

挑战：S是固定的数据字节数，节点中的所有m个GPU数量，n是S个字节的数据分配到的GPU数量，每个gpu占s/n个，节点内延迟由于GPU上有n/m次非连续的内存访问，随着n的扩展而增加。

算法：（2DH）全对全算法，将多个连接上的小信息合并成一个单一连接上的大信息。

为了避免由于非连续的内存访问造成的速度下降，2DH All-to-All包括一些额外的阶段，这些阶段进行有效的跨步内存拷贝，将非连续的块对齐到一个连续的地址空间。

首先通过stride内存拷贝来对齐共享相同的本地目标GPU的块（阶段1），然后进行节点内的All-to-All（阶段2）。在接下来的阶段，我们再次对齐共享相同的远程目标GPU的块（阶段3），然后最后进行节点间All-to-All（阶段4）。通过利用stride内存拷贝，2DH All-to-All实现了较高的内存带宽利用率，无论前三个阶段的n如何，都能保持恒定的低延迟。与现有算法相比，2DH All-to-All的优势随着S/n的变小而增加（数据大小S较小或GPU数量n较大）。

具有SIMT效率的快速编码和解码

SIMT：二者折中方案，单指令多线程，线程内部执行相同指令，但比SIMD更灵活，比SMT效率更高.

稀疏版本的时间复杂度低，但是稀疏运算无发利用各种tensor加速器。

→3种快速编译码运算器，通过调用各种只适用于密集计算的优化实现稀疏计算优化

4 实现

4.1 特征

动态TOP-K MOE门控网络：

TUTEL支持大多数现有框架不支持的topANY路由。

k值也可以在运行时动态更改，即一个MoE层的不同迭代可以使用它们首选的top-k设置，而不是使用相同的k值。用户可以利用这个特性动态地微调MoE层的稀疏性。

动态容量因子：

TUTEL支持在每次迭代中动态调整容量因子。如图16所示，调整行为由传递给我们的MoE层API的capacity_ f actor = x参数控制。

如果x为正数，则直接将其作为MoE层的容量因子。如果x为零，TUTEL自动将容量因子调整为在每次迭代中不删除任何令牌的最小值。当x为负数时，TUTEL进行与x为零时相同的自动自适应，只是将- x设置为容量因子的上限，即任何超过的值都将自适应为- x。

并行分布的 控制：

TUTEL自动适应第3.2节中描述的并行性，它还支持用户对并行性的透明控制。

TUTEL为控件count_per_node = x提供了一个易于使用的接口，它是传递给MoE层API的参数。

即使只有一个参数，TUTEL也会自动相应地构造一个所需的专家分布，它涵盖了不同的专家分布，包括单个GPU中的多个专家或每个专家横跨多个GPU。如图17所示，如果x是正整数，每个GPU管理x个本地专家。否则，每个专家都是并行在−x个GPU上，每个GPU处理单个专家的全部输入数据的1/(−x)。Count_per_node只处理吞吐量，同时保持训练算法不变。

**不同的设备和数据类型。**TUTEL支持不同的计算设备，包括NVIDA GPU, AMD GPU和CPU。它还支持不同设备上的所有原生数据类型:FP64/FP32用于CPU后端，FP64/FP32/FP16/BF16用于GPU后端。

兼容任何DNN架构。虽然本文仅使用基于transformer的模型架构来评估TUTEL，但它可以与任何其他DNN架构一起使用，如MLP、CNN或LSTM。

5.3.4 TUTEL支持的一种新的余弦路由器

通过TUTEL，我们提供了更多的MoE基线来丰富算法选择，并举例说明如何利用这个框架进行算法创新。一个尝试是新的余弦路由器，希望通过增加模型尺寸来提高数值稳定性，每一列代表一位专家;τ是一个可学习的温度，它被设置为最低0.01，以避免温度太小;P为选择专家的路由得分。

我们在表13中的初步实验表明，当使用32位专家时，余弦路由器在图像分类中与普通线性路由器一样准确。虽然目前它在图像分类方面并不优越，但我们仍然鼓励TUTEL用户在他们的问题中尝试这个选项.