PopFetcher: Towards Accelerated Mixture-of-Experts Training Via Popularity Based Expert-Wise Prefetch

作者/机构: Junyi Zhang, Chuanhu Ma, Xiong Wang, and Yuntao Nie, 华中科技大学; Yuqing Li, 武汉大学; Yuedong Xu, 复旦大学; Xiaofei Liao, 华中科技大学; Bo Li, 香港科技大学; Hai Jin, 华中科技大学

A1 主要贡献

本文旨在解决混合专家（Mixture-of-Experts, MoE）模型训练中的通信和计算瓶颈。随着模型参数规模扩展至万亿级别，MoE架构虽然通过稀疏激活专家（expert）实现了亚线性的训练计算增长，但其训练效率受到严重的All-to-All通信开销和计算负载不均衡问题的制约。

核心问题：
1. 通信瓶颈：MoE训练中，输入令牌（token）需要在不同GPU工作节点（worker）之间通过两次All-to-All通信进行分发和结果收集，这在模型规模较大时极其耗时。
2. 计算负载不均衡：由于门控网络（gate network）的动态路由，部分“热门”专家会接收远超其他专家的令牌，导致托管这些专家的GPU负载过重，形成计算瓶颈。
3. 现有方案的局限性：现有方法如复制部分专家或在训练前静态决策，要么引入额外的参数同步开销，要么无法适应动态变化的专家负载，且专家调度与All-to-All通信并发执行，可能相互干扰。

研究目标与创新点：
为了解决上述问题，本文提出了一个名为PopFetcher的可扩展MoE训练系统，其核心思想是基于专家流行度的预取（prefetching）。
1. 发现并利用专家选择的规律：本文揭示了MoE模型中专家选择存在的偏斜（skewed）和跨层相关（correlated）模式。基于此，开发了一种轻量级的滑动窗口预测方案，能够准确预测下一层中即将成为热点的专家。
2. 提出混合拉取-推送（pull-push）与预取机制：PopFetcher设计了一种专家拉取和令牌推送的混合模式。它在执行当前非MoE层计算时，利用空闲的网络链接，预先拉取（prefetch）下一层中可能的热门专家。这减少了即将到来的All-to-All通信中需要分发的令牌数量，从而缓解通信瓶颈。
3. 建立端到端延迟模型与最优预取策略：PopFetcher严格地形式化了端到端的训练延迟，综合考虑了计算和通信成本。通过一个定制的剪枝策略，系统能够推导出全局最优的专家预取方案，该方案能根据底层网络设施恢复通信和计算的平衡，同时满足GPU的内存限制。
4. 优化反向传播中的网络流调度：在反向传播过程中，PopFetcher优先处理All-to-All数据流，而不是预取专家的All-Reduce梯度聚合流，并通过流水线化调度非MoE层的All-Reduce任务与MoE的反向计算，显著减轻了网络争用和通信阻塞。

图1：MoE训练期间的All-to-All通信。
图1：MoE训练期间的All-to-All通信。

A3 背景与动机

2.1 背景知识

2.1.1 基于Transformer的预训练语言模型（PLM）

Transformer模型架构：Transformer模型【索引38，Attention is all you need，2017，Advances in Neural Information Processing Systems】已成为序列处理的基础架构。一个典型的Transformer块包含一个注意力（Attention）层和一个多层感知机（MLP）。MLP通常由两个全连接（FC）层及其间的ReLU激活函数组成。注意力网络将输入令牌转换为查询（query）、键（key）和值（value）矩阵，经过缩放点积注意力处理后，结果被送入MLP。MLP的输出会与该块的输入相加，并通过层归一化处理。通过堆叠多个相同的Transformer块，可以构建出强大的PLM，如BERT【索引8，BERT: pre-training of deep bidirectional transformers for language understanding，2019，NAACL】和GPT【索引5，Language models are few-shot learners，2020，NeurIPS；索引29，Language models are unsupervised multitask learners，2019，OpenAI blog】。

图2：MoE块的结构。
图2：MoE块的结构。

2.1.2 MoE架构与模型训练

MoE架构解决扩展性挑战：尽管Transformer架构取得了成功，但扩展PLM容量因训练计算需求的增加而充满挑战。MoE模型【索引31，Outrageously large neural networks: The sparselygated mixture-of-experts layer，2017，ICLR】通过将密集的MLP划分为多个稀疏的专家来解决此问题，如图2所示。这种架构利用Transformer的内在稀疏性，以亚线性的计算成本增加来扩展模型规模。训练期间，每个输入令牌由一个门控网络路由到一个或多个专家，通常采用top-k选择（k常为1或2）。令牌到专家的分发及结果返回通过All-to-All通信原语管理。

