作者/机构: Shigang Li (shigangli.cs@gmail.com), Torsten Hoefler (htor@inf.ethz.ch)
Department of Computer Science, ETH Zurich Switzerland

A1 主要贡献

核心问题：训练大规模深度学习模型极具挑战性。随着模型规模的增长（例如，GPT-3拥有1750亿参数），简单的并行化方案（如数据并行）已不再适用，因为模型无法装入单个加速器内存。为了解决这个问题，模型必须被分布在数百个加速器上。模型分发主要有两种方式：（1）算子并行，将一个层内的算子拆分到多个加速器上；（2）流水线并行，将模型按层分发。算子并行通常通信开销巨大，而流水线并行虽然通信成本较低（只需点对点通信），但会遭受“气泡”（bubbles，即设备空闲）或权重陈旧（staleness）问题的困扰。此外，现有的流水线方法（如DAPPLE、PipeDream）存在严重的激活内存消耗不均衡问题，流水线的第一个加速器需要存储比最后一个加速器多得多的激活值，这限制了微批次大小，从而影响了整体性能。

研究目标：本文旨在提出一种新的流水线并行方案，以解决现有方法的效率瓶颈。该方案的目标是：
1. 保持训练过程的完全同步，避免权重陈旧问题。
2. 相比现有同步方法，显著提高流水线利用率（即减少气泡）。
3. 实现更均衡的激活内存消耗，提高内存资源利用率。
4. 易于配置，并能通过精确的性能模型指导其在不同神经网络和系统架构上的应用。

创新点：本文提出了Chimera，一种新颖的流水线并行方案，其核心思想是在同一组加速器上运行完全填充的双向流水线。Chimera的主要贡献和特点如下：
* 完全同步：Chimera是一种同步方法，不依赖陈旧的权重进行训练，保证了模型的收敛性和准确性。
* 更高的流水线利用率：与现有的同步方法相比，Chimera能将气泡数量减少高达50%，从而显著提升性能。
* 均衡的内存消耗：通过精巧的双向流水线调度，Chimera实现了比现有技术更均衡的激活内存消耗，同时峰值内存消耗与最先进方法相当。
* 易于配置和性能建模：Chimera可以通过一个精确的性能模型进行配置，以适应不同的深度神经网络和系统架构。

例如，在Piz Daint超级计算机的2048个GPU节点上训练一个13亿参数的GPT-2模型时，Chimera相比最先进的同步和异步流水线方法，训练吞吐量提升了1.16倍至2.34倍。

图1：在2048个GPU节点上，针对mini-batch size=2048的GPT-2模型。我们展示了每种方法的气泡率（流水线开销）、内存成本（OOM需要激活重计算，用R表示），以及最佳吞吐量（p是对应的流水线阶段数）。细节将在第2节和第4.2节讨论。
图1：在2048个GPU节点上，针对mini-batch size=2048的GPT-2模型。我们展示了每种方法的气泡率（流水线开销）、内存成本（OOM需要激活重计算，用R表示），以及最佳吞吐量（p是对应的流水线阶段数）。细节将在第2节和第4.2节讨论。

A3 背景知识与相关工作

小批量随机梯度下降（SGD）：小批量SGD【索引6，Optimization methods for large-scale machine learning，2018，SIAM Rev.】是训练深度神经网络的主流方法。设$M$为小批量大小，$W_t$为第$t$步的神经网络权重，$(x_i, y_i)$为小批量中的一个样本，$\ell$为损失函数。训练过程中，它为每个样本计算前向传播的损失$\ell(f(x_i, W_t), y_i)$，然后在反向传播中计算随机梯度：

公式：SGD梯度计算
公式：SGD梯度计算

模型通过迭代进行训练，使得 $W_{t+1} = W_t - \eta \nabla L_t$。更广义的一阶随机梯度更新规则可以有不同形式（例如，通过添加动量项），表示为 $W_{t+1} = W_t + \Delta(W_t, L(0,...,M), \eta)$。

并行化训练：为了将训练过程扩展到并行机器上，数据并行【索引20，Accurate, large minibatch sgd: Training imagenet in 1 hour，2017，arXiv；索引33，Taming unbalanced training workloads in deep learning with partial collective operations，2020，PPoPP；索引49，Horovod: fast and easy distributed deep learning in TensorFlow，2018，arXiv；索引60，Imagenet training in minutes，2018，ICPP】是常用方法，即将小批量数据在$P$个工作节点间划分，每个节点维护一份完整的模型副本。梯度同步通常通过allreduce实现。然而，对于参数量从百万级扩展到十亿级的现代深度学习模型【索引7，Language models are few-shot learners，2020，arXiv；索引13，Bert: Pre-training of deep bidirectional transformers for language understanding，2018，arXiv；索引43，Language models are unsupervised multitask learners，2019，OpenAI blog；索引45，Regularized evolution for image classifier architecture search，2019，AAAI】，纯数据并行可能因同步整个模型梯度而效率低下，或者模型太大而无法装入单个设备。

