作者/机构: Aaron Harlap† ∗, Deepak Narayanan‡ ∗, Amar Phanishayee, Vivek Seshadri, Nikhil Devanur?, Greg Ganger†, Phil Gibbons†

?Microsoft Research, †Carnegie Mellon University, ‡Stanford University

A1 主要贡献

本文针对深度神经网络（DNN）训练中通信开销成为瓶颈的问题，提出了一种名为 PipeDream 的新型分布式训练系统。随着模型规模增大和GPU算力提升，传统的数据并行方法因其与模型大小成正比的通信量而效率下降，而传统的模型并行方法则存在严重的计算资源闲置问题。PipeDream 旨在解决这些问题，其核心思想是流水线并行（pipeline parallelism），一种结合了模型并行、主动流水线化和数据并行的混合并行策略。

核心问题与研究目标:
- 核心问题: 数据并行训练在模型较大或网络带宽有限时，通信开销会超过计算时间，限制了扩展性。传统模型并行则导致GPU利用率低下。
- 研究目标: 设计一个能够有效解决通信瓶颈、实现高硬件利用率，并且能够自动化处理模型划分和任务调度的分布式DNN训练系统。

主要贡献与创新点:
1. 提出流水线并行方法: 创新性地将模型并行与主动流水线技术相结合，并在必要时融入数据并行，以应对DNN训练中的通信瓶颈。该方法通过将模型分层到不同工作节点，并流水化处理多个微批次（minibatch），从而大幅减少节点间通信量并实现计算与通信的完美重叠。
2. 解决流水线并行的关键挑战:
- 自动分区: 设计并实现了一种基于动态规划的算法，该算法通过对模型进行简短的性能剖析，自动将DNN层划分为计算负载均衡且通信量最小的“阶段”（stage）。
- 工作调度: 提出了一种“一前一后”（1F1B）调度机制，在启动阶段填满流水线后，使每个阶段交替执行前向和后向传播，从而保证所有工作节点持续繁忙，避免流水线停顿。
- 保证有效学习: 提出了“权重暂存”（Weight Stashing）机制，为每个处理中的微批次保留其前向传播时所使用的权重版本，并在后向传播时使用同一版本计算梯度，解决了因流水线引入的权重版本不一致导致的模型不收敛问题。
3. 实现PipeDream系统: 开发了一个高效实现流水线并行训练的系统，能够自动化地完成模型分区、资源分配和执行调度。
4. 实验验证: 通过在两个不同集群上对五种DNN模型（包括CNN和RNN）进行的实验，证明了PipeDream的有效性。与数据并行BSP相比，PipeDream在“达到目标精度时间”上最高可提速5倍，尤其是在通信与计算比较高的场景下优势更为明显。

A3 背景知识与相关工作

DNN训练基础

DNN模型与训练过程。一个DNN模型由一系列不同类型的层（如卷积层、全连接层、池化层）组成。训练过程通常包含多个轮次（epoch），每个轮次会遍历整个数据集。在每个步骤中，模型首先对一小批训练样本（即微批次）进行预测，这个过程称为前向传播。输入数据被送入模型的第一层，每层计算其输入的函数（通常使用学习到的参数或权重），并将输出传递给下一层，最后一层的输出即为类别预测。然后，根据模型的预测标签和真实标签，输出层计算一个损失（或误差）。在随后的后向传播中，每一层计算两部分内容：1）前一层的误差，以及2）所有相关层的权重更新（损失的梯度），这些更新会使模型的预测更接近期望的输出。

训练目标与并行化方法。DNN训练的目标是在尽可能短的时间内获得一个高精度的模型。这个目标可以由两个指标来衡量：1）统计效率，即达到期望精度所需的轮次数；2）硬件效率，即完成单个轮次所需的时间。总训练时间是这两个指标的乘积【15，Omnivore: An optimizer for multi-device deep learning on cpus and gpus，2016】。为了在合理时间内训练大型模型，通常采用分布式训练，主要有两种方法：数据并行或模型并行。

数据并行及其局限性

数据并行机制。在使用数据并行时，输入数据集被划分到多个GPU上。每个GPU都持有一份完整的模型副本，并在各自的数据分区上进行训练，同时定期与其他GPU同步权重，这通常通过集体通信原语【14，Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour，2017】或参数服务器【28，Scaling distributed machine learning with the parameter server，2014；9，Geeps: Scalable deep learning on distributed gpus with a gpu-specialized parameter server，2016】完成。参数同步的频率会影响统计效率和硬件效率。

同步并行（BSP）的通信瓶颈。一种极端情况是在每个微批次结束后进行同步（称为批量同步并行或BSP【42，A bridging model for parallel computation，1990】），这能减少用于计算梯度的权重陈旧度，保证良好的统计效率。然而，如图2所示，BSP要求每个GPU等待或停顿，以接收来自其他GPU的梯度，从而显著降低硬件效率。尽管存在诸如“无等待反向传播”（Wait-free Backpropagation）【47，Poseidon: An efficient communication architecture for distributed deep learning on GPU clusters，2017】等优化，但由于DNN计算的结构特性以及通信时间常能主导总执行时间，数据并行训练中的通信停顿是不可避免的。此外，计算速度的快速增长进一步将训练瓶颈推向通信。