专家并行（EP）训练：大规模MoE模型的内存需求远超单个GPU容量。为方便训练，专家并行（EP）被用于将每个MoE层的不同专家分布到不同的GPU工作节点上【索引21，Gshard: Scaling giant models with conditional computation and automatic sharding，2021，ICLR】。通常，位于MoE层之间的非MoE层（如注意力层）会像数据并行【索引19，A unified architecture for accelerating distributed DNN training in heterogeneous GPU/CPU clusters，2020，OSDI】一样被复制到每个工作节点上。

2.2 动机

专家激活偏斜导致热点专家：在MoE架构中，门控模块路由输入令牌的可变性常常导致专家间工作负载不均衡。如图3所示的MoE-GPT模型【索引41，Mpmoe: Memory efficient moe for pre-trained models with adaptive pipeline parallelism，2024，IEEE Trans. Parallel Distributed Syst.】中一个层内的专家选择分布，令牌分发模式在训练早期演变显著，但很快呈现出局部化活动的趋势，且相邻迭代间的变化微小。尽管有平衡门控策略（如对每个专家施加容量限制【索引21，Gshard: Scaling giant models with conditional computation and automatic sharding，2021，ICLR】），某些专家仍会持续接收更多令牌，给托管这些热点专家的GPU带来更重的计算负担。这种负载不均衡在MoE架构的前沿层尤为明显，而在深层则趋于稳定。热点专家的出现要求有效的负载分布管理以避免GPU工作节点过载。

图3：一个MoE层内偏斜的专家选择（y轴上的k代表1000个单位）。
图3：一个MoE层内偏斜的专家选择（y轴上的k代表1000个单位）。

MoE训练中的All-to-All通信瓶颈：在MoE模型训练期间，每个令牌在每个MoE层都要经历两次All-to-All通信（令牌分发和输出收集），这常导致显著延迟。表1显示，在八个工作节点上训练MoE-GPT时，All-to-All通信可占总时间的50%-60%。主要的通信瓶颈源于每个MoE层中强制的数据同步，即专家执行和后续计算必须等待所有必需的令牌及其处理输出都接收完毕后才能开始。当存在热点专家时，这种同步问题尤其严重，因为它们增加的负载会进一步加剧网络争用。因此，精心设计的专家并行（EP）对于解决同步执行挑战和减少通信开销至关重要。

表1：一个MoE层内的时间消耗。

表1：一个MoE层内的时间消耗。
表1：一个MoE层内的时间消耗。

粗粒度的专家调度是不够的：为提升MoE训练效率，一种方法是在多个GPU上复制部分专家，以减轻过载工作节点的计算偏斜，使得特定令牌可以由“影子”专家在本地处理，而非远程分发【索引16，Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models，2022，PPoPP】。虽然有效，但向所有GPU广播专家参数可能因数据同步而引入意外延迟。另一种方法是尝试将必要的专家参数拉取到本地GPU，以减少将令牌推送到远程专家的需求【索引24，Janus: A unified distributed training framework for sparse mixture-of-experts models，2023，SIGCOMM】。然而，当专家参数远大于输入令牌时，这种方法会遇到显著限制，使得拉取专家变得更加通信密集。尽管有这些努力，现有的粗粒度MoE方案仍然难以解决计算负载不均衡和持续的All-to-All通信瓶颈问题。

图4：不同专家调度方案的执行时间线。
图4：不同专家调度方案的执行时间线。

2.3 机会与挑战

解决通信瓶颈同时恢复计算平衡对于提高MoE模型训练效率至关重要。接下来我们将探讨实现这一目标的机会，并指出潜在的挑战。

2.3.1 机会

提升MoE效率的努力通常沿着两个正交方向进行。
* 模型设计：文献中尝试改进MoE模型架构本身。例如，GShard【索引21，Gshard: Scaling giant models with conditional computation and automatic sharding，2021，ICLR】引入了负载均衡损失，以在不同GPU工作节点间平均分配工作负载。类似地，专家路由策略【索引42，Mixture-of-experts with expert choice routing，2022，NeurIPS】被重新定义，允许门控模块选择top-k个最匹配的令牌，以增强负载均衡。虽然这些方法有助于实现更公平的工作负载分布，但它们也可能牺牲模型精度，且通信瓶颈依然存在。
* 系统设计：面向系统的方法，如FasterMoE【索引16，Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models，2022，PPoPP】和FlexMoE【索引28，Flexmoe: Scaling large-scale sparse pre-trained model training via dynamic device placement，2023，Proceedings of the ACM on Management of Data】，通过“影子”或复制热门专家来减轻负载不均衡的不利影响。然而，它们可能为参数同步或系统调整引入额外开销，从而可能抵消改善工作负载分布带来的好处。

尽管存在这些挑战，MoE层之间的非MoE组件为优化提供了机会。图4展示了MoE层内的数据流，表明在非MoE部分（如注意力层）的计算期间，一个工作节点仅使用本地输入数据，导致网络链接未被充分利用。随着训练的进行，专家负载变化趋于稳定，这为利用历史令牌分布进行专家预取提供了机会。也就是说，我们可以在预期下一MoE层中的热门专家被选择之前，主动地从远程GPU将它们拉取到本地工作节点，从而利用当前空闲的链接来减少未来的All-to-All通信需求。

