FlexMoE: Scaling Large-scale Sparse Pre-trained Model Training via Dynamic Device Placement
FlexMoE：通过动态设备配置实现大规模稀疏预训练模型的扩展训练

The Transformer architecture (Vaswani et al., 2017) has demonstrated its superior performance in various tasks (Devlin et al., 2019; Brown et al., 2020; Raffel et al., 2020; Dosovitskiy et al., 2021; Liu et al., 2021; Wei et al., 2022; Maruta and Kato, 2022), which mainly consists of attention networks and feed-forward networks. Each attention network first linearly transforms the input tokens into corresponding queries (Q), keys (K) and values (V) and then performs the scaled dot-product on them as Equation 1, where $d$ is the dimension of queries and keys. Meanwhile, this attention network benefits from capturing the dependencies between tokens within the sequence, and thus is effective in sequence transduction tasks.
变压器架构（Vaswani 等人，2017 年）在各种任务中展示了其卓越的性能（Devlin 等人，2019 年；Brown 等人，2020 年；Raffel 等人，2020 年；Dosovitskiy 等人，2021 年；Liu 等人，2021 年；Wei 等人，2022 年；Maruta 和 Kato，2022 年），主要由注意力网络和前馈网络组成。每个注意力网络首先将输入的令牌线性变换为相应的查询（Q）、键（K）和值（V），然后对它们执行缩放的点积运算，如方程 1 所示，其中 $d$ 是查询和键的维度。同时，这个注意力网络通过捕捉序列内令牌之间的依赖关系而受益，因此在序列转换任务中效果显著。

Each feed-forward network (FFN) is composed by two fully connected layers and an activation function, i.e., ReLU, formulated as Equation 2. Different from the attention network, FFNs model the relation of different feature dimensions within a single token by scaling it into larger dimension space. For example, the output dimension of $W_{1}$ is always set as $4\times$ of its input dimension. Existing famous pre-trained models are usually stacked by a series of transformer layers to improve its model capacity as well as model quality.
每个前馈网络（FFN）由两个全连接层和一个激活函数（例如 ReLU）组成，如方程 2 所公式化。与注意力网络不同，FFN 通过将单个令牌扩展到更大的维度空间来模拟不同特征维度之间的关系。例如， $W_{1}$ 的输出维度总是设置为其输入维度的 $4\times$ 倍。现有的著名预训练模型通常由一系列变压器层堆叠而成，以提高其模型容量以及模型质量。

Recently, scaling with more data and more parameters has driven significant model quality improvement (Dosovitskiy et al., 2021; Brown et al., 2020) while requires large amounts of computing resources for training. To scale models efficiently, researchers have adopted the sparse-gated Mixture-of-Experts (MoE) paradigm (Shazeer et al., 2017; Fedus et al., 2022) by replacing the FFN network with an MoE layer, where each input token activates only a subset of model parameters, and thus introducing model sparsity.
近期，通过增加数据量和模型参数，显著提升了模型的质量（Dosovitskiy 等人，2021 年；Brown 等人，2020 年），但这同时也需要大量的计算资源来进行训练。为了高效地扩展模型，研究人员采用了稀疏门控的专家混合（MoE）范式（Shazeer 等人，2017 年；Fedus 等人，2022 年），通过用 MoE 层替换 FFN 网络，其中每个输入令牌只激活模型参数的一个子集，从而引入了模型的稀疏性。

The key components of a MoE layer include a data-dependent sparse gate network $g(x)$ and a series of experts $\mathcal{E}$, as shown in Figure 1(a). For each input token $x$, the gate network first produces the probability of $x$ with respective to all experts and then routes $x$ to its corresponding experts. The Top-K gate (Shazeer et al., 2017) is formulated in Equation 3, which keeps only the top k values before the softmax function.
MoE 层的关键组成部分包括一个数据依赖的稀疏门控网络和一系列专家，如图 1(a) 所示。对于每个输入令牌，门控网络首先产生相对于所有专家的概率，然后将令牌路由到其对应的专家。Top-K 门控（Shazeer 等人，2017 年）在方程 3 中被定义，它在 softmax 函数之前只保留前 k 个值。

As soon as each expert $e_{i}$ receives its input token $x$ ($e_{i}\in\mathcal{E}$), it produces its corresponding output $e_{i}(x)$ and then the final output $y$ is obtained by the linear combination of $e_{i}(x)$ weighted by the output of the gate network $g(x)_{i}$ as follows:
一旦每个专家接收到其输入令牌，它就会产生相应的输出，然后通过门控网络的输出加权线性组合这些输出，得到最终的输出

When the gate network’s sparsity is static, e.g., Top-1 (Fedus et al., 2022) or Top-2 (Lepikhin et al., 2021), the computation and communication costs of given inputs nearly remain constant as the number of experts increases, allowing models to be scaled effectively with MoE.
当门控网络的稀疏性是静态的，例如，Top-1（Fedus 等，2022）或 Top-2（Lepikhin 等，2021），随着专家数量的增加，给定输入的计算和通信成本几乎保持不变，允许模型通过 MoE 有效地扩展。