图2：一个包含4台机器的数据并行设置示例。其中一台机器的时间轴显示了模型参数交换期间的通信停顿。
图2：一个包含4台机器的数据并行设置示例。其中一台机器的时间轴显示了模型参数交换期间的通信停顿。

通信开销的量化分析。图1量化展示了在“商用”公有云服务器上，使用三代不同的NVIDIA GPU（Kepler (K80)、Pascal (Titan X)和Volta (V100)），通过10Gbps网络运行时，五种不同DNN模型的训练时间中通信停顿所占的比例。我们关注三个要点：首先，即使在像K80这样较慢的GPU上，一些CNN（VGG16和AlexNet）和用于视频转录的序列到序列模型（S2VT）也花费了大量时间在通信上。像ResNet50和Inception-v3这样的网络通信开销相对较低。其次，随着数据并行工作节点数量的增加，所有模型的通信开销都会增加。第三，随着GPU计算速度的提升（从K80到V100），所有五个模型的通信停顿也会增加。

图1：不同硬件配置下，通信开销占总训练时间的百分比。许多模型（AlexNet、VGG16、S2VT）即使在相对较慢的K80上也有很高的通信开销。两个因素导致所有模型的通信开销增加：（i）数据并行工作节点数量的增加，以及（ii）GPU计算能力的提升。
图1：不同硬件配置下，通信开销占总训练时间的百分比。许多模型（AlexNet、VGG16、S2VT）即使在相对较慢的K80上也有很高的通信开销。两个因素导致所有模型的通信开销增加：（i）数据并行工作节点数量的增加，以及（ii）GPU计算能力的提升。

异步并行（ASP）的权衡。之前的工作提出了更宽松的同步模型，我们称之为异步并行或ASP，其中每个GPU在不等待其他GPU梯度的情况下继续计算下一个微批次。这种方法减少了GPU的空闲时间，从而提高了硬件效率。然而，这个过程可能导致梯度是基于陈旧的权重计算的，从而降低了统计效率。我们的实验结果证实了近期的发现，即ASP并不能减少端到端的训练时间。

模型并行及其局限性

模型并行机制。在模型并行中，模型被划分到多个GPU上，每个GPU只负责模型的一部分。对于矩阵分解、主题建模和线性回归等机器学习问题，先前的工作【29，Distributed graphlab: A framework for machine learning in the cloud，2012；27，On model parallelization and scheduling strategies for distributed machine learning，2014；45，Lightlda: Big topic models on modest computer clusters，2015；23，STRADS: a distributed framework for scheduled model parallel machine learning，2016】已经表明，由于避免了使用极大批量大小带来的统计效率提升，模型并行通常能比数据并行实现更快的训练时间；STRADS框架【23】表明，对多个微批次进行流水线处理可以进一步缩短这些机器学习问题的训练时间。模型并行也已用于DNN，但传统上仅在模型训练的工作集太大无法装入单个工作节点的内存或缓存时作为最后手段【25，One weird trick for parallelizing convolutional neural networks，2014；7，Project adam: Building an efficient and scalable deep learning training system，2014；10，Large scale distributed deep networks，2012】（这使得数据并行不可行）。这是因为传统的模型并行DNN训练存在两个主要限制。

传统模型并行的低效性。首先，模型并行的DNN训练导致GPU资源利用率严重不足，如图3所示。该图展示了将DNN层划分到四台机器上的情况，每台机器负责一组连续的层；在这种模式下，这些组之间的层间值（激活值和梯度）是唯一需要在机器间通信的参数。对于每个微批次，任何时刻只有一个阶段处于活动状态。虽然背靠背地流水线处理多个微批次可以提高利用率，但传统上不这么做，因为：1）DNN的双向性（前向传播后是反向通过相同层的后向传播）使流水线处理变得困难；更重要的是，2）一个简单的流水线机制会引入基于陈旧权重的权重更新计算，导致最终模型的精度低于数据并行训练。

图3：使用4台机器的模型并行训练。数字表示微批次ID。为简单起见，这里我们假设每个阶段的前向和后向工作都花费一个时间单位，并且跨机器通信激活值没有开销。
图3：使用4台机器的模型并行训练。数字表示微批次ID。为简单起见，这里我们假设每个阶段的前向和后向工作都花费一个时间单位，并且跨机器通信激活值没有开销。

模型划分的困难。其次，将模型划分到多个GPU的负担留给了程序员【25】，这导致了许多特定的点解决方案。最近有工作探索使用强化学习来自动确定模型并行的设备布局【30，Device placement optimization with reinforcement learning，2017】。不幸的是，这种在线决策技术耗时且资源密集；它们也无法无缝地结合流水线、数据并行和模型并行。

其他相关工作

高性能集群与专用硬件。近期关于快速训练卷积神经网络（CNN）的工作利用了高度优化且昂贵的集群，这些集群具有高速的机器内和机器间互连【14，Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour，2017；12，NVIDIA DGX-1】。许多公有云提供商尚未提供这类优化的服务器SKU，而且即使提供，其价格也可能高得令人望而却步。相比之下，我们的工作研究的是使用公有云产品中可用的商用SKU，这是大众可以随时访问的训练基础设施。

