Netmoe: Accelerating Moe Training Through Dynamic Sample Placement

作者/机构: Xinyi Liu1 Yujie Wang1 Fangcheng Fu1 Xupeng Miao2 Shenhan Zhu1 Xiaonan Nie1 Bin Cui1,3
1School of CS & Key Lab of High Confidence Software Technologies (MOE), Peking University 2Purdue University 3Institute of Computational Social Science, Peking University (Qingdao) {xy.liu, alfredwang, ccchengff}@pku.edu.cn
xupeng@purdue.edu, {shenhan.zhu, xiaonan.nie, bin.cui}@pku.edu.cn

A1 主要贡献

本文旨在解决专家混合（MoE）模型分布式训练中的核心瓶颈：All-to-All通信开销巨大。由于MoE模型中专家（Experts）分布在不同GPU上，每次路由后都需要通过All-to-All通信在GPU间交换训练令牌（tokens），这一过程频繁且数据量大，通信时间可占总训练时间的80%。

核心问题与研究目标：
现有的优化方法主要从模型角度入手（如动态调整专家位置或修改模型定义），但这些方法或引入巨大的额外通信开销，或可能损害模型性能。本文发现，从训练样本（data sample）的角度进行优化是一个被忽视的方向。研究观察到，同一训练样本中的令牌在专家路由上存在一定的局部性（locality），即它们倾向于被路由到相似的专家。同时，深度学习集群具有固有的网络局部性，即节点内通信（如NVLINK）远快于节点间通信（如InfiniBand）。本文的研究目标是利用数据路由的局部性和集群网络的局部性，通过动态调整训练样本的放置位置来加速All-to-All通信。

创新点 (NetMoE)：
1. 提出动态样本放置策略：本文首次提出NetMoE框架，通过在训练过程中动态调整数据样本的放置，将原本需要跨节点（inter-node）的慢速通信转化为节点内（intra-node）的快速通信，从而降低All-to-All通信的整体时间成本。如图2所示，通过交换样本0和样本3的位置，节点间的通信量从5个token减少到2个token。

图 1: 一个专家并行应用于MoE模型的示例，该模型有 J 个设备和 E = 2J 个专家（每个设备有两个不同的专家）。 **(a) 一个MoE层示例概览。** 图 2: 样本交换示例。该图展示了一个MoE层中的All-to-All gather操作，该层有两个节点，每个节点包含两个设备，每个设备有一个专家。不同颜色代表发送给不同专家的令牌，i-j表示第i个样本中的第j个令牌。图2(a)展示了前向传播期间MoE层的完整过程。图2(b)显示了在不调整样本放置的情况下MoE层中的All-to-All gather操作，其中每个节点的节点间通信量为5个令牌。图2(c)显示了启用样本放置调整后的All-to-All gather操作——设备上的样本位置发生变化（样本0和3被交换），每个节点的节点间通信量减少到2个令牌。

1. 建立优化问题模型：将动态样本放置问题形式化为一个组合优化问题。该问题旨在根据专家的路由结果，寻找能够最大化通信效率（即最小化All-to-All通信时间）的最佳样本放置方案。

2. 设计高效的多项式时间求解器：为了在训练过程中实时、高效地求解该优化问题，本文将其分解为两个阶段，并设计了一个多项式时间的算法来快速推导出样本放置方案，确保求解过程不会成为新的性能瓶颈。

实验成果：
在32个NVIDIA A800 GPU上的实验表明，与现有的MoE训练框架相比，NetMoE最高可实现1.67倍的训练效率提升。

A3 背景知识

2.1 分布式训练中的并行策略

数据并行与模型并行。在数据并行中【索引21，Pytorch distributed: Experiences on accelerating data parallel training，2020，Proc. VLDB Endow.】、【索引36，Horovod: fast and easy distributed deep learning in tensorflow，2018，CoRR】、【索引44，A communication efficient admmbased distributed algorithm using two-dimensional torus grouping allreduce，2023，Data Sci. Eng.】、【索引53，Ai computing systems for llms training: a review，2024b，J. Comput. Sci. Technol.】，每个设备都持有一份完整的模型参数副本，但分配到不同的训练样本。反向计算完成后，所有设备的模型梯度会被聚合，然后用于更新模型参数。在模型并行中【索引26，Efficient large-scale language model training on GPU clusters using megatron-lm，2021b，SC】、【索引13，Gpipe: Efficient training of giant neural networks using pipeline parallelism，2019，NeurIPS】、【索引25，Memory-efficient pipeline-parallel DNN training，2021a，ICML】、【索引10，Advances of pipeline model parallelism for deep learning training: An overview，2024，J. Comput. Sci. Technol.】，模型参数被分布到多个设备上，每个设备只负责模型的一部分。为了完成前向和反向传播，需要进行通信操作来传输中间结果（即前向激活及其反向梯度）。

