[Paper Review] Memory-Efficient Pipeline-Parallel DNN Training

Deepak Narayanan, Amar Phanishayee+3 more|arXiv (Cornell University)|Jun 16, 2020

Topic ModelingComputer Science22 references62 citations

TL;DR

PipeDream-2BW introduces memory-efficient pipeline parallelism with double-buffered weight updates and automatic partitioning, enabling faster training of huge transformer models with similar accuracy.

ABSTRACT

Many state-of-the-art ML results have been obtained by scaling up the number of parameters in existing models. However, parameters and activations for such large models often do not fit in the memory of a single accelerator device; this means that it is necessary to distribute training of large models over multiple accelerators. In this work, we propose PipeDream-2BW, a system that supports memory-efficient pipeline parallelism. PipeDream-2BW uses a novel pipelining and weight gradient coalescing strategy, combined with the double buffering of weights, to ensure high throughput, low memory footprint, and weight update semantics similar to data parallelism. In addition, PipeDream-2BW automatically partitions the model over the available hardware resources, while respecting hardware constraints such as memory capacities of accelerators and interconnect topologies. PipeDream-2BW can accelerate the training of large GPT and BERT language models by up to 20$ imes$ with similar final model accuracy.

Motivation & Objective

• Motivate the challenge of training billion-parameter DNNs under accelerator memory limits.
• Propose a memory-efficient pipeline-parallel system to improve throughput without flush-induced stalls.
• Automatically partition models across heterogeneous interconnects while honoring device memory constraints.
• Demonstrate convergence and throughput benefits on large transformer models (GPT/BERT) compared to baselines.

Proposed method

• Introduce 2BW (double-buffered weight updates) to reduce memory footprint by maintaining two weight versions and accumulating gradients over microbatches.
• Use a planning algorithm to partition DNN operators into equally replicated stages (equi-replicated) and assign microbatches and gradient accumulation to balance throughput and memory.
• Provide a variant PipeDream-Flush with single weight version and pipeline flushes for even lower memory footprint at some throughput cost.
• Formalize weight update semantics with a 1-step delay: W(t+1)=W(t)−η·∇f(W(t−1)) and discuss convergence similarity to vanilla updates.
• Incorporate activation recomputation to reduce activation storage and increase feasible per-GPU microbatch sizes.
• Evaluate on large GPT/BERT models using PyTorch on AWS with NVLink, comparing against Megatron-style tensor model parallelism and GPipe.

Experimental results

Research questions

• RQ1 How can memory and throughput be balanced in pipeline-parallel training for billion-parameter models?
• RQ2 Can a double-buffered weight update scheme achieve high throughput without pipeline flushes and with manageable memory Footprint?
• RQ3 What automatic partitioning strategies best map transformer-like architectures across heterogeneous interconnects while preserving convergence?
• RQ4 Do 2BW and PipeDream-Flush semantics preserve convergence/accuracy comparable to vanilla optimizers on large NLP models?

Key findings

• End-to-end speedups of 1.3x to 20x over optimized model-parallel baselines for various GPT models.
• Up to 3.2x faster than GPipe due to absence of pipeline flushes and lower memory footprint.
• Can train large transformer models up to almost 30 billion parameters using 64 eight- or sixteen-GPU setups, outperforming Megatron-style configurations in memory-limited scenarios.
• Adam optimizer with 2BW shows similar training loss trajectories and downstream finetuning accuracy to vanilla Adam across BERT and GPT pretraining tasks.
• PipeDream-Flush offers lower memory footprint at the cost of throughput, while 2BW achieves higher throughput with two weight versions and activation recomputation.
• Memory footprint comparisons show GPipe memory challenges, PipeDream-Flush moderate footprint, and PipeDream-2BW intermediate footprint.