2.3.2 挑战

在非MoE计算阶段预取热门专家有望最小化通信开销和平衡GPU工作负载，但也带来了一些需要解决的独特挑战。
* 如何准确预测下一MoE层的热门专家？ 随着训练的进行，专家间的负载分布动态变化。为了进行有效的专家调度，开发一种能够准确预测哪些专家在后续层中需求量大且轻量级的方法至关重要，但这因其动态性而具有挑战性。
* 如何在工作节点容量限制下预取合适的专家？ 鉴于GPU工作节点的内存容量有限，且本地GPU上已存在多个专家，决定预取哪些以及多少专家是一个难题。这个决策必须在专家预取的好处与可用内存资源之间取得平衡。
* 如何管理与专家预取相关的额外成本？ 虽然预取可以安排与非MoE计算重叠，但它必须遵守MoE工作流的数据依赖性。确保预取的专家在下一层计算开始前准备就绪，可能会引入额外的同步开销。此外，预取专家间的All-Reduce操作可能与All-to-All通信竞争，可能阻碍反向传播中的梯度计算。

A2 方法细节

3 PopFetcher 概览

本文提出了PopFetcher，一个基于流行度的专家级预取高级MoE训练系统，旨在减少All-to-All通信开销并确保GPU工作负载均衡。PopFetcher在每个工作节点上维护一个专家池，用于存储本地专家参数和从远程工作节点预取来的参数。图5展示了PopFetcher的系统架构，主要由三个关键模块组成：路由信息收集器、预取决策器和异步调度执行器。

图5：PopFetcher的系统架构。
图5：PopFetcher的系统架构。

3.1 路由信息收集器

功能与实现：PopFetcher中的每个工作节点都配备了一个路由信息收集器，用于跟踪和监控专家的运行时流行度。该模块记录每个令牌在每个MoE层的门控选择细节，并更新路由到每个专家的令牌分布。基于这些数据，收集器更新每个MoE层内的专家流行度，并将其传递给预取决策器以进行更明智的专家检索。每个工作节点编译的专家流行度通过torch.distributed.all_reduce迅速同步，以保持所有工作节点间的一致性。由于流行度向量很小，同步开销可以忽略不计。

3.2 预取决策器

决策过程：预取模块定期聚合所有工作节点的路由信息收集器数据，以识别每个专家的流行度。考虑到每个GPU工作节点已托管本地专家和前向传播的相应激活值，因内存限制，获取所有剩余专家是不切实际的。根据更新的流行度统计数据，决策器在有限的GPU内存限制内为每个工作节点制定量身定制的预取策略。PopFetcher精心编排决策过程以避免干扰常规训练活动，因为路由信息可以在GPU进行训练操作的同时，在CPU上进行并发分析。此外，PopFetcher决定为即将到来的MoE层预取专家，以便在当前层进展时最小化端到端的训练延迟。

图6：一次迭代中MoE层之间的专家相关性（使用每层四个专家的MoE-GPT进行训练）。
图6：一次迭代中MoE层之间的专家相关性（使用每层四个专家的MoE-GPT进行训练）。

3.3 异步调度执行器

执行策略：调度执行器旨在通过抢先并异步地从远程工作节点获取下一层的专家，来利用MoE训练期间的空闲网络链接。这一策略减少了All-to-All通信中传输的令牌量，从而有效解决了通信瓶颈。理想情况下，专家调度应与当前Transformer非MoE部分（即注意力层）的计算同时进行，这确保了最小的开销，并避免了对当前MoE层执行的干扰。

具体操作与优化：通常，执行器根据预取决策器的指令，从全局池中优先选择热门专家，为每个工作节点智能地调度专家。因此，PopFetcher不仅减轻了All-to-All通信量，还确保了均衡的GPU工作负载。考虑到由于专家预取，多个专家可能存在于不同的工作节点上，PopFetcher在反向传播期间并发执行非MoE层的All-Reduce操作和专家的梯度计算，以减少通信开销。为进一步最小化对反向All-to-All通信的影响并避免计算阻塞，我们还优先处理All-to-All数据流，使其优于All-Reduce流。

4 轻量级流行度预测

本节将深入探讨一种动态的专家流行度预测方法，该方法被设计为既轻量级又准确，以增强专家预取效果。在标准场景中，专家激活是由门控网络动态预先确定的，这可能限制了系统级优化的机会。

4.1 MoE层之间的专家相关性

专家选择的规律性：如前所述，输入令牌对专家的选择表现出可辨别的规律性。在单个MoE层内，某些专家持续吸引更多令牌，显示其流行度，而其他专家则较少被选中。图3的观察也表明，一度流行的专家倾向于在时间上保持其流行度。