专家并行。如图1所示，专家并行【索引17，Gshard: Scaling giant models with conditional computation and automatic sharding，2021，ICLR】、【索引9，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，22，J. Mach. Learn. Res.】可以看作是模型并行和数据并行的结合。它像模型并行一样将专家参数分布在不同设备上，同时像数据并行一样在所有设备上复制其他参数。在每个MoE层中，每个令牌（token）由门控网络路由到top K个不同的专家进行处理，其中K是一个超参数，通常是一个较小的值，如1或2，这有助于降低计算复杂性。MoE层获得门控路由后，令牌会根据路由结果被发送到相应专家所在的设备。专家计算的结果随后被发送回令牌所在的原始设备。由于专家分布在不同的设备上，此过程中的通信涉及所有设备之间互相发送和接收消息，从而导致所谓的All-to-All通信。

2.2 MoE模型的分布式训练加速技术

动态专家放置。MoE模型的效率受到训练期间所需的大量且频繁的All-to-All通信的限制。为解决此问题，一些研究观察到数据在训练过程中倾向于偏好某些专家。基于此观察，他们进一步提出动态调整专家的放置以减少通信量【索引12，Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models，2022，PPoPP ’22】、【索引27，Flexmoe: Scaling large-scale sparse pre-trained model training via dynamic device placement，2023，Proc. ACM Manag. Data】、【索引51，Smartmoe: Efficiently training sparsely-activated models through combining offline and online parallelization，2023，USENIX ATC 2023】。例如，可以将热门专家以数据并行的方式放置在更多设备上，从而减少与它们相关的通信量。然而，由于专家规模日益增大，这些方法在设备间传输专家参数时会产生巨大的开销，因此无法在每次迭代中都调整专家放置，导致次优性。相比之下，我们的工作试图从一个不同的角度减少通信量：我们在每次迭代中动态调整样本的放置来加速All-to-All通信。具体来说，我们构建了一个优化问题来推导出最小化All-to-All通信时间成本的最佳样本放置方案。正如我们将在第4节中评估的，我们的工作在训练MoE模型时优于基于动态专家放置的现有工作。

模型定义修改。为了在MoE训练中实现更好的工作负载平衡，现有许多工作通过修改模型定义（例如，路由机制、模型架构）来实现。一些方法修改路由机制以平衡专家间的负载，这有助于减少设备间的同步时间【索引18，BASE layers: Simplifying training of large, sparse models，2021，ICML】。认识到分布式训练中的网络局部性，一些工作引入了路由拓扑损失，以优先在同一节点内路由令牌，从而减少节点间通信【索引20，Locmoe: A low-overhead moe for large language model training，2024，CoRR】、【索引4，Ta-moe: Topology-aware large scale mixture-of-expert training，2022，NeurIPS】。其他方法【索引50，Scomoe: Efficient mixtures of experts with structured communication，2023，ICLR】在节点间通信前将令牌映射到更小的隐藏层维度，进一步降低了通信负载。SCMoE【索引3，Shortcut-connected expert parallelism for accelerating mixture-of-experts，2024，CoRR】提出将当前注意力层的输出直接输入到下一个MoE层，从而实现与当前MLP层的前向传播并行，以完全重叠All-to-All通信与计算。尽管这些方法提高了训练效率，但它们不可避免地会影响模型收敛。在应用这些方法时，我们通常需要进行大量试验来调整超参数，例如调整拓扑感知路由损失的权重【索引4，Ta-moe: Topology-aware large scale mixture-of-expert training，2022，NeurIPS】或为不同通信信道调整超参数【索引50，Scomoe: Efficient mixtures of experts with structured communication，2023，ICLR】。鉴于LLM的每次训练可能需要数天甚至数月，它们的实用性不可避免地受到限制。相比之下，我们的工作专注于如何在不影响模型收敛的情况下加速All-to-All通信。

A2 方法细节

图 3: NetMoE方法的概览。

在本节中，我们介绍NetMoE，一个旨在通过同时考虑数据和网络局部性来优化MoE模型分布式训练的新颖框架。给定一个目标MoE模型和硬件环境，NetMoE旨在最小化All-to-All通信成本。其核心创新在于优化每个MoE层内样本的放置，以最大化利用更快的节点内带宽，从而减少通过较慢的节点间连接的通信量。具体来说，NetMoE在每个MoE层期间跨设备交换样本，使得更多令牌能够在All-to-All通信期间在节点内进行通信。

