作者/机构: Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Mia Xu Chen, Dehao Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V. Le, Yonghui Wu, Zhifeng Chen (Google)

A1 主要贡献

本文的核心问题在于，扩大深度神经网络的容量是提升模型质量的有效方法，但这常常受限于单个加速器（如GPU或TPU）的内存，需要开发专用的算法或基础设施，而这些解决方案通常与特定架构绑定，难以迁移。为了解决这一挑战，研究的目标是设计一种高效、任务无关且易于使用的模型并行化方法。

为实现这一目标，本文提出了GPipe，一个灵活的流水线并行库，其核心创新点如下：
1. 提出GPipe库：该库通过将模型划分为多个分区并置于不同加速器上，允许用户扩展任何可以表示为层序列的网络，从而突破单加速器的内存限制。
2. 新颖的批次分割流水线算法：GPipe将一个训练小批量（mini-batch）分割成多个更小的微批次（micro-batches），并在不同加速器上以流水线方式执行这些微批次。这种设计使得不同加速器可以同时处理不同的数据块，从而提高了硬件利用率。
3. 同步梯度更新与训练一致性：在每个小批次结束时，GPipe会累加所有微批次的梯度，并进行一次同步的参数更新。这种同步更新机制确保了训练过程与分区数量无关，保持了与单设备训练的数学一致性，使得研究人员可以通过增加加速器数量来轻松扩展模型规模。
4. 支持重计算（Re-materialization）：为了进一步减少内存占用，GPipe支持重计算技术，在前向传播时只存储分区边界的激活值，在反向传播时重新计算中间层的激活，显著降低了峰值内存需求。
5. 有效性验证：本文通过两个大规模任务验证了GPipe的有效性和灵活性：
* 图像分类：训练了一个包含5.57亿参数的AmoebaNet模型，在ImageNet-2012数据集上达到了84.4%的Top-1准确率。
* 多语言神经机器翻译：训练了一个包含60亿参数、128层的Transformer模型，该模型在一个涵盖100多种语言的语料库上进行训练，其翻译质量超过了所有对应的双语模型。

A2 方法细节

2.1 接口

GPipe的接口设计。任何深度神经网络都可以被定义为一个包含L层的序列。其中，每一层 $L_i$ 由一个前向计算函数 $f_i$ 和相应的参数集 $w_i$ 组成。GPipe还允许用户为每层指定一个可选的计算成本估算函数 $c_i$。给定K个分区，这个L层的序列可以被划分成K个复合层，也称为单元（cell）。如果一个单元 $p_k$ 由从第i层到第j层的连续层组成，那么该单元对应的参数集就是 $w_i, w_{i+1}, \dots, w_j$ 的并集，其前向计算函数为 $F_k = f_j \circ \dots \circ f_{i+1} \circ f_i$。相应的反向传播函数 $B_k$ 可以通过自动符号微分从 $F_k$ 计算得出。该单元的成本估算值 $C_k$ 则设为 $\sum_{l=i}^{j}c_l$。GPipe的接口极其简单直观，用户只需指定：(i) 模型的份区数K，(ii) 微批次的数量M，以及(iii) 定义模型的L个层的序列和具体定义。


图 2: (a) 一个包含顺序层的示例神经网络被划分到四个加速器上。Fk是第k个单元的复合前向计算函数。Bk是反向传播函数，它依赖于来自上层的Bk+1和Fk。(b) 朴素的模型并行策略由于网络的顺序依赖性导致严重的利用率不足。(c) 流水线并行将输入的小批量（mini-batch）划分为更小的微批量（micro-batch），使不同的加速器能同时处理不同的微批量。梯度在最后同步应用。

2.2 算法

GPipe的并行算法与梯度更新机制。用户定义了网络的层序列（包括模型参数 $w_i$、前向计算函数 $f_i$ 和成本估算函数 $c_i$）之后，GPipe会将网络划分为K个单元，并将第k个单元放置在第k个加速器上。在分区边界，系统会自动插入通信原语，以允许相邻分区之间的数据传输。分区算法的目标是最小化所有单元估算成本的方差，从而通过同步所有分区的计算时间来最大化流水线的效率。在前向传播过程中，GPipe首先将每个大小为N的小批量（mini-batch）分成M个大小相等的微批次（micro-batch），这些微批次通过K个加速器以流水线方式处理。在反向传播过程中，每个微批次的梯度是基于与前向传播相同的模型参数计算的。在每个小批次结束时，来自所有M个微批次的梯度被累加起来，并用于更新所有加速器上的模型参数。图2c展示了这一系列操作的流程。

