微软提出DNN并行训练新方法PipeDream，比传统方法快四倍

本文介绍了微软新提出的PipeDream，旨在使深度学习网络训练并行化水平更高，进而提高训练效率。

深度神经网络（DNNs：Deep Neural Networks）已经在大量应用中取得了巨大进展，这些应用包括图像分类、翻译、语言建模以及视频字幕等。但DNN训练极其耗时，需要多个加速器高效地并行化。

在论文“PipeDream: Greneralized Pipeline Parallelism for DNN Training”（该论文发表于第27届ACM操作系统原理研讨会：SOSP 2019）中，微软系统研究小组的研究员，与来自卡内基梅隆大学和斯坦福大学的同事和学生们一起提出了一种DNN并行化训练的新方法。正如论文里在大量模型上展示的那样，这套被称为PipeDream的系统比传统的训练方法快了最多4.3倍。

DNN训练是在前向和后向通道计算中迭代进行的。在每轮迭代中，训练过程循环处理输入数据的一个minibatch，并且更新模型参数。最常见的DNN并行化训练方法是一种被称为数据并行化的方法（见下图1），这种方法把输入数据分散到各个workers（加速器）中运行。

不幸的是，尽管在数据并行化加速方面有一些性能优化的进展，但若要放在云基础设施上训练就会产生很大的通信开销。而且，GPU计算速度的飞快提升，更进一步地把所有模型训练的耗时瓶颈转向了通信环节。

不那么常见的并行化训练形式是模型的并行化（见下图1），是把算子分散到各个worker上计算的，这在以往通常用于训练大型DNN模型。模型并行化也遇到了挑战：它不能高效地利用硬件资源，并且需要程序员决定怎样按照给定的硬件资源来分割特定的模型，这其实无形中加重了程序员的负担。

PipeDream是微软研究院Fiddle项目开发的一个系统，它引入了流水线并行化，这是一种DNN并行化训练的新方法，结合了传统的batch内并行化（模型并行化和数据并行化）和batch间并行化（流水线）。

图1：传统的batch内并行化训练方式（如数据并行化和模型并行化）对硬件利用率太低。左图中，数据并行化中的单个worker在交换梯度数据时不得不进行通信等待。右图是模型并行化，worker之间只能同时处理一个minibatch，这大大限制了硬件利用率。

使用流水线并行化训练来解决batch内的并行化限制

PipeDream重新审视了模型的并行化，希望以此来优化性能，这与以往的动机不同，以前模型并行化是因为训练大型模型时，训练过程受限于数据集的大小。PipeDream使用多个输入数据的流水线作业，来克服模型并行训练时硬件效率限制的问题。典型的流水线并行化设置涉及到不同stage之间layer的分割，每个stage都会被复制并且并行运行数据。流水线会被注入多个batch，以使流水线满负荷运行在稳定状态。在大部分情况下，流水线并行化训练比数据并行化训练需要通信的数据要少很多，因为它只需要在两个stage边界之间传输activation和梯度。在稳定状态下，所有的workers时刻都在运转，不像模型并行化训练中会有停下来等待的时候（如下图所示）。

图2：左图中，我们展示了一个流水线并行化的例子，8块GPU被分配到4个stage中。通信只用于stage边界的actiavtion和梯度上。stage 1,2和3通过对其各自的stage进行复制来保证流水线的负载平衡。右图展示了具有4个worker的一个流水线，展示了启动阶段和稳定阶段。在这个例子中，后向处理花费的时间是前向处理的两倍。

在PipeDream中克服流水线并行化训练的挑战

为了获得流水线并行化训练的潜在收益，PipeDream必须克服三个主要挑战：

• 首先，PipeDream必须在不同的输入数据间，协调双向流水线的工作。
• 然后，PipeDream必须管理后向通道里的权重版本，从而在数值上能够正确计算梯度，并且在后向通道里使用的权重版本必须和前向通道里使用的相同。
• 最后，PipeDream需要流水线里的所有stage都花费大致相同的计算时间，这是为了使流水线得到最大的通量（因为最慢的stage会成为流水线通量的瓶颈）。

图3：PipeDream工作流程概览

图3展示了PipeDream工作流程的顶层概览。给定一个模型和硬件部署方式，PipeDream在单个GPU上进行短暂的运行时性能分析后，可以自动决定怎样分割这些DNN算子，如何平衡不同stage之间的计算负载，而同时尽可能减少目标平台上的通信量。PipeDream即使是在不同的模型（不同点体现在计算和通信方面）和不同的平台上（不同点体现在互联的网络拓扑和分层带宽上）也能够有效的均衡负载。

由于DNN并不总是在可用的worker间进行均等分割，所以PipeDream可能在某些stage上使用数据并行化——多个worker会被分配到给定的stage上，并行化地处理不同的minibatch。PipeDream使用称作1F1B的调度算法来使硬件保持满负荷运转，同时还能达到类似数据并行化的语义。

在1F1B的稳定状态下，每个worker为它所在的stage严格地切换前向和后向通道，保证资源的高利用率（可忽略的流水线暂停，没有流水线flush），即使在常见情况下，后向通道花费的时间多于前向通道时也是如此。上面的图2已经通过例子展示了这一点。与数据并行化不同，1F1B还使用不同版本的权重来维持统计有效性。最后，PipeDream扩展了1F1B，在数据并行stage上引入了循环调度策略，保证了后向通道计算的梯度被引流至前向通道对应的worker上。

PipeDream是基于PyTorch（PipeDream的早期版本使用Caffe）构建出来的。我们的评估围绕着DNN模型、数据集和硬件配置的不同组合进行，证实了PipeDream流水线并行化所带来的训练时间上的收益。相比于数据并行化训练，PipeDream在多GPU机器上达到了很高的精确度，性能方面，在图像分类上有5.3倍提升，机器翻译上3.1倍提升，语言建模任务有4.3倍提升，视频字幕模型则有3倍提升。PipeDream比模型并行化有2.6至15倍的性能提升，相比混合并行化有1.9倍提升。

如果你对PipeDream更多细节感兴趣，可以在GitHub上找到源码。

原文链接：

https://www.microsoft.com/en-us/research/blog/pipedream-a-more-effective-way-to-train-deep-neural-networks-using-pipeline-parallelism/