[论文解读] GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

Yanping Huang, Youlong Cheng+9 more|arXiv (Cornell University)|Nov 16, 2018

Digital Imaging for Blood Diseases被引用 881

一句话总结

GPipe 引入批量拆分流水线并行，以跨越多个加速器在同步梯度下训练巨型神经网络，实现近线性加速和可扩展的内存效率，适用于可表示为层序列的架构。

ABSTRACT

Scaling up deep neural network capacity has been known as an effective approach to improving model quality for several different machine learning tasks. In many cases, increasing model capacity beyond the memory limit of a single accelerator has required developing special algorithms or infrastructure. These solutions are often architecture-specific and do not transfer to other tasks. To address the need for efficient and task-independent model parallelism, we introduce GPipe, a pipeline parallelism library that allows scaling any network that can be expressed as a sequence of layers. By pipelining different sub-sequences of layers on separate accelerators, GPipe provides the flexibility of scaling a variety of different networks to gigantic sizes efficiently. Moreover, GPipe utilizes a novel batch-splitting pipelining algorithm, resulting in almost linear speedup when a model is partitioned across multiple accelerators. We demonstrate the advantages of GPipe by training large-scale neural networks on two different tasks with distinct network architectures: (i) Image Classification: We train a 557-million-parameter AmoebaNet model and attain a top-1 accuracy of 84.4% on ImageNet-2012, (ii) Multilingual Neural Machine Translation: We train a single 6-billion-parameter, 128-layer Transformer model on a corpus spanning over 100 languages and achieve better quality than all bilingual models.

研究动机与目标

• 使比单个加速器内存更大的网络的可扩展训练成为可能。
• 提供一个灵活的、与任务和架构无关的模型并行库。
• 开发一个具有同步梯度更新、用于稳定性的批量拆分流水线算法。
• 通过再材料化减少激活内存以容纳巨型模型。
• 在图像分类和多语言机器翻译任务上展示可扩展性。

提出的方法

• 将由 L 层组成的串行网络划分为 K 个单元，每个单元放置在一个独立的加速器上。
• 将大小为 N 的每个小批分成 M 个微批，并在 K 台加速器上进行流水线处理。
• 在更新参数之前，使用跨越 M 个微批的梯度累积的同步小批量梯度下降。
• 在反向传播期间应用再材料化以重新计算前向激活，峰值内存降至 O(N + (L/K) * (N/M))。
• 最小化估计单元成本的方差，以平衡流水线效率并降低气泡开销 (K-1)/(M+K-1)。
• 可选地将 GPipe 与数据并行结合以实现进一步扩展。

实验结果

研究问题

• RQ1 GPipe 是否能够扩展到表达为层序列的任意深度网络，且超出单个加速器的内存限制？
• RQ2 微批拆分（M）和分区计数（K）对训练效率、内存和可扩展性的影响是什么？
• RQ3 再材料化如何影响巨型模型的内存使用和训练稳定性？
• RQ4 在卷积网络（AmoebaNet）和基于 transformer 的（multilingual Transformer）架构上可实现的性能与吞吐提升是多少？
• RQ5 层分布不均衡的网络架构在 GPipe 下如何影响加速比？

主要发现

• GPipe 使在 8 个加速器上训练高达 1.8B AmoebaNet 参数成为可能，相较于单一加速器的限制有显著提升。
• 一个 6B 参数的多语言 Transformer 可以在 128 台加速器上进行训练，在 100 种语言上实现相对于双语基线的改进。
• 当 M 至少为 4×K 时，气泡开销变得可忽略，在 Transformer 工作负载中实现近线性加速。
• GPipe 通过再材料化实现显著的内存降低，使得在给定硬件上能够容纳更大的模型。
• 实验在 ImageNet 上显示出强劲结果，使用 557M-parameter AmoebaNet，并在数据集间展现有竞争力的迁移学习。
• 在多语言机器翻译中，深度与宽度存在权衡；更深的模型在低资源语言上表现出色，而极大模型则提升高资源语言。