模型并行策略：算子并行是一种通过在多个工作节点间划分层的算子来训练大模型的方法，但如第1节所述，它可能会遭受高通信量的困扰。结合算子并行和数据并行的混合方法【索引29，Beyond data and model parallelism for deep neural networks，2019，SysML；索引30，One weird trick for parallelizing convolutional neural networks，2014，arXiv】也面临与纯算子并行类似的问题。为减少算子并行的通信量，流水线并行被广泛研究【索引16，DAPPLE: a pipelined data parallel approach for training large models，2021，PPoPP；索引19，AMPNet: Asynchronous model-parallel training for dynamic neural networks，2017，arXiv；索引26，Gpipe: Efficient training of giant neural networks using pipeline parallelism，2019，NeurIPS；索引28，GEMS: GPU-Enabled Memory-Aware Model-Parallelism System for Distributed DNN Training，2020，SC20；索引38，PipeDream: generalized pipeline parallelism for DNN training，2019，SOSP；索引39，Memory-efficient pipeline-parallel DNN training，2020，arXiv；索引55，FPDeep: Scalable acceleration of CNN training on deeply-pipelined FPGA clusters，2020，IEEE Trans. Comput.；索引57，PipeMare: Asynchronous Pipeline Parallel DNN Training，2019，arXiv】。其关键思想是按层划分模型，并将每个工作节点（及其上的层）视为一个流水线阶段。小批量被划分为多个微批次，这些微批次在各阶段间流水化执行以提高资源利用率。然而，高效地流水化深度神经网络训练是具有挑战性的，原因在于：（1）一个训练步骤包含一次前向传播和一次反向传播；（2）反向传播中的梯度计算依赖于前向传播的中间结果；（3）为获得良好的收敛精度，小批量大小通常不会非常大。本文将分析现有流水线方法如何处理这些挑战，并展示我们的方法如何在各方面实现最佳平衡。

流水线中的气泡：为了获得更好的收敛质量，同步方法在每次训练迭代结束时同步梯度并清空流水线，如图2所示。因此，同步方法会导致流水线气泡。为了利用流水线，GPipe【索引26，Gpipe: Efficient training of giant neural networks using pipeline parallelism，2019，NeurIPS】同时向流水线中注入$m$个微批次；DAPPLE【索引16，DAPPLE: a pipelined data parallel approach for training large models，2021，PPoPP】使用带有周期性清空的“一前一后”（1F1B）【索引38，PipeDream: generalized pipeline parallelism for DNN training，2019，SOSP；索引39，Memory-efficient pipeline-parallel DNN training，2020，arXiv】调度。GPipe和DAPPLE都会产生$2(p-1)$个气泡（即前向传播有$p-1$个气泡，反向传播有$p-1$个气泡）。相比之下，Chimera（本文工作）产生$p-2$个气泡（即前向传播有$p/2-1$个气泡，反向传播有$p/2-1$个气泡），与DAPPLE和GPipe相比减少了约50%。我们将气泡率定义为气泡开销除以流水线的总运行时间。实践中，反向传播的工作量大约是前向传播的两倍，这导致Chimera在流水线中间产生$p/2-1$个气泡（气泡率 = $(p-2)/(3m/2+p-2)$），如图2所示。我们将在3.5节讨论如何移除中间的气泡（气泡率 = $(p-2)/(2m+p-2)$）。

表1：本文中常用的符号列表。

表1：本文中常用的符号列表。
表1：本文中常用的符号列表。

表2：不同流水线方案的比较。

表2：不同流水线方案的比较。
表2：不同流水线方案的比较。

图2：流水线并行方案，在一个训练迭代内有四个流水线阶段（p=4）和四个微批次（m=4），除了PipeDream在每个微批次的后向传播后更新模型。
图2：流水线并行方案，在一个训练迭代内有四个流水线阶段（p=4）和四个微批次（m=4），除了PipeDream在每个微批次的后向传播后更新模型。

微批次数量的限制：表2展示了考虑前向和反向传播典型工作量下所有方法的气泡率。虽然GPipe和DAPPLE的气泡率随着$m$（每个工作节点在一个训练迭代中执行的微批次数）的增加而减少，但通常无法获得足够大的$m$（如【索引26，Gpipe: Efficient training of giant neural networks using pipeline parallelism，2019，NeurIPS】建议的$m \geq 4p$）而不损害效率，原因有三：（1）对于一个模型，通常存在一个经验上的最大$M̂$（小批量大小），超过这个值会影响模型收敛【索引5，Demystifying parallel and distributed deep learning: An in-depth concurrency analysis，2019，ACM Computing Surveys；索引12，Why Dataset Properties Bound the Scalability of Parallel Machine Learning Training Algorithms，2021，IEEE Transactions on Parallel and Distributed Systems；索引58，Large-batch training for LSTM and beyond，2019，SC；索引59，Large batch optimization for deep learning: Training bert in 76 minutes，2019，arXiv；索引60，Imagenet training in minutes，2018，ICPP】。（2）对于给定的$M̂$，增加$m$意味着减小$b$（微批次大小）。然而，现代加速器需要足够大的$b$才能实现高计算效率。（3）通过与数据并行结合来扩展到大规模机器（这已被证明是有效的方法【索引16，DAPPLE: a pipelined data parallel approach for training large models，2021，PPoPP；索引39，Memory-efficient pipeline-parallel DNN training，2020，arXiv】）会对于给定的$M̂$减小$m$。GEMS【索引28，GEMS: GPU-Enabled Memory-Aware Model-Parallelism System for Distributed DNN Training，2020，SC20】是一种内存高效的流水线方法，它在两个模型副本之间调度微批次。由于GEMS主要为小的$M̂$设计，并且同时最多有两个活跃的微批次，其气泡率远高于其他方法，并且不能通过更大的$m$来缓解。异步方法（如PipeDream【索引38，PipeDream: generalized pipeline parallelism for DNN training，2019，SOSP】和PipeDream-2BW【索引39，Memory-efficient pipeline-parallel DNN training，2020，arXiv】）没有周期性的流水线清空，因此没有气泡问题，但存在权重陈旧。