跨层相关性分析：除了单个MoE层，连续层之间的专家选择也存在相关性。假设一个令牌在第$i$层MoE层选择的专家索引为$x$，在随后的第$(i+1)$层为$x'$。元组$(x, x')$代表一组相关的选择。图6展示了在给定前一层选择了专家$x$的情况下，当前层选择专家$x'$的条件概率。结果揭示了相邻MoE层之间专家选择的明显趋势，某些相关性比其他更普遍。本质上，下一层的热门专家通常可以根据前一层中的专家偏好来预测。

规律性的稳定性：尽管专家的流行度及其相关性随着训练的进行而不断演变，但这些变化随时间趋于平缓，这支持了在非MoE计算期间进行早期专家预取的有效性。

4.2 专家流行度预测

基于滑动窗口的预测方法：实现工作节点间计算和通信的最佳平衡取决于准确预测专家流行度，这项任务因其动态性而变得复杂。为了解决这个问题同时最小化对正常训练的干扰，我们提出了一种基于滑动窗口的预测方法来精确定位热门专家【索引6，Prediction is all moe needs: Expert load distribution goes from fluctuating to stabilizing，2024，CoRR】。具体来说，路由信息收集器在前向传播期间定期从门控网络收集运行时选择数据，然后分析这些数据以更新我们当前对专家流行度的理解。

表2：主要符号说明

表2：主要符号说明
表2：主要符号说明

历史令牌分布分析：首先，我们分析最近几次迭代直到当前MoE层的令牌分布：

公式1
公式1

这里$s$代表滑动窗口的大小，根据我们的测试，最佳设置为10次迭代。

结合跨层相关性进行预测：我们的目标是准确预测即将到来的MoE层的专家流行度，这涉及到整合专家相关性，如图6所示。让$M$表示令牌总数，每个令牌表示为$T_m, m \in \{1,2,···,M\}$。我们用$E_{i,j}$表示第$j$层（即Transformer块）的第$i$个专家，每层总共有$K'$个专家。那么，任何专家$E_{i,j}$和$E_{h,j+1}$之间的相关性可以通过它们的条件概率来评估：

公式2
公式2

基于相关性的专家流行度由下式给出：
公式3
公式3

公式(3)通过利用前一层及时更新的选择数据，实现了对下一层专家流行度的准确、轻量级预测。因此，PopFetcher可以动态适应专家需求的变化，而不会给系统带来过重负担。

5 专家级预取与调度

本节详细介绍了我们在有限GPU内存下进行专家预取和调度的策略。为了方便后续介绍，我们在表2中总结了主要符号。

图7：混合推拉范式。
图7：混合推拉范式。

5.1 混合推拉范式

现有方法的局限性：由于专家选择偏斜导致的通信和计算不平衡，严重延长了MoE训练时间。以往的尝试通过发送（推送）令牌给专家（以专家为中心）【索引12，Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2022，JMLR；索引16，Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models，2022，PPoPP；索引21，Gshard: Scaling giant models with conditional computation and automatic sharding，2021，ICLR】或将专家拉取到本地GPU（以数据为中心）【索引24，Janus: A unified distributed training framework for sparse mixture-of-experts models，2023，SIGCOMM】来解决此问题，然而，这在管理MoE架构固有的不均匀令牌分发方面被证明是不足的，如图7所示。由于这些方法只发送令牌或拉取专家，它们无法处理大数据传输期间的阻塞，因为所有数据块都被同等对待并同时调度。

推拉决策的量化分析：具体来说，每个专家模块通常包含一个前馈网络，其中有两个线性层，矩阵维度分别为$H \times \alpha H$和$\alpha H \times H$，其中$\alpha$通常设为2。因此，单个专家层的参数量为$4H^2$。使用float32数据类型，当$H=1024$时，参数达到16MB。此外，在两次All-to-All通信期间，每个令牌的跨机数据传输量为$2H$。因此，我们推断当令牌传输量超过2048个时，拉取专家变得更可行。否则，发送令牌的传统方法开销更低。

PopFetcher的混合策略：我们引入了一种新的令牌和专家混合推拉策略，它结合了以专家为中心和以数据为中心的框架的优点。根据图7，令牌通过门控模块后，其到不同工作节点的分发差异显著。通过采用同时推送令牌和拉取专家的混合范式，我们可以根据当前专家间的令牌分布来优化数据传输。具体来说，当令牌的体积超过专家参数的体积时，我们选择将专家拉取到数据所在位置；否则，我们默认发送令牌，因为专家参数负担较小。

预取计算模式：我们的混合方法还包括专家预取，以实现远程专家的提前检索。为此，我们需要适应两种MoE计算模式：本地计算和预取计算。在本地计算中，令牌被发送到其特定专家进行处理，结果返回到始发节点。预取计算涉及被预取并可在本地使用的专家，从而消除了通过网络发送令牌的需要。然而，预取计算引入了一个额外的步骤来确保模型更新的逻辑完整性。即，在反向传播期间，复制专家的梯度必须发送回托管主专家的工作节点进行全局归约。梯度同步虽然至关重要，但不可避免地会消耗网络带宽，并且必须在预取决策中加以考虑。