批归一化处理。如果网络中使用了批归一化（Batch Normalization）【20，Batch normalization: Accelerating deep network training by reducing internal covariate shift，2015，ICML】，训练期间输入的充分统计量会在每个微批次上计算，并在必要时跨副本计算【21，Megdet: A large mini-batch object detector，2017，CVPR】。同时，我们还跟踪整个小批次上充分统计量的移动平均值，以便在评估阶段使用。

2.3 性能优化

通过重计算降低激活内存。为了减少激活值所占用的内存，GPipe支持重计算（re-materialization）【14，Training deep nets with sublinear memory cost，2016，arXiv】。在前向计算过程中，每个加速器只存储位于分区边界的输出激活值。在反向传播过程中，第k个加速器会重新计算其复合前向函数 $F_k$。这样一来，峰值激活内存需求被降低到 $O(N + L_K \times N_M)$，其中 $N_M$ 是微批次的大小，$L_K$ 是每个分区的层数。相比之下，在没有重计算和分区的情况下，内存需求为 $O(N \times L)$，因为计算梯度 $b_i$ 需要上层梯度 $b_{i+1}$ 和缓存的激活值 $f_i(x)$。

流水线气泡开销分析。如图2c所示，分区操作会在每个加速器上引入一些空闲时间，我们称之为气泡开销（bubble overhead）。这个气泡时间在M个微步数上摊销后为 $O(\frac{K-1}{M+K-1})$。在我们的实验中，当 $M \geq 4 \times K$ 时，我们发现气泡开销可以忽略不计。这部分原因也是因为反向传播过程中的重计算可以提前安排，无需等待来自前几层的梯度。

通信开销与负载均衡。GPipe引入的通信开销很低，因为我们只需要在加速器之间的分区边界传递激活张量。因此，即使在没有高速互连的加速器上，我们也能实现高效的扩展性能。图2c假设分区是均匀平衡的。然而，不同层的内存需求和计算浮点数操作（flops）通常是不平衡的。在这种情况下，不完善的分区算法可能会导致负载不平衡。更好的分区算法有可能在我们目前使用的启发式方法之上进一步提升性能。

A3 性能分析

我们使用两种截然不同的模型架构来评估GPipe的性能：AmoebaNet【12，Regularized evolution for image classifier architecture search，2018，arXiv】卷积模型和Transformer【15，Attention is all you need，2017，Neurips】序列到序列模型。我们通过实验研究了它们的可扩展性、效率和通信成本。

模型大小扩展性分析。我们预期重计算和流水线并行都能有效利用内存，从而使训练巨大模型成为可能。表1报告了在合理的大输入尺寸下，GPipe支持的最大模型尺寸。对于AmoebaNet，我们在每个加速器拥有8GB内存的Cloud TPUv2上进行了实验。我们使用了固定的224×224输入图像尺寸和128的小批量大小。在不使用GPipe的情况下，单个加速器最多可以训练一个8200万参数的AmoebaNet，这受限于设备内存。由于在反向传播中采用了重计算和批次分割，GPipe将中间激活内存需求从6.26GB减少到3.46GB，使得单个加速器上可以容纳一个3.18亿参数的模型。通过模型并行，我们能够在8个加速器上将AmoebaNet扩展到18亿参数，比不使用GPipe时增加了25倍。在这种情况下，最大模型尺寸没有实现完全线性扩展，这是因为AmoebaNet中模型参数在不同层上的分布不均衡。


表 1: 在不同场景下GPipe支持的AmoebaNet最大模型尺寸。Naive-1指没有GPipe的顺序版本。Pipeline-k指在k个加速器上使用GPipe进行k个分区。AmoebaNet-D (L, D)：具有L个正常单元层和滤波器大小D的AmoebaNet模型。Transformer-L：具有L层、2048模型维度和8192隐藏维度的Transformer模型。由于我们使用RMSProp进行训练，每个模型参数需要12字节。