内存消耗：内存开销主要来自两个方面：权重参数和激活值（前向传播的中间结果，在反向传播中计算梯度时需要）。对于GPipe和DAPPLE，每个工作节点维护一个流水线阶段的权重。对于GEMS和Chimera（默认设置），每个工作节点维护两个流水线阶段的权重，因为有两个方向的流水线。PipeDream在每个微批次的反向传播后更新模型（$m=1$）；因此，为了确保前向和反向传播之间的权重版本一致性，它需要在单个工作节点上存储多达$p$个版本的权重，这与纯数据并行的内存成本相同。通过使用梯度累积（$m \geq p$），PipeDream-2BW将需要存储的权重版本数减少到2。GEMS在流水线开始时只注入一个微批次，因此只存储一个微批次前向传播的激活值。然而，这导致了前面讨论的低流水线效率。由于GPipe同时向流水线中注入$m$个微批次（为充分利用流水线$m \geq p$），激活内存消耗与$m$成正比，这对于大的小批量大小扩展性不佳。使用经典（或变体）的1F1B【索引38，PipeDream: generalized pipeline parallelism for DNN training，2019，SOSP；索引39，Memory-efficient pipeline-parallel DNN training，2020，arXiv】调度，PipeDream、PipeDream-2BW、DAPPLE和Chimera在流水线开始时最多注入$p$个微批次，这对于大的小批量大小扩展性很好。通过计算图2中Chimera每个工作节点上注入的微批次数，我们可以观察到Chimera在工作节点间的激活内存消耗更均衡（通用分析见表2），优于PipeDream、PipeDream-2BW和DAPPLE，因此内存资源利用率更高。请注意，激活内存消耗可以通过激活重计算【索引11，Training deep nets with sublinear memory cost，2016，arXiv】技术来减少，但这会增加约1/3的计算开销【索引16，DAPPLE: a pipelined data parallel approach for training large models，2021，PPoPP；索引39，Memory-efficient pipeline-parallel DNN training，2020，arXiv】。ZeRO【索引44，Zero: Memory optimizations toward training trillion parameter models，2020，SC20；索引46，ZeRO-Offload: Democratizing Billion-Scale Model Training，2021，arXiv】通过在数据并行进程间划分三个模型状态（即优化器状态、梯度和参数）来消除内存冗余，但通信量略有增加。我们的流水线方法与ZeRO是正交的。进一步减少我们流水线方法的内存消耗是一个有趣的未来工作。

收敛友好性：通过周期性的流水线清空，同步方法确保在一次训练迭代中所有阶段和所有微批次都使用相同版本的权重，不会引入陈旧性。从算法角度看，同步方法等同于标准且经过充分验证的小批量SGD，因此保证收敛。为了消除流水线清空，异步方法要么激进地导致前向和反向传播之间的权重版本不匹配（如AMPNet【索引19，AMPNet: Asynchronous model-parallel training for dynamic neural networks，2017，arXiv】和PipeMare【索引57，PipeMare: Asynchronous Pipeline Parallel DNN Training，2019，arXiv】），要么保守地引入权重陈旧性，同时确保权重版本的一致性（如PipeDream-2BW和PipeDream）。尽管它们在经验上显示出有希望的收敛结果，但缺乏通用性的证明。最近的研究【索引4，Stochastic gradient push for distributed deep learning，2019，ICML；索引34，Breaking (Global) Barriers in Parallel Stochastic Optimization with Wait-Avoiding Group Averaging，2020，IEEE Transactions on Parallel and Distributed Systems；索引36，Asynchronous Decentralized Parallel Stochastic Gradient Descent，2018，ICML；索引37，SwarmSGD: Scalable decentralized SGD with local updates，2019，arXiv；索引52，Communication-efficient distributed deep learning: A comprehensive survey，2020，arXiv】观察到异步训练算法可能导致较低的收敛性能。

模型准确性：对于模型准确性，所有同步流水线方法（如Chimera、DAPPLE、GPipe和GEMS）都保证能达到与标准小批量SGD算法相同的准确性。对于异步方法（如PipeDream-2BW和PipeDream），由于引入了权重陈旧性，不能保证达到与标准算法理想的准确性，并且收敛质量可能在不同的神经网络和任务上表现出差异。

总结：总的来说，如表2所示，Chimera在所有方面都实现了最佳平衡。我们将在下一节讨论Chimera的实现细节。

A2 方法细节

3.1 双向流水线

