torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models 文章翻译

在细读文章前，首先感谢 罗西的思考 博主，大家可以去看下这个博主写的流水线并行的相关文章，介绍得真的很好，前段时间的组会，全靠他的文章救我一命。

---

摘要：我们在PyTorch中设计并实现了一个现成的库来执行 micro-batch 流水线并行，并使用了GPipe[11]提出的检查点checkpointing。特别是，我们开发了一组设计组件，以在PyTorch的按运行定义和急切执行的环境（？）中实现流水线并行的梯度计算。我们证明了在这样的环境中，每个组件都是充分受益于流水线并行性所必需的，并通过将其应用于包括AmoebaNet-D[23]和U-Net[24]在内的各种网络体系结构来展示该库的效率。我们的库在https://github.com/kakaobrain/torchgpipe.

1 介绍

近年来，深度学习在几种方法的推动下有了显著的增长，这些方法使得深度神经网络(DNN)的训练能够以可扩展的方式进行，并通过开发更强大的硬件来推动。可以看出，DNN容量的增加有效地改善了性能。例如，AmoebaNet-B[23]使用GPipe[11]进行了缩放，拥有5.57亿个参数，达到了TOP-1的84.4%的准确率(这是当时最先进的结果)，而GPT-2[22]是一个基于Transformer的[28]语言模型，拥有15亿个参数(有关模型缩放的效果，请参阅[11]的图1)。然而，训练如此庞大的模型是非常耗费资源的。人们可以通过削减模型[8，1]、设计更高效的体系结构[10、27]、在资源约束下进行体系结构搜索[3]等等来减小模型的大小而不损失性能。

我们可能会想，有没有可能采取一种相当直接的方法：在拥有大量设备的情况下，我们能足够快地训练出一个庞大的模型吗？一个障碍是，训练神经网络的常见优化技术本质上是连续的。这些算法一次重复计算相对于给定mini-batch的损失的梯度，并使用该梯度更新模型参数。由于具有丰富的计算资源，通常使用data parallelism[17]来加速整体优化过程，方法是将mini-batch划分为micro-batch，并将每个micro-batch的计算委托给可用的设备。通过仔细的超参数调整，这可以有效地将训练时间减少到一定大小的mini-batch，这可能取决于模型、优化算法和数据[6，25]。data- parallel训练的一个缺点是设备拥有它们自己的网络版本来执行细分的任务，并且必须在每次参数更新后同步网络参数。当有许多参数需要同步时，这可能会导致很大的通信负载。

请注意，当模型太大以至于无法计算梯度时，数据并行性不适用，即使在单个数据点被馈入网络时也是如此。Model parallelism[5]是一种训练庞大的模型的方法，它将模型划分为几个部分，并将它们放置在不同的设备上。每个设备只计算模型的一小部分，并且只更新该部分中的参数。然而，模型并行性受到其未充分利用行为的影响。由于大多数神经网络由层序列组成，因此保存模型较后部分的设备必须等待，直到保存模型较早部分的设备进行计算。

另一种可能的解决方案是使用梯度检查点[4] gradient checkpointing，它只存储激活图的子集，并在必要时重新计算被丢弃的激活图，从而节省内存。显然，这需要对模型的某些部分进行两次计算，并且会增加整体训练时间.

将不同类型的并行化策略[16、14、26、12、9、11、7]和最近的一系列研究问题结合起来是有益的，如何找到最优策略[15, 19, 18, 29]。在这些方法中，流水线并行pipeline parallelism是一种通过将模型并行性与数据流水线相结合来加速神经网络训练的方法，既可以像GPipe[11]那样以同步方式进行，也可以像[12]、PipeDream[9]和XTube[7]那样以异步方式进行。我们注意到，梯度检查点gradient checkpointing(也称为re-materialization)进一步结合在GPipe中，以允许训练更大的模型。