To help readers better understand the routine of MoE layers, we also provide an illustration in Figure 1(b), where each GPU is associated with one expert and 9 tokens are fed to the gate network and routed to their target experts. For example, GPU 1 routes 6 and 3 tokens to the $1st$ and $Nth$ experts, while GPU N routes 4 and 5 tokens to these two experts. This assignment would cause the workload imbalance problem under current expert placement (i.e., 10 tokens for GPU 1 and 8 tokens for GPU N). Since there is a reverse routing step after the expert computation, all GPUs have to wait for each other before executing the following layers. Such an imbalanced workload inevitably results in GPU under-utilization and low training efficiency.
为了帮助读者更好地理解 MoE 层的常规流程，我们还在图 1(b)中提供了一个示意图，其中每个 GPU 与一个专家相对应，9 个令牌被送入门控网络并被路由到它们的目标专家。例如，GPU 1 将 6 个和 3 个令牌路由到 $1st$ 和 $Nth$ 专家，而 GPU N 将 4 个和 5 个令牌路由到这两个专家。这种分配会在当前的专家布局下引起工作负载不平衡问题（即 GPU 1 有 10 个令牌，而 GPU N 有 8 个令牌）。由于在专家计算后有一个反向路由步骤，所有 GPU 必须等待彼此完成后才能执行后续层。这种不平衡的工作负载不可避免地导致 GPU 利用率低下和训练效率低。

As models are sparsely scaled with MoE, multiple GPUs would be involved for model management and training acceleration (Narayanan et al., 2021; Rajbhandari et al., 2022; Nie et al., 2023). We will analyze the existing popular parallelism strategies in the following:
随着 MoE 的稀疏扩展模型，多个 GPU 将被用于模型管理和训练加速（Narayanan 等，2021 年；Rajbhandari 等，2022 年；Nie 等，2023 年）。我们将在下文中分析现有的流行并行策略：

Data Parallelism. Data parallelism (DP) is usually used to scale training when the model can fit in the device’s available GPU memory, as the communication primitives (e.g., AllReduce) of DP achieve good scalability performance. In DP, training samples are partitioned while model parameters are duplicated for each device. Given a batch of training samples, each worker executes the forward and backward computation, synchronizes gradients globally (i.e., averaged), and updates local parameters based on the synchronized gradients. However, each device has to maintain a full replica of the model and thus DP can’t be used to scale up to large models.
数据并行。数据并行（DP）通常用于在模型可以适应设备可用 GPU 内存的情况下扩展训练，因为 DP 的通信原语（例如，AllReduce）实现了良好的可扩展性性能。在 DP 中，训练样本被分割，而模型参数为每个设备复制一份。给定一批训练样本，每个工作器执行前向和后向计算，全局同步梯度（即，平均化），并根据同步的梯度更新本地参数。然而，每个设备必须维护模型的完整副本，因此 DP 不能用于扩展到大型模型。

Model Parallelism. If the memory requirement of the given model exceeds the GPU memory, it should be partitioned across multiple devices. In tensor parallelism (TP), every single tensor can be partitioned over multiple devices to reduce the memory footprint of each GPU. For example, Megatron-LM (Narayanan et al., 2021) proposed to partition the attention network by exploiting its inherent parallelism in the multi-head attention operation where the queries ($Q$), keys ($K$) and values ($V$) matrices can be partitioned in a column-parallel fashion. Pipeline parallelism (PP) is an alternative partitioning method, where different groups of layers are placed on different devices. For example, GPipe (Huang et al., 2019) partitions the input batch into a number of micro batches and pipelines each device’s computation across micro batches to improve the resource utilization of devices.
模型并行性。如果给定模型的内存需求超过了 GPU 内存，那么它应该被分割到多个设备上。在张量并行性（TP）中，每一个单独的张量都可以被分割到多个设备上，以减少每个 GPU 的内存占用。例如，Megatron-LM（Narayanan 等人，2021 年）提出通过利用多头注意力操作中固有的并行性来分割注意力网络，其中查询（ $Q$ ）、键（ $K$ ）和值（ $V$ ）矩阵可以以列并行的方式被分割。流水线并行性（PP）是另一种分割方法，不同组的层被放置在不同的设备上。例如，GPipe（Huang 等人，2019 年）将输入批次分割成多个微批次，并在微批次之间对每个设备的计算进行流水线处理，以提高设备的资源利用率。

Expert Parallelism. GShard (Lepikhin et al., 2021) first proposed expert parallelism (EP) as a specific model parallelism method for MoE models, where experts within an MoE layer are placed on different devices and non-MoE layers are replicated on devices as DP. The workflow of distributed training a Top-1 gate MoE layer is illustrated in Figure 1(b). After getting the target expert for each token, an All-to-All communication is involved to send tokens to target experts for processing. And another All-to-All communication is needed to send data back for the execution of data-parallel non-MoE layers. As MoE models often have numerous experts, e.g., 1024 Experts in Switch Transformer (Fedus et al., 2022), EP can scale up with model size better than model parallelism.
专家并行性。GShard（Lepikhin 等人，2021 年）首次提出专家并行性（EP）作为 MoE 模型的一种特定模型并行方法，其中 MoE 层内的专家被放置在不同的设备上，而非 MoE 层则作为 DP 在设备上被复制。分布式训练 Top-1 门 MoE 层的工作流程如图 1(b)所示。在为每个令牌获取目标专家后，需要进行一次全到全通信，以将令牌发送给目标专家进行处理。然后，需要另一次全到全通信将数据发送回来，以执行数据并行的非 MoE 层。由于 MoE 模型通常拥有大量的专家，例如 Switch Transformer（Fedus 等人，2022 年）中的 1024 个专家，EP 可以随着模型大小的增加而比模型并行性更好地扩展。