核心思想：双向流水线：我们考虑具有重复结构的大规模模型（即同一模块重复多次），如Bert【索引13，Bert: Pre-training of deep bidirectional transformers for language understanding，2018，arXiv】和GPT-2/3【索引7，Language models are few-shot learners，2020，arXiv；索引43，Language models are unsupervised multitask learners，2019，OpenAI blog】，这些模型可以被划分为具有相等数量模块的平衡阶段。PipeDream-2BW【索引39，Memory-efficient pipeline-parallel DNN training，2020，arXiv】也利用了这一重复结构特性。如何将任何模型高效地划分为多个阶段不是本文的主题，已在近期工作中得到充分研究【索引16，DAPPLE: a pipelined data parallel approach for training large models，2021，PPoPP；索引38，PipeDream: generalized pipeline parallelism for DNN training，2019，SOSP】。Chimera的核心思想是将两个不同方向的流水线（我们分别称之为下行和上行流水线）结合在一起。图3展示了一个有四个流水线阶段（即$p=4$）的例子。这里我们假设每个工作节点在一个训练迭代中执行$m$个微批次，即$m=p$，这是保持所有阶段活跃的最低要求。如何扩展到超过$p$个微批次（即$m>p$）将在3.5节讨论。在下行流水线中，stage0到stage3线性地映射到P0到P3，而在上行流水线中，阶段的映射顺序完全相反。$m$个微批次（假设为偶数）被平均分配给两个流水线。每个流水线使用1F1B【索引38，PipeDream: generalized pipeline parallelism for DNN training，2019，SOSP】策略调度$m/2$个微批次，如图3左侧所示。然后，通过将这两个流水线合并在一起，我们得到了Chimera的流水线调度（图3右上）。给定偶数个阶段$p$（这在实践中很容易满足），可以保证在合并过程中没有冲突（即，每个工作节点的同一时间槽最多只有一个微批次）。我们可以看到，前向和反向传播中的气泡数量分别减少到$p/2-1$。考虑到前向和反向传播之间工作量的不均衡，我们得到了一个更实际的Chimera调度（图3右下）。

图3：Chimera的模型副本和双向流水线调度。
图3：Chimera的模型副本和双向流水线调度。

小批量情况下的适应性：对于那些为保证收敛必须使用小$M̂$的模型，一个训练迭代中可能少于$p$个微批次（即$m

推广到多条流水线：请注意，Chimera可以推广到结合两个以上的流水线（将在3.6节讨论），这会进一步减少气泡并平衡激活内存消耗，但代价是更高的通信开销和权重内存消耗。

3.2 通信方案

通信类型与朴素同步：Chimera使用p2p（点对点）通信在前向传播和反向传播中分别传输流水线阶段间的中间激活值和梯度（相对于输入）。由于Chimera结合了双向流水线，因此在下一次训练迭代之前，使用集体通信（即allreduce）来同步阶段副本间的权重梯度。图4(a)展示了一种简单的梯度同步方式，即在当前迭代的所有本地计算完成后，同步工作节点维护的每个阶段的梯度。请注意，中间阶段的梯度同步被开始和结束阶段的计算部分重叠了。

图4：通过重叠隐藏梯度同步开销。
图4：通过重叠隐藏梯度同步开销。

利用气泡重叠通信：为了实现更深度的通信重叠，我们利用流水线中的气泡来提前启动allreduce。以图4(b)中的P0和P3为例，在这两个工作节点分别完成微批次3和微批次1的反向传播后，stage3的权重梯度计算已经完成；因此，P0和P3可以在梯度计算完成后立即使用非阻塞集体通信【索引23，A case for standard non-blocking collective operations，2007，EuroPVM/MPI；索引25，Implementation and Performance Analysis of Non-Blocking Collective Operations for MPI，2007，SC07】启动异步allreduce来同步stage3的梯度，并在所有本地计算之后调用一个等待操作以确保allreduce完成。通过这种方式，stage3的梯度同步被气泡和后续计算所重叠。然而，与P0和P3不同，我们选择不对P1和P2上的stage2（一个中间阶段）进行提前梯度同步，因为从stage2的梯度完成到本地计算结束之间没有气泡。尽管异步通信可能在计算进行时发生，但它可能会导致额外的开销（初始化、线程等【索引24，Message Progression in Parallel Computing - To Thread or not to Thread?，2008，IEEE Cluster】），这可能会延长流水线的关键路径并危及整体性能。通信方案的性能建模将在3.4节中介绍。

3.3 流水线与数据并行的混合

混合并行策略：Chimera支持流水线并行和数据并行的混合。Chimera的双向流水线被复制$d$次以扩展到$p \cdot d$个工作节点。由于我们考虑的是可以轻松划分为平衡阶段的大型模型，因此在混合并行中，所有$p$个阶段都被同等地复制$d$次。当扩展到配备高性能互连网络（如Infiniband【索引50，InfiniBand network architecture，2003，Addison-Wesley Professional】、Cray Aries【索引2，Cray XC series network，2012，Cray Inc.】或Slingshot【索引48，An In-Depth Analysis of the Slingshot Interconnect，2020，SC20】以及NVLink【索引18，Ultra-performance Pascal GPU and NVLink Interconnect，2017，IEEE Micro】）的并行机器时，混合并行通常比纯流水线并行实现更好的性能【索引16，DAPPLE: a pipelined data parallel approach for training large models，2021，PPoPP；索引39，Memory-efficient pipeline-parallel DNN training，2020，arXiv】。这是因为纯流水线并行在流水线中有$p \cdot d$个阶段，而混合并行有$p$个阶段（少了$d$倍），这有助于减少阶段间的p2p通信开销并增加每个阶段的计算工作量。尽管混合并行导致阶段副本间的梯度同步，但其开销可以通过前述的高性能互连网络来缓解。然而，随着$d$增加（$p$减少），流水线阶段变得更粗粒度，直到某个点增加的梯度同步开销无法被减少的p2p通信开销所摊销。因此，找到最佳平衡点以实现最佳性能非常重要。

