GPipe 简介：用于高效训练大型神经网络模型的开源库

文 / Yanping Huang，Google AI 软件工程师

深度神经网络 (DNN) 已推动诸多机器学习任务取得进步，其中包括语音识别、视觉识别和语言处理。BigGan、Bert 以及 GPT2.0 取得的最新进展表明，随着 DNN 模型规模逐渐增大，任务处理的表现也逐渐提升。此外，视觉识别任务过往取得的进步也表明模型大小和分类准确度之间存在很强的关联性。例如，2014 年 ImageNet 视觉识别挑战赛的获胜者 GoogleNet 使用 400 万参数实现 74.8% 的 top-1 准确率，而短短三年之后，2017 年 ImageNet 挑战赛的获胜者 Squeeze-and-Excitation Networks 就使用 1.458 亿（超过 36 倍）参数实现 82.7% 的 top-1 准确度。然而，在此期间，GPU 内存仅扩大了 3 倍左右，而目前最先进的图像模型已达到 Cloud TPUv2 可用内存的上限。因此，我们急需一款可扩容的高效基础架构来实现大规模深度学习，并克服当前加速器上的内存限制问题。

最近开发的具有代表性的图像分类模型呈现出 ImageNet 准确率与模型大小之间的高度关联性

在  "GPipe:Efficient Training of Giant Neural Networks using Pipeline Parallelism"（使用流水线并行技术有效培训巨型神经网络）一文中，我们展示了如何使用流水线并行技术扩展 DNN 训练，从而克服这一限制。GPipe 是一个分布式机器学习库，使用同步随机 梯度下降 和流水线并行技术进行训练，适用于任何由多个序列层组成的 DNN。重要的是，GPipe 让研究人员无需调整超参数，即可轻松部署更多加速器，从而训练更大的模型并扩展性能。为了证明 GPipe 的有效性，我们在 Google Cloud TPUv2 上训练了一个 AmoebaNet-B，它拥有 5.57 亿个模型参数，输入图像大小为 480 x 480。该模型在多个广泛使用的数据集上表现良好，包括将单裁剪的 ImageNet 准确率提升到 84.3%，将 CIFAR-10 的准确率提升到 99%，以及将 CIFAR-100 的准确率提升到 91.3%。核心 GPipe 库 已在 Lingvo 框架下开源。

注：GPipe 库 链接

https://github.com/tensorflow/lingvo/blob/master/lingvo/core/gpipe.py

从小批次到微批次

有两种标准方法可以加速中等规模的 DNN 模型。数据并行方法可以利用更多的机器，并将输入数据从中分离出来。另一种方法是将模型移到诸如 GPU 或 TPU 的加速器上。这些加速器拥有能够加速模型训练的特殊硬件，只是内存有限，而且与主机的通信带宽也有限。因此，我们需要在加速器上使用模型并行技术训练规模较大的 DNN 模型，方法是对模型进行分区，然后为不同的加速器分配不同的分区。但由于 DNN 具有有序性，这种 “天真的” 策略可能导致计算期间只有一个加速器处于活跃状态，从而严重浪费加速器的计算能力。另一方面，标准的数据并行方法允许在多个加速器上同时训练拥有不同输入数据的同一个模型，但无法扩大加速器能够支持的最大模型大小。 

为了在多个加速器间实现高效训练，GPipe 在不同的加速器中对模型进行划分，然后自动将小批次训练示例拆分为更小的微批次示例。通过在微批次中采取流水线式执行操作，加速器可以并行运行。此外，微批次中的梯度不断累积，以免分区数量影响模型质量。

上图：由于网络具有有序性，“天真的” 模型并行策略导致计算能力遭到严重浪费。每次只有一个加速器处于活跃状态。下图：GPipe 将输入的小批次示例拆分成更小的微批次，使不同的加速器可以同时在单独的微批次上运行

最大化内存和效率