通信优化技术。CNTK的1比特量化技术【33，CNTK: Microsoft’s open-source deep-learning toolkit，2016】解决了数据并行训练中的通信瓶颈问题【35，1-bit stochastic gradient descent and its application to data-parallel distributed training of speech dnns，2014】。这种近似策略缺乏通用性，且仅在有限场景下有效；它对某些语音模型的收敛没有损害【34，On parallelizability of stochastic gradient descent for speech dnns，2014】，但在许多其他模型中会因梯度带噪而损害统计性能【9，Geeps: Scalable deep learning on distributed gpus with a gpu-specialized parameter server，2016；1，Tensorflow: A system for large-scale machine learning，2016】。

高效集合通信。Goyal等人【14】使用了更高效的all_reduce实现，如递归减半加倍算法和桶算法，以减少网络上传输的数据量【39，Optimization of collective communication operations in mpich，2005】。其他人则探索了来自HPC领域的技术来减少通信开销【2，Bringing HPC Techniques to Deep Learning，2017；41，Meet Horovod: Uber’s Open Source Distributed Deep Learning Framework for TensorFlow，2017】。但所有这些归约方法仍然涉及同步通信模式，导致的网络停顿虽然变小，但只能轻微缓解因模型规模不断增长和计算能力更快提升而引入的通信瓶颈。

早期流水线探索。Chen等人【6，Pipelined Back-propagation for Context-dependent Deep Neural Networks，2012】简要探讨了在模型并行训练中对微批次进行流水线处理的潜在好处，但没有解决良好统计效率、规模化和普适性等适用于大型真实世界模型的问题。事实上，我们的工作表明，对于真实世界的模型，简单地进行计算流水线处理是不够的。在我们提出的解决方案（第3节）中，我们解决了先前工作中被忽略的关键问题，并提供了一个通用且自动化的解决方案。

A2 方法细节

PipeDream结合了传统的数据并行与通过流水线增强的模型并行。我们称这种方案为流水线并行（Pipeline Parallelism, PP）。本节首先描述PP，然后介绍PipeDream为有效实现流水线并行训练而设计的解决方案，以应对相关挑战。

流水线并行机制

流水线并行（PP）的核心思想。流水线并行训练将待训练模型的层划分为多个阶段（stage），每个阶段包含模型中一组连续的层。每个阶段被映射到一个独立的GPU上，该GPU负责执行该阶段所有层的前向传播和后向传播。我们将包含输入层的阶段称为输入阶段，包含输出层的阶段称为输出阶段。图4展示了一个简单的流水线并行分配示例，其中DNN被分割到四台机器上。

图4：一个包含四台机器的流水线并行分配示例，以及其中一台机器的时间轴，突显了计算与激活/梯度通信在时间上的重叠。
图4：一个包含四台机器的流水线并行分配示例，以及其中一台机器的时间轴，突显了计算与激活/梯度通信在时间上的重叠。

单微批次与多微批次流水线。在最简单的情况下，系统中只有一个活跃的微批次，就像传统的模型并行训练一样。图3展示了一个包含四台机器和一个活跃微批次的示例配置的计算时间线。在前向阶段，每个阶段为该微批次执行其所在层的前向传播，并将结果发送到下一个阶段。输出阶段在完成其前向传播后，计算该微批次的损失。在后向阶段，每个阶段执行后向传播并将损失传播到前一个阶段。当只有一个活跃微批次时，任何给定时间点最多只有一个GPU处于活动状态。

流水线化的优势。为了确保任何时间点都没有GPU空闲，我们一个接一个地向流水线中注入多个微批次，从而通过流水线技术增强模型并行训练。在完成一个微批次的前向传播后，每个阶段会异步地将输出激活值发送到下一个阶段，同时开始处理另一个微批次。类似地，在完成后向传播后，每个阶段异步地将梯度发送到前一个阶段，同时开始计算另一个微批次。与数据并行训练相比，流水线化有两个主要优点。

优势一：通信量更少。与BSP相比，PP所需的通信量要少得多。图5比较了VGG16中每一层输出的大小与模型参数的总体大小。在PP执行中，每台机器只需通信其中一层的输出数据，而无需像BSP那样通信所有参数。这通常能大幅减少通信量（例如，VGG16减少了90%）。

图5：在ImageNet1K数据上，VGG16使用32的微批次大小时各层输出数据的大小。黑色虚线表示模型参数的大小。
图5：在ImageNet1K数据上，VGG16使用32的微批次大小时各层输出数据的大小。黑色虚线表示模型参数的大小。

优势二：计算与通信重叠。跨阶段异步通信前向输出激活值和后向梯度，可以使通信与后续微批次的计算显著重叠，从而实现比BSP更高的硬件效率（如图4所示）。