图5：Chimera中结合流水线并行与数据并行（d = 2, p = 4）。
图5：Chimera中结合流水线并行与数据并行（d = 2, p = 4）。

混合并行下的通信：图5展示了一个$d=2$和$p=4$的例子。注意，与数据并行结合后，需要同步的本地梯度大小不变，但参与梯度同步的进程数量增加了$d$倍。我们还使用与3.2节讨论的相同通信方案来重叠混合并行中的梯度同步开销。在下一节中，我们将讨论如何基于性能建模找到$p$和$d$的最佳配置。

3.4 基于性能建模的配置选择

配置参数与优化目标：给定小批量大小$M̂$和工作节点数$N$，配置$p, d$和$b$对训练吞吐量有很大影响。更大的微批次大小（$b$）通常能提高加速器的计算效率。由于Chimera极大地缓解了气泡问题，它贪婪地选择适应设备内存的最大微批次大小。与现有的必须在气泡和计算效率之间权衡的同步流水线方法相比，Chimera的贪婪策略显著减小了调优空间。

性能模型：计算开销：为了选择$p$和$d$的最佳配置，我们建立一个性能模型来预测每个可用配置的单次训练迭代的运行时间（用$T_{iter}$表示）。对于计算开销，我们使用微基准测试来测量单个流水线阶段的前向传播运行时间（用$T_f$表示）。反向传播的运行时间（用$T_b$表示）在不使用激活重计算时建模为前向传播的两倍，否则为三倍。我们将关键路径定义为一系列具有依赖关系的流水线阶段，它们决定了单次训练迭代的整体计算开销。图6展示了一个关键路径的例子。设$n_f$和$n_b$分别表示流水线关键路径上的前向传播和反向传播次数。对于图6中所示的例子，$n_f=6$和$n_b=10$。总计算开销为$n_f T_f + n_b T_b$。

图6：Chimera中用于重叠梯度同步的空闲区域，其中 p = m = 6。
图6：Chimera中用于重叠梯度同步的空闲区域，其中 p = m = 6。

性能模型：通信开销：为了建模通信开销，我们假设计算节点之间存在双向直接的点对点通信，并使用经典的延迟-带宽（$\alpha - \beta$）成本模型。发送大小为$S$的消息的成本是$\alpha + S\beta$，其中$\alpha$（延迟）和$\beta$（每字传输时间）都可以通过微基准测试测量。如3.2节所述，Chimera有两种类型的通信：阶段间的p2p通信（$T_{p2p}$）和用于梯度同步的allreduce（$T_{allreduce}$）。$T_{p2p}$可以简单地用$\alpha - \beta$模型建模。总p2p通信成本为$(n_f + n_b)T_{p2p}$。注意，如果存在中间气泡，$T_{p2p}$可以被部分重叠，但为简化模型，我们不考虑这种效应。对于$T_{allreduce}$，我们假设其实现利用了Rabenseifner算法【索引42，Optimization of collective reduction operations，2004，ICCS；索引53，Optimization of collective communication operations in MPICH，2005，Int. J. High Perform. Comput. Appl.】，其成本为：

公式：Allreduce通信成本
公式：Allreduce通信成本

其中$S$是待同步的梯度大小，$d$是阶段副本的数量。Rabenseifner算法在带宽项上达到了基于主机的allreduce的下界，因此对大模型效果最好。我们为$T_{allreduce}$建模了通信重叠（在3.2节讨论）的效果。图6展示了Chimera中用于重叠梯度同步的空闲区域（即超出该区域会增加总运行时间）的示例。注意每个工作节点上有两个阶段副本。区域(a)和(b)可用于重叠第一个阶段副本（阶段ID较大的那个）的梯度同步，区域(c)可用于重叠两个阶段副本的梯度同步。设$T_{allreduce}^{non-overlap}(i)$表示工作节点$i$上无法被空闲区域重叠的$T_{allreduce}$部分，那么$p$个工作节点中$T_{allreduce}^{non-overlap}(i)$的最大值贡献给总运行时间。

整体运行时间模型：总的来说，单次训练迭代的运行时间建模为：

公式：单次迭代运行时间
公式：单次迭代运行时间

我们使用这个模型来选择$p$和$d$的最佳配置（见4.2.2节）。

3.5 扩展到更多微批次

问题背景与基本单元：对于一个大的$M̂$，特别是在计算资源有限时，每个工作节点在一个训练迭代中可能会有超过$p$个微批次（即$m > p$）。为了扩展到大的$M̂$，我们首先选择$m$个微批次的最大$b$来饱和设备内存，并使用前面讨论的双向流水线调度这$m$个微批次。然后，我们使用$m$个微批次的调度作为一个基本调度单元，并通过将$k$（$k=m/p$且$m=\hat{M}/b/d$）个基本单元连接在一起来扩展到大的$M̂$。图7(a)展示了一个例子，其中每个工作节点在一个训练迭代中有$m=2p$个微批次，这有两个基本调度单元（即$k=2$）。

图7：Chimera扩展到一次训练迭代中超过p个微批次（即m > p）。
图7：Chimera扩展到一次训练迭代中超过p个微批次（即m > p）。

a. m=2p微批次，带有流水线清空
a. m=2p微批次，带有流水线清空

