Memory-Efficient Pipeline-Parallel DNN Training

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$\times$ with similar final model accuracy.

Accepted to ICML 2021

References in corpus (5)

• Language Models are Few-Shot Learners
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• PipeMare: Asynchronous Pipeline Parallel DNN Training

Cited by in corpus (11)

• Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines
• ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
• Near-Optimal Sparse Allreduce for Distributed Deep Learning
• HongTu: Scalable Full-Graph GNN Training on Multiple GPUs (via communication-optimized CPU data offloading)
• Efficient Pipeline Planning for Expedited Distributed DNN Training
• Reducing Energy Bloat in Large Model Training
• Automatic Graph Partitioning for Very Large-scale Deep Learning
• Tensor Relational Algebra for Machine Learning System Design
• Varuna: Scalable, Low-cost Training of Massive Deep Learning Models
• Scheduling Optimization Techniques for Neural Network Training
• DistIR: An Intermediate Representation and Simulator for Efficient Neural Network Distribution