GPipe 会最大化模型参数的内存分配。我们在 Cloud TPUv2 上进行实验，每个 TPUv2 都有 8 个加速器核心和 64 GB 内存（每个加速器 8 GB）。如果没有 GPipe，由于内存限制，单个加速器最多可以训练 8200 万个模型参数。得益于通过反向传播和批量分割进行的重新计算，GPipe 将中间激活内存从 6.26 GB 减少到 3.46GB，从而实现了在单个加速器上训练 3.18 亿个参数的成果。我们还发现，通过采用流水线并行技术，最大模型大小与分区数量成正比，这与我们的预期相符。借助 GPipe，AmoebaNet 能够在 Cloud TPUv2 的 8 个加速器上并入 18 亿个参数，比没有 GPipe 时多出 25 倍。 

为了测试效率，我们衡量了 GPipe 对 AmoebaNet-D 模型吞吐量的影响。鉴于训练需要至少两个加速器来适应模型大小，我们只对有两个分区但没有采用流水线并行技术的“天真”案例进行了衡量。我们观察到训练过程中存在近乎线性的加速效果。与使用两个分区的“天真”方法相比，将模型分布在四倍数量的加速器上实现了 3.5 倍的加速。虽然论文中的所有实验都使用 Cloud TPUv2，但我们发现使用当前可用的 Cloud TPUv3（均拥有 16 个加速器核心和 256 GB 内存，每个加速器 16 G）会获得更好的表现。当我们将模型分配给全部 16 个加速器时，GPipe 可以让针对 1024 个令牌句子的 80 亿参数 Transformer 语言模型加速 11 倍。

使用 GPipe 对 AmoebaNet-D 进行加速。此模型不适合单个加速器。naive-2 基线表示在将模型拆分为两个分区时，采用 “天真” 分区方法的加速器性能。Pipeline-k 表示 GPipe 将模型拆分为包含 k 个加速器的 k 个分区时的性能

GPipe 还可以在不更改超参数的情况下，利用更多加速器来扩大训练规模。因此，它可以与数据并行技术结合，使用更多加速器，以一种互补的方式扩大神经网络训练的规模。 

测试准确率

我们使用 GPipe 来验证假设，即扩展现有的神经网络可以提升模型质量。我们在 ImageNet ILSVRC-2012 数据集上训练了拥有 5.57 亿个模型参数，输入图像大小为 480 x 480 的 AmoebaNet-B。我们将该网络分成 4 个分区，并同时对模型和数据进行并行训练。该巨型模型在未采用任何外部数据的情况下，实现了目前最高的 84.3% top-1 / 97% top-5 单裁剪验证准确率。大型神经网络不仅适用于 ImageNet 等数据集，还可以通过迁移学习与其他数据集相关联。实验证明，ImageNet 模型越好，迁移效果也越好。我们在 CIFAR10 和 CIFAR100 数据集上进行了迁移学习实验。我们的巨型模型将已知的最高 CIFAR-10 准确率提升到 99%，将 CIFAR-100 准确率提升到 91.3%。

结论

许多实用机器学习应用（例如自动驾驶和医学成像）之所以能够不断发展并取得成功，原因在于尽可能地实现了最高的准确度。然而，这也意味着我们需要构建更大、更复杂的模型。因此，我们很乐意将 GPipe 提供给更广泛的研究社区，并且希望其成为有助于高效训练大型 DNN 的基础架构。 

致谢

特别感谢论文的联合作者：Youlong Cheng、Dehao Che、HyoukJoong Lee、Jiquan Ngiam、Quoc V. Le 和 Zhifeng Chen。我们要感谢 Esteban Real、Alok Aggarwal、Xiaodan Song、Naveen Kumar、Mark Heffernan、Rajat Monga、Megan Kacholia、Samy Bengio 和 Jeff Dean 提供的支持和宝贵意见；感谢 Noam Shazeer、Patrick Nguyen、Xiaoqiang Zheng、Yonghui Wu、Barret Zoph、Ekin Cubuk、Jonathan Shen、Tianqi Chen 和 Vijay Vasudevan 与我们开展有益的讨论并激发我们的灵感；最后要感谢整个 Google Brain 团队。

更多 AI 相关阅读：