b. 直接拼接（中间有气泡）
b. 直接拼接（中间有气泡）

方法一：直接拼接：我们提出了三种方法来连接多个基本单元。（1）直接拼接，如图7(b)所示。第一个基本单元末尾的气泡可以被第二个基本单元开始的前向传播所占据。如果反向传播和前向传播的工作量相同，基本单元可以无缝连接。然而，反向传播的工作量大约是前向传播的两倍，这会导致中间产生气泡。

方法二和三：前向加倍与后向减半：为了消除直接拼接的中间气泡，我们提出了（2）前向加倍（如图7(d)所示）和（3）后向减半，其关键思想是均衡前向和反向传播的工作量。在前向加倍中，我们将每次前向传播的微批次数增加到两个，然后连接两个连续的反向传播块，每个块只有一个微批次，如图7(c)所示。然后，我们微调调度以消除流水线开始时50%的气泡，如图7(d)所示。前向加倍消除了中间气泡，但导致两倍的激活内存消耗，因此可能超过设备内存容量。我们采用激活重计算来减少内存开销。注意，重计算增加了反向传播的工作量，但由于输出两个微批次，前向传播中的p2p通信开销也加倍；因此，我们仍然认为前向传播（两个微批次）和反向传播（一个微批次并重计算）的工作量大致相等。前向加倍偏好于即使$b=1$也超过设备内存容量的大模型，因为在这种情况下必须使用激活重计算。对于$b$较大的较小模型，我们建议使用后向减半，它使用与前向加倍相同的调度，但不是在前向传播中执行两个微批次，而是将反向传播的微批次大小减半。后向减半不增加激活内存（因此不需要激活重计算），但由于使用了次优的$b$，可能会降低计算效率。选择这三种方法中最好的一种并非先验的，我们依赖于经验结果。注意，前向加倍和后向减半总共有$p-2$个气泡（前向传播$p/2-1$个，反向传播$p/2-1$个），如图7(d)所示，与DAPPLE和GPipe相比减少了约50%。对于$k>2$，我们使用$2p$个微批次的调度作为一个基本调度单元（如图7(c)所示）用于前向加倍和后向减半，并连接$\lfloor k/2 \rfloor$个基本单元和剩余的$p$个微批次（如果$k$是奇数）。

前向加倍与后向减半的额外优势：前向加倍和后向减半的另一个好处是，它们比经典的1F1B调度【索引38，PipeDream: generalized pipeline parallelism for DNN training，2019，SOSP；索引39，Memory-efficient pipeline-parallel DNN training，2020，arXiv】有更多空间来重叠p2p通信（在前向传播中）。在图7(d)中，以微批次1的前向传播为例，从P1到P2的p2p通信可以被中间的前向传播计算所重叠，而对于1F1B，可能没有足够的计算来重叠p2p通信。

3.6 推广到两个以上的流水线

推广到2k条流水线：到目前为止，我们只讨论了Chimera中组合两条流水线（一条下行和一条上行）的情况。然而，对于偶数个流水线阶段（即$p$为偶数），Chimera可以推广到组合两个以上的流水线。对于$p'=p/2$，设$D$为$p'$的所有约数的集合，包括1和$p'$本身。对于任何$k \in D$，我们可以为Chimera生成一个方案，该方案将$k$条下行流水线和$k$条上行流水线组合在一起，每条流水线有$p/(2k)$个微批次，由1F1B策略调度。图8给出了一个例子，其中Chimera组合了四条流水线，有八个阶段（即$p=8$和$k=2$）。对于下行流水线$i$（$i \in [0, k-1]$），$p$个阶段依次映射到$p$个工作节点上，第一个阶段（即stage0）映射到工作节点$i*(p/k)$。例如，对于图8中的下行流水线1，阶段[0,1,2,3,4,5,6,7]分别映射到工作节点[4,5,6,7,0,1,2,3]。对于$k$条上行流水线，$p$个阶段以与相应下行流水线完全相反的顺序映射到$p$个工作节点上。可以很容易地证明，这$2k$条流水线的调度可以无冲突地叠加。

图8：Chimera与四个8阶段流水线的组合。
图8：Chimera与四个8阶段流水线的组合。

多流水线的权衡：对于任何$k \in D$，Chimera可以使用3.5节讨论的方法扩展到更多的微批次（即$m > p$）。对于给定的$p$，Chimera会产生$2(p/(2k)-1)$个气泡，但必须在每个工作节点上维护$2k$个模型副本并同步$2k$个阶段的梯度。$k$的值越大，气泡越少（激活内存消耗越均衡），但梯度同步开销和权重内存消耗也越高。当$k=p/2$时，Chimera退化为纯数据并行。4.4节的经验结果表明，对于评估中使用的模型，$k>1$很少带来更多性能优势。因此，除非另有说明，本文中Chimera的默认设置为$k=1$（即两条流水线的组合）。我们预计，其特性总结在表3中的$k>1$方案，将能为未来具有更深流水线和每个阶段更高计算密度的深度模型进一步提升性能。

表3：具有2k条流水线的Chimera。

表3：具有2k条流水线的Chimera。
表3：具有2k条流水线的Chimera。