Transformer模型扩展性分析。接下来，我们使用每个加速器核心拥有16GB内存的Cloud TPUv3来训练Transformer模型。我们使用了固定的32k词汇表大小、1024序列长度和32的批量大小。每个Transformer层具有2048的模型维度、8192的前馈隐藏层维度和32个注意力头。我们通过改变层数来扩展模型。重计算使得在单个加速器上可以训练2.7倍大的模型。通过128个分区，GPipe允许将Transformer扩展到高达839亿参数，比单个加速器所能支持的增加了298倍。与AmoebaNet不同，Transformer的最大模型尺寸随加速器数量线性扩展，因为每个层具有相同的参数数量和输入大小。


表 2: 在TPUs上使用不同分区数K和不同微批量数M的GPipe的归一化训练吞吐量。性能随着微批量数的增加而提升。当M >> K时，Transformer模型的加速比与加速器数量几乎成线性关系。必要时调整批量大小以适应内存。

训练效率评估。为了评估效率，我们在表2中报告了AmoebaNet-D (18, 256)和Transformer-48使用不同分区数和微批量数的GPipe归一化训练吞吐量。每个分区被分配到一个独立的加速器上。我们观察到，当微批次的数量M至少是分区数K的4倍时，气泡开销几乎可以忽略不计。对于Transformer模型，当分区到4倍多的加速器上时，获得了3.5倍的加速。此外，训练吞吐量几乎与设备数量成线性扩展，这得益于计算在Transformer各层之间的均匀分布。相比之下，AmoebaNet模型由于其计算分布不均衡，实现了次线性的加速。当M相对较小时，气泡开销变得不可忽略。当M为1时，实际上没有流水线并行。我们观察到无论使用多少加速器，吞吐量都相对恒定，这表明在任何给定时间只有一个设备在进行有效计算。


表 3: 在没有高速互连的GPU上使用GPipe的归一化训练吞吐量。

通信开销影响。为了测量GPipe的通信开销影响，我们在单个主机上运行实验，该主机配备了多个NVIDIA P100 GPU，但没有NVLink。因此，GPU之间的数据传输必须通过相对较慢的PCI-E进行设备到主机和主机到设备的传输。微批次的数量固定为32。如表3所示，当我们将AmoebaNet-D (18, 128)的分区数从2增加到8时，我们观察到2.7倍的加速。对于24层的Transformer，加速比为3.3倍。这与我们在配备了高速互连的TPU上观察到的线性加速相似。设备之间的通信带宽不再是模型并行的瓶颈，因为GPipe只在分区的边界传输激活张量。

A6 设计特性与权衡

已有的多种方法被提出来实现高效的大规模模型并行。然而，每种方法都有其自身的权衡，使其适用于在特定硬件约束下扩展特定架构。在此，我们重点介绍几种模型并行方法的设计选择和权衡，并比较它们在各种硬件约束和架构变体下与GPipe在灵活性、可扩展性和效率方面的差异。

现有模型并行方法的局限性。模型并行的核心思想是将网络划分为不同的计算单元，然后将它们放置在不同的设备上【42，One weird trick for parallelizing convolutional neural networks，2014，arXiv；43，On model parallelization and scheduling strategies for distributed machine learning，2014，Neurips；44，Device placement optimization with reinforcement learning，2017，arXiv；45，Large scale distributed deep networks，2012，Neurips】。从概念上讲，这支持将各种模型扩展到巨大的容量。然而，这些方法通常存在硬件利用率低和设备通信瓶颈的问题。为了应对这些挑战，人们提出了单程序多数据（SPMD）和流水线并行。

与Mesh-Tensorflow (SPMD)的对比。Mesh-Tensorflow【34，Mesh-tensorflow: Deep learning for supercomputers，2018，Neurips】遵循SPMD范式，它将用于数据并行的单指令多数据（SIMD）方法扩展到其他张量维度。SPMD允许将每个计算拆分到多个设备上，使用户能够随加速器数量线性扩展单个矩阵乘法的大小（从而扩展单个层的模型参数）。然而，这也引入了加速器之间的高通信开销，因为需要大量类似AllReduce的操作来合并每个并行化矩阵乘法的输出。这限制了该方法的适用场景，通常需要加速器之间有高速互连。此外，SPMD限制了可以有效扩展的操作类型，使其用途局限于特定的网络架构和机器学习任务。例如，在这种范式下，沿卷积层的通道维度进行拆分效率不高，因为通道实际上是全连接的；而沿空间维度拆分则需要复杂的技术来处理“光晕区域”（halo regions）。虽然SPMD通过使每个操作变小来支持模型深度的扩展，但它需要将每个层拆分到更多的加速器上，这反过来又进一步增加了设备间的通信开销。