To model the training cost, we consider a single MoE layer, which consists of a gate network and a set of experts $\mathcal{E}$ ($e\in\mathcal{E}$). Given current expert-to-device mapping $\mathcal{P}$ ($(e,g)\in\mathcal{P}$ represents GPU $g$ has expert $e$), the assignment of tokens $\mathcal{I}$ ($I_{e,g}$ represents the number of tokens for expert $e$ to GPU $g$) at the current step, the training cost can be formulated as $T(\mathcal{I},\mathcal{P})$. Our objective is to minimize this training cost, which can be expressed as follows:
为了建模训练成本，我们考虑一个单独的 MoE 层，它由一个门控网络和一组专家 $\mathcal{E}$ （ $e\in\mathcal{E}$ ）组成。给定当前的专家到设备的映射 $\mathcal{P}$ （ $(e,g)\in\mathcal{P}$ 表示 GPU $g$ 拥有专家 $e$ ），在当前步骤下令牌的分配 $\mathcal{I}$ （ $I_{e,g}$ 表示专家 $e$ 到 GPU $g$ 的令牌数量），训练成本可以被表述为 $T(\mathcal{I},\mathcal{P})$ 。我们的目标是最小化这个训练成本，可以如下表达：

Existing System-Friendly Optimizations. As discussed in Section 1, existing works focus on addressing the workload imbalance problem by modifying the definitions of models (e.g., balance loss and capacity), which doesn’t need users to make modifications on the underlying DL systems and therefore is system-friendly.
现有的系统友好优化。如第 1 节所讨论的，现有工作通过修改模型的定义（例如，平衡损失和容量）来解决工作负载不平衡问题，这不需要用户对底层深度学习系统进行修改，因此是系统友好的。

• •

  The balance loss depicts the imbalance level of $\mathcal{I}$ and is widely used in MoE training (Fedus et al., 2022; Lepikhin et al., 2021). Once the gate network produces an imbalanced token assignment, it would be penalized by a large balance loss. In practice, the training process aims to minimize the weighted sum of balance loss and training loss. Thus, there exists a trade-off — emphasizing the balance loss would drive the assignment $\mathcal{I}$ to be even but harms the model quality since the training loss would be higher, and vice versa.

  
  • 平衡损失描述了 $\mathcal{I}$ 的不平衡水平，并在 MoE 训练中广泛使用（Fedus 等，2022；Lepikhin 等，2021）。一旦门控网络产生了不平衡的令牌分配，它将会因较大的平衡损失而受到惩罚。实际上，训练过程旨在最小化平衡损失和训练损失的加权和。因此，存在一个权衡——强调平衡损失会使分配 $\mathcal{I}$ 变得更加均匀，但会损害模型质量，因为训练损失会更高，反之亦然。
• •

  The capacity is a threshold that limits the number of input tokens for each expert (i.e., ${I}_{e,g}$). Tokens that exceeds the capacity will be skipped by the expert and directly forwarded to the next layer by the residual connection. In short, the capacity upper bounds the training cost of the heaviest expert to improve the overall efficiency, however, at the price of degraded model quality since a certain amount of tokens cannot be fully trained.

  
  • 容量是限制每个专家（即 ${I}_{e,g}$ ）的输入令牌数量的阈值。超出容量的令牌将被专家跳过，并通过残差连接直接转发到下一层。简而言之，容量上限限制了最重专家的训练成本，以提高整体效率，然而，代价是模型质量下降，因为一定数量的令牌无法得到充分训练。

Observation 1: Limitations of the existing techniques. We first studied the effect of the existing techniques empirically and took Swin-MoE (Liu et al., 2021) as an example, where we varied the balance loss coefficient and did not restrict the capacity of each expert to ensure no token was absent. The results are summarized in Figure 2. By increasing the balance loss coefficient from 0 to 0.05, the GPU utilization is improved from $18.77\%$ to $63.30\%$ while the top-5 accuracy is decreased from $94.588\%$ to $93.981\%$. It demonstrates that enforcing workload balance by adjusting the assignment of tokens $\mathcal{I}$ would inevitably lead to the trade-off between system efficiency and model quality.
观察 1：现有技术的局限性。我们首先从经验上研究了现有技术的效果，并以 Swin-MoE（Liu 等，2021）为例，我们调整了平衡损失系数，并且没有限制每个专家的容量以确保没有令牌缺失。结果总结在图 2 中。通过将平衡损失系数从 0 增加到 0.05，GPU 利用率从 $18.77\%$ 提高到 $63.30\%$ ，而 top-5 准确率则从 $94.588\%$ 降低到 $93.981\%$ 。这表明，通过调整令牌的分配来强制工作负载平衡 $\mathcal{I}$ ，不可避免地会导致系统效率与模型质量之间的权衡。

Observation 2: Dynamic imbalanced workloads. We also recorded the trace of training GPT-MoE models (64 expert per MoE layer) and studied the loads of experts during training. Results are summarized in Figure 3 and we have identified two key characteristics:
观察 2：动态不平衡的工作负载。我们还记录了训练 GPT-MoE 模型（每个 MoE 层 64 个专家）的轨迹，并研究了训练期间专家的负载。结果总结在图 3 中，我们确定了两个关键特征：

• •

  Skewness. As shown in Figure 3(a), we visualize the computational loads of experts as the cumulative distribution function (CDF), where we observe that the Top-10 experts (10 out of 64) receive almost $75\%$ tokens, leading to routing imbalance in MoE training. If the experts are evenly distributed among GPUs as in expert parallelism (Lepikhin et al., 2021), such imbalanced workloads would result in severe resource under-utilization, as most experts need to wait for the slowest.

  
  • 偏斜性。如图 3(a)所示，我们将专家的计算负载可视化为累积分布函数（CDF），其中我们观察到前 10 名专家（64 个中的 10 个）接收到几乎 $75\%$ 令牌，导致 MoE 训练中的路由不平衡。如果专家像在专家并行中那样均匀分布在 GPU 上（Lepikhin 等，2021），这种不平衡的工作负载将导致严重的资源未充分利用，因为大多数专家需要等待最慢的专家。