图3展示了本节的概览。我们首先在§3.1中介绍MoE训练中All-to-All通信的建模并构建我们的优化问题。然后在§3.2中说明如何解决该问题，详细算法如算法1所示。我们还在§3.3中介绍了我们的实现细节。为清晰起见，常用的符号列在表1中。

表 1: 本文使用的符号。我们假设 I 可被 J 整除，J 可被 N 整除，这在分布式训练中很常见。

表 2: 我们实验中使用的NVIDIA A800 GPU集群各通道的带宽。

3.1 问题构建

通信建模。我们首先讨论All-to-All通信的数学建模，这是NetMoE的优化目标。我们使用α-β模型【索引34，Connection-level analysis and modeling of network traffic，2001，IMW】来分析All-to-All通信，其中α代表延迟成本，β代表带宽成本。具体来说，我们将通信分为三类：设备内、节点内和节点间通信，每种通信使用不同的通道。表2列出了我们实验中使用的每个通道的带宽。由于设备内通信通常通过内存复制实现，其速度远快于另外两类，因此在我们的建模中不予考虑。因此，通信时间由跨节点内和节点间通道数据传输所需的最大时间决定。这些通道的带宽分别用$v_{intra}$和$v_{inter}$表示。因此，对于每次All-to-All通信，其时间成本可以用以下公式表示，其中$s·$代表相应通道的通信量。

$$\begin{aligned} \begin{aligned} t = \max(t_{intra}, t_{inter}), \quad \text{where} \quad & t_{intra} = \alpha_{intra} + \beta_{intra}s_{intra}, \beta_{intra} = 1/v_{intra}, \\ & t_{inter} = \alpha_{inter} + \beta_{inter}s_{inter}, \beta_{inter} = 1/v_{inter} \end{aligned} \end{aligned}$$

通信量计算。带宽（$v·$）和延迟（$α·$）可以在训练前通过分析硬件环境获得，而通信量（$s·$）需要根据MoE层内的路由结果动态确定。我们接着分析如何计算通信量。设$route \in \mathbb{N}^{I \times L \times K}$为门控网络的令牌路由结果，表示每个令牌将被发送到的K个专家。那么，第i个样本需要发送给第e个专家的令牌数量可以计算为：

$$num_{i,e} = \sum_{l,k} \mathbb{I}[route_{i,l,k} = e] \text{ for } i \in \lbrack\lbrack I \rbrack\rbrack, e \in \lbrack\lbrack E \rbrack\rbrack$$

通信量建模。接下来，可以使用$num \in \mathbb{N}^{I \times E}$来建模跨不同通道的通信量。设$ExpDev(e)$为第e个专家的设备索引，$SmpDev(i)$为第i个样本应该被路由到的设备索引，$Node(j)$为第j个设备的节点索引。通过将通信量视为需要传输的令牌数量，我们有：

$$s_{intra} = \sum_{(i, e) \in \mathcal{S}_{intra}} num_{i,e}, \quad s_{inter} = \sum_{(i, e) \in \mathcal{S}_{inter}} num_{i,e}$$

其中$S_{intra}$和$S_{inter}$可以通过专家和样本的设备索引计算得出：

$$\begin{aligned} \begin{aligned} S_{i n t r a} & =\{(i, e) \mid \operatorname{Node}(\operatorname{SmpDev}(i))=\operatorname{Node}(\operatorname{ExpDev}(e)) \wedge \operatorname{SmpDev}(i) \neq \operatorname{ExpDev}(e)\} \\ S_{i n t e r} & =\{(i, e) \mid \operatorname{Node}(\operatorname{SmpDev}(i)) \neq \operatorname{Node}(\operatorname{ExpDev}(e))\} \end{aligned} \end{aligned}$$