5.2 问题形式化

我们的目标是设计一种包含专家预取的最佳调度策略，以最小化每个MoE层的训练延迟。在此背景下，通信和计算延迟主要由最慢工作节点的性能决定。

5.2.1 无预取时的训练延迟

延迟分析：首先，我们考虑不包括专家预取的常规情况。设$B_w$为工作节点$w$接收的令牌数，则有$B_w = \sum_{n=1}^{N} T_{n,w}$，其中$T_{n,w} = I_n \sum_{i=1}^{K} p_{i,w}$表示基于专家流行度预测$p_i$从工作节点$n$来的$I_n$个令牌。对于Transformer中的每个MoE层，专家处理令牌涉及前向传播中的一次GeMM操作，以及反向传播中的两次（因为需要计算输入和专家参数的梯度）。此外，FC层之间的中间嵌入向量大小为$\alpha H$，其中$H$是每个令牌的嵌入大小。因此，MLP中两个FC层的总操作数为$4B_w \alpha H^2$。假设数据类型为float32，每个GeMM需要$4B_w \alpha H^2 / P_w$的时间，导致计算延迟为$3 \times 4B_w \alpha H^2 / P_w$。此外，前向和反向传播共包含四轮All-to-All通信，通信延迟为$4H \sum_{n=1}^{N} \sum_{i=1}^{K} B_{i,w}^{n} / W$，其中$B_{i,w}^{n}$表示工作节点$n$发送给专家$E_{i,w}$的令牌数。

延迟公式：因此，工作节点$w$的训练延迟为：

公式4
公式4

由于总延迟取决于最慢的工作节点，一个MoE层的端到端训练时间变为：
公式5
公式5

5.2.2 有预取时的训练延迟

预取引入的延迟：基于公式(5)的推导，我们继续探讨预取热门专家对训练延迟的影响。具体来说，专家预取会产生两种延迟：
1. 预取专家计算延迟：这源于将专家参数从远程工作节点取到本地节点进行计算，受专家流行度及应路由到它的本地令牌数量影响。
2. 模型梯度归约延迟：将预取专家的梯度与主专家进行归约所需的时间。

详细延迟分项：
* 前向传播：(1) 计算时间：本地令牌的计算时间为 $Comp_{f}^l = 4B_w\alpha H^2 / P_w$，预取专家的计算时间为 $Comp_{f}^p = 4\alpha H^2 / P_w \sum_{n=1}^{N} \sum_{i=1}^{K} B_{i,w}^{n} \delta_{i,w}^{n}$，其中指示符$\delta_{i,w}^{n}$表示是否预取。(2) 通信时间：令牌传输时间为 $Comm_{f}^t = 2H \sum_{n=1}^{N} \sum_{i=1}^{K} B_{i,w}^{n} (1-\delta_{i,w}^{n}) / W_{n,w}$。
* 反向传播：(1) 计算时间：反向传播中，本地令牌计算时间加倍为 $Comp_{b}^l = 2 \times 4B_w\alpha H^2 / P_w$，预取专家计算需要 $Comp_{b}^p = 2 \times 4\alpha H^2 / P_w \sum_{n=1}^{N} \sum_{i=1}^{K} B_{i,w}^{n} \delta_{i,w}^{n}$。(2) 通信时间：令牌传输$Comm_{b}^t$与前向传播的$Comm_{f}^t$相同，梯度归约时间为 $Comm_{b}^r = 2\alpha H^2 \sum_{n=1}^{N} \sum_{i=1}^{K} \delta_{i,w}^{n} / W_{n,w}$。

总延迟公式：综合以上各项，工作节点$w$的总延迟为：

公式6
公式6

我们的目标是最小化端到端延迟，即：
公式7
公式7

考虑到全局专家的数量庞大（例如Azure上的256个GPU，每个工作节点128个专家【索引30，Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation AI scale，2022，ICML】），决定是否预取一个专家的开销是巨大的，需要进一步优化。

5.3 专家预取决策

我们的目标转向解决公式(7)中概述的问题，重点是以高效的方式优化预取专家到GPU工作节点的分配$\{\delta_{i,w}^{n}\}$。

5.3.1 专家预取剪枝

挑战与剪枝策略：设计专家预取的主要挑战是决定下一层专家的巨大搜索空间。为了克服这个问题，我们通过专家剪枝来缩小搜索空间，有效地将需要考虑的专家数量减少到最多$k \times N$，其中$k$表示门控模块的top-k专家选择。

剪枝约束条件：
* GPU内存限制：由于每个GPU已经托管了本地专家参数和中间激活，可以预取的专家总大小不得超过其可用内存。

公式8
公式8

* 传输时间约束：预取应理想地与非MoE层计算重叠，这意味着获取专家所需的时间必须满足：
公式9
公式9

这些约束为可行的预取专家数量建立了一个上限$r_w$：
公式10
公式10

