[Paper Review] 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 Diseases881 citations

TL;DR

GPipe introduces batch-splitting pipeline parallelism to train giant neural networks across multiple accelerators with synchronous gradients, enabling near-linear speedups and scalable memory efficiency for architectures expressible as layer sequences.

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.

Motivation & Objective

• Enable scalable training of networks larger than a single accelerator's memory.
• Provide a flexible, task- and architecture-agnostic model-parallelism library.
• Develop a batch-splitting pipeline algorithm with synchronous gradient updates for stability.
• Reduce activation memory with re-materialization to fit giant models.
• Demonstrate scalability on image classification and multilingual machine translation tasks.

Proposed method

• Partition a sequential network of L layers into K cells, each placed on a separate accelerator.
• Split each mini-batch of size N into M micro-batches and pipeline them across K accelerators.
• Use synchronous mini-batch gradient descent with gradient accumulation across M micro-batches before updating parameters.
• Apply re-materialization to recompute forward activations during backward pass, reducing peak memory to O(N + (L/K) * (N/M)).
• Minimize variance in estimated cell costs to balance pipeline efficiency and reduce bubble overhead (K-1)/(M+K-1).
• Optionally combine GPipe with data parallelism for further scaling.

Experimental results

Research questions

• RQ1 Can GPipe scale arbitrary deep networks that are expressible as sequences of layers beyond a single accelerator's memory limit?
• RQ2 What is the impact of micro-batch splitting (M) and partition count (K) on training efficiency, memory, and scalability?
• RQ3 How does re-materialization affect memory usage and training stability in giant models?
• RQ4 What performance and throughput gains can be achieved on convolutional (AmoebaNet) and transformer-based (multilingual Transformer) architectures?
• RQ5 How do network architectures with imbalanced layer distributions affect speedups under GPipe?

Key findings

• GPipe enables training up to 1.8B AmoebaNet parameters across 8 accelerators, a substantial increase over single-accelerator limits.
• A 6B-parameter multilingual Transformer can be trained across 128 accelerators, achieving improvements over bilingual baselines across 100 languages.
• With M at least 4×K, bubble overhead becomes negligible, yielding near-linear speedups in Transformer workloads.
• GPipe achieves substantial memory reductions via re-materialization, enabling far larger models on given hardware.
• Experiments show strong ImageNet results with 557M-parameter AmoebaNet and competitive transfer learning across datasets.
• In multilingual MT, depth–width trade-offs exist; deeper models excel on low-resource languages, while very large models improve high-resource languages.