动态样本放置的合理性。根据上述建模，毫无疑问，All-to-All通信的时间成本与专家和样本的放置高度相关。在实践中，动态调整放置位置不会影响训练结果，因为All-to-All通信仍然被正确执行。结合网络局部性这一普遍事实，即$v_{intra} > v_{inter}$，我们可以调整样本和/或专家的放置，以减少节点间通信量，即使节点内通信量略有增加。事实上，出于类似的目标，一些现有工作已提出根据专家的受欢迎程度动态调整其放置【索引12，Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models，2022，PPoPP ’22】、【索引27，Flexmoe: Scaling large-scale sparse pre-trained model training via dynamic device placement，2023，Proc. ACM Manag. Data】、【索引51，Smartmoe: Efficiently training sparsely-activated models through combining offline and online parallelization，2023，USENIX ATC 2023】，如§2所述。然而，所有这些工作都忽略了数据局部性——同一样本的令牌通常被路由到同一专家【索引47，Openmoe: An early effort on open mixture-of-experts language models，2024，CoRR】、【索引15，Mixtral of experts，2024，CoRR】，从而错失了动态样本放置的优化机会。更重要的是，专家参数的大小通常远大于样本的大小。这使得先前的工作无法在每次迭代中调整专家放置。相比之下，样本放置的调整可以与All-to-All通信自然地融合（详见下文），需要零额外通信。因此，本工作专注于尚未探索的方面，旨在通过动态调整样本放置来加速MoE训练。

示例说明。为了帮助读者更好地理解动态样本放置的优势，我们以图2为例，其中$I = 4, L = 4, E = 4, K = 1$，专家和样本都按顺序放置，即$ExpDev = [0, 1, 2, 3], SmpDev = [0, 1, 2, 3]$。图2(b)显示了不改变样本放置的通信情况。根据公式3和公式4，如果我们只考虑节点0的发送量，那么$S_{inter} = \{(0, 2), (0, 3), (1, 2), (1, 3)\}$，表明$s_{inter} = 5$。然而，在优化样本放置后，如图2(c)所示，即$SmpDev = [3, 1, 2, 0]$，相应的节点间通信量变为$S_{inter} = \{(3, 2), (3, 3), (1, 2), (1, 3)\}$，这使得$s_{inter} = 2$。此外，值得注意的是，样本放置调整可以与All-to-All gather操作相结合。具体来说，令牌不是恢复到其原始位置，而是根据改变后的样本放置直接放置到其新位置。这种方法直接优化了当前的通信操作，而没有引入任何额外的通信。

问题公式化。在确定样本放置调整后，可以看出改变SmpDev会影响两次All-to-All操作：当前MoE层的gather操作和下一个MoE层的scatter操作。因此，我们的优化目标是这两次操作。对于第l层，优化问题可以写成如下形式。

$$\begin{aligned} \begin{aligned} \underset{\text{SmpDev}(i) \in \lbrack\lbrack J \rbrack\rbrack \text{ for } i \in \lbrack\lbrack I \rbrack\rbrack}{\arg \min } & t^{(l, \text {gather})} + t^{(l+1, \text {scatter})} \\ & =\max \left(t_{i n t r a}^{(l, \text {gather})}, t_{i n t e r}^{(l, \text {gather})}\right)+\max \left(t_{i n t r a}^{(l+1, \text {scatter})}, t_{i n t e r}^{(l+1, \text {scatter})}\right) \\ \text { s.t. } \quad & \sum_{i \in \lbrack\lbrack I \rbrack\rbrack} \mathbb{I}[\operatorname{SmpDev}(i)=j]=I / J \text { for } j \in \lbrack\lbrack J \rbrack\rbrack \end{aligned} \end{aligned}$$

由于一次样本放置的改变会影响两次All-to-All操作，因此优化目标中包含了两次通信时间。此外，为了确保设备间的计算和内存平衡，每个设备在样本放置调整前后应保留相同数量的样本。这构成了我们优化中的约束条件。

3.2 问题求解

公式5是一个复杂的组合优化问题，无法在多项式时间内得到最优解。随着集群规模的增加，即使找到一个近似解也可能需要大量时间。由于这个问题需要在每次gather操作之前解决，直接求解会导致无法承受的额外时间成本。为了解决这个问题，我们设计了一种高效的方法来获得近似解。具体来说，我们首先将优化问题分解为两个阶段，并开发了一个多项式时间算法来实现求解，如下所述。

两阶段分解。尽管公式1取两种通信成本的最大值，但在实践中，由于节点间和节点内连接的带宽差异巨大，最耗时的通常是节点间通信项。因此，我们提出了一种两阶段求解策略：第一阶段在全局范围内优化$t_{inter}$，而第二阶段在每个节点内最小化$t_{intra}$，而不影响$t_{inter}$。形式上，假设有N个节点，每个节点由J/N个设备组成，那么第一阶段的优化公式可以写成如下的整数线性规划（ILP）问题：