在本文中，我们设计并实现了一个在PyTorch[21]中为GPipe提供的现成的库--torchgpipe库。特别是，我们开发了一组设计组件，用于在PyTorch的按运行定义和急切执行的环境中进行优化的流水线并行计算。我们证明了在这样的环境中，每个组件都是充分受益于流水线并行性所必需的，并通过在AmoebaNet-D[23]和U-Net[24]上进行速度和内存基准测试来演示Torchgtube的效率。

论文的其余部分组织如下。在第二节中，我们讨论了如何将前向和后向遍分解为子任务(在一定的假设下)，描述了micro-batch流水线并行的设备布局策略，并演示了每个设备的期望执行顺序。在第3节中，我们讨论了在PyTorch中实现流水线并行性的最佳时间线的复杂性，并解释了torchgipe如何解决这些问题。此外，我们放宽了模型是按顺序组合的假设，并提供了一种表示具有长跳跃连接的模型的方法，从而在不牺牲效率的情况下仍然适用流水线并行。然后，我们论证了本文提出的优化组件对性能是必不可少的，并在第四节对提出的库进行了性能评估。

2 流水线并行

假设我们有一个神经网络，它被表示为子网络序列的组合。让我们用

表示这些子网络，参数  ，完整网络：

参数化为为清楚起见，我们将f^j称为f的第j个划分，并假设划分的参数是相互不相交的。 

当训练网络时，基于梯度的方法例如随机梯度下降需要在给定一个小批量x的训练数据和相应的损失的情况下计算网络的结果f(X)，以及损失相对于网络参数θ的梯度g。这两个阶段分别称为向前和向后传播。

由于f是按顺序组成的，因此在前向传递中，f(X)可以通过设并对 j = 1,···L顺序地应用分区。此外，如果x由称为微批次micro-batch的m个较小批次组成，则计算f(x)分解为任务

假设 f 不涉及任何批次内计算。这方面的一个显著例外是批处理标准化[13]。损失是通过聚合并计算其上的损失函数来获得的。

以类似的方式，向后传递被分解成任务，其中是损失相对于的梯度，并且

是一个通过分区 进行反向传播（也称为向量雅可比积）的函数，的定义类似。因此，我们通过对和求和得到了关于的损失的梯度。

请注意，任务之间存在数据依赖性。例如，which is only available可用的 after（completed 完成的),图 1 显示了 m = 4 和 n = 3 情况下的完整依赖图。

 给定一组任务和可以并行工作的设备池，不同的并行化策略有自己的规则来将任务分配给设备。一旦依赖关系得到解决，每个设备就会计算一个或多个分配的任务。在上面的设置中，所有依赖项都位于具有相同微批次索引 i 的任务之间。因此，通过将具有不同微批次索引的任务分配给不同的设备，可以有效地并行化任务——这就是数据并行。

2.1  GPipe的依赖图

流水线并行的策略是根据分区索引 j 分配任务，使得第 j 个分区完全位于第 j 个设备中。除此之外，还强制规定 必须在执行 之前完成，必须在执行 之前完成。 

除了微批量micro-batch流水线之外，GPipe[11]通过对每个使用梯度检查点gradient checkpointing进一步减少了内存需求。由于第 j 个设备一次执行一个，因此只需要从 获得的激活映射即可完成 。通过在执行 之前重新计算前向传递 ，内存消耗减少了 m 倍。此外，重新计算可以在设备等待  完成时进行。图 2 对此进行了总结，其中虚线箭头表示由微批量顺序引起的独立任务之间的执行顺序，表示 的重新计算。

我们注意到最后一个微批次的重新计算，即是没必要，无用的。这是因为在第 j 个设备上，前向传递中的最后一个任务是 ，因此，在前向传递中丢弃中间激活并在后向传递开始时重新计算它们并不会减少内存，只会减慢流水线速度。因为这个原因，图中省略了。

2.2  设备方面的执行顺序