• •

  Smoothness and continuousness. In Figure 3(b), we present evolution of expert loads throughout the entire training process, where different intervals represent different experts. We observe that the load of each expert is continuously changing during training, for example, from less to more, from more to less, and from more to less and then more again, etc, which poses routing fluctuation when training MoE models. Fortunately, the load of each expert does not change dramatically in a short period of time, which means a smooth and continuous change.

  
  • 平滑性与连续性。在图 3(b)中，我们展示了整个训练过程中专家负载的演变，其中不同的区间代表不同的专家。我们观察到，每个专家的负载在训练过程中持续变化，例如，从少到多，从多到少，再从多到少然后再多等等，这在训练 MoE 模型时引起了路由波动。幸运的是，每个专家的负载在短时间内不会发生剧烈变化，这意味着变化是平滑且连续的。

Challenges. Motivated by these observations, our work explores how to develop a system that achieves workload balance by adjusting the expert-to-device mapping (i.e., $\mathcal{P}$) for better system efficiency while maintaining the model quality. Moreover, the system should dynamically adapt to the varying routing distribution during the training process. However, since both token routing and expert loads are decided by the data-dependent gating mechanism, we can not determine the optimal mapping ahead of its execution. The main challenge lies in designing and modifying the expert-to-device mapping efficiently under the fast GPU computation and rapid change in workloads. Another challenge is how to efficiently implement these irregular MoE operations.
挑战。受到这些观察的启发，我们的工作探索了如何开发一个系统，通过调整专家到设备的映射（即 $\mathcal{P}$ ），在保持模型质量的同时，实现工作负载平衡，提高系统效率。此外，系统应能够动态适应训练过程中路由分布的变化。然而，由于令牌路由和专家负载都由数据依赖的门控机制决定，我们无法在执行之前确定最优映射。主要挑战在于，在快速的 GPU 计算和工作负载迅速变化的情况下，高效地设计和修改专家到设备的映射。另一个挑战是如何高效实现这些不规则的 MoE 操作。

Opportunities. Fortunately, as previously mentioned, the distribution changes smoothly and continuously. This means that the optimal expert-to-device mapping would not shift significantly in a short period of time. Therefore, it is feasible to refine the mapping based on the ad-hoc routing determination, without the need to predict the optimum for the next few training steps. Furthermore, the cost of computation and communication can be estimated before the actual execution. Based on the cost models, we can predict the benefits and overheads of different mapping candidates to find the best one.
机遇。幸运的是，如前所述，分布变化平滑且连续。这意味着最优的专家到设备映射在短时间内不会发生重大变化。因此，基于特定情况下的路由决定来细化映射是可行的，无需预测接下来几个训练步骤的最优解。此外，计算和通信的成本可以在实际执行之前估算。基于成本模型，我们可以预测不同映射候选方案的利益和开销，以找到最佳方案。

The workflow of our FlexMoE is illustrated in Figure 4, where (1) input tokens are first processed by the gate network to determine their corresponding target experts and then (2) the router leverages a greedy algorithm to distribute tokens to different replicas of each expert. Finally, (3) each replica performs computations on its assigned tokens and delivers the output results back.
我们的 FlexMoE 工作流程如图 4 所示，其中（1）输入令牌首先由门控网络处理，以确定它们对应的目标专家，然后（2）路由器利用贪心算法将令牌分配给每个专家的不同副本。最后，（3）每个副本对其分配的令牌进行计算，并将输出结果返回。

To handle the dynamic workload imbalance, we have designed a dynamic expert management mechanism that adaptively adjusts the expert-to-device mapping $\mathcal{P}$ and the assignment of tokens $\mathcal{I}$ during training. In addition to the regular workflow, we also introduce two new components: Scheduler and Policy Maker. (4) Scheduler monitors the real-time loads of experts and sends them to Policy Maker if the current imbalance metric (i.e., balance ratio in Equation 6) exceeds the predetermined threshold. Based on the received loads and the current placement of experts, Policy Maker produces modification instructions and sends them back to Scheduler. Scheduler then interacts with the executor of current DL frameworks and triggers modifications of expert placement at runtime.
为了处理动态工作负载不平衡，我们设计了一种动态专家管理机制，该机制在训练过程中自适应调整专家到设备的映射 $\mathcal{P}$ 和令牌的分配 $\mathcal{I}$ 。除了常规工作流程外，我们还引入了两个新组件：调度器和策略制定者。（4）调度器监控专家的实时负载，并在当前不平衡指标（即，方程 6 中的平衡比率）超过预定阈值时，将其发送给策略制定者。基于接收到的负载和专家当前的位置，策略制定者产生修改指令并将其发送回调度器。调度器随后与当前深度学习框架的执行器交互，并在运行时触发专家位置的修改。

To cope with the large decision space of flexible expert-to-device mapping and its dynamic modifications, FlexMoE proposes the abstraction of vExpert to represent the minimum unit for scheduling, as detailed in Section 3.2. Moreover, FlexMoE decouples the expert placement modifications using three placement modification primitives, including Expand, Shrink and Migrate, as explained in Section 3.3. Finally, we design a greedy heuristic algorithm to generate a sequential combination of these modification primitives to produce the adjustment plan of expert-to-device mapping, as described in Section 3.4.
为了应对灵活的专家到设备映射的大决策空间及其动态修改，FlexMoE 提出了 vExpert 的抽象概念，以表示调度的最小单元，详见 3.2 节。此外，FlexMoE 通过使用三种位置修改原语（包括扩展、收缩和迁移）来解耦专家位置的修改，如 3.3 节所述。最后，我们设计了一个贪心启发式算法，以生成这些修改原语的顺序组合，以产生专家到设备映射的调整计划，如 3.4 节所述。