$$\begin{aligned} \begin{aligned} \underset{\text{Node}(\text{SmpDev}(i)) \in \lbrack\lbrack N \rbrack\rbrack \text{ for } i \in \lbrack\lbrack I \rbrack\rbrack}{\arg \min} \quad & t_{inter}^{(l, gather)} + t_{inter}^{(l+1, scatter)} \\ \text{s.t.} \quad & \sum_{i \in \lbrack\lbrack I \rbrack\rbrack} \mathbb{I} [\text{Node}(\text{SmpDev}(i)) = n] = I / N \text{ for } n \in \lbrack\lbrack N \rbrack\rbrack \end{aligned} \end{aligned}$$

第二阶段优化。在公式5中，设备间的平衡约束在第一阶段被转化为节点间的平衡约束，因为我们专注于节点间通信。在获得第一阶段的最优解后，第二阶段考虑在每个节点内单独重新排列样本。对于第n个节点，记$I^*_n \subseteq I$为求解公式6后分配给它的样本集合。并令$J_n = \{j|j \in [J] \land Node(j) = n\}$为驻留在其上的专家集合（$J_n$由设备放置决定，而不是通过公式6获得）。然后，为了优化第n个节点，我们应该解决以下ILP问题：

$$\mathop{\arg\min}_{\mathrm{SmpDev}(i) \in \lbrack\lbrack J \rbrack\rbrack_{n} \text{ for } i \in \lbrack\lbrack I \rbrack\rbrack_{n}^{*}} t_{intra}^{(l,gather)} + t_{intra}^{(l+1,scatter)} \quad \text{s.t. } \sum_{i \in \lbrack\lbrack I \rbrack\rbrack_{n}^{*}} \mathbb{I}[\mathrm{SmpDev}(i) = j] = I / J \text{ for } j \in \lbrack\lbrack J \rbrack\rbrack_{n}$$

分解示例。具体来说，图2可以被看作是第一阶段的优化，而图4则展示了在其基础上进行的第二阶段优化。尽管第二阶段包含N个ILP问题，每个节点一个，但它们是独立的，可以并行求解。

图 4: 第二阶段优化示例。图4(a)显示了在图2(c)中第一阶段优化后Node1中的MoE层。通过在节点内应用第二阶段优化，节点内通信可以减少1个令牌（通过交换sample0和sample2），如图4(b)所示。

多项式时间求解器。通过分解原始的组合优化问题，我们得到了N+1个ILP问题，这些问题可以通过现有的库如PuLP【索引24，Pulp: a linear programming toolkit for python，2011，The University of Auckland】来解决。然而，回想一下，我们需要在每个训练迭代的每一层中解决这些问题，问题求解的效率至关重要。不幸的是，由于ILP问题是NP-hard的，当我们尝试用PuLP解决它们时，求解的时间成本超过了scatter通信和专家计算的时间成本（如§4.4中的评估），使其不切实际。考虑到每个样本必须分配给一个设备，我们通过将ILP问题转化为加权二分图匹配问题，重新考虑它们为分配问题，并随后开发了一种基于广泛使用的Kuhn-Munkres（KM）算法的多项式时间求解器。

问题转化。我们首先介绍如何将ILP问题转化为加权二分图匹配问题。令$c_{i,n}$和$c'_{i,j}$分别表示将第i个样本放置在第n个节点的第j个设备上时的节点间和节点内通信量。它们可以用以下公式计算：

$$\begin{aligned} \begin{gathered} c_{i, n}=\sum_{e \in S} n u m_{i, e}, \quad c_{i, j}^{\prime}=\sum_{e \in S^{\prime}} n u m_{i, e}, \quad \text { where } \\ S=\{e \mid \operatorname{Node}(\operatorname{ExpDev}(e)) \neq n\}, S^{\prime}=\{e \mid \operatorname{Node}(\operatorname{ExpDev}(e))=\operatorname{Node}(j) \wedge \operatorname{ExpDev}(e) \neq j\} . \end{gathered} \end{aligned}$$

目标函数重构。为了使表达式更清晰，令$p_{i,n}, p'_{i,j} \in \{0, 1\}$分别表示第i个样本是否被放置在第n个节点和第j个设备上。那么，优化目标可以表示为：

$$\begin{aligned} \begin{gathered} t_{inter} = \alpha_{inter} + \beta_{inter} \sum_{i \in \lbrack\lbrack I \rbrack\rbrack, n \in \lbrack\lbrack N \rbrack\rbrack} c_{i,n} p_{i,n}, \quad t_{intra} = \alpha_{intra} + \beta_{intra} \sum_{i \in \lbrack\lbrack I \rbrack\rbrack, j \in \lbrack\lbrack J \rbrack\rbrack} c'_{i,j} p'_{i,j}, \quad \text{where} \\ p_{i,n} = \mathbb{I} [\text{Node}(\text{SmpDev}(i)) = n], \quad p'_{i,j} = \mathbb{I} [\text{SmpDev}(i) = j] \quad \text{for } i \in \lbrack\lbrack I \rbrack\rbrack, n \in \lbrack\lbrack N \rbrack\rbrack, j \in \lbrack\lbrack J \rbrack\rbrack \end{gathered} \end{aligned}$$

