FasterMoE: Modeling and Optimizing Training of Large-Scale Dynamic Pre-Trained Models

文章标题: FasterMoE: 建模与优化大规模动态预训练模型的训练
作者: Jiaao He, Jidong Zhai, Tiago Antunes, Haojie Wang, Fuwen Luo, Shangfeng Shi, Qin Li
机构: 清华大学

A1 主要贡献

本文旨在解决训练大规模Mixture-of-Expert (MoE)模型时面临的效率挑战，这些挑战主要源于动态专家选择和灵活的MoE结构。
核心问题:
1. 动态负载不均衡: 由于训练样本对专家的选择存在偏斜且动态变化，导致部分Worker负载过重而其他Worker空闲，严重影响硬件利用率和训练效率。
2. 低效的同步操作: 训练中耗时的all-to-all通信操作通常是同步的，在计算和通信负载不均衡的情况下，这种模式会导致大量资源浪费。将all-to-all操作分解并异步执行存在数据依赖和死锁风险。
3. 模型设计与网络拓扑不匹配: 现有的专家选择策略主要关注计算负载均衡，忽略了通信开销，而专家分配直接决定了通信流量模式。在当前广泛使用的网络拓扑中，MoE复杂的通信模式常引发网络拥塞。

研究目标与创新点:
为了应对上述挑战，本文提出了FasterMoE，一个高效的分布式系统，用于训练大型动态预训练模型。其主要贡献如下：
* 设计了一个性能模型：该模型能准确估计给定MoE模型在特定并行策略下的性能。
* 提出了一个DDL-Roofline模型：这是一个类似Roofline的可视化模型，用于分析不同并行策略的性能和理论上限，并展示本文优化措施带来的改进。
* 发明了一种动态影分身（dynamic shadowing）方法：通过在运行时有选择地复制热门专家模型参数而非传输其输入数据，来应对专家受欢迎程度偏斜导致的负载不均衡问题。
* 创建了一个智能的细粒度调度策略：将通信和计算任务分解，并进行异步调度，以减少它们的总延迟。
* 设计了一种调整后的专家选择策略：该策略能在运行时避免网络拥塞，从而以更快的通信速度换取更快的迭代，同时保证模型损失以理想的斜率下降。
* 实现并集成了一个端到端的MoE训练系统FasterMoE：该系统集成了上述技术，并在与现有最先进系统的对比中取得了高达17.87倍的加速比。

A3 背景知识与关键挑战

2.1 Transformer：预训练模型的骨干

Transformer的基本结构。Transformer【索引34，Attention is all you need，2017，Advances in neural information processing systems】是处理序列的先进结构，构成了许多预训练模型的基础，这些模型处理文本、蛋白质甚至图像像素等序列数据。如图2所示，一个Transformer块由两部分组成：注意力（attention）层和多层感知机（MLP）。

注意力层与MLP层的功能。注意力层通过在特定线性空间中对每对令牌（token）进行点积运算来提取序列中令牌的关系，结果形成的注意力矩阵用于加权求和不同令牌的嵌入向量。之后，结果被送入一个MLP层，该层通常由两个巨大的全连接（FC）层组成。Transformer块中计算最耗时的部分是MLP层中的通用矩阵乘法（GeMM）。

2.2 MoE结构

MoE的核心思想。MoE（Mixture-of-Experts）【索引17，Mixture of experts: a literature survey，2014，Artificial Intelligence Review】在现代大规模预训练模型中显示出强大的能力。其核心思想是，一个大模型由许多小的模型（即专家）构成，直觉上不同的小模型是不同领域的专家，只有当其领域内的数据输入时才被激活。

MoE在Transformer中的应用。在Transformer中，MLP层通常通过MoE进行扩展。处理一个令牌时，只有少数最适合其领域的专家被激活。在非MoE模型中，MLP中的两个相邻FC层在模型规模扩大时会变得巨大，导致GeMM计算过于繁重。而在MoE模型中，GeMM的权重矩阵沿特定维度被分割，每个部分仍然产生相同大小的输出，但GeMM计算量保持较小。这意味着MoE允许在不增加计算量的情况下增加模型参数，使其成为目前生产万亿级及以上规模预训练模型最可行的方法。

门控网络（Gate）的作用。对于给定的输入，一个名为“门”（gate）的附加模块被引入，以决定哪些专家应该被激活。门通常是一个小型的FC层，用于为每个专家计算一个拟合分数，并选择得分最高的k个专家。

2.3 并行策略

三种常用的并行策略。数据并行、模型并行和专家并行是分布式训练中三种常用的并行策略。

数据并行。数据并行在所有工作节点（worker）上复制模型参数，每个工作节点处理不同批次的训练样本。迭代后，工作节点全局同步梯度并更新模型。虽然迭代内没有通信，但模型大小不能超过单个工作节点的容量，因此无法扩展到大型模型。

模型并行。模型并行沿特定维度划分权重张量，即将模型分割成多个分区放置在不同工作节点上。所有工作节点共同处理全局批次，并使用其对应的权重分区进行计算。每层之后，嵌入向量需要被聚合和重新分发。然而，由于分区维度的限制和层间巨大的通信开销，模型并行也难以高效地扩展到非常大的模型。