总而言之，在流水线并行性（带有检查点）中，每个设备都被分配了一组具有规定顺序的任务。一旦满足跨设备依赖性，每个设备就会一一执行给定的任务。然而，这张图中缺少一个组件——设备之间的数据传输。为了便于说明，图 3 显示了设备 j 必须遵循的完整执行顺序。这里为了强调，数据传输操作明确表示为“接收”和“发送”。

 3 torchgpipe: GPipe的PyTorch库

torchgpipe 是一个 PyTorch 库，用于具有检查点的微批量流水线并行性，称为 GPipe。该库提供了一种将 GPipe 应用于用 PyTorch 编写的通用顺序模块的简单方法。 torchgpipe 的用法类似于 PyTorch 的数据并行模块 - 只需用包装器包装您的模型即可。

用户必须指定微批次的数量 m 以及连续层如何形成 n 个分区。在这里我们指出，尽管我们将模型简化为模型是一系列分区，但在 torchgpipe 中严格要求模型是一顺序层，以便为用户提供如何拆分模型的灵活性。 torchgpipe 将假设每一层都是不可分割的、黑盒的、引用透明的算法。

为了方便起见，该库提供了子模块 torchgpipe.balance，它计算成对资源差异较小的分区，其中资源消耗是通过分析计算的。具体来说，我们使用[2]中的算法。

由于 torchgpipe 是在配备 CUDA 后端的 PyTorch 上构建的，因此我们在本节中通常会假设设备是 NVIDIA GPU。尽管如此，该库的基本原理通常适用于在任何急切执行环境中实现管道并行性。

3.1  PyTorch 中的并发性

我们最关心的是效率。正如我们在 2.2 节中讨论的，为了使流水线并行性按需要工作，必须以正确的顺序将任务分配给每个设备。在 PyTorch 中实现这一点有几个复杂的问题。

首先，由于 PyTorch 的“按运行定义”风格及其急切的执行行为（与“构造并运行”类型框架相反），内核会即时发布到每个设备。因此，必须仔细设计主机代码，以便不仅设备绑定任务在每个设备内以正确的顺序发出，而且设备上任务的执行（与 CPU 异步）不会因 Python 解释器失败而延迟提前请求。当某些任务是 CPU 密集型任务或涉及大量廉价内核调用时，可能会发生这种延迟。作为一种解决方案，torchgpipe 引入了确定性时钟周期，它给出了任务的总排序。

其次，后向传递的计算图是在 PyTorch 中前向传递期间动态构建的。换句话说，“它避免了具体化“前向图”，只记录区分计算所需的内容。” [21] 由于 PyTorch 不记录前向计算图，也不维护梯度带，因此 PyTorch 的自动微分（autograd）引擎仅针对图进行反向传播。这意味着 autograd 引擎可能不会完全按照与前向传递相反的执行顺序运行，除非由图的结构强制执行。为了解决这个问题，我们开发了一对称为“fork”和“join”的原始函数，以在向后计算图中动态创建显式依赖关系。

第三，如果不仔细管理，多个设备之间的通信可能会导致双向同步。这可能会导致利用率不足，因为即使副本和队列中的下一个任务之间没有显式依赖关系，发送方也可能会等待与接收方同步，反之亦然。 torchgpipe 通过使用非默认 CUDA 流来避免此问题，这样副本就不会阻塞计算，除非计算必须等待数据。

最后，torchgpipe 试图放宽微批量流水线并行性的限制，即模型必须是顺序的。尽管原则上任何神经网络都可以以顺序形式编写，但这需要提前了解整个计算图，而 PyTorch 中的情况并非如此。特别是，如果有一个张量从设备 中的一层跳到设备  中的另一层，则该张量将被复制到其间的所有设备，因为 torchgpipe 无法提前知道它。为了解决这个问题，我们设计了一个接口来表示跳过哪些中间张量以及哪些层使用它们。

3.2  优化组件

在本节的其余部分中，将解释 torchgpipe 的组件是如何设计的以及为什么每个组件对于性能至关重要。

