摘要

扩大深度神经网络的容量已被证明是提高多种不同机器学习任务模型质量的有效方法。在许多情况下，当模型容量超过单个加速器的内存限制时，需要开发特殊的算法或基础设施。这些解决方案通常具有架构特定性，无法迁移到其他任务中。为了解决高效且与任务无关的模型并行需求，我们引入了GPipe，这是一种管道并行库，允许对任何可以表示为层序列的网络进行扩展。通过将不同的层子序列分配到单独的加速器上，GPipe 提供了灵活地高效扩展各种不同网络至巨大规模的能力。此外，GPipe 还利用了一种新颖的批量分割管道算法，在模型跨多个加速器分区时几乎实现了线性的加速效果。我们通过在两个具有不同网络架构的任务上训练大规模神经网络来展示 GPipe 的优势：(i) 图像分类：我们在 ImageNet-2012 数据集上训练了一个包含 5.57 亿参数的 AmoebaNet 模型，达到了 84.4% 的 top-1 准确率；(ii) 多语言神经机器翻译：我们在涵盖超过 100 种语言的语料库上训练了一个包含 60 亿参数、128 层的 Transformer 模型，并且其质量优于所有双语模型。