专家并行。专家并行是GShard【索引11，Gshard: Scaling giant models with conditional computation and automatic sharding，2020，arXiv preprint arXiv:2006.16668】首次为MoE模型提出的一种特定并行方法。如图3所示，专家被放置在不同的工作节点上，每个工作节点处理不同批次的训练样本。对于非MoE层，专家并行的行为与数据并行相同。在MoE层中，序列中的令牌被发送到其所需专家所在的节点。与模型并行类似，每个MoE层的输出需要再次交换，以便重新组织成原始序列以供下一层计算。由于MoE模型通常有大量专家，专家并行比模型并行更能随模型规模扩展。

2.4 挑战与观察

专家并行训练中的主要挑战。在使用专家并行训练Transformer时，一系列挑战极大地影响了训练效率。

偏斜的专家选择导致动态负载不均衡。如图3所示，专家0接收3个令牌，是专家2工作量的3倍，导致工作节点2在下一次通信开始前长时间空闲，未能充分利用其计算能力。训练数据天然的偏斜分布使得一些专家比其他专家更容易被选中。

专家选择分布的实际观察。我们收集了训练两个包含16个专家的真实模型时每个令牌的专家选择情况。图4展示了训练过程中部分迭代的采样。在前500次迭代中观察到快速变化的非均匀分布。在图4a所示的MoE层中，专家的受欢迎程度在整个训练过程中不断变化。图4b展示了另一个模型中的不同层，其受欢迎程度更稳定，但仍有许多不受欢迎的专家。实际上，放大图可以看到许多微小的条纹，表明这些专家不受欢迎，但仍在处理其领域特定的数据。同时，16个专家中有4个处理了约20%的令牌，是平均水平的3.2倍。

负载不均衡的后果。更受欢迎的专家接收到比不受欢迎的专家更多的令牌，导致它们所在的节点负载更重。这种动态行为影响硬件利用率并降低模型训练效率。因此，MoE训练系统面临的第一个挑战是处理由偏斜的专家选择引起的动态负载不均衡。

同步执行模式效率低下。专家并行中的all-to-all操作通常由通信库（如MPI【索引6，Using MPI: portable parallel programming with the message-passing interface，1999，MIT press】或NCCL【索引8，Optimized inter-GPU collective operations with NCCL 2，2017】）提供的同步操作符实现。考虑到非均匀的专家选择导致计算和通信的不均衡，这种同步执行方法会造成更大的资源浪费。当执行通信或计算时，另一部分硬件处于未充分利用状态。然而，由于不同通信和计算任务之间存在依赖关系，分解all-to-all通信并不容易，若数据传输顺序设计不当，极易引入死锁。因此，第二个挑战是如何高效地组织通信和计算任务以并行执行。

专家并行导致严重的网络拥塞。最后，我们强调专家分配与网络拓扑之间的不兼容性。在每次迭代中，多个通信操作同时执行，可能因少数链路饱和而导致严重的性能下降。由于令牌的专家分配决定了负载均衡和通信路径，执行智能的令牌分配有助于降低端到端的训练延迟，而不影响模型质量。因此，第三个挑战是如何设计一种感知网络拓扑的令牌分配策略，以避免严重的网络拥塞。

A2 方法细节

3 性能建模

为了估计和分析训练任务的性能，我们首先分别为计算和通信建立模型，然后引入一个类似Roofline的模型来研究通信延迟和计算延迟如何共同决定整体训练效率。模型中使用的符号定义在表1中。

表1. 模型中使用的符号

3.1 感知负载的计算建模

GeMM作为主要计算。训练Transformer时的主要计算是GeMM。现代大规模计算设备（如GPU）对GeMM等规则计算进行了高度优化，能达到很高的性能。根据我们的测量，NVIDIA Tesla V100 GPU在运行Transformer中典型模型尺寸和批次大小时，其GeMM性能可达到峰值吞吐量的90%以上。因此，我们通过以下公式预测Transformer块中MLP层前向传播的计算延迟。

公式解释与负载不均衡的体现。公式中，$B_i$ 是工作节点 $i$ 上的批次大小，因为在专家并行中不同节点的模块批次大小可能不同。$h$ 是令牌嵌入向量的长度，$h_{ff}$ 是MLP中FC层之间的中间嵌入长度。由于单个FMA操作计为2次运算，每个FC执行需要 $2Bh h_{ff}$ 次操作。共有2个FC层，因此常数因子为4。$P_i$ 是工作节点 $i$ 执行GeMM的平均吞吞吐量。端到端延迟是每个单个工作节点延迟的最大值，因为所有工作节点在计算后必须交换特征。因此，该公式反映了计算中的负载不均衡。

小批量计算的潜在问题。一个潜在问题是，对于 $B_i$ 非常小的工作节点，其计算设备可能无法达到良好的利用率，导致延迟估计不准确。然而，尽管未达到峰值性能，小 $B_i$ 的计算延迟通常不大于大 $B_i$ 的计算延迟。由于大的 $B_i$ 主导了整体计算延迟，这种预测上的不准确性并不会使其有效性失效。

3.2 拓扑感知的通信建模

通信延迟模型。根据LogP模型【索引2，LogP: Towards a realistic model of parallel computation，1993，Proceedings of the fourth ACM SIGPLAN symposium on Principles and practice of parallel programming】，通信的总延迟由开销（overhead）和延迟（latency）组成。每个令牌的特征向量通常大于1024，意味着数据传输的最小粒度超过4KB。因此，我们简化模型，认为通信开销可忽略不计。假设没有拥塞，互连带宽可以被充分利用。鉴于一个节点内通常有多个加速器（每个都是一个工作节点），我们不仅要考虑节点间连接，还要考虑节点内连接，如PCIe、UPI和NVLink。