与PipeDream的对比。其他方法尝试利用基于流水线并行的方法来扩展神经网络【46，Performance analysis of a pipelined backpropagation parallel algorithm，1993，IEEE Transactions on Neural Networks；47，Google’s neural machine translation system: Bridging the gap between human and machine translation，2017，Transactions of the Association for Computational Linguistics】。应用于神经网络训练的最新流水线并行迭代是PipeDream【48，Pipedream: Fast and efficient pipeline parallel dnn training，2018，arXiv】，其目标是减少参数服务器【49，Scaling distributed machine learning with the parameter server，2014，OSDI】的通信开销。PipeDream将前向传播的执行流水线化，并与反向传播交错进行，以最大化硬件利用率。这种设计存在由异步反向更新引入的权重陈旧性（weight staleness）问题。为了避免由权重陈旧性引起的优化问题，PipeDream需要在每个加速器上维护模型参数的多个版本化副本，以便准确计算梯度更新，这阻止了用户扩展到更大的模型。

GPipe的设计优势与局限性。GPipe引入了一种新型的流水线并行，它在应用整个小批次的单个同步梯度更新之前，对微批次的执行进行流水线化。我们新颖的批次分割流水线并行算法，结合重计算，允许扩展到大量的微批次。这最小化了气泡开销，而无需异步梯度更新。GPipe使用户能够随所用加速器数量线性扩展模型大小。与SPMD不同，流水线并行在扩展模型时几乎不引入额外的通信开销。设备间通信仅在每个微批次的分区边界发生，引入的通信开销很小，这使得GPipe的实用性扩展到没有高速设备互连的场景。然而，GPipe目前假设单个层可以容纳在单个加速器的内存需求内。此外，微批次分割需要复杂的策略来支持那些需要在整个批次上进行计算的层（例如，BatchNorm在训练时使用微批次的统计数据，但在评估时累积小批次的统计数据）。

A4 实验环境

  • 数据集:
    • 图像分类: ImageNet 2012 数据集用于预训练。CIFAR-10, CIFAR-100, a variety of other target datasets用于迁移学习微调。
    • 机器翻译: 一个大规模多语言平行语料库，包含102种语言到英语的翻译对，总计250亿个训练样本。
  • 模型架构:
    • AmoebaNet: AmoebaNet-B(18, 512)，包含5.57亿参数。输入图像尺寸为480×480。
    • Transformer: 多种尺寸的Transformer模型，从4亿参数到60亿参数。最大的模型有128层（64个编码器层和64个解码器层），模型维度2048，前馈隐藏层维度16384，32个注意力头。词汇表大小为32k或64k，序列长度为1024。
  • 硬件配置:
    • TPU: Cloud TPUv2（每个核心8GB内存）和Cloud TPUv3（每个核心16GB内存）。
    • GPU: 多个NVIDIA P100 GPU，无NVLink高速互连，通过PCI-E连接。
  • 软件配置:
    • 框架: GPipe库在Lingvo框架【16，Lingvo: a modular and scalable framework for sequence-to-sequence modeling，2019，arXiv】下实现。
    • 优化器: 训练中使用RMSProp。

A5 实验结果

图像分类

  • 实验内容: 使用GPipe训练一个包含5.57亿参数的AmoebaNet-B(18, 512)模型，该模型在ImageNet 2012数据集上进行训练，并将输入图像尺寸扩大到480×480。网络被划分为4个分区。随后，使用这个预训练模型在多个其他图像分类数据集上进行迁移学习微调。
  • 实验结果:
    • 在ImageNet 2012验证集上，该模型单次裁剪（single-crop）的Top-1准确率达到84.4%，Top-5准确率达到97%（见表4）。
    • 在迁移学习任务中，该模型在所有目标数据集上都取得了有竞争力的结果。例如，在CIFAR-10上的错误率降至1.0%，在CIFAR-100上的错误率降至8.7%（见表4）。
  • 分析结论: 实验结果证明了GPipe能够成功训练巨大的卷积神经网络。同时，这些结果也证实了Kornblith等人【25，Do better imagenet models transfer better?，2018，CoRR】的发现，即在ImageNet上表现更好的模型在迁移学习任务中也表现更佳。