预取收益分析：此外，对任何工作节点的专家预取都应该减少端到端训练延迟。根据公式(4)和(6)，这意味着以下差值大于0：

公式11
公式11

具体来说，当以下条件成立时，$Lat_{origin,w} > Lat_{prefetch,w}$：
公式12
公式12

这里，为简洁起见，我们用$\epsilon$作为计算带宽比$W_{n,w}/P_w$的简写。观察到只有当$\epsilon$超过$3\alpha H$时，预取才能提高训练效率，这表明了工作节点间带宽有限但GPU计算能力强大的场景。例如，NVIDIA DGX B200配备了72 PFLOPS的计算能力和400 Gb/s的NVLink连接。给定典型的嵌入大小$H \in \{768, 1024, 2048, 4096\}$，很明显$\epsilon > 3\alpha H$，即专家预取是可行的。在这种情况下，如果专家$E_{i,n}$被预取，它接收的令牌必须满足：
公式13
公式13

通过结合公式(10)和公式(13)，我们根据专家的流行度来优先预取专家，直到达到GPU内存限制。如前所述，专家流行度随着训练的进行趋于稳定。因此，在MoE训练的中后期，可能不需要在每次迭代中评估每个专家的预取需求。相反，可以采用固定的预取策略，或减少专家重新规划的频率，以进一步降低操作开销。

PopFetcher机制概述：总结来说，PopFetcher的运行机制如下。以公式(6)中的端到端延迟为条件，我们旨在通过搜索最优的专家预取方案来最小化延迟。鉴于随着机器和GPU数量的增加，搜索空间的复杂性呈指数级增长，我们实施了两个关键约束来管理这种复杂性：公式(10)中定义的GPU内存限制和公式(13)中指定的专家传输时间约束。这些约束有效地修剪了搜索空间，使我们能够制定一个全面的全局专家预取策略。具体来说，对于我们策略中确定的每组$\delta_{i,w}^{n}$，我们在非MoE层计算期间执行将工作节点$n$上的第$i$个专家$E_{i,n}$预取到工作节点$w$的操作。

图8：不同类型的GPU机器拓扑结构。
图8：不同类型的GPU机器拓扑结构。

5.3.2 本地GPU间的内部专家共享

利用异构网络拓扑：在GPU集群中，连接各种组件的链路通常具有异构带宽，如图8所示。通常，同一台机器内的GPU通过NVLink通信，带宽可达1800GB/s，而通过PCIe与机器主CPU互连的GPU共享64GB/s的带宽。对于服务器间通信，机器通常通过GDR NIC连接，带宽可达400Gb/s，专家必须先通过NIC拉取到本地CPU内存，然后通过PCIe移动到GPU内存。鉴于这种多样化的网络结构，当一个工作节点希望从其他节点获取专家时，结合异构连接显然可以提高带宽利用率。具体来说，工作节点将优先选择通过更高带宽链路连接的源来检索所需的专家参数。

CPU内存作为共享缓存：此外，内部专家可以通过CPU内存在所有本地工作节点间共享。一旦某个专家在同一台机器上可用，PopFetcher就绕过外部获取的需要，即工作节点直接从本地CPU内存访问该专家，以减轻专家预取的开销。为此，我们为每台服务器维护一个缓存管理器，负责在本地工作节点间共享远程专家参数，并利用CPU内存作为中介来加快预取过程。

最优调度策略：总的来说，我们的调度方案旨在通过最快的可用链路优化地检索专家，无论具体的机器内部和机器间带宽如何。通过利用CPU内存，PopFetcher有效地避免了专家参数的冗余传输，从而提高了MoE训练系统的整体效率。

图9：反向传播中数据流的流水线调度。
图9：反向传播中数据流的流水线调度。

6 反向传播中的流调度

多流竞争问题：由于专家预取，一个专家的多个副本可能存在于不同的工作节点上，这需要在反向传播期间进行额外的AllReduce操作。通常，EP的All-to-All通信、非MoE层的All-Reduce梯度聚合以及预取专家间的All-Reduce都在不同的进程组中处理，这将启动三个不同的CUDA流并发运行。如果没有适当的调度，这些流可能会竞争有限的网络带宽和GPU资源，可能扰乱反向计算的正常流程。

理想调度与现实限制：理想情况下，应为All-to-All通信分配严格的优先级以最小化带宽争用。该策略涉及在All-Reduce任务就绪时立即启动以防止延迟，同时允许All-to-All任务在启动时抢占GPU资源。然而，使用像NCCL这样的高度优化的多GPU通信库实现如此精确的流调度是不可行的，因为NCCL在调用点将通信原语锁定到CUDA流中。因此，传输必须预先确定，排除了与其他操作结合进行核心和带宽资源的任何实时调整。

流水线化微操作解决方案：为了克服这些挑战，我们建议将All-Reduce和All-to-All通信分解为微操作，并以流水线方式执行。实际上，专家预取导致反向传播期间All-to-All令牌传输量减少，从而减少了它们与专家梯度计算的重叠。通过策略性地交错All-Reduce和All-to-All通信的微操作，如图9生动展示的那样，我们可以充分利用GPU工作节点在其他任务上忙碌时的网络空闲期。