采用拓扑感知模型预测延迟。我们采用拓扑感知模型来预测集体通信操作的延迟。假设链路 $e$ 的单向带宽为 $BW_e$，流经该链路的流量大小为 $V_e$。通信的端到端延迟计算如下：

模型参数的获取。$BW_e$ 可以通过硬件规格和点对点带宽基准测试来确定。我们强调网络拓扑图是有向的。为了获得 $V_e$，我们将每个链路建模为图中的一条边。使用两条有向边来表示一个双工链路，分别考虑双向的流量，因为在负载不均衡的情况下，一个链路上两个方向的流量可能差异很大。链路的有效带宽不直接等于两个方向都繁忙时的带宽。

不同通信操作的流量建模。每个链路上的流量取决于算法和路由策略。我们展示了如何对三种常见的通信类型进行建模：
* All-to-all-v: 用于将令牌从其在序列中的位置路由到其所需的专家。由于专家选择的灵活性，每对工作节点之间的流量变化很大。我们假设all-to-all操作在所有工作节点对之间简单地创建链路，并同时传输数据。每对工作节点之间的路径由一个根据拓扑类型计算的算法得出。对于每对工作节点，它们之间的流量会累加到路径上的所有有向边上。
* All-reduce: 广泛用于同步数据，包括数据并行中重复模型参数的梯度和模型并行中的嵌入向量。在一个大小为 $S$ 的张量上对 $N$ 个工作节点应用环形all-reduce【索引30，Horovod: fast and easy distributed deep learning in TensorFlow，2018，arXiv preprint arXiv:1802.05799】会导致每个节点在流水线中向其邻居总共发送 $2\frac{N-1}{N}S$ 的数据。
* Broadcast和reduce: 与all-reduce一样规则，利用环形连接和流水线来降低延迟。但与all-reduce不同，它们在每个链路上只发送总大小为 $S$ 的消息。

3.3 DDL-Roofline模型

DDL-Roofline模型简介。我们提出了一个分布式深度学习（DDL）Roofline模型，用于描述特定训练任务在给定集群上的性能。

X轴：计算-通信比。计算和通信是并行MoE模型中的两个关键因素。因此，我们定义计算-通信比 $\rho$，呈现在DDL-Roofline的X轴上，如下所示：

$\rho$的含义。$T_{comp}$ 和 $T_{comm}$ 分别是我们预测器估计的计算和通信延迟。$\rho$ 表示任务是由计算还是通信所限制。当 $\rho > 1$ 时，计算时间主导端到端延迟；否则，通信占据大部分延迟。这个比率指明了应用不同优化的方向。

Y轴：平均计算吞吐量。Y轴的变量是 $\Theta$，即所有工作节点的平均计算吞吐量。在训练MoE MLP层时，可以按如下方式计算：

$\Theta$的含义。$12h h_{ff} \sum_{i} B_i$ 表示所有专家为所有令牌处理的总计算量，$N$ 是工作节点数量。$T_{e2e}$ 是通过估计或测量得到的单次迭代的端到端延迟。例如，在同步专家并行中，我们估计 $T_{e2e} = 3T_{comp} + 4T_{comm}$，因为前向和后向传播总共有3轮计算和4轮通信。$\Theta$ 直观地反映了所有工作节点设备的平均利用率，也可以直接指示系统的可扩展性。

理论上限（理想情况）。理想情况下，通信和计算同时进行，我们得到一条类似屋顶线的多段线作为理论上限，如图5中的实线所示。其计算方式如下：

半理想情况。我们还用虚线标出了一条半理想曲线，它指的是当训练以同步方式执行时，硬件达到完全利用率的情况。

半理想情况的解释。在半理想情况下，端到端延迟是通信延迟和计算延迟之和。与原始的Roofline模型【索引37，Roofline: an insightful visual performance model for multicore architectures，2009，Commun. ACM】描述单个设备上程序（内存访问和计算自然同时执行）不同，分布式训练程序通常需要对系统进行重大优化才能同时执行它们。

DDL-Roofline的应用示例。给定一个训练任务及其并行配置，DDL-Roofline有助于更好地理解模型的训练吞吐量。下面，我们通过图5中一个特定Transformer模型的训练示例，展示不同并行策略在DDL-Roofline中的反映：
* 数据并行: 在理想多段线的左侧显示为2个点。由于为MoE MLP层同步梯度涉及对 $2h h_{ff}$ 个元素执行all-reduce，代价过高，导致 $\rho$ 值很差。但是，由于all-reduce可能与反向计算重叠，它可以略微移动到半理想曲线之上。
* 模型并行: 具有更大的 $\rho$，因为它引入的通信较少。它对一个包含 $B$ 个令牌的嵌入矩阵执行2次all-reduce，总大小为 $2B h$。与数据并行相比，它减少了通信，使得 $\rho > 1$。但在同步嵌入向量时，无法执行其他计算。这一特性迫使模型并行以同步方式执行，并阻止该点移动到半理想曲线之上。
* 专家并行: 由于负载不均衡，其计算延迟大于通信延迟，因此具有较大的 $\rho$ 但较差的 $\Theta$，远低于半理想曲线。FasterMoE中的优化也在图5中展示，我们将在下一节中说明它们在DDL-Roofline中的特性。

4 模型指导的优化方法

4.1 轻量级动态影分身策略