To tackle the optimization problem related to expert-to-device mapping and its corresponding token routing, we introduce a novel abstraction called vExpert, which defines the minimum unit for scheduling GPU computations for an expert and enables dynamic expert management. The vExpert abstraction helps in determining how to duplicate experts, when to increase or decrease replicas, and how to partition tokens between replicas.
为了解决与专家到设备映射及其相应令牌路由相关的优化问题，我们引入了一个名为虚拟专家（vExpert）的新概念抽象，它定义了为专家安排 GPU 计算的最小单元，并实现了动态专家管理。虚拟专家抽象有助于确定如何复制专家，何时增减副本，以及如何在副本之间分配令牌。

vExpert has the following characteristics:
虚拟专家具有以下特点：

• •

  Each vExpert can only be assigned as the replica of exactly one expert and process part of tokens for the master expert.

  
  • 每个虚拟专家只能被指定为一个专家的副本，并处理该主专家的部分令牌。
• •

  Each vExpert shares weights with other vExperts for the same expert on the same GPU, which means packing the same vExperts of the same GPU.

  
  • 每个虚拟专家（vExpert）与同一 GPU 上的同一专家的其他虚拟专家共享权重，这意味着将同一 GPU 上的相同虚拟专家进行打包。
• •

  Evenly workloads (i.e., tokens) partitioning among all vExperts for the same expert.

  
  • 在所有虚拟专家（vExperts）之间均匀分配工作负载（即，令牌）。

With this abstraction, we can significantly reduce the large search space of the optimization problem by making decisions at the vExpert level. Besides, we have designed three placement modification primitives to manage the dynamic expert management at the vExpert level, including Expand, Shrink and Migrate (see Section 3.3 for details). These primitives allow for arbitrary modifications of expert placement by composing them in different ways. Moreover, because the amount of data in each iteration is fixed ($\sum B_{i}=B$), we can calculate the ideal capacity of each vExpert ($B/(G*E)$) as a reference for decision-making, where $B$ is the total number of input tokens, $G$ is the number of GPUs and $E$ is the number of vExperts on each device.
通过这种抽象，我们可以通过在虚拟专家（vExpert）层面上做决策，显著减少优化问题的大型搜索空间。此外，我们设计了三种位置修改原语，以在虚拟专家层面管理动态专家管理，包括扩展、收缩和迁移（详见第 3.3 节）。这些原语允许通过以不同方式组合它们来对专家位置进行任意修改。而且，因为每次迭代中的数据量是固定的（ $\sum B_{i}=B$ ），我们可以计算每个虚拟专家的理想容量（ $B/(G*E)$ ）作为决策参考，其中 $B$ 是输入令牌的总数， $G$ 是 GPU 的数量， $E$ 是每个设备上的虚拟专家数量。

Figure 4 illustrates the role of the Scheduler in connecting the Runtime Executor and the Policy Maker in FlexMoE. Our design includes a metric, the balance ratio, which measures real-time workloads and determines when to trigger the Policy Maker for adjustment decisions. Additionally, we propose three atomic primitives for the Runtime Executor to use in scheduling experts. These primitives will be further explained in the following section.
图 4 展示了调度器在 FlexMoE 中连接运行时执行器和策略制定者的角色。我们的设计包括一个度量标准，平衡比率，用于测量实时工作负载并确定何时触发策略制定者进行调整决策。此外，我们提出了三个原子原语供运行时执行器在调度专家时使用。这些原语将在下一节中进一步解释。

Balance ratio. Owing to the synchronous execution mode in MoE layer, the slowest GPU will dominate the finish time of all-to-all communication as well as the global training. Thus, we design the balance ratio as Equation 6, which adds up the loads of all vExperts on a single GPU and finds the max value among all GPUs to get the final results. While there may be other metrics that can be used to measure loads, we compare the balance ratio in Equation 6 with variance as a metric and provide a detailed analysis of our findings in Section 5.3.
平衡比率。由于 MoE 层中的同步执行模式，最慢的 GPU 将决定全体通信以及全局训练的完成时间。因此，我们将平衡比率设计为方程 6，它将单个 GPU 上所有 vExpert 的负载相加，并找出所有 GPU 中的最大值以得到最终结果。虽然可能有其他度量标准可以用来测量负载，我们将方程 6 中的平衡比率与方差作为度量进行比较，并在第 5.3 节提供我们发现的详细分析。

Expand. When an expert receives increasing workloads, previously allocated vExpert resources may handle them slower than others and thus produce stragglers (over-utilized). In this scenario, FlexMoE will call the Expand primitive to allocate an extra new vExpert resource for this expert at a time. The Expand primitive copies the expert parameter as well as the corresponding optimizer states from the source vExpert to the target vExpert via parameter sharing for intra-GPU communication and point-to-point communication of inter-GPU communication.
扩展。当一个专家接收到增加的工作负载时，之前分配的 vExpert 资源可能比其他资源处理得更慢，从而产生拖后腿者（过度利用）。在这种情况下，FlexMoE 将调用扩展原语，一次为该专家分配一个额外的新 vExpert 资源。扩展原语通过参数共享以及 GPU 间通信的点对点通信，将专家参数及相应的优化器状态从源 vExpert 复制到目标 vExpert。

Shrink. In the opposite, when previously allocated vExpert resources are not fully utilized due to decreasing workloads (under-utilized), FlexMoE will call the Shrink primitive to release an vExpert resource at a time. The Shrink primitive is executed by a tag without any communication, and thus introduces no overheads.
缩减。相反，当之前分配的 vExpert 资源由于工作负载减少（未充分利用）而没有被充分利用时，FlexMoE 将调用缩减原语，一次释放一个 vExpert 资源。缩减原语通过一个标签执行，不涉及任何通信，因此不引入任何开销。