流水线并行的挑战。虽然流水线本身能比数据并行减少训练时间，但我们观察到模型并行和数据并行对不同类型的层效果最佳【25】。因此，PipeDream旨在以最小化总训练时间的方式结合流水线模型并行和数据并行。图6展示了流水线并行如何将一个假设模型的层划分到8台机器上的不同阶段。然而，要使这种方法对大型真实世界的DNN模型有效，需要解决三个挑战：
1. 在可用计算资源间自动划分工作。
2. 调度计算以最大化吞吐量，同时确保学习任务取得进展。
3. 在流水线引入的异步性面前，确保学习的有效性。
本节余下部分将描述这些挑战以及PipeDream的解决方法。

图6：PipeDream中的流水线并行训练结合了流水线、模型并行和数据并行训练。
图6：PipeDream中的流水线并行训练结合了流水线、模型并行和数据并行训练。

跨机器的层划分

自动划分的挑战与工作流。给定一个模型和一组机器，PipeDream的首要挑战是自动地将模型的层划分到可用的机器上，以最小化总训练时间。图7展示了PipeDream用于在可用机器间划分DNN层的工作流程。在将层划分为不同阶段时，PipeDream的分区算法必须确保每个阶段大致执行相同总量的工作。同时，分区算法还必须确保跨阶段通信的数据量尽可能小，以避免通信停顿。机器间的负载不平衡或过多的通信会降低硬件效率（吞吐量）。

图7：PipeDream将DNN层划分为阶段的自动化机制。PipeDream首先对输入的DNN进行性能分析，以获取每层的计算时间和输出大小的估计值。利用这些估计值，PipeDream的优化器将层划分到可用机器上。
图7：PipeDream将DNN层划分为阶段的自动化机制。PipeDream首先对输入的DNN进行性能分析，以获取每层的计算时间和输出大小的估计值。利用这些估计值，PipeDream的优化器将层划分到可用机器上。

分区算法流程。考虑到这些因素，对于一个有N层和M台可用机器的DNN，PipeDream首先在一台机器上对模型进行性能分析，然后运行一个分区算法，该算法将层分组为阶段，同时确定每个阶段的复制因子，以最小化模型的总训练时间。

DNN模型性能剖析。我们的性能剖析机制利用了DNN训练在不同微批次间的计算和通信时间变化很小这一事实。PipeDream为每层l记录三个量：1) $T_l$，该层前向和后向传播的总计算时间；2) $a_l$，该层的输出激活值大小（也即后向传播中的输入梯度大小）；3) $w_l$，第l层的参数大小。

计算时间$T_l$的确定。为了确定所有层的$T_l$，PipeDream在一台机器上使用1000个微批次对DNN模型进行一次简短的运行分析。利用这个分析结果，PipeDream将$T_l$计算为第l层的前向和后向计算时间之和。

通信时间的估算。所有通信都分三步进行：1) 将数据从发送方的GPU移动到CPU；2) 通过网络将数据从发送方发送到接收方；3) 将数据从接收方的CPU移动到GPU。PipeDream将通信所需时间估算为需要传输的数据量除以通信链路上的网络带宽。在流水线中，从第l层到第l+1层通信激活值所需的时间$C_l$，是使用$a_l$来估算的。在有m台机器的数据并行配置中，每个工作节点通信的数据量为 $4 \times (m-1) \times |w_l|/m$；这被用来估算当使用分布式参数服务器时，该层权重同步所需的时间$W^m_l$。

PipeDream的分区算法。我们的分区算法接收性能剖析步骤的输出，并计算：1) 将层划分为阶段的分区方案；2) 每个阶段的复制因子；3) 保持训练流水线繁忙所需的最优微批次数。

算法目标与动态规划。分区算法试图最小化模型的总训练时间。对于一个流水线系统，这个问题等同于最小化流水线中最慢阶段所需的时间。这个问题具有最优子结构特性；即在给定机器数量下最大化吞吐量的流水线，是由在更少机器数量下最大化吞吐量的子流水线组成的。因此，我们可以使用动态规划找到这个问题的最优解。

动态规划定义。我们用$A(j, m)$表示使用m台机器在第1层到第j层之间构建的最优流水线中，最慢阶段所花费的时间。算法的目标是找到$A(N, M)$以及相应的分区方案。令$T(i \to j, m)$表示一个跨越第i层到第j层、并在m台机器上复制的单个阶段所花费的时间。

阶段时间计算公式。

公式
公式

其中，max函数内的左项是该阶段所有层的总计算时间，右项是该阶段所有层的总通信时间。

递归关系：情况一。由第1层到第j层、使用m台机器的最优流水线，要么是一个被复制m次的单一阶段，要么由多个阶段组成。情况1：最优流水线只包含一个阶段，被复制m次。在这种情况下：

公式
公式

递归关系：情况二。情况2：最优流水线包含多个阶段。在这种情况下，它可以被分解为一个由第1层到第i层、使用$m-m_0$台机器的最优子流水线，后面跟着一个由第$i+1$层到第j层、在$m_0$台机器上复制的单一阶段。然后，利用最优子结构特性，我们有：

公式
公式

其中，max函数内的第一项是使用$m-m_0$台机器在第1层和第i层之间的最优子流水线中最慢阶段的时间，第二项是第i层和第i+1层之间通信激活值和梯度的时间，第三项是包含剩余层的单一阶段在$m_0$台机器的数据并行配置下所需的时间。