负载不均衡问题的根源与权衡。在MoE模型中，热门专家可能被超过一半的输入令牌选中，导致严重的负载不均衡，如图4所示。尽管单个输入令牌的嵌入远小于模型参数，但来自所有其他工作节点的成批输入可能等于甚至大于模型参数。因此，这里存在一个权衡：是否可以用复制热门专家的模型参数来替代传输大量嵌入向量的延迟。

动态影分身（shadowing）的直观想法。如图6所示，一些专家被复制到所有工作节点上，称为“影分身专家”（shadowed experts），这样它们的参数（而不是输入令牌）通过网络传输。相关的计算在本地执行，直观上减轻了包含热门专家的工作节点的负载。

动态影分身的挑战。然而，动态影分身具有挑战性，因为专家的受欢迎程度随训练过程而变化，每次迭代的决策可能不同。此外，参数不能像普通的分布式内存系统那样被缓存，因为它们在每次迭代中都会更新，并且需要全局收集梯度作为更新过程的一部分。如果将它们缓存在每个工作节点上，就必须在每次迭代中更新它们，这会引入显著的额外开销。

基于性能模型的决策。为了解决这个挑战，我们利用我们的性能模型在运行时分析一个专家是否应该被影分身。我们预测训练迭代的端到端延迟来检查性能增益，并据此采取行动。我们的分析阐述如下。

原始不均衡情况下的延迟建模。在原始的不均衡情况下，通信和计算由一组热门专家主导。我们通过计算其批次大小 $B_i = \sum_{j=1}^{N} B_{ji}$ 来为一个工作节点 $i$（总共 $N$ 个）建模其工作负载，其中 $B_{ji}$ 个令牌从工作节点 $j$ 发送到工作节点 $i$。

原始延迟公式。训练一个MLP层的单次迭代包含前向传播中的1次GeMM和反向传播中计算输入梯度的2次GeMM。共有4轮all-to-all通信，前向和后向各2轮。在一个简化的网络带宽固定的情况下，训练延迟计算如下：

影分身策略下的操作流程与延迟建模。要对一个热门专家进行影分身，我们必须首先将其参数广播到所有工作节点，然后使用获取的模型对本地令牌进行计算。在反向传播阶段，每个工作节点分别计算其获取的专家的梯度，然后进行一次reduce操作。最后，参数更新操作在热门专家最初所在的那个工作节点上执行。

影分身延迟公式。在这种情况下，执行不均衡计算的开销被替换为对2个大小均为 $h h_{ff}$ 的参数进行的2次集体通信操作。由于多个热门专家被发送到许多其他工作节点，负载不均衡发生的可能性较小。对 $K$ 个模型进行影分身的延迟计算如下：

启用影分身的条件。如果影分身策略更快，预计满足以下任一条件：

或

条件解读与DDL-Roofline中的表现。第一个条件表明，传输输入的总开销高于传输模型的开销。第二个条件表明，减少的计算延迟超过了增加的通信开销。在任何一种情况下，启用动态影分身以减少端到端延迟。否则，通信成本太高，模型交换对减少延迟没有好处，因此不执行。这种情况通常发生在工作负载在不同工作节点之间均衡时。如图5中箭头(1)所示，由于减少了空闲时间，计算延迟缩短，导致 $\rho$ 降低且 $\Theta$ 提高。

运行时选择影分身专家的算法。我们在每次迭代中运行时选择要影分身的专家。一个轻量级算法，如算法1所示，在每个工作节点上执行。由于矩阵 $B$ 必须始终在所有工作节点上可用，因此不会引入额外的通信。它根据上述公式返回一组要影分身的专家。

算法1 选择要影分身的专家

4.2 异步细粒度智能调度

调度优化的动机。如我们的DDL-Roofline所示，当通信和计算分开执行时，程序无法超越半理想曲线。此外，由于固有的通信量大，增加 $\rho$ 很困难。因此，我们提出一种智能调度方法，将任务划分为更小的部分，并以效率为目标重新调度细粒度的通信和计算操作。细粒度调度允许计算和通信异步执行，从而更好地利用硬件，使性能跃过半理想曲线，如图5中箭头(2)所示。

细粒度任务划分。在专家并行中，通信遵循复杂的all-to-all模式。我们首先通过将工作节点划分为细粒度的组来分解all-to-all通信。我们使用分组配对交换算法【索引33，Automatically tuned collective communications，2000，SC’00: Proceedings of the 2000 ACM/IEEE Conference on Supercomputing】来执行all-to-all。这些组形成一个大小为 $G$ 的环，并以从0到 $G-1$ 递增的步长向其他组发送数据。对于分组分配，我们遵循一个启发式规则，即连接紧密的工作节点被放在同一组中，从而使同组工作节点之间的连接更快。组的大小由连接拓扑和相关的计算粒度共同决定。

操作分解示例。从工作节点0的视角看，在 $S_{0,1}$ 中，工作节点0在从组1接收数据的同时向组2发送数据。它在 $S_{0,1}$、 $C_{0,2}$ 和 $S_{0,2}$ 中的行为类似。由于空间限制，它们未被绘制出来。

分解后的操作序列。在MoE层的前向或后向阶段，涉及2次对称的all-to-all通信，中间夹着计算。我们根据配对交换来分解计算，为重新组织它们留出空间。在 $G$ 个步骤中，所有工作节点在 $G$ 个组中执行 $3G$ 个操作。在步骤 $s$ 中，组 $g$ 中的工作节点执行以下3个操作，如图7中的示例所示：
* $S_{g,s}$：向组 $g_{send} = (g-s)$ 发送令牌，并从 $g_{recv} = (g+s)$ 接收令牌（均为模 $G$）。
* $C_{g,s}$：使用本地专家对来自 $g_{recv}$ 的令牌进行计算。
* $S'_{g,s}$：从 $g_{send}$ 接收本地令牌的输出，并将输出发送回 $g_{recv}$。

