介绍 GPipe，一个用于高效训练大规模神经网络模型的开源库

lixia01

深度神经网络(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 TPUv2s上的可用内存。因此，迫切需要一种高效、可扩展的基础设施，以实现大规模深度学习并克服当前加速器的内存限制。
在“ GPipe：使用流水线并行高效训练巨型神经网络”中，我们展示了如何使用流水线并行来扩展 DNN 训练以克服这一限制。GPipe 是一个分布式机器学习库，它使用同步随机梯度下降和流水线并行进行训练，适用于任何由多个连续层组成的 DNN。重要的是，GPipe 允许研究人员轻松部署更多加速器来训练更大的模型，并在不调整超参数的情况下扩展性能。为了证明 GPipe 的有效性，我们在 Google Cloud TPUv2s 上训练了一个具有 5.57 亿个模型参数和 480 x 480 输入图像大小的AmoebaNet -B。该模型在多个流行数据集上表现良好，包括将单裁剪 ImageNet 准确率提高到 84.3%，将CIFAR -10 准确率提高到 99%，将 CIFAR-100 准确率提高到 91.3%。核心GPipe 库已在Lingvo框架下开源。
从小批量到微批量
有两种标准方法可以加速中等规模的 DNN 模型。数据并行方法使用更多机器并在它们之间分割输入数据。另一种方法是将模型移动到加速器上，例如 GPU 或 TPU，它们具有特殊的硬件来加速模型训练。但是，加速器的内存有限，与主机的通信带宽也有限。因此，需要模型并行来在加速器上训练更大的 DNN 模型，方法是将模型划分为多个分区，并将不同的分区分配给不同的加速器。但由于 DNN 的顺序性，这种简单的策略可能导致在计算过程中只有一个加速器处于活动状态，从而严重未充分利用加速器的计算能力。另一方面，标准数据并行方法允许在多个加速器上使用不同的输入数据同时训练同一模型，但不能增加加速器可以支持的最大模型大小。
为了实现跨多个加速器的有效训练，GPipe 将模型划分到不同的加速器上，并自动将小批量的训练示例分割成更小的微批量。通过跨微批次流水线执行，加速器可以并行运行。此外，梯度会在微批次之间持续累积，因此分区数量不会影响模型质量。
最大化内存和效率
GPipe 最大化了模型参数的内存分配。我们在 Cloud 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 GB）。当将模型分布在所有 16 个加速器上时，GPipe 在 1024 个 token 句子上启用了 80 亿参数Transformer语言模型，速度提高了 11 倍。
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 团队。