算法1 NetMoE优化

(0,1)-ILP问题转化。在修改了相应的约束条件后，我们将ILP问题转化为(0,1)-ILP问题。例如，下面展示了第一阶段（即公式6）的转化问题：

$$\begin{aligned} \begin{aligned} \underset{p_{i, n} \text { for } i \in \lbrack\lbrack I \rbrack\rbrack, n \in \lbrack\lbrack N \rbrack\rbrack}{\arg \min } & \alpha_{\text {inter }}+\beta_{\text {inter }} \sum_{i \in \lbrack\lbrack I \rbrack\rbrack, n \in \lbrack\lbrack N \rbrack\rbrack}\left(c_{i, n}^{(l, g a t h e r)}+c_{i, n}^{(l+1, s c a t t e r)}\right) p_{i, n} \\ \text { s.t. } & \sum_{i \in \lbrack\lbrack I \rbrack\rbrack, n \in \lbrack\lbrack N \rbrack\rbrack} p_{i, n}=I / N \text { for } n \in \lbrack\lbrack N \rbrack\rbrack, \quad \sum_{n \in \lbrack\lbrack N \rbrack\rbrack} p_{i, n}=1 \text { for } i \in \lbrack\lbrack I \rbrack\rbrack \end{aligned} \end{aligned}$$

二分图匹配建模。这个(0,1)-ILP问题可以被建模为一个加权二分图匹配问题。具体来说，考虑一个具有两组图节点P和Q的二分图。集合P代表所有训练样本，且$|P| = I$。集合Q代表所有训练节点（机器），其中每个训练节点可以处理 $B := I/N$ 个训练样本。为了对此建模，Q中的每个图节点被复制B次，从而得到$|Q| = I$。在P和Q的每对图节点之间都存在一条加权边。令$P_i$表示第i个训练样本，$Q_n$表示第$\lfloor n/B \rfloor$个训练节点。$P_i$和$Q_n$之间的边权重表示为$W_{i,n} = c^{(l,gather)}_{i,\lfloor n/B \rfloor} + c^{(l+1,scatter)}_{i,\lfloor n/B \rfloor}$。这种转换将问题简化为在这个二分图中寻找一个最小权重完美匹配，这可以使用Kuhn-Munkres（KM）算法在多项式时间内高效地求得最优解。

二分图示例。图5展示了在图2第一阶段构建二分图的一个例子。左侧的图节点代表集合P，右侧的图节点代表集合Q。每对图节点之间都由一条带权重的边连接，用虚线表示。红色的边表示最终的匹配方案，其中所有匹配边的总权重最小。

图 5: 二分图示例。

3.3 实现

NetMoE是在PyTorch【索引30，Pytorch: An imperative style, highperformance deep learning library，2019，NeurIPS】之上实现的，自定义操作（例如，num, c, c′的计算和KM算法）用C++和CUDA实现。NetMoE的完整工作流程如算法1所示。除了在§3.2中介绍的问题求解外，NetMoE还通过以下方式进行了优化。

专家残差内联。在经典的MoE模型中，残差连接独立于MoE层。然而，在NetMoE中，训练数据的位置在All-to-All gather操作后发生了变化，而残差连接中的样本仍保留在原始位置。为确保模型的正确性，我们将残差连接内联到专家计算中，如算法1的第12行所示。这种优化确保了应用该算法前后模型精度的一致性。关于内联的更多细节在附录A.1中阐述。

表 3: 评估模型的配置。

表 4: 不同求解器的时间成本与All-to-All scatter和专家计算的总时间成本（毫秒）的对比（MoE-GPT-S, J = 16）。

卸载求解过程。KM算法难以并行化，因此不适合像GPU这样的高度并行化加速器，所以我们在CPU上执行求解过程。如算法1的第9行所示，在获得当前层的路由结果后，每个设备计算并将num传输到CPU内存。优化算法所需的下一层的路由结果可以通过直接将当前层的输入传递给下一层的路由器来预测【索引8，Fast inference of mixture-of-experts language models with offloading，2023，CoRR】、【索引39，HOBBIT: A mixed precision expert offloading system for fast moe inference，2024，CoRR】。求解过程只需要在All-to-All gather操作之前提供新的样本位置。通过这种方式，求解过程可以与All-to-All scatter和专家计算重叠。正如我们将在§4.4中展示的，求解时间被完全隐藏，因此引入了零开销。关于算法选择和重叠潜力的更多讨论在附录A.2中描述。

A4 实验环境

实验设置：
* 对比方法：NetMoE与基于动态专家放置的最新方法进行比较，包括FasterMoE 【索引12，Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models，2022，PPoPP ’22】和SmartMoE 【索引51，Smartmoe: Efficiently training sparsely-activated models through combining offline and online parallelization，2023，USENIX ATC 2023】。同时，将FastMoE 【索引11，Fastmoe: A fast mixture-of-expert training system，2021，CoRR】作为不调整专家或样本位置的基线。
* 硬件配置：实验在一个由4个节点组成的集群上进行，每个节点配备8个NVIDIA A800-SXM4-40GB GPU。根据表2，节点内的GPU通过NVLink连接，带宽为400 GB/s；节点间通过InfiniBand互连，带宽为100 GB/s。
* 模型架构：实验采用GPT模型架构【索引32，Language models are unsupervised multitask learners，2019，OpenAI blog】、【索引2，Language models are few-shot learners，2020，NeurIPS】作为骨干网络，并将模型中的所有FFN层替换为MoE层。具体模型配置如表3所示。由于SmartMoE要求每个设备至少有2个专家，因此设置专家数量E = 2 × J（J为GPU数量），每个token选择的专家数量K = 2。
* 软件配置：默认情况下，每个节点使用8个GPU进行实验。所有报告的结果均为50次迭代的平均值。

A4 实验结果

端到端性能

如图6所示，NetMoE相较于FastMoE、FasterMoE和SmartMoE，分别取得了最高1.67倍、1.37倍和1.33倍的加速比。
* 与FasterMoE对比：FasterMoE通过重叠专家计算和支持动态专家放置实现了显著优化。但随着模型隐藏维度的增加，与专家通信的成本上升，使其难以维持同等级别的加速效果，导致其与NetMoE之间存在性能差距。
* 与SmartMoE对比：SmartMoE在FasterMoE的优化基础上，通过调整专家位置来确保负载均衡，因此性能优于FasterMoE。然而，SmartMoE主要关注计算负载的均衡，而没有重点优化通信效率。当通信成为主要瓶颈时，负载均衡的优势就不那么明显。
* NetMoE的优势：NetMoE通过动态调整样本放置，始终优于这些最先进的系统。值得注意的是，NetMoE的方法与动态专家放置兼容，未来可以通过结合两者实现更高的效率。

图 6: 不同方法的端到端加速比（均值和标准差）。

All-to-All性能

如图7所示，实验对比了应用NetMoE前后All-to-All通信的性能差异，并与求解器提供的理论优化值进行了比较。
* 实验结果：可以看出，All-to-All通信的实际加速比略低于理论值。
* 分析结论：这种差异是合理的，因为对All-to-All通信的建模是基于理想条件，并未考虑潜在的路由冲突或硬件引起的错误。附录C提供了更多关于All-to-All通信加速的实验结果分析。

图 7: All-to-All通信成本的实际和理论加速比。

求解器性能

为了验证求解器的效率，实验比较了不同规模下求解器的时间与All-to-All scatter和专家计算的总时间，结果如表4所示。
* 实验内容：对比了NetMoE中使用的KM算法和通用线性规划工具包PuLP【索引24，Pulp: a linear programming toolkit for python，2011，The University of Auckland】的求解时间。
* 实验结果：尽管KM算法的求解时间随着样本数I的增加呈超线性增长，但在各种场景下，其求解过程始终能被All-to-All scatter和专家计算的时间所覆盖（隐藏）。相比之下，PuLP的求解时间很难被完全覆盖。
* 分析结论：这凸显了在有高实时性要求的场景中设计专门优化方法的必要性。

A5 结论

本文提出了NetMoE，旨在优化MoE模型训练中的主要瓶颈——All-to-All通信。通过利用数据和网络局部性，我们的方法在训练期间动态调整训练样本的放置，将节点间通信转化为节点内通信，从而提升All-to-All通信效率。我们将All-to-All通信时间和样本放置问题建模为一个优化问题，并设计了一种多项式时间的方法来解决它。实验结果表明，NetMoE在训练效率上比现有的MoE训练系统高出最多1.67倍。

A6 附录

A.1 专家残差内联细节