调度策略。这样，通信和计算都被分解为具有特定顺序要求的细粒度任务。图8a展示了一个忠实地顺序执行这些操作的调度，其延迟等于使用粗粒度操作符的原始同步执行模式的延迟。只要满足依赖关系，这些操作就可以乱序执行。

双流并行调度。调度的目标是并行执行操作。我们为每个工作节点创建一个通信流和一个计算流来执行不同类型的操作符。如图8b所示，在其通信流中，它首先执行 $S_{g,0}, S_{g,1}, ..., S_{g,G-1}$，然后执行从 $S'_{g,0}$ 到 $S'_{g,G-1}$。其计算流执行从 $C_{g,0}$ 到 $C_{g,G-1}$。通过并行执行这些操作，端到端延迟显著减少。然而，所有操作必须遵守其数据依赖关系，并在开始前等待前序任务执行完毕。

最小化开销的智能调度。我们用一个双流调度来说明我们最小化开销的方法。我们假设计算流大部分时间是繁忙的。我们强调，在相反的情况下，即通信占据大部分时间，根据DDL-Roofline，优化效果会很小。由于计算流已完全被占用，优化的主要机会是减少第一个 $S$ 和最后一个 $S'$ 的延迟。图8c给出了一个例子。由于组2比组3引入的开销更低，将其放在调度的最后位置可以降低端到端延迟。注意，$S_{g,0}$ 从本地组接收令牌，这预计是最快的操作，因为没有涉及上层连接。$S_{g,G-1}$ 只与组 $g$ 的邻居交换数据。从全局看，步骤 $G-1$ 中的所有组被组织成一个环，并沿环交换数据。这使得网络带宽在除步骤0外的所有步骤中得到最佳利用。因此，最快的两个操作，即 $S_{g,0}$ 和 $S'_{g,G-1}$，被放置在智能调度的最前和最后，从而最小化了开销。

4.3 避免拥塞的专家选择策略

专家选择的灵活性与目标。在MoE模型中，最终目标是用足够多的输入样本训练所有专家，而不是用它们最期望的专家来处理每个令牌。由于拟合分数被用作权重来聚合每个专家的输出，改变专家的选择不会引入数值上的不正确性。GShard【索引11，Gshard: Scaling giant models with conditional computation and automatic sharding，2020，arXiv preprint arXiv:2006.16668】和BASE Layer【索引12，Base layers: Simplifying training of large, sparse models，2021，arXiv preprint arXiv:2103.16716】都改变了专家选择策略以实现特定目的。我们观察到，专家选择策略可以与训练系统协同设计以获得更好的效率。然而，模型的准确性可能会受到选择策略的影响。专家可能会被喂给与其专业知识不太相关的令牌，使其能力下降。因此，除了吞吐量，令牌和专家之间更好的拟合也是值得追求的。

拓扑感知的门控设计。我们设计了一个拓扑感知的门（gate），将输入导向延迟较低的专家。通过考虑特定硬件的网络拓扑，可以提高训练吞吐量。在常见的具有树状拓扑的集群中，上层连接的带宽通常低于本地连接。与其他常规的集体通信不同，all-to-all会导致这些连接上更高的拥塞。

流量控制机制。假设一个交换机连接 $M$ 个节点，每个节点上有 $N_w$ 个工作节点。一个工作节点与主机之间的流量大约为 $V_{h} = \frac{N_w-1}{N_w}B_i$。同时，每个节点的网络接口上的流量为 $V_{n} = \frac{M(N_w-1)}{N_w}B_i$，大约比 $V_{h}$ 大 $M$ 倍。

拥塞避免策略。为了减少拥塞，我们最多允许 $k = \frac{B_{local}}{B_{net}}$ 的令牌被导向另一个节点。这里，$B_{net}$ 和 $B_{local}$ 分别表示节点间和节点内的通信带宽。具体来说，如果有超过 $k$ 个令牌的最佳匹配选择在另一个节点上，其中得分最高的 $k$ 个被允许过去。其余的令牌与其它令牌一起，在本地节点内重新选择它们所需的专家。跨网络的流量减少到 $\frac{B_{net}}{B} V_{n}$，与本地通信花费的时间相同。结果，上层链路的拥塞可能性降低，通信开销也减少了。随着通信开销的降低，模型可以在相同时间内进行更多的迭代训练。此外，专家和令牌的最佳匹配对被保留下来，减少了对其他令牌选择空间限制的影响。

方法论的推广。请注意，对于其他类型的拓扑，应设计另一种专门的拓扑感知门来提高性能。通过在特定的树状拓扑上展示这种拓扑感知门作为一个实例，我们倡导一种协同设计的方法论。在DDL-Roofline模型的指导下，可以轻松设计出高吞吐量的门，并更好地理解它们的性能模式。

A4 实验

实验环境