A4 实验环境

7 实现

PopFetcher是基于PyTorch开发的，使用Python、C++和CUDA构建了MoE训练框架，代码量超过8000行。
* 软件配置：
* 路由信息收集器：使用Python实现，通过All-Gather算子聚合运行时数据。
* 专家调度执行器：利用torch.distributed的通信能力，核心专家预取逻辑由C++和CUDA编写。
* 并行化：通过torch.cuda.Stream实现计算和通信任务的并行处理，并管理一个专用的预取流。
* 集成：PopFetcher实现为PyTorch插件，可独立使用或集成到Megatron-LM框架【索引34，Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，CoRR】中。它使用torch.autograd.Function类自定义MoE算子的前向和反向行为，无缝集成到计算图中。
* 流水线调度：计算和通信流的流水线调度在C++和CUDA中实现以优化性能。
* 依赖库：PyTorch 2.3，CUDA 12.4，通信后端基于NCCL。

8.1 实验设置

硬件配置

数据集与模型

表3：不同MoE模型的配置

表3：不同MoE模型的配置
表3：不同MoE模型的配置

基线系统

PopFetcher与以下框架进行了比较：
* DeepSpeed【索引30，Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation AI scale，2022，ICML】：支持大规模MoE模型的高性能系统。
* FasterMoE【索引16，Fastermoe: modeling and optimizing training of large-scale dynamic pre-trained models，2022，PPoPP】：通过在GPU工作节点间“影子”专家来简化专家管理的系统。
* Megablocks【索引13，Megablocks: Efficient sparse training with mixture-of-experts，2023，MLSys】：一个轻量级的MoE训练库，特点是块稀疏GPU核。
* Tutel【索引17，Tutel: Adaptive mixture-of-experts at scale，2023，MLSys】：实现了自适应专家执行和流水线调度的MoE系统。
* Janus【索引24，Janus: A unified distributed training framework for sparse mixture-of-experts models，2023，SIGCOMM】：一个以数据为中心的MoE训练系统。由于其源码未公开，实验中通过预取所有专家来复现其功能。

所有基线系统都可集成到Megatron-LM【索引27，Efficient large-scale language model training on GPU clusters using megatron-lm，2021，SC】中，并在相同的软硬件环境下进行公平比较。

评估指标

A5 实验结果

8.2 统计训练等效性

实验内容：为验证PopFetcher在加速训练的同时不影响模型收敛性，在集群A上使用MoE-GPT模型，对比了传统FasterMoE和PopFetcher的训练损失曲线。实验覆盖了朴素top-k和GShard两种专家路由机制。
实验结果：如图10所示，无论采用哪种门控机制，PopFetcher的损失曲线与基线（FasterMoE）完全重合。
分析结论：PopFetcher通过额外的梯度归约操作，确保了预取专家和原始专家之间的梯度更新同步，从而在不牺牲模型准确性的前提下实现了训练加速，达到了统计上的等效性。

图10：训练迭代过程中的损失值。
图10：训练迭代过程中的损失值。

8.3 整体端到端性能

实验内容：在两个GPU集群上，使用GShard门控机制，评估了PopFetcher与多个基线系统在每次迭代训练时间上的性能。
实验结果：如图11所示，PopFetcher在集群A和集群B上分别实现了1.28-2.4倍和1.18-18.3倍的训练速度提升。Janus由于其“拉取所有专家”的策略对GPU内存要求极高，在实验设置下频繁导致OOM错误，故未包含在图中。
分析结论：PopFetcher的性能提升得益于其专家预取设计和网络流优先级调度，它能在非MoE层计算期间并发地获取热门专家参数，且不阻塞反向传播。尤其在像集群B这样的低带宽环境中，PopFetcher表现依然稳健，证明了其在高计算带宽比（ε）场景下的有效性。

图11：相较于基线的加速比。
图11：相较于基线的加速比。

8.4 消融研究

8.4.1 专家流行度与流调度

实验内容：在集群A上使用MoE-GPT和MoE-BERT模型，对比了PopFetcher与两种变体：（1）随机预取专家；（2）不使用优化的流调度策略。
实验结果：如图12所示，基于流行度的专家调度相比随机预取，为MoE-GPT和MoE-BERT模型分别带来了1.30倍和1.26倍的训练加速。此外，PopFetcher的流水线化流调度策略，为MoE-GPT和MoE-BERT分别减少了10.9%和10%的平均迭代时间。
分析结论：准确的专家流行度预测和精细的反向传播流调度是PopFetcher性能提升的关键。前者确保了预取的高效性，后者避免了网络争用和计算阻塞。

图12：专家与流调度效果。
图12：专家与流调度效果。

图13：令牌传输量。
图13：令牌传输量。

8.4.2 专家预取对令牌传输的影响