动态规划初始化。$A(1, m) := T(1 \to 1, m)$，其中$T()$如上定义，m从1变化到M（总机器数）。$A(i, 1) := T(1 \to i, 1)$，其中i从1变化到N（模型总层数）。

运行时分析。子问题的总数是$O(N \cdot M)$。每个子问题的时间复杂度也是$O(N \cdot M)$，导致总时间复杂度为$O(N^2 M^2)$。

最优活跃微批次数（NOAM）。基于我们算法生成的分区，为使流水线在稳态下保持满负荷，每个输入阶段应接纳的最优微批次数由 ceil((# machines) / (# machines in the input stage)) 给出。我们称这个量为NUM_OPT_ACTIVE_MINIBATCHES (NOAM)。

工作调度

双向流水线的调度挑战。与传统的单向流水线不同，流水线化的DNN训练涉及一个双向流水线。一个微批次的前向传播从输入层开始，而后向传播在输入层结束。因此，流水线中每个活跃的微批次可能处于不同的层，或者在前向传播阶段，或者在后向传播阶段。结果是，系统中的每台机器都必须在两个选项之间做出选择：i) 为一个微批次执行前向传播，从而将该微批次推向流水线下游的机器；ii) 为另一个不同的微批次执行后向传播，从而确保学习过程取得进展。

朴素调度的弊端。一个总是优先处理前向工作的简单调度机制会阻碍整体学习进展，因为权重更新只有在后向传播完成后才能应用。同样，总是优先处理后向工作可能会周期性地导致机器因无可用工作而空闲。我们提出了一种避免这些问题的调度机制。

1F1B调度机制。在启动阶段，输入阶段接纳NOAM个微批次，以使流水线在稳态下保持满负荷。一旦进入稳态，每个阶段会交替为一个微批次执行前向传播和为另一个微批次执行后向传播。我们称这种机制为一前一后（one-forward-one-backward, 1F1B）。在一个平衡的流水线中，1F1B确保在稳态下没有GPU空闲，并且我们能从每个微批次中取得学习进展。

1F1B工作流程图解。图8展示了一个包含4个阶段的流水线对应的计算时间线，每个阶段运行在一台机器上。此配置的NOAM为4。在启动阶段，输入阶段正好接纳四个微批次，它们依次传播到输出阶段。一旦输出阶段完成了第一个微批次的前向传播，它会立即为同一个微批次执行后向传播，然后开始为后续的微批次交替执行前向和后向传播。随着后向传播开始向流水线的前期阶段传播，每个阶段都开始为不同的微批次交替执行前向和后向传播。如图所示，在稳态下，每台机器都忙于为某个微批次执行前向或后向传播。为了使1F1B有效，前向传播所需时间不一定要与后向传播相同。事实上，我们观察到在实践中，后向传播总是比前向传播耗时更长，而1F1B仍然是一种有效的调度机制。

图8：一个包含4台机器的流水线示例，展示了启动和稳态。
图8：一个包含4台机器的流水线示例，展示了启动和稳态。

数据并行阶段的调度。当阶段以数据并行配置运行，即被复制到多个GPU上时，我们使用确定性的轮询负载均衡（minibatchID mod stageReplicaID）来将来自前一阶段的工作分散到各个副本上。这种确定性的负载均衡确保了一个微批次的后向传播会在负责其前向传播的同一台机器上执行。

静态策略的优势。无论是流水线中各阶段的1F1B调度策略，还是跨复制阶段进行负载均衡的轮询调度策略，都是静态策略。因此，它们可以由每台机器独立执行，无需昂贵的分布式协调。

有效学习

朴素流水线中的权重版本问题。在一个朴素的流水线系统中，每个微批次的前向传播使用一个版本的参数执行，而后向传播则使用另一个不同版本的参数。图8以一个无数据并行的分区为例说明了这一点。如果我们观察阶段1（机器1），微批次5的前向传播是在应用了微批次1的更新之后执行的，而微批次5的后向传播则是在应用了微批次2、3和4的更新之后执行的。结果是，在阶段1为微批次5进行后向传播时，梯度是使用与相应前向传播中使用的不同权重集计算的；这种权重版本的不一致性可能阻止模型收敛。

跨阶段的陈旧度不对称性。此外，DNN模型中的不同阶段遭受不同程度的陈旧性。例如，在第三个阶段，每个微批次的前向和后向传播之间只有一个交错的更新，而输出阶段则没有交错的更新。这种跨层的不对称性会进一步影响模型收敛。我们的实验结果表明，朴素的流水线处理无法达到与数据并行训练相同的准确度。为解决此问题，PipeDream使用了两种技术。

权重暂存（Weight Stashing）。权重暂存维护多个版本的权重，每个活跃的微批次一个版本。在执行前向传播时，每个阶段使用可用的最新版本权重处理一个微批次。完成前向传播后，PipeDream将所使用的权重作为该微批次中间状态的一部分存储起来。在执行该微批次的后向传播时，会使用相同版本的权重来计算权重梯度。