我们在两个代表性的集群上评估FasterMoE：
* johnny: 一个拥有16个GPU的集群，分布在2个工作节点上。每个节点有8个NVIDIA Tesla V100-PCIE GPU，通过PCIe交换机连接到2个CPU插槽。虽然配备了Infiniband EDR，但由于主板上缺少×16 PCIe插槽，带宽降至50Gb/s。该集群代表了深度学习训练中广泛使用的一类常见硬件。
* trevor: 一个超级计算机的分区。每个节点有4个NVIDIA V100-SXM2 GPU，通过NVLink互连形成一个异构环，其中一半的边通过两条链路的绑定具有双倍带宽。使用100Gb/s的Infiniband EDR进行通信。该集群代表了用于深度学习和其他传统HPC任务的超级计算机。我们在trevor上使用了16个节点的64个GPU进行实验。

评估模型:
评估中使用的模型规格如表2所示。MoE-GPT和MoE-BERT-Deep在johnny上训练。MoE-BERT-Wide在trevor上训练。

表2. 评估模型规格

软件与基线:
* 实现: FasterMoE基于FastMoE【索引7，FastMoE: A Fast Mixture-of-Expert Training System，2021，arXiv preprint arXiv:2103.13262】实现，通过扩展其功能来支持动态影分身和智能调度。
* 基线系统:
* ZeRO Optimizer【索引25，Zero: Memory optimizations toward training trillion parameter models，2020，SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】: 使用DeepSpeed【索引27，Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters，2020，Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining】的实现，特别是ZeRO Stage 3作为主要数据并行基线。
* FastMoE: 作为专家并行的基线实现。
* GShard【索引11，Gshard: Scaling giant models with conditional computation and automatic sharding，22020，arXiv preprint arXiv:2006.16668】: 其负载均衡策略作为自定义门在FastMoE中实现。
* BASE Layer【索引12，Base layers: Simplifying training of large, sparse models，2021，arXiv preprint arXiv:2103.16716】: 由FairSeq【索引21，fairseq: A Fast, Extensible Toolkit for Sequence Modeling，2019，Proceedings of NAACLHLT 2019: Demonstrations】实现，作为插件层使用。
* 模型训练框架: 使用Megatron-LM【索引20，Efficient large-scale language model training on gpu clusters，2021，arXiv preprint arXiv:2104.04473】作为基线，修改其MLP模块以进行MoE训练。

实验结果

性能模型准确性:
* 实验内容: 在两个集群上使用不同的模型和层，以及不同的$d$和$h$值来验证性能预测器的准确性。将预测延迟与实际延迟进行比较。
* 实验结果: 如图9所示，预测器能准确预测通信和计算的延迟。在johnny上，端到端延迟预测的$R^2$分数为0.987；在trevor上为0.967。计算延迟在负载较高时略有低估，因为模型假设为90%峰值性能，而实际中大矩阵计算效率更高。trevor上的一些异常值是由于大型共享网络的性能波动。

不同标记代表不同模型中的不同层。

总体加速比:
* 实验内容: 将FasterMoE与ZeRO optimizer的stage 3进行端到端性能比较。
* 实验结果: 如图10所示，FasterMoE在johnny和trevor集群上分别实现了6.63倍和17.87倍的加速比。与ZeRO stage 2相比，FasterMoE仍实现了高达3.94倍的加速。未优化的FastMoE（专家并行基线）也优于ZeRO stage 3，表明数据并行在MoE场景下固有地效率低下。

动态影分身和智能调度分析:
* 实验内容: 分别测试动态影分身和智能调度两项优化的性能增益。
* 实验结果:
* 动态影分身: 如图11所示，专家受欢迎程度是高度动态的，平均有19%的专家被影分身。该策略成功地降低了迭代延迟，当有1个专家被影分身时，实现了最高1.97倍的加速。
* 智能调度: 如图12所示，与理论上限（最高1.71倍）相比，智能调度在实际中实现了高达1.42倍的加速，在某些层中达到了理论加速的99%。在更大的模型和更多工作节点上，实际加速更接近理论上限。
* 联合优化: 如图13所示，动态影分身在johnny上带来1.95倍加速，在trevor上带来4.74倍加速。智能调度带来1.40倍加速。两者结合使用时，在johnny和trevor上分别观察到2.20倍和5.72倍的加速。

拓扑感知门加速比:
* 实验内容: 训练MoE-GPT模型，比较FasterMoE（带和不带拓扑感知门）与GShard、BASE Layer的训练损失随时间的变化。
* 实验结果: 如图14和表3所示，GShard虽然单次迭代时间最短，但达到目标损失（LM-loss=4.0）所需的步数是忠实top-2门的2.38倍。BASE Layer需要更多步数且单次迭代时间更长。
* FasterMoE（不带拓扑门）：与FastMoE的损失曲线几乎相同，但每次迭代加速1.33倍。这表明优化不改变专家选择，因此不增加收敛所需的步数。
* FasterMoE（带拓扑门）：迭代速度比不带拓扑门的版本快9.4%，但需要多18%的步数才能收敛，这与其他修改专家选择的基线类似。
* 结论: 总体而言，FasterMoE的收敛时间比GShard快1.37倍，比BASE Layer快2.19倍。

表3. 达到LM-loss=4.0时的时间和迭代次数

A7 补充细节：相关工作