3.2.1  前向依赖性：确定性时钟周期

正如我们在 3.1 节中讨论的，任务的总排序由前向传递中的主机代码确定。每个设备根据 CPU 分配任务的顺序隐式地了解任务之间的依赖关系。理想情况下，如果可以免费将任务分配给设备，那么只要设备内的顺序正确，CPU 就可以按任意顺序将任务分配给设备。然而，这种假设不够现实，因为在 GPU 上启动内核对于 CPU 来说并不是免费的，GPU 之间的内存传输可能需要同步，或者任务是 CPU 密集型的。因此，我们通过按到  的距离对所有任务进行排序来最小化来自 CPU 的延迟。（？GPU运算速度很快，但CPU读写慢）

我们将此称为确定性时钟周期（算法 1）。在该算法中，CPU 执行从计数器 k = 1 到 k = m + n − 1 的时钟周期。 在第 k 个时钟周期，首先发出执行任务 所需数据的所有复制内核，其中 i + j − 1 = k，然后将执行任务的计算内核注册到相应的设备（可以安全地多线程，因为同一时钟周期内的任务是独立的）。

3.2.2  向后依赖：分叉和加入

现在假设我们根据确定性时钟周期运行前向传递。即使设备 j 上的前向任务  按顺序执行，后向所得的计算图看起来也更像 图1 而不是 图2。从这样的图中，PyTorch 的 autograd 引擎永远不会知道  必须在 之前执行，这会打乱向后传递的时间线。因此，必须在前向传递期间显式绘制虚拟依赖关系（图 2 中的虚线箭头）。

我们设计了一对称为 Fork 和 Join 的原始简单函数（或者应该翻译为 简单函数）来表达这种依赖关系。基本上，Fork 是将张量 x 映射到 对的 autograd 函数，其中 是一个空张量 ，Join 是将一对 映射到张量 x 的 autograd 函数。现在， 对 的依赖关系（在后向计算图中转化为 对 的依赖关系）可以表示为

请参见图 4 进行说明。 

3.2.3  并发复制和计算：流

PyTorch 将每个设备绑定的内核发布到默认流，除非另有指定。流是按顺序执行的设备绑定内核序列。同一流中的内核保证按规定顺序执行，但不同流中的内核可以交错，甚至在可能的情况下可以重叠。特别是，几乎所有具有计算能力 1.1 及更高版本的 CUDA 设备都支持并发复制和执行：设备之间的数据传输始终可以与内核执行重叠（请参阅[20]的第 4.5.1.5 节）。

torchgpipe 将每个复制内核注册到非默认流，同时将计算内核保留在默认流上。这允许设备 j 在向设备 j + 1 发送  和/或从设备 j − 1 接​​收 ​​​​​​​的同时处理。此外，每个设备对每个微批次使用不同的流。由于不同微批次之间不存在真正的依赖性，因此这种流的使用是安全的，并且允许尽可能快地进行复制。请参见图 5 进行说明。

3.2.4  具有共享内存的 Autograd 函数

到目前为止，在本节中，我们没有讨论在使用梯度检查点时如何调度重新计算任务 。它必须在 完成后安排在反向传播任务 之前。对于 autograd 引擎，这也必须编码在计算图中。事实上，PyTorch 通过用于检查点的内部 autograd 函数支持此类功能。 

PyTorch 中的检查点是通过定义一个 autograd 函数来实现的，该函数在前向传递中像平常的函数一样计算，不存储中间激活图，只存储输入。在向后传递中，该函数通过使用存储的输入重新计算函数来构造向后的局部计算图，并通过局部图反向传播来计算梯度。然而，这将 和 紧密地结合在一起。最终，我们希望在 和 之间插入等待 的结果 从设备 j + 1 复制到设备 j 的指令，以允许 和复制同时发生。

 对于这种细粒度的顺序控制，torchgpipe 使用两个独立的自动分级函数 Checkpoint 和 Recompute 来实现检查点。在任务 执行时，会生成一对具有共享内存的 Checkpoint 和 Recompute。该共享内存在向后传递中用于将执行重新计算生成的本地计算图传输到检查点以进行反向传播。通过安排函数使得 、接收 的同步、在向后传递期间按顺序执行，确保了重新计算和复制可以同时发生。