Migrate. If an expert retains multiple replicas on different GPUs, the training is slowed down due to the need for extra communications (i.e., gradients all-reduce) for synchronizations. The efficiency of this communication depends on the number of GPUs involved and their locality. To reduce the number of GPUs holding expert replicas and lower the synchronization cost, FlexMoE will call the Migrate primitive exchanges the model states between two vExperts.
迁移。如果一个专家在不同的 GPU 上保留了多个副本，由于需要额外的通信（即，梯度全归约）来进行同步，训练速度会因此变慢。这种通信的效率取决于参与的 GPU 数量及其位置关系。为了减少持有专家副本的 GPU 数量并降低同步成本，FlexMoE 将调用迁移原语在两个虚拟专家之间交换模型状态。

We present the workflow of Scheduler in Algorithm 1, which monitors the real-time workloads, i.e., the assignment of tokens $\mathcal{I}$, and measures the balance ratio under current placements $\mathcal{P}$ (Line 1-2). If current balance ratio exceeds the pre-defined $threshold$, Scheduler will repeatedly ask the Policy Maker for modification primitives until no beneficial modifications exist (Line 3-8). Then, Scheduler turns to the Migrate operation to reduce the synchronization cost and continuously optimizes it at backend (Line 9). Intuitively, the vExpert-based Expand and Shrink primitives tackle the complicated expert-device mapping problem under the dynamic changing workloads, and the Migrate primitive continuously optimizes the placement of replicas for each expert.
我们在算法 1 中展示了调度器的工作流程，它监控实时工作负载，即令牌的分配 $\mathcal{I}$ ，并测量当前放置下的平衡比率 $\mathcal{P}$ （第 1-2 行）。如果当前平衡比率超过预定义的 $threshold$ ，调度器将反复请求策略制定者进行修改原语，直到不存在有益的修改为止（第 3-8 行）。然后，调度器转向迁移操作以减少同步成本，并在后端持续优化它（第 9 行）。直观地说，基于虚拟专家的扩展和收缩原语解决了在动态变化的工作负载下复杂的专家-设备映射问题，而迁移原语则持续优化每个专家副本的放置。

Using the modification primitives described above, FlexMoE employs an efficient vExpert-based scheduling algorithm for generating sequential modification operations, as shown in Algorithm 2. Specifically, we leverage a cost-model driven search planning approach, which makes decisions based on feed-backs from simulating the training time.
在上文描述的修改原语基础上，FlexMoE 采用了一种高效的 vExpert 为基础的调度算法来生成序列化的修改操作，如算法 2 所示。具体来说，我们采用了一种基于成本模型驱动的搜索规划方法，该方法根据模拟训练时间的反馈来做出决策。

To model the training cost of an MoE layer, we decompose it into three parts as Equation 5, including computation cost $T_{C}$, All-To-All communication cost $T_{A2A}$ and expert synchronization cost $T_{Sync}$. For each part, we build cost models that take into account input variables (e.g., $\mathcal{I}$: the assignment of tokens and $\mathcal{P}$: the expert-to-device mapping) as well as environmental variables (e.g., TPS: tokens-per-second for an expert). By leveraging a profiling-based approach, we first profile the function’s running time under different input sizes and then estimate the corresponding environmental variables. Besides, we also consider the cost of expert adjustments, which could be executed concurrently with model training.
为了建模 MoE 层的训练成本，我们将其分解为三部分，如方程 5 所示，包括计算成本 $T_{C}$ 、全对全通信成本 $T_{A2A}$ 和专家同步成本 $T_{Sync}$ 。对于每个部分，我们构建了考虑输入变量（例如， $\mathcal{I}$ ：令牌的分配和 $\mathcal{P}$ ：专家到设备的映射）以及环境变量（例如，TPS：每秒令牌数）的成本模型。通过采用基于分析的方法，我们首先在不同输入大小下对函数的运行时间进行分析，然后估计相应的环境变量。此外，我们还考虑了专家调整的成本，这可以与模型训练同时执行。

Computation Cost. The computation cost refers to the time taken for an expert to perform the forward and backward computation of experts during training, which can be formulated as:
计算成本。计算成本指的是专家在训练期间执行前向和后向计算所需的时间，可以表述为：

All-To-All Cost. The All-To-All operation is involved to send tokens to target experts and send back their results after processing, which will be called for totally 4 times in each training step. We predict the All-To-All cost by a topology-aware model as Equation 8, where the expert $e$ on GPU $g$ has received $I_{e,g}.count(g^{\prime})$ tokens from GPU $g^{\prime}$ and $Bw_{g,g^{\prime}}$ is the profiled bandwidth between GPU $g$ and GPU $g^{\prime}$. Our cost model considers the intra-node bandwidth (e.g., PCIe, NvLink) and inter-node bandwidth (e.g., IB, NIC) separately.
全对全成本。全对全操作涉及向目标专家发送令牌并在处理后发送回其结果，每个训练步骤中将调用总共 4 次。我们通过一个拓扑感知模型预测全对全成本，如方程 8 所示，其中 GPU $g$ 上的专家 $e$ 已从 GPU $g^{\prime}$ 接收到 $I_{e,g}.count(g^{\prime})$ 个令牌，而 $Bw_{g,g^{\prime}}$ 是 GPU $g$ 与 GPU $g^{\prime}$ 之间的已测量带宽。我们的成本模型分别考虑了节点内带宽（例如，PCIe、NvLink）和节点间带宽（例如，IB、NIC）。