A4 实验环境

  • 硬件配置：

    1. Piz Daint超级计算机：主要实验平台。每个Cray XC50计算节点包含一个Intel Xeon E5-2690 CPU和一块NVIDIA P100 GPU（16GB全局内存）。节点通过Cray Aries互联网络以Dragonfly拓扑连接。
    2. V100 GPU集群：一个小型集群，配备32块NVIDIA V100 GPU（32GB全局内存）。集群由4台GPU服务器组成，每台服务器有8块V100 GPU，通过NVLink进行节点内连接，通过Infiniband进行节点间连接。

  • 软件配置：

    • 实现框架：所有方案均在PyTorch【索引41，Pytorch: An imperative style, high-performance deep learning library，2019，NeurIPS】中实现。
    • 分布式后端：使用GLOO【索引15，Gloo，2018，GitHub】作为分布式后端，用于流水线阶段间的点对点（p2p）通信和阶段副本间的梯度同步（allreduce）。虽然PyTorch的NCCL【索引40，NVIDIA Collective Communications Library，2020，NVIDIA】后端在GPU节点间的allreduce性能更好，但它不支持p2p通信，同时使用两者会失败。
    • 操作系统与编译器：SUSE SLES 11.3，gcc 9.3.0（见附录）。

  • 模型与数据集：

    • 模型架构：
      1. Bert-48：一个包含48个Transformer层的BERT模型，参数量为11亿。
      2. GPT-2：一个包含48个Transformer层的GPT-2模型，参数量为13亿。
    • 数据集：
      1. Wikipedia dataset 和 WikiText-2 dataset（见附录）。
    • 关键参数：Bert-48的最大序列长度设为128，GPT-2设为632。小批量大小和序列长度与机器学习社区的常用设置保持一致。

表4：用于评估的神经网络。

表4：用于评估的神经网络。
表4：用于评估的神经网络。

A4 实验结果

4.1 内存消耗

实验内容：在Piz Daint的32个GPU节点上，比较了不同配置下Bert-48和GPT-2模型的内存消耗（包括激活值和权重）分布。

实验结果与分析：
* GPipe：由于一次性注入所有微批次，其高昂的激活内存成本导致在所有配置中都出现内存溢出（OOM）。
* PipeDream：内存消耗第二高，因为它需要缓存多达$p$个版本的权重。
* PipeDream-2BW和DAPPLE：由于流水线调度和语言模型第一层（包含嵌入层）权重更大的双重不平衡，第一个工作节点通常是内存消耗的峰值。PipeDream-2BW因需要缓存两个权重版本，在流水线阶段变粗时也出现OOM。
* Chimera：得益于双向流水线调度，Chimera实现了更均衡的内存消耗分布。其激活内存成本在第一个（也是最后一个）工作节点上最低，这可以摊销第一层额外的权重开销。因此，尽管Chimera维护了两个模型副本，但在六个配置中的四个中，其峰值内存消耗仍略低于只维护一个副本的DAPPLE。
* GEMS：内存消耗最低且最均衡，但这是以牺牲并行性为代价的。

结论：Chimera在峰值内存消耗方面与最先进的方法相当，同时在工作节点之间实现了更均衡的内存使用，这与理论分析一致。

图9：Piz Daint上32个GPU节点的内存消耗分布。红色和蓝色圆圈分别表示最大和最小内存消耗。
图9：Piz Daint上32个GPU节点的内存消耗分布。红色和蓝色圆圈分别表示最大和最小内存消耗。

4.2 并行可扩展性

实验内容：首先为每个基准方法搜索最佳配置（$b, p, d$），然后在弱扩展性测试中比较它们的最佳性能。

实验结果与分析：
1. 基准方法性能优化（图10、11）：针对Bert-48和GPT-2，通过搜索参数空间为所有基准方法（GPipe, DAPPLE, GEMS, PipeDream-2BW, PipeDream）找到了最佳性能配置。这确保了比较的公平性。
2. Chimera性能建模与通信策略（图12、13）：
* 实验证明，论文中提出的优化版提前同步策略（eager-sync-opt）比简单的提前同步（eager-sync）吞吐量更高（例如在64个节点上高出1.09倍）。
* 性能模型预测的吞吐量与实际吞吐量误差在10%以内，并且能够准确地选择出接近最优的（$p, d$）配置。
3. 最佳性能对比（图14、15、16）：
* Piz Daint (P100)：在Bert-48和GPT-2的弱扩展性测试中，Chimera在所有规模上均优于所有基准方法。
* 在64个节点上运行Bert-48，Chimera比异步方法快1.17x-1.94x，比同步方法快1.19x-2.41x。
* 在2048个节点上运行GPT-2，Chimera比异步方法快1.16x-2.01x，比同步方法快1.38x-2.34x。Chimera的主要优势在于低气泡率和均衡的内存消耗（避免了激活重计算），并实现了91.4%的并行效率。
* V100集群：在V100集群上的实验也得出了相同的结论，Chimera的吞吐量比同步和异步方法分别提升了1.10x-2.39x和1.05x-1.89x，证明了其方法的普适性。

结论：Chimera展现了卓越的并行扩展能力，在不同规模和硬件平台上均显著优于现有最先进的流水线并行方法。

图14：Bert-48在Piz Daint上的弱扩展性。当N从16扩展到64时，M̂从256扩展到1024，除了PipeDream的M̂从24扩展到96。
图14：Bert-48在Piz Daint上的弱扩展性。当N从16扩展到64时，M̂从256扩展到1024，除了PipeDream的M̂从24扩展到96。