3.3 处理非序列模型

在第 2 节中，我们假设模型 f 由按顺序的分区 组成。原则上，任何神经网络都可以通过按拓扑顺序对 f 的前向计算图中的所有节点进行排序来表示为这种形式。因此，流水线并行性适用于任何模型。 

但是，考虑一种症状情况，即除了第一个和最后一个分区之外的所有分区都是并行的，即

在这种情况下，以其本机形式使用流水线并行性的效率非常低，因为在设备 j − 1 和 j 的边界处，必须复制元组 而不是单个张量，这是计算第 j 个分区所需的唯一数据。

torchgpipe 提供了一个子模块，允许用户指示从哪一层跳到哪一层张量：torchgpipe.skip。使用装饰器@skippable，用户定义的层可以存储一个张量供以后使用，或者通过Python中的yield运算符弹出存储的张量而不返回它。这尤其不会改变层的输入和输出签名。因此，将跳跃连接添加到预先存在的顺序模型中只需付出最少的努力。

3.3.1  在图中隐藏跳过张量：门户

将跳过连接添加到依赖关系图中（图 2）相当简单。事实上，无论添加多少个跳过连接，都不会引入额外的依赖关系，因此只有用于跳过连接的复制内核需要额外小心。在torchgpipe中，这是由门户网站负责的，门户网站由三个自动分级功能PortalBlue、PortalOrange和PortalCopy共享内存组成，如第3.2.4节中的检查点和重新计算。每个都分别执行保存跳过张量、加载张量以及将保存的张量移动到跳过的设备的工作（在向后传递中反之亦然）。该机制如图 6 所示。

4 实验 

每个实验均使用配备 CUDA 10.1.243 的 NVIDIA Tesla P40 GPU 进行，每个 GPU 具有 22 GiB 内存。为了可重复性，本节中提供的所有基准测试的代码均在存储库中提供。

4.1 优化组件的效果

我们进行了一项实验，证明 torchgpipe 的每个组件都是实现最大效率所必需的。从仅具有确定性时钟周期而没有其他时钟周期的基线开始，逐步添加每个组件（通过 Fork 和 Join 的向后依赖、用于复制内核的非默认流以及用于跳过连接的门户）。我们报告每个设置下的吞吐量、GPU 利用率和内存使用情况，以衡量每个组件对 torchgpipe 性能的贡献。我们发现添加每个组件都会加快速度，并且（添加）所有组件使 torchgpipe 的运行速度几乎是基线的两倍。结果见表1。

我们使用U-Net进行实验。该架构的详细信息可以在第 4.2.2 节中找到，我们将 (B，C) 设置为 (5，64)，如速度基准中所示。在没有门户的设置中，模型被实现为完全顺序的版本，其中跳跃连接被编码为它们所经过的层的输入和输出，如第 3.3 节的症状示例中所述。对于所有组件的设置，都是通过torchgpipe.skip实现的，但架构是相同的。

我们还可视化了每个 GPU 时间线，以帮助理解每个组件的角色，如图 7 所示。每张图片的说明总结如下。

(a) 通过确定性时钟周期，在前向传递期间所有内核都以正确的顺序发出。它由时间线的左侧部分说明。然而，如果没有在计算图中编码显式的依赖关系，autograd 引擎就会以不可控的顺序处理微批次，因此时间线会变得混乱。

(b) 通过向后依赖，现在可以在向后传递中以正确的、确定性的顺序发布内核。

(c) 通过使用非默认复制流，复制和计算现在是并发的，如重叠的蓝色和红色条所示。

(d) 门户删除由于将跳跃张量传输到其间的所有设备而导致的不必要的副本。与 (c) 相比，红色条的长度减少了，这说明了这一点。