权重暂存的作用。权重暂存确保了在一个阶段内，给定微批次的前向和后向传播使用相同版本的模型参数。例如，在图8中，微批次5在机器1上使用了来自批次1的参数更新，在机器2上使用了来自批次2的更新。权重暂存并未解决一个给定微批次在不同阶段间使用的参数版本的一致性问题。

纵向同步（Vertical Sync）。纵向同步消除了跨阶段的潜在不一致性。例如，在图8中，使用纵向同步，微批次5在其前向和后向传播中，在所有机器上都使用由微批次1更新的参数。每个进入流水线的微批次（$m_i$）都与在输入阶段看到的最新的权重版本（$w^{(i-x)}$）相关联。当微批次$m_i$在前向方向流过流水线时，这个信息会与激活值和梯度一起传播。在所有阶段中，对$m_i$的前向传播都使用暂存的权重$w^{(i-x)}$，而不是最新的权重更新。在为$m_i$执行后向传播后（使用暂存的权重$w^{(i-x)}$），每个阶段独立地应用权重更新来创建最新的权重（$w^{(i)}$），然后可以删除$w^{(i-x)}$。这种跨阶段的协调是异步的。

陈旧度（Staleness）的数学形式化。现在我们可以为每种技术形式化权重更新的陈旧度。在此讨论中，我们假设一个模型被分成n个阶段的直通流水线；每个阶段的权重表示为$w_1, w_2$等。此外，我们用$w_1^{(t)}$表示经过t个微批次后阶段1的权重。

标准SGD更新。在每个微批次之后，我们计算在微批次中所有样本上平均的梯度$\nabla f(w_1, w_2, \ldots, w_n)$。标准的微批次SGD（$f$是我们试图优化的损失函数，$\eta$是学习率）具有以下梯度更新：

公式
公式

权重暂存下的更新。通过权重暂存，阶段1中的梯度是使用延迟了n步的权重计算的，阶段2的梯度是使用延迟了$n-1$步的权重计算的，依此类推。在数学上，这意味着我们的权重更新看起来像：

公式
公式

无权重暂存的问题。若不使用权重暂存，权重更新对于任何权重向量$w_1, w_2, \ldots, w_n$都不是损失函数$f$的有效梯度。

纵向同步下的更新。加入纵向同步后，权重更新变为：

公式
公式

这在语义上与在n台机器上进行BSP同步的数据并行相同（每台机器上的原始微批次大小相同）。

PipeDream的默认策略。权重暂存对于有意义的学习至关重要。PipeDream的默认语义（有权重暂存但无纵向同步）介于单机上常规的微批次SGD和带BSP同步的数据并行之间【8，Exploiting bounded staleness to speed up big data analytics，2014；18，More effective distributed ml via a stale synchronous parallel parameter server，2013】。我们的评估证明了它在多种模型、数据集和硬件配置上的有效性。

GPU内存管理

内存管理的挑战。随着微批次进入和离开流水线，系统必须确保GPU计算所需的输入、权重和其他中间状态存在于GPU内存中。如果管理不善，GPU中动态内存分配以及GPU和CPU内存之间数据传输的开销会极大地降低硬件效率。

PipeDream的内存管理策略。PipeDream从DNN模型中提取层参数，并计算在流水线中存在的活跃微批次上，每个阶段需要存储的激活值、参数和中间状态的大小。每个阶段需要维护中间状态的微批次数是不同的。输出阶段只需为一个活跃微批次维护中间状态，而输入阶段则需要为NOAM个微批次维护。PipeDream在训练开始时分配所有必需的GPU内存，并根据需要重用已分配的内存。这显著减少了动态管理GPU内存的开销。

实现细节

PipeDream的顶层工作流。图7展示了PipeDream的高层工作流。我们系统的输入是模型架构、训练数据集以及将用于训练的GPU数量。PipeDream首先在一台机器上使用训练数据集的一个子集对模型进行性能剖析。然后，它运行第3.2节中描述的优化算法，将DNN模型划分为k个阶段，其中一些阶段可能被复制。之后，PipeDream运行时将每个阶段分配给一个单独的GPU。

阶段运行时的架构。图9展示了PipeDream中阶段运行时的高层架构。PipeDream的接口实现为一个C++库，它为在GPU上运行的机器学习（ML）工作线程管理参数和中间数据。在我们当前的实现中，我们使用Caffe【20，Caffe: Convolutional architecture for fast feature embedding，2014】作为ML工作线程。然而，PipeDream是可扩展的，也可以与其他ML框架如Tensorflow【1】、MXNet【5】和CNTK【33】一起工作。

图9：阶段运行时架构及与Caffe的集成。PipeDream向ML工作线程（Caffe）提供指向包含层输入数据、参数以及用于记录层输出和参数更新的缓冲区的GPU内存指针。PipeDream管理缓冲池，并处理所有机器内和机器间的通信。
图9：阶段运行时架构及与Caffe的集成。PipeDream向ML工作线程（Caffe）提供指向包含层输入数据、参数以及用于记录层输出和参数更新的缓冲区的GPU内存指针。PipeDream管理缓冲池，并处理所有机器内和机器间的通信。