参数服务器与数据并行。参数服务器【索引10，A Unified Architecture for Accelerating Distributed DNN Training in Heterogeneous GPU/CPU Clusters，2020，14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20)】、【索引13，Scaling Distributed Machine Learning with the Parameter Server，2014，Proceedings of the 11th USENIX Conference on Operating Systems Design and Implementation (OSDI’14)】、【索引14，Scaling Distributed Machine Learning with the Parameter Server，2014，Proceedings of the 11th USENIX Conference on Operating Systems Design and Implementation (OSDI’14)】是支持数据并行的最早系统，很快被使用all-reduce以获得更好性能的Horovod【索引30，Horovod: fast and easy distributed deep learning in TensorFlow，2018，arXiv preprint arXiv:1802.05799】所取代。异步部分模型更新方法【索引15，Taming unbalanced training workloads in deep learning with partial collective operations，2020，Proceedings of the 25th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming】、【索引16，Prague: High-performance heterogeneity-aware asynchronous decentralized training，2020，Proceedings of the Twenty-Fifth International Conference on Architectural Support for Programming Languages and Operating Systems】、【索引18，SwarmSGD: Scalable decentralized SGD with local updates，2019，arXiv preprint arXiv:1910.12308】、【索引32，ASYNC: A Cloud Engine with Asynchrony and History for Distributed Machine Learning，2020，2020 IEEE International Parallel and Distributed Processing Symposium (IPDPS)】被引入以加速异构环境中数据并行的收敛速度。

大规模模型训练系统。SuperNeurons【索引35，Superneurons: Dynamic GPU memory management for training deep neural networks，2018，Proceedings of the 23rd ACM SIGPLAN symposium on principles and practice of parallel programming】通过细粒度的内存管理方法将大型模型放置在单个GPU上。ZeRO Offload【索引28，Zero-offload: Democratizing billion-scale model training，2021，arXiv preprint arXiv:2101.06840】通过将数据交换到主机内存来减少数据并行的内存消耗，数据还可以进一步卸载到磁盘【索引26，ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning，2021，arXiv preprint arXiv:2104.07857】。Megatron-LM【索引20，Efficient large-scale language model training on gpu clusters，2021，arXiv preprint arXiv:2104.04473】是专为预训练设计的训练系统，为Transformer块提供了一种富有洞察力的模型并行方法。TOFU【索引36，Supporting very large models using automatic dataflow graph partitioning，2019，Proceedings of the Fourteenth EuroSys Conference 2019】和FlexFlow【索引9，Beyond data and model parallelism for deep neural networks，2018，arXiv preprint arXiv:1807.05358】是通用系统，通过执行模拟器和搜索提供数据和模型并行的最优混合策略。另一种节省内存的方法是流水线并行【索引19，PipeDream: generalized pipeline parallelism for DNN training，2019，Proceedings of the 27th ACM Symposium on Operating Systems Principles】，它可以与数据并行混合使用【索引4，DAPPLE: A pipelined data parallel approach for training large models，2021，Proceedings of the 26th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming】。

预训练模型与MoE系统。BERT【索引3，Bert: Pre-training of deep bidirectional transformers for language understanding，2018，arXiv preprint arXiv:1810.04805】和GPT【索引1，Language models are few-shot learners，2020，arXiv preprint arXiv:2005.14165】、【索引24，Language models are unsupervised multitask learners，2019，OpenAI blog】分别是用于语言理解和生成的流行预训练模型。基于Mesh TensorFlow【索引31，Mesh-tensorflow: Deep learning for supercomputers，2018，arXiv preprint arXiv:1811.02084】的GShard【索引11，Gshard: Scaling giant models with conditional computation and automatic sharding，2020，arXiv preprint arXiv:2006.16668】首次引入了专家并行。作为FairSeq【索引21，fairseq: A Fast, Extensible Toolkit for Sequence Modeling，2019，Proceedings of NAACLHLT 2019: Demonstrations】一部分的BASE layers【索引12，Base layers: Simplifying training of large, sparse models，2021，arXiv preprint arXiv:2103.16716】是来自PyTorch【索引23，Pytorch: An imperative style, high-performance deep learning library，2019，Advances in neural information processing systems】社区的另一个MoE训练系统，使用匹配算法进行专家分配。

A5 结论

本文提出了FasterMoE，通过一个性能模型和多项优化来解决MoE模型分布式训练中的挑战。性能通过一个精确的通信和计算预测器以及一个新颖的DDL-Roofline模型进行建模，该模型展示了特定训练任务在给定平台上的理论上限。在性能模型的指导下，我们提出了一种动态影分身方法，可以减少负载不均衡引入的开销。同步操作符被分解为工作节点组之间更小的任务，并被智能调度以并发执行，从而最小化通信开销。我们还设计了一种专家选择方法来避免网络拥塞，以有希望的收敛速度实现了更高的吞吐量。FasterMoE使训练大型动态MoE模型的效率提高了高达17.87倍。

A6 附录

A.1 访问

代码库。本文实验的脚本可在 https://github.com/laekov/fastermoe-ae 获取以供复现。

A.2 先决条件

A.2.1 硬件。脚本的目标是在16个NVIDIA V100 GPU上获得结果。虽然可以使用其他型号的GPU，但结果可能会有所不同。推荐使用Infiniband网络连接，因为MoE模型是高度通信密集型的。

A.2.2 FasterMoE的安装。FasterMoE是基于FastMoE【索引7，FastMoE: A Fast Mixture-of-Expert Training System，2021，arXiv preprint arXiv:2103.13262】实现的。它需要CUDA、NCCL和PyTorch (v1.10.0)。PyTorch可以使用以下命令安装：

pip install --user -f https://download.pytorch.org/whl/cu113/torch_stable.html torch==1.10.0+cu113

NCCL依赖。需要NCCL的开发者包（≥2.9.9），其版本应与PyTorch使用的版本相同（例如，上述命令安装的PyTorch v1.10.0附带NCCL 2.10.3）。可以从 https://developer.nvidia.com/nccl/nccl-legacy-downloads 下载。此外，需要matplotlib来绘制图形。
编译安装FasterMoE。在fastermoe目录中使用以下命令进行编译和安装：