表 4: 使用AmoebaNet-B (18, 512) 先在ImageNet 2012上训练，然后在其他数据集上微调的图像分类准确率。微调结果是5次运行的平均值。基线结果来自Real等人[12]和Cubuk等人[26]，他们是从头开始直接训练的。*Mahajan等人的模型[27]达到了85.4%的top-1准确率，但它是在非公开的Instagram数据上预训练的。Ngiam等人[28]通过在私有数据集（JFT-300M）上进行预训练取得了更好的结果。

大规模多语言机器翻译

  • 实验内容: 在一个包含102种语言到英语的大规模语料库上，训练一个单一的、巨大的多语言Transformer模型，并与专门训练的双语基线模型进行比较。通过增加模型的深度（层数）和宽度（隐藏层维度、注意力头数）来扩展模型，最大的模型达到60亿参数、128层。
  • 实验结果:
    • 模型容量与性能: 如图3所示，将模型容量从4亿参数增加到13亿，再到60亿，所有语言的翻译质量都得到了显著提升。对于高资源语言的提升尤为明显，而低资源语言也从多语言训练中获得了巨大的迁移学习增益。
    • 深度与宽度的权衡: 对比同样是13亿参数的深模型 T(24, 8192, 16) 和宽模型 T(12, 16384, 32)，发现虽然它们在高资源语言上表现相似，但深模型在低资源语言上表现明显更优，这表明增加模型深度可能更有利于泛化和知识迁移。
    • 深度模型的训练挑战: 在训练深度模型时遇到了严重的训练不稳定性问题，表现为激活值变得尖锐（正峰度）和对噪声敏感，导致梯度爆炸。通过两种方法解决了这个问题：(1) 按照Zhang等人【38，Fixup initialization: Residual learning without normalization，2019，arXiv】的方法，按层数比例缩小前馈层的初始化；(2) 对logits（softmax的输入）进行裁剪，限制其最大值。
    • 大批量训练: 测试了将批量大小从26万个token增加到400万个token的效果。如表5所示，随着批量大小的增加，验证损失和BLEU分数都得到了显著改善。
  • 分析结论: GPipe的灵活性使其能够成功扩展Transformer模型。实验证明，一个足够大的单一多语言NMT模型可以同时学习超过100个语言对的映射，并且在所有语言上都优于双语模型。这凸显了拥有高效灵活的模型并行工具的重要性。


图 3: 随着多语言模型容量的增加，所有语言的翻译质量变化。语言按训练数据集大小从左到右降序排列。T (L, H, A)表示一个Transformer模型，具有L个编码器和L个解码器层，前馈隐藏维度为H，注意力头数为A。我们注意到，将模型容量从4亿参数（T(6, 8192, 16)）增加到13亿（T(24, 8192, 16)），再到60亿（T(64, 16384, 32)），所有语言的质量都显著提高。我们还注意到，与双语基线相比，低资源语言（图右侧）的质量有巨大提升，突显了训练多语言模型带来的显著迁移增益。


表 5: 批量大小的影响

A5 结论

本文介绍了GPipe，一个用于训练巨型神经网络的可扩展模型并行库。我们提出并实现了一种新颖的批次分割流水线并行算法，该算法使用同步梯度更新，从而在实现高硬件利用率的同时保证了训练的稳定性。我们利用GPipe训练了大规模的卷积模型和基于Transformer的模型，并在图像分类和多语言机器翻译任务上展示了强大的实证结果。GPipe库具有三个关键特性：1) 高效性：通过新颖的批次分割流水线算法，GPipe实现了几乎与设备数量成线性的加速比；2) 灵活性：GPipe支持任何可以表示为层序列的深度网络；3) 可靠性：GPipe采用同步梯度下降，确保无论分区数量多少，训练结果都保持一致。