4.2 性能基准

为了证明 torchgpipe 的效率，我们报告了类似于 GPipe [11] 进行的性能基准。

4.2.1  AmoebaNet-D 速度基准

我们测量了 AmoebaNet-D 使用不同数量设备的吞吐量。为此，我们测量了应用 torchgpipe 时模型的吞吐量，其中包含 n 个分区和 m 个微批次。这里的吞吐量是指每秒处理的样本数。 

对每对 (m , n) 进行实验，其中 。当 m = 1 时，我们对所有微批次使用检查点，以对检查点造成的损失与[11]进行公平比较。我们使用的模型是我们在 PyTorch 中实现的 AmoebaNet-D 的顺序版本。

该模型通过普通 SGD 训练 10 个时期，并报告除第一个时期之外的各个时期的平均吞吐量。为了排除数据加载造成的开销，我们使用了一个由 10,000 张尺寸为 3 × 224 × 224 的图像组成的合成数据集。对于每个设置，选择批量大小和微批量数量以最大化吞吐量。相对加速比是根据基线情况 (m , n) = (1 , 2) 计算的，并在表 2 中报告。我们将 GPipe 的加速比纳入其中进行比较。

torchgpipe 的相对加速表现出与 GPipe 相似的趋势。我们注意到表 2 中报告的性能差异可能是由于许多未知因素造成的，例如分区的平衡、实现之间的差异、设备的差异等。 

4.2.2  U-Net 内存基准测试

为了评估 torchgpipe 对于具有长跳跃连接的模型的有效性，我们使用 U-Net [24] 进行二维分割。我们使用的 U-Net 版本有五个下采样层和五个上采样层，以及两个决定模型大小的超参数 B 和 C。这里B代表下采样层之间的卷积块的数量，C代表第一个卷积的输出通道的数量。通道在每个下采样层之后加倍（或在每个上采样层之后分别减半）。我们的 U-Net 实现比 [24] 中提出的原始模型更加对称，以实现有效平衡。

我们进行了一项实验来衡量 torchgpipe 训练更大模型的能力。对于 1、2、4 和 8 个 GPU，我们发现占用每个数量的设备的最大值（B , C）。在所有设置中，输入大小设置为 3 × 192 × 192，输出大小设置为 1 × 192 × 192，批量大小设置为 32。表 3 报告了训练每个模型的总内存使用量。这里参数消耗每个 8 个字节用于其自身及其梯度。

4.2.3  U-Net 速度基准

我们还测量了 U-Net 与不同数量设备的吞吐量。 Naive-1 表示没有管道并行性和检查点的基线，Pipeline-1、-2、-4、-8 表示模型使用 torchgpipe 和相应数量的分区进行训练。 在本实验中，决定 U-Net 大小的超参数设置为 (B , C) = (5 , 64)。选择批次大小、微批次数量 (m) 以及分区平衡以最大化吞吐量。对于每个设置，吞吐量的测量如第 4.2.1 节所示，不同之处在于本实验中图像大小为 3 × 192 × 192。结果总结于表 4 中。

5 总结

在本文中，我们介绍了 torchgpipe，它是 PyTorch 中的一个现成库，用于由 GPipe [11] 提出的带有检查点的微批量管道并行性。该库是在 PyTorch 的运行时定义和热切执行环境中设计和实现的。第 4 节中介绍的消融研究和性能基准表明，torchgpipe 的所有组件对于在急切执行环境中通过检查点实现管道并行性和检查点的所需优势至关重要。我们相信，我们在本文中建立的一般原则适用于任何其他具有急切执行环境的框架。

我们试图避免过于深入地讨论 torchgpipe 所涉及的技术细节。我们的代码可在 https://github.com/kakaobrain/torchgpipe 上找到，对于那些对更多细节感兴趣的人以及那些想要将管道并行性应用于 PyTorch 中的模型的人。