初始化与执行流程。作为一个初始化步骤，每台机器中的PipeDream库会初始化分配给该机器的阶段所对应的GPU数据结构。这包括：1) 用该阶段必须执行的层来初始化ML工作线程，以及 2) 在GPU中为激活值、权重、梯度和中间状态（包括每个活跃微批次的输入激活值和暂存权重）静态分配内存。一旦一台机器初始化完成，ML工作线程就从PipeDream获取其下一个工作任务；PipeDream的运行时为ML工作线程提供指向输入数据的指针。输入阶段通过为其ML工作线程创建一个前向工作的微批次来启动流水线。从那时起，每台机器都遵循1F1B调度算法，同时将活跃微批次的数量限制在NOAM以内。

ML工作线程的交互。对于分配的微批次（前向或后向），ML工作线程遍历该阶段的每一层，并为该微批次执行相关工作。ML工作线程使用适当的PipeDream API调用来获取指向每层输入、参数、输出、梯度和中间状态的指针。一旦微批次处理完毕，ML工作线程会向PipeDream指示工作完成，并拉取其下一个工作项。

参数状态管理。对于每个阶段，PipeDream直接在GPU内存中维护分配给该阶段的层的所有相关参数。每层的参数分开存储，并被分配一个唯一的ID。如果阶段没有被复制，当权重更新在提供的GPU缓冲区中可用时，PipeDream会将更新应用到存储在GPU内存中的最新版本的参数数据上。如果阶段被复制，权重更新会被复制到主机内存，然后发送到参数服务器。当有新版本的参数可用时，作为权重暂存方案的一部分，旧版本不会立即被丢弃。参数数据只有在执行了使用更新鲜参数的后向传播后才会被丢弃。

中间状态管理。每一层的中间数据也被分配一个唯一的blob ID。从前一个阶段（或对于输入阶段，从磁盘）接收到中间数据后，PipeDream将中间数据复制到GPU内存，并将指向相关缓冲区的指针放入一个工作队列中。前向传播的中间数据在相关微批次完成该阶段的后向传播之前不会被丢-弃。后向传播的中间数据在ML工作线程使用完毕后，以及在必要时发送到下一阶段后，会立即被释放。由于前向和后向传播对中间数据的需求不同，PipeDream中的阶段通常会管理多个版本的前向传播中间数据，而只管理当前运行的后向传播的一个版本的中间数据。

数据并行支持。PipeDream使用一个类似于GeePS【9】的分布式参数服务器来同步数据并行阶段各层的参数。通过使用无等待反向传播，权重梯度在计算出来后会立即被传送到服务器，而不是等待所有层的计算完成。每个工作节点包含一个参数服务器分片的实例，该实例存储参数的一个唯一子集。服务器分片在聚合了所有阶段副本的更新后，会立即将最新版本的参数推送到其他分片。

通信与检查点。PipeDream阶段之间的所有机器间通信，无论是在数据并行还是流水线并行设置中，都使用ZeroMQ【46，Distributed Messaging】和一个带有快速自定义序列化的高效通信栈。PipeDream支持为容错而定期对模型参数进行检查点，默认在每个轮次结束时跨阶段进行检查点。检查点不需要昂贵的全局协调；每个阶段在为轮次中的最后一个微批次执行后向传播时，会本地决定转储其模型参数。因阶段失败而重启失败的训练运行时，需要从所有阶段都成功创建检查点的最后一个轮次开始。

A4 实验环境

  • 数据集:

    • ILSVRC12 (ImageNet 1K): 用于图像分类，包含约130万张训练图片和5万张验证图片，共1000个类别【32】。
    • MSVD (Microsoft Video description corpus): 用于视频描述生成，包含1970个视频片段和12594个词的词汇表【3】。

  • 模型架构:

    • VGG16: 模型大小550MB，用于ILSVRC12数据集【36】。
    • Inception-v3: 模型大小157MB，用于ILSVRC12数据集【19】。
    • S2VT: 序列到序列的视频转录模型，大小349MB，用于MSVD数据集【43】。

  • 硬件配置:

    • Cluster-A (私有集群):
      • GPU: NVIDIA Titan X (12GB显存)
      • CPU: E5-2698Bv3 Xeon
      • 内存: 64GB RAM
      • 网络: 25 Gbps 以太网
    • Cluster-B (公有云 AWS p3.2xlarge):
      • GPU: NVIDIA V100 (16GB显存)
      • CPU: E5-2690 Xeon
      • 内存: 64GB RAM
      • 网络: 10 Gbps 以太网
    • 对比: Cluster-B的GPU计算能力更强，但网络带宽低于Cluster-A，这导致其计算通信比更低。

  • 软件与训练配置:

    • 软件栈: Ubuntu 16.04 (64位), CUDA 8.0, cuDNN v6。PipeDream使用Caffe【20】作为ML工作后端。
    • 训练方法:
      • VGG16/S2VT: 使用带动量的SGD优化器（动量0.9），初始学习率0.01。
      • Inception-v3: 使用RMSProp优化器【40】，初始学习率0.045，每两个epoch以0.94的指数率衰减。
    • 微批次大小: VGG16和Inception-v3为每台机器32，S2VT为每台机器80。
    • 评估标准: 测量模型达到其公开宣称的验证准确度所需的时间。VGG16为top-1准确率68%，Inception-v3为top-1准确率67%，S2VT为METEOR【11】得分0.294。
    • 基线对比: 将PipeDream与数据并行（BSP）和单机配置进行比较。PipeDream的数据并行实现与高效的开源系统GeePS【9】进行了对比，性能相当或更优。