图15：GPT-2在Piz Daint上的弱扩展性。当N从512扩展到2048时，M̂从512扩展到2048，除了PipeDream的M̂从128扩展到512。
图15：GPT-2在Piz Daint上的弱扩展性。当N从512扩展到2048时，M̂从512扩展到2048，除了PipeDream的M̂从128扩展到512。

4.3 在固定机器数量上扩展到大批量

实验内容：在固定的机器数量上（Bert-48为32个GPU，GPT-2为512个GPU），评估当微批次数量远大于流水线深度时（$m \gg p$），各方法的训练吞吐量。

实验结果与分析：
* Bert-48（图17）：Chimera的三种扩展方法中，“直接拼接”性能最佳，因为 uneven workloads 产生的中间气泡可以用来容纳p2p通信。在这种情况下，Chimera的性能非常接近使用陈旧权重的异步方法PipeDream-2BW，并显著优于其他同步方法（平均比GPipe快1.13x，比GEMS快2.07x，比DAPPLE快1.06x）。
* GPT-2（图18）：对于GPT-2这样的大模型，激活重计算是必需的。Chimera的“前向加倍”方法因能消除中间气泡而性能最佳。得益于其精巧的调度（更少气泡和更多通信重叠），Chimera优于所有基准方法（平均比PipeDream-2BW快1.13x，比GPipe快1.18x，比GEMS快2.60x，比DAPPLE快1.34x）。

结论：Chimera能高效地扩展到大批量训练场景，无论是否需要激活重计算，都能通过合适的调度策略取得领先性能。

图17：Bert-48在Piz Daint的32个GPU节点上扩展到大mini-batch size。PipeDream使用p=8，其他使用p=4。
图17：Bert-48在Piz Daint的32个GPU节点上扩展到大mini-batch size。PipeDream使用p=8，其他使用p=4。

图18：GPT-2在Piz Daint的512个GPU节点上扩展到大mini-batch size。当M̂ >= 1,024时，DAPPLE和PipeDream-2BW从p=16切换到p=8，其他使用p=8。
图18：GPT-2在Piz Daint的512个GPU节点上扩展到大mini-batch size。当M̂ >= 1,024时，DAPPLE和PipeDream-2BW从p=16切换到p=8，其他使用p=8。

4.4 Chimera与超过两条流水线

实验内容：在一个32层的GPT-2模型上，测试Chimera组合超过两条流水线（即超过两个模型副本）时的性能。

实验结果与分析（图19）：
* 两条流水线（Chimera的默认设置）在大多数配置下都能实现最高的性能。
* 在特定配置下（例如$p=32$），四条流水线性能略好，因为它在减少气泡和增加allreduce开销之间达到了最佳平衡点。
* 然而，当流水线阶段变粗（例如$p$减少到16），四条流水线的性能会因allreduce开销增加而劣于两条流水线。

结论：两条流水线的默认设置是一个稳健且高效的选择。对于未来更深、计算密度更高的模型，组合更多流水线可能会带来进一步的性能提升。

图19：在Piz Daint的64个GPU节点上，针对32层GPT-2模型（M̂=64），Chimera使用更多流水线时的吞吐量。
图19：在Piz Daint的64个GPU节点上，针对32层GPT-2模型（M̂=64），Chimera使用更多流水线时的吞吐量。

A5 结论

总结：Chimera为大规模神经网络的高效流水线训练提供了新的见解。与最先进的流水线方法相比，Chimera在流水线效率、内存成本和收敛友好性之间实现了最佳平衡。在多达2048个GPU节点上对大型语言模型进行的实验结果表明，Chimera的训练吞吐量显著优于其对应方法。

未来工作：作者预见，Chimera将成为大规模扩展深度学习训练的主要解决方案之一。下一步工作将是通过利用深度学习训练中的稀疏化【索引22，Sparsity in Deep Learning: Pruning and growth for efficient inference and training in neural networks，2021，arXiv；索引47，SparCML: High-performance sparse communication for machine learning，2019，SC】和量化【索引1，QSGD: Communication-efficient SGD via gradient quantization and encoding，2017，NeurIPS】技术来降低梯度同步的通信成本。

A6 附录

A.1 实验总结

本文在CSCS Piz Daint超级计算机上对Chimera进行了评估。每个Cray XC50计算节点包含一个12核Intel Xeon E5-2690 CPU和64GB RAM，以及一块NVIDIA Tesla P100（16GB内存）。节点由Cray Aries互连网络以Dragonfly拓扑连接。我们使用PyTorch中的GLOO作为分布式后端。所有实验都利用GPU进行加速。Chimera的源代码如下：
* 项目名称: Chimera
* 永久ID: https://github.com/Shigangli/Chimera

A.2 基准实验设置

  • 相关硬件细节: CSCS Piz Daint超级计算机。每个Cray XC50计算节点包含一个12核Intel Xeon E5-2690 CPU和一块NVIDIA Tesla P100 GPU。文件系统为Lustre。
  • 操作系统与版本: SUSE SLES 11.3
  • 编译器与版本: gcc 9.3.0
  • 应用与版本: Bert, GPT-2
  • 库与版本: PyTorch 1.6
  • 关键算法: 随机梯度下降
  • 输入数据集与版本: Wikipedia dataset, WikiText-2 dataset
  • 执行环境信息收集脚本输出URL: https://www.dropbox.com/s/md1jlcn3sm5bl9x/environment.txt?dl=0