Synchronization Cost. When expert $e$ holds multiple replicas on different GPUs, their gradients must be synchronized via AllReduce communication, which is determined by the message size, the number of involved devices and their connected network. We enumerate different device groups and profile them before training to get their BPS (bytes-per-second). We predict the synchronization cost for the expert $e$ by finding the device group that includes $e$ (i.e., $\mathcal{P}.index(e)$) and then obtaining its corresponding $BPS$ from profiling data. ${size}(e)$ represents the size of gradients for an expert.
同步成本。当专家 $e$ 在不同 GPU 上持有多个副本时，必须通过 AllReduce 通信同步它们的梯度，这取决于消息大小、涉及设备的数量及其连接的网络。我们在训练前枚举不同的设备组并对它们进行评估，以获得它们的 BPS（每秒字节数）。我们通过找到包含 $e$ （即， $\mathcal{P}.index(e)$ ）的设备组，然后从评估数据中获取其相应的 $BPS$ 来预测专家 $e$ 的同步成本。 ${size}(e)$ 代表一个专家的梯度大小。

Expert Adjustment Cost. FlexMoE proposes three modification primitives, including Expand, Shrink and Migrate. The Expand and Migrate primitives involve transferring model states from the source GPU to the target GPU via NCCL Point-to-Point communication. On the other hand, the Shrink primitive is executed with no overheads by marking a tag. We predict the cost of transferring model states as $\frac{{size(e.model\_states)}}{Bw_{g,g^{\prime}}}$.
专家调整成本。FlexMoE 提出了三种修改原语，包括扩展、收缩和迁移。扩展和迁移原语涉及通过 NCCL 点对点通信将模型状态从源 GPU 传输到目标 GPU。另一方面，收缩原语通过标记一个标签而无需开销地执行。我们预测传输模型状态的成本为 $\frac{{size(e.model\_states)}}{Bw_{g,g^{\prime}}}$ 。

The scheduling policy of Policy Maker is summarized in Algorithm 2. Firstly, the algorithm gets the time cost $t_{0}$ of current placement as the baseline (Line 2), and then finds the expert id with the maximum workload and the expert id with the minimum workload (Line 3-6). After that, our Policy Maker estimates the time cost $t_{1}$ after applying the Expand and Shrink primitives (Line 7) and decides whether to return the modification by comparing $t_{0}$ and $t_{1}$ (Line 8-10).
策略制定者的调度策略在算法 2 中总结。首先，算法获取当前布局的时间成本 $t_{0}$ 作为基线（第 2 行），然后找到工作量最大的专家 id 和工作量最小的专家 id（第 3-6 行）。之后，我们的策略制定者估计在应用扩展和收缩原语后的时间成本 $t_{1}$ （第 7 行），并通过比较 $t_{0}$ 和 $t_{1}$ 来决定是否返回修改（第 8-10 行）。

Machine environment. We conduct experiments on Azure VMs (azu, 2014), each equipped with 192-core AMD CPUs and 8 NVIDIA Ampere A100 GPUs. The GPUs are connected via NVLink 3.0 intra-node and the servers are connected via 8 InfiniBand NICs (8*200 Gbps totally). RDMA is used by default and the PyTorch version is 1.11.
机器环境。我们在 Azure VMs（azu, 2014）上进行实验，每台虚拟机配备了 192 核的 AMD CPU 和 8 块 NVIDIA Ampere A100 GPU。GPU 通过 NVLink 3.0 内节点连接，服务器之间通过 8 个 InfiniBand NICs（总共 8*200 Gbps）连接。默认使用 RDMA，PyTorch 版本为 1.11。

Baselines. We compare FlexMoE with other popular frameworks, including DeepSpeed (Rajbhandari et al., 2022) and FasterMoE (He et al., 2022). Deepspeed used expert parallelism, which was first proposed by GShard (Lepikhin et al., 2021). FasterMoE proposed the shadowing strategy to replicate the popular expert among all GPUs.
基准。我们将 FlexMoE 与其他流行框架进行比较，包括 DeepSpeed（Rajbhandari 等，2022）和 FasterMoE（He 等，2022）。DeepSpeed 使用了 GShard（Lepikhin 等，2021）首次提出的专家并行机制。FasterMoE 提出了影子策略，以在所有 GPU 中复制流行的专家。

Benchmarks and datasets. Our evaluations are conducted by scaling representative transformer models in different application domains with the MoE architecture, including BERT (Devlin et al., 2019) and GPT (Radford et al., 2019) in NLP and Swin (Liu et al., 2021) in CV, as shown in Table 1. We adopt Top-2 Gate for each model, which is adopted by widely used sparse MoE models (GShard (Lepikhin et al., 2021), GLaM (Du et al., 2022), V-MoE (Riquelme et al., 2021)), and set the capacity factor as $1.0$ for each expert and balance loss as $0.001$. We pretrain BERT-MoE with masked language modeling (MLM) and next-sentence-prediction (NSP) (Devlin et al., 2019) tasks, GPT-MoE with language modeling task (LM) (Radford et al., 2019), Swin-MoE with image classification task (Liu et al., 2021; Hwang et al., 2022). NLP models are trained on wikipedia (wik, 2014) and vision models are trained on the ImageNet-1K (ima, 2014). All the hyper-parameters (e.g., learning rate) are fixed as for the same model.
基准和数据集。我们通过扩展不同应用领域中具有代表性的变压器模型与 MoE 架构的结合进行评估，包括 NLP 领域的 BERT（Devlin 等，2019 年）和 GPT（Radford 等，2019 年），以及 CV 领域的 Swin（Liu 等，2021 年），如表 1 所示。我们为每个模型采用 Top-2 门控，这一做法被广泛应用于稀疏 MoE 模型中（GShard（Lepikhin 等，2021 年），GLaM（Du 等，2022 年），V-MoE（Riquelme 等，2021 年）），并将每个专家的容量因子设置为 $1.0$ ，平衡损失设置为 $0.001$ 。我们使用遮蔽语言建模（MLM）和下一句预测（NSP）（Devlin 等，2019 年）任务对 BERT-MoE 进行预训练，使用语言建模任务（LM）（Radford 等，2019 年）对 GPT-MoE 进行预训练，使用图像分类任务（Liu 等，2021 年；Hwang 等，2022 年）对 Swin-MoE 进行预训练。NLP 模型在维基百科（wik，2014 年）上训练，视觉模型在 ImageNet-1K（ima，2014 年）上训练。所有超参数（例如，学习率）对相同模型保持不变。