USE_NCCL=1 python setup.py install --user

GPU架构设置。如果节点中安装了用于编译和测试运行的不同类型的GPU，则应提供环境变量TORCH_CUDA_ARCH_LIST。具体解释可参考 https://pytorch.org/docs/stable/cpp_extension.html#torch.utils.cpp_extension.CUDAExtension。

A.2.3 基线系统的安装。Megatron-LM和FairSeq【索引12，Base layers: Simplifying training of large, sparse models，2021，arXiv preprint arXiv:2103.16716】需要NVIDIA Apex。PyPI上的apex包已损坏，因此必须从源码安装。从 https://github.com/NVIDIA/apex 克隆，并使用以下命令安装：

python3 setup.py install --user --cuda_ext --cpp_ext

DeepSpeed安装。使用DeepSpeed【索引25，Zero: Memory optimizations toward training trillion parameter models，2020，SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】 v0.4.4作为基线。它可以在 https://github.com/microsoft/DeepSpeed 找到并直接从PyPI安装：

pip install --user deepspeed==0.4.4

FairSeq安装。FairSeq用作BASE Layers的基线。由于它没有在PyPI或GitHub上正式发布新版本，因此使用主分支。在fairseq目录中使用以下命令进行安装：

python3 setup.py build_ext 
python3 setup.py install --user

如果安装因未知的pyx文件类型失败，请尝试在安装命令中添加--editable标志。

A.2.4 训练和基准测试数据。对于单层测试，专家选择数据集可以从 https://pacman.cs.tsinghua.edu.cn/laekov/fastermoe-data/dumps.tgz 下载。该数据集是为16个专家生成的。下载并解压数据集，以使用我们的基准测试脚本测试FasterMoE和基线的性能。
对于训练，数据集是预处理过的，可以从 https://pacman.cs.tsinghua.edu.cn/laekov/fastermoe-data/wikidataset.tgz 下载。

A.3 运行实验

复现脚本。您可以简单地运行AE仓库中提供的runme.sh来复现下述所有产物。脚本开头的数据目录必须根据其实际位置进行修改。
环境假设与修改。假设您正在使用由SLURM管理的集群，每个节点安装了8个GPU。对于不同的设置，您可能需要修改runme.sh中的srun命令和scripts目录中的文件。还应注意，不同系统中的网络带宽不同。性能模型中的默认参数在其他系统上可能无效。因此，FasterMoE中的部分Python代码应该被修改。有关如何修改性能模型的详细信息，请参阅附录A.5。

A.4 复现结果

结果生成。实验结果（即图形和表格）在results目录中生成。以下是可用结果的列表及其一些讨论。

A.4.1 图9：预测准确性。results/fig9a.pdf和results/fig9b.pdf是图9中所示的性能预测器预测准确性的散点图。为减少复现时间，数据点的数量有所减少。然而，它们仍应接近虚线，但通信时间有时会因不确定的网络流量而被低估。

A.4.2 图10：相对于ZeRO的加速比。图10中将FasterMoE与ZeRO Optimizer【索引25，Zero: Memory optimizations toward training trillion parameter models，2020，SC20: International Conference for High Performance Computing, Networking, Storage and Analysis】进行了比较。结果可在results/fig10.pdf中找到。出于类似原因，测试用例的数量被减少以缩短运行时间。因此，加速比可能与论文中的原始图有所不同。

A.4.3 图11：影分身的效果。图11中详细分析了动态影分身。62次迭代中的影分身专家显示为黑色块。复现的图在results/fig11.pdf中。如果性能模型不够准确（例如，使用了与您实际系统不符的网络带宽），影分身的有效性可能会下降。

A.4.4 图12：智能调度的加速比。计算了理论加速比，并与FasterMoE实现的实际加速比进行了比较。结果可在results/fig12.pdf中看到。

A.4.5 图13：加速比分解。分别启用动态影分身和智能调度以检查它们各自的效果。图13的复现版本可在results/fig13.pdf中找到。

A.4.6 表3：拓扑感知门的迭代时间。由于将模型训练到我们的目标损失需要数天时间，我们只展示了前500次迭代的平均迭代时间，仅与表3的最后一列进行比较。复现结果可在results/table3.txt中找到。论文中呈现的数据是从数万次迭代中获得的。由于我们只运行前500次迭代，不同系统的平均迭代时间可能会有所不同，这是由于专家选择负载均衡的差异，如前文所述。

A.5 拓展

脚本与参数调整。这些图由不同的脚本生成，并在runme.sh中收集。请检查该文件以获取详细信息。改变模型几何形状（如论文中提到的$d$和$h$）并观察性能如何变化会很有趣。这可以通过为基准测试脚本指定不同的输入来实现，这应该很简单。这些变化也可以由我们的DDL-Roofline模型预测。
优化开关。此外，由于FasterMoE是基于FastMoE实现的，我们使用环境变量作为我们优化的开关。默认情况下，所有优化都关闭，其行为与原始FastMoE相同。
可调参数。一些参数，如动态影分身中的通信带宽和智能调度中的组大小，是可调的。可以对它们进行调整，性能可能会有所不同。具体来说，动态影分身的公式位于fastermoe/fmoe/transformer.py:34。分组的粒度由环境变量FMOE_FUSE_GRAN控制。