残差连接的调整。如图8所示，原始的残差加法方法是将注意力（attention）输出与gather操作得到的结果相加。然而，在NetMoE中，它被调整到在scatter操作之后、gather操作之前进行。这种内联方式有助于样本放置的调整，同时确保了计算的正确性。

图 8: 包含和不包含专家残差内联的Transformer层示意图。

A.2 算法选择与重叠潜力的讨论

KM算法的选择依据。我们的设计采用KM算法是基于两个实际因素：（1）尽管KM算法的时间复杂度为$O(I^3)$，但由于GPU内存有限，当前的训练过程通常采用梯度累积【索引40，Gradient accumulation tensorflow，2019，Tensorflow】、【索引31，Gradient accumulation pytorch，2019，Pytorch】。因此，$I$的值通常被限制在一个可接受的大小，确保求解时间可以被有效重叠；（2）该算法的运行时间与通信阶段完全重叠，使得为隐藏求解器开销而进行的进一步加速变得不必要。虽然存在更快的近似求解器【索引29，New scaling algorithms for the assignment and minimum mean cycle problems，1992，Math. Program.】、【索引6，Linear-time approximation for maximum weight matching，2014，J. ACM】，但在当前计算-通信重叠已经掩盖了优化时间的训练配置中，它们的好处将是微乎其微的。

B 每节点使用较少GPU的端到端性能

不同GPU配置下的性能。图9展示了每个节点配置2个或4个GPU时的端到端加速比。结果表明，NetMoE在各种实验设置下均取得了最佳性能，这与每个节点有8个GPU时的结果（如图6所示）是一致的。

实验配置的合理性。值得注意的是，标准服务器配置通常每个节点最多容纳8个NVIDIA GPU。因此，每节点8个GPU代表了大型语言模型分布式训练的标准设置【索引7，The llama 3 herd of models，2024，CoRR】、【索引1，Nemotron-4 340b technical report，2024，CoRR】、【索引5，Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models，2024，ACL】、【索引35，BLOOM: A 176bparameter open-access multilingual language model，2022，CoRR】。尽管像NVIDIA GB200 NVL72这样的超级计算机支持超过8个GPU之间的高速连接（例如NVLink），但它们依赖于定制硬件且价格高昂。在超级计算机上的训练场景很少见，并且与GPU集群或云中的典型场景有显著不同。因此，本文选择使用每节点最多8个GPU的配置进行实验。

图 9: 不同总设备数（表示为J）和节点数（表示为N）下的端到端加速比（均值和标准差）。

C All-to-All通信优化的详细分析

为了更深入地了解NetMoE优化的来源，我们评估了两种统计数据：
* NetMoE分别在节点间或设备间交换的训练样本比例。比例越高，表示在节点间/设备间调整的样本越多。
* 应用NetMoE前后的节点内和节点间通信量。

通信量与样本调整比例分析。首先，表5总结了所有迭代中的均值和标准差。应用NetMoE后，很大一部分训练样本在节点间被交换，从而减少了节点间的通信量。值得注意的是，尽管应用NetMoE后，节点内通信量占总通信量（$s_{intra}+s_{inter}$）的比例很大（即$s_{intra}$增加），但这不会成为性能瓶颈，因为节点间通信带宽要低得多。因此，由于样本放置调整带来的节点间通信量减少，All-to-All通信得以加速。

动态路由下的鲁棒性分析。其次，由于MoE模型的训练过程中路由结果是动态变化的，为了探究路由分布的影响，我们在图10中绘制了（1）节点间通信的减少量，和（2）跨节点交换的样本比例，这两个指标在不同迭代中的变化。同时，我们遵循先前的工作【索引12，Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models，2022，PPoPP ’22】、【索引27，Flexmoe: Scaling large-scale sparse pre-trained model training via dynamic device placement，2023，Proc. ACM Manag. Data】记录了不同迭代中专家选择的分布，以描述路由分布。可以观察到，在模型训练过程中路由分布是变化的。然而，NetMoE在给定动态分布的情况下，通过调整样本放置，始终能减少节点间通信。因此，NetMoE的有效性对路由分布具有鲁棒性。

表5: 通信量和序列调整比例摘要。对于通信量，我们提供了应用NetMoE前后的节点内和节点间通信量，括号内给出了增加或减少的比例。对于序列调整比例，“跨节点”表示跨节点交换的序列比例，“全部”表示所有被调整的序列比例。

(a) 2 节点, 16 GPUs

(b) 4 节点, 32 GPUs

图 10: 左：节点间通信量的减少量。中：跨节点交换的样本比例。右：专家选择的分布（第0层）。