To evaluate the end-to-end efficiency of FlexMoE, we compare it with DeepSpeed and FasterMoE. For DeepSpeed, we use the traditional expert parallelism approach and set the number of experts per GPU as 1 in every MoE layer. We evaluate two different width of models, X-MoE-S with 32 experts and X-MoE-L with 64 experts, on 32 GPUs and 64 GPUs, respectively.
为了评估 FlexMoE 的端到端效率，我们将其与 DeepSpeed 和 FasterMoE 进行了比较。对于 DeepSpeed，我们使用传统的专家并行方法，并在每个 MoE 层中将每个 GPU 的专家数量设置为 1。我们在 32 个 GPU 和 64 个 GPU 上分别评估了两种不同宽度的模型，X-MoE-S 拥有 32 个专家，而 X-MoE-L 拥有 64 个专家。

Model Quality. To demonstrate the importance of not dropping tokens, we compare the model quality of FlexMoE with DeepSpeed on various models and various tasks. We use the validation perplexity for language models (e.g., BERT-MoE and GPT-MoE), which is the lower the better, and top-1/top-5 accuracy of Imagenet-1K for vision models (e.g., Swin-MoE). As shown in Table 2, FlexMoE outperforms DeepSpeed in almost all tasks, indicating that the limited capacity in existing MoE systems (e.g., DeepSpeed) can lead to model quality degradation. Considering the training cost, the Swin-MoE models (scaled based on the Swin-B model) are trained for 100 epochs (less than 300 epochs in the standard configuration) and thus the accuracy is slightly worse than the benchmark, where $94.04\%$ of FlexMoE v.s. $96.5\%$ of the benchmark as for top-5 accuracy. And we believe it’s enough to show the benefits of no dropping tokens. What’s more, Swin-MoE-L performs slightly worse than Swin-MoE-S as the size of training data is small and thus it may lead to overfitting.
模型质量。为了展示不丢弃令牌的重要性，我们比较了 FlexMoE 与 DeepSpeed 在各种模型和任务上的模型质量。我们使用语言模型（例如，BERT-MoE 和 GPT-MoE）的验证困惑度，其值越低越好，以及视觉模型（例如，Swin-MoE）的 Imagenet-1K 的 top-1/top-5 准确率。如表 2 所示，FlexMoE 在几乎所有任务中都优于 DeepSpeed，这表明现有 MoE 系统（例如，DeepSpeed）的有限容量可能导致模型质量下降。考虑到训练成本，Swin-MoE 模型（基于 Swin-B 模型的缩放）训练了 100 个周期（标准配置下少于 300 个周期），因此准确率略低于基准，其中 FlexMoE 的 top-5 准确率与基准的 $94.04\%$ 相比为 $96.5\%$ 。我们认为这足以展示不丢弃令牌的好处。更重要的是，由于训练数据量小，Swin-MoE-L 的表现略逊于 Swin-MoE-S，因此可能导致过拟合。

System Efficiency. We also evaluate FlexMoE against other SOTA systems on efficiency. To measure efficiency, we record the required training time to achieve the target model quality, and the results are illustrated in Figure 5. Our experiments show that FlexMoE achieves the best training efficiency, outperforming DeepSpeed by $1.70\times$ on average and up to $2.10\times$, and FasterMoE by $1.30\times$ on average and up to $1.45\times$. As mentioned above, although DeepSpeed obtains the smallest iteration time thanks to its limited capacity, it drops tokens to skip the expert network and thus requires more iterations to converge. FasterMoE proposes the dynamic shadowing strategy for loading balance, which replicates the popular expert on all GPUs. However, due to its coarse-grained expert management (i.e., on 1 GPU or on all GPUs), it falls back to a sub-optimal solution. As the number of GPUs increases, FasterMoE suffers from the global synchronization of expert replicas.
系统效率。我们还评估了 FlexMoE 与其他 SOTA 系统在效率上的表现。为了衡量效率，我们记录了达到目标模型质量所需的训练时间，结果如图 5 所示。我们的实验表明，FlexMoE 在训练效率上表现最佳，平均优于 DeepSpeed $1.70\times$ ，最高可达 $2.10\times$ ，平均优于 FasterMoE $1.30\times$ ，最高可达 $1.45\times$ 。如上所述，尽管 DeepSpeed 由于其有限的容量获得了最小的迭代时间，但它通过跳过专家网络丢弃令牌，因此需要更多的迭代才能收敛。FasterMoE 提出了动态阴影策略来实现负载平衡，该策略在所有 GPU 上复制热门专家。然而，由于其粗粒度的专家管理（即，在 1 个 GPU 上或在所有 GPU 上），它回落到了次优解。随着 GPU 数量的增加，FasterMoE 受到专家副本全局同步的影响。