A5 实验结果

PipeDream与数据并行（BSP）的性能对比

实验结果表明，PipeDream在通信开销大的场景下显著优于数据并行BSP。

  • PipeDream自动配置（表1）: PipeDream的优化器能自动生成结合了模型并行、流水线和数据并行的混合配置（如"2-1-1"表示3个阶段，第一阶段有2个数据并行副本）。结果显示，最优配置通常不是纯粹的数据并行或模型并行，而是混合模式，且性能远超BSP。

    表1：PipeDream与数据并行配置（BSP）在训练模型达到其宣称的最终准确度时的结果摘要。“PipeDream配置”代表由我们的分区算法生成的配置——例如，“2-1-1”是一个将模型分为三个阶段，其中第一个阶段在2台机器上复制的配置。
    表1：PipeDream与数据并行配置（BSP）在训练模型达到其宣称的最终准确度时的结果摘要。“PipeDream配置”代表由我们的分区算法生成的配置——例如，“2-1-1”是一个将模型分为三个阶段，其中第一个阶段在2台机器上复制的配置。

  • 基准测试（Cluster-A, 8台机器，图10）:

    • VGG16: 通信开销高达72%，导致BSP仅比单机快2.35倍。PipeDream通过减少95%的通信，实现了比单机快7.04倍的加速（比BSP快2.99倍）。
    • Inception-v3: 通信开销仅5%，BSP实现了接近线性的7.66倍加速。此时，PipeDream的优化器自动选择了数据并行配置，性能与BSP持平。
      图10：在Cluster-A上使用8台机器时，VGG16和Inception-v3的准确率与时间关系图
      图10：在Cluster-A上使用8台机器时，VGG16和Inception-v3的准确率与时间关系图

  • 硬件影响（Cluster-B, 8台机器，图11）:

    • Cluster-B的GPU更快但网络更慢，计算通信比更低，放大了PipeDream的优势。
    • VGG16: PipeDream相对于BSP的加速比从Cluster-A的2.99倍提升到Cluster-B的5.12倍。
    • Inception-v3: 即使对于计算密集型模型，PipeDream在Cluster-B上也比BSP快45%。
    • 其他模型: 在Cluster-B上，PipeDream对AlexNet和ResNet-50的吞吐量分别提升了6.78倍和1.21倍。
      图11：在Cluster-B上使用8台机器时，VGG16和Inception-v3的准确率与时间关系图
      图11：在Cluster-B上使用8台机器时，VGG16和Inception-v3的准确率与时间关系图

  • 扩展性（VGG16, Cluster-A, 图12）:

    • 随着机器数量增加，BSP扩展性很差（4、8、16台机器下加速比仅为1.47x, 2.35x, 3.28x）。
    • PipeDream扩展性良好（相应加速比为3.14x, 7.04x, 9.86x）。值得注意的是，4台机器的PipeDream性能几乎与16台机器的BSP相当。
      图12：在Cluster-A上使用4和16个工作节点时，VGG16不同配置的准确率与时间关系图
      图12：在Cluster-A上使用4和16个工作节点时，VGG16不同配置的准确率与时间关系图

  • 与异步并行（ASP）的比较（图12）: 尽管ASP没有通信开销，但由于其统计效率低下，PipeDream达到48%准确率的速度比4机ASP快7.4倍。

  • RNN模型（S2VT）: 对于S2VT模型，BSP几乎无加速（1.1倍）。PipeDream通过减少95%的通信，实现了比单机快3.34倍的加速（比BSP快3.01倍）。

混合并行的价值分析

通过对VGG16在Cluster-A上的实验（图13），分析了不同并行策略的贡献：
- 纯模型并行: 由于资源闲置，性能甚至低于单机。
- 直通流水线（无数据并行）: 性能优于BSP，4机和8机下分别实现2.56倍和3.49倍加速。
- PipeDream（混合并行）: 性能最佳，4机和8机下分别实现3.14倍和7.04倍加速。这证明了将流水线、模型并行和数据并行相结合的策略是最优的。

图13：在Cluster-A上，VGG16的模型并行、流水线并行与PipeDream的性能对比
图13：在Cluster-A上，VGG16的模型并行、流水线并行与PipeDream的性能对比

A6 结论

流水线并行DNN训练解决了数据并行训练在处理超大型DNN时遇到的通信开销瓶颈问题。PipeDream系统能够自动地对DNN训练进行分区，并在工作节点之间进行主动的流水线处理。通过在两个不同集群上对五种不同的DNN模型进行的实验，结果表明，与当前最先进的方法相比，PipeDream在“达到目标精度时间”上最高可提速5倍。