实验内容：评估了PopFetcher的预取策略对All-to-All通信中令牌传输量的减少效果，并可视化了哪些专家被频繁预取。
实验结果：如图13所示，在集群A上，全局预取策略为MoE-GPT和MoE-BERT分别减少了14.85%和13.46%的令牌传输量。图14展示了训练过程中被预取的专家（深色单元格表示更频繁），显示出某些专家在连续迭代中被持续请求，呈现出“局部性”。
分析结论：专家预取有效减少了通信量。专家流行度的“局部性”验证了预测和预取策略的有效性，即集中预取少数热门专家就能显著提升性能。

图14：PopFetcher在训练迭代中预取专家的展示。
图14：PopFetcher在训练迭代中预取专家的展示。

8.5 混合推拉范式分析

实验内容：在集群A上使用MoE-GPT和MoE-BERT模型，对比了PopFetcher的混合推拉模式与纯数据中心（Janus）和纯专家中心（无预取）模式的训练吞吐量和迭代时间。
实验结果：如图15所示，PopFetcher在两个指标上均优于单一模式的系统。
分析结论：PopFetcher的混合通信方法结合了发送令牌和拉取专家的优点，优化了All-to-All通信中的数据传输量，有效解决了负载不均衡的粒度问题，从而提升了整体训练效率。即使发生“坏预取”（预取的专家接收令牌较少），由于预取操作与非MoE计算重叠，PopFetcher也能无缝退回到传统训练模式，不会产生额外开销。

图15：PopFetcher混合令牌发送和专家拉取的效果。
图15：PopFetcher混合令牌发送和专家拉取的效果。

8.6 参数敏感性分析

实验内容：分析了滑动窗口大小$s$对专家流行度预测准确率的影响。由于GPU内存限制，通常每层预取的专家不超过4个，因此评估了预测top-5热门专家与实际top-5的准确率。
实验结果：如表4所示，在评估了不同配置下的预测准确率后，选择滑动窗口大小为10，以在预测性能和内存消耗之间达到最佳平衡。
分析结论：滑动窗口大小$s=10$是实现高预测准确率和资源高效利用的最佳选择。

表4：不同滑动窗口大小下的预测准确率

表4：不同滑动窗口大小下的预测准确率
表4：不同滑动窗口大小下的预测准确率

8.7 平衡GPU工作负载

实验内容：比较了使用PopFetcher前后，负载最轻和最重的工作节点接收令牌数量的差异。
实验结果：如图16所示，PopFetcher显著减少了工作节点间的负载差异：MoE-GPT减少了43.1%，MoE-BERT减少了57.1%。
分析结论：通过预取热门专家并由多个工作节点共享其负载，PopFetcher实现了更均衡的工作负载分布，避免了传统MoE系统中因专家选择偏斜导致的负载不均和网络拥塞。

图16：负载最轻和最重工作节点接收令牌数量的差异（每层16个专家）。
图16：负载最轻和最重工作节点接收令牌数量的差异（每层16个专家）。

8.8 GPU内存消耗

实验内容：通过调整每层MoE的专家数量，比较了PopFetcher、FasterMoE和Janus在集群A上能够容纳的最大模型参数规模。
实验结果：如表5所示，与FasterMoE相比，PopFetcher在MoE-GPT和MoE-BERT上分别提升了12.3%和20.1%的模型规模。与Janus相比，提升幅度分别为49.0%和58.2%。
分析结论：PopFetcher的混合推拉范式能更有效地利用GPU内存，从而支持更大规模模型的训练。

表5：可容纳的最大模型规模

表5：可容纳的最大模型规模
表5：可容纳的最大模型规模

8.9 运行时开销

实验内容：评估PopFetcher引入的额外运行时开销。
实验结果：在实际测试中，流行度预测的开销小于100毫秒。
分析结论：PopFetcher的运行时开销极小。因为专家流行度预测和预取方案搜索是在CPU上异步执行的，不干扰主训练流程。同时，剪枝策略显著减小了搜索空间，使得计算成本可以忽略不计。这点开销完全被预取热门专家带来的巨大收益所覆盖。

A6 结论

本文提出了PopFetcher，一个有效且可扩展的MoE训练框架，它通过基于流行度的专家级预取来缓解MoE训练中的All-to-All通信瓶颈。PopFetcher识别出专家选择中的偏斜和相关模式，并采用基于滑动窗口的预测来确定未来层的专家流行度。基于此统计数据，PopFetcher精确分析了每个MoE层的端到端训练延迟，并为GPU工作节点建立了最优的预取专家映射。通过这样做，PopFetcher不仅最小化了通信时间，还确保了工作节点间的负载均衡。此外，PopFetcher利用网络空闲期将All-Reduce流与非MoE计算进行流水线化处理，从而加速了反向传播过程。在真实世界数据集上的实验评估表明，PopFetcher优于现有的最先进MoE训练系统，能够将训练时间减少15%-94.5%。