Memory-Efficient Pipeline-Parallel DNN Training

By D. Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, Matei A. Zaharia. 2020

299 citations recorded by the source corpus.

Many state-of-the-art results in domains such as NLP and computer vision have been obtained by scaling up the number of parameters in existing models. However, the weight parameters and intermediate outputs of these large models often do not fit in the main memory of a single accelerator device; this means that it is necessary to use multiple accelerators to train large models, which is challenging to do in a time-efficient way. In this work, we propose PipeDream-2BW, a system that performs memory-efficient pipeline parallelism, a hybrid form of parallelism that combines data and model parallelism with input pipelining. Our system 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 being cognizant of constraints such as compute capabilities, memory capacities, and interconnect topologies, and determines when to employ existing memory-savings techniques, such as activation recomputation, that trade off extra computation for lower memory footprint. PipeDream-2BW is able to accelerate the training of large language models with up to 2.5 billion parameters by up to 6.9x compared to optimized baselines.

Subfields: Artificial Intelligence, Hardware and Architecture, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Hardware and Architecture

In Efficiency & Compression

Citations in the corpus

PaLM: Scaling Language Modeling with Pathways (2022, methodology): “Many other works aim to increase of the scale of models, while limiting communication overheads (Rajbhandari et al., 2020; Lepikhin et al., 2020; Li et al., 2020; Rasley et al., 2020; Rajbhandari et al., 2021; Ren et al., 2021; Narayanan et al., 2021a).”
ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep learning (2021, methodology): “…the growth in model size has been made possible mainly through advances in system technology for training large DL models, with parallelism technologies such as model parallelism [7], pipeline parallelism [8–10], and ZeRO [11, 12] creating a path to training larger and more powerful models.”
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021, methodology): “Previous work like PipeDream-2BW [23] has looked at the performance impact of using activation recomputation.”
Reducing Activation Recomputation in Large Transformer Models (2022, methodology): “To quantify this, we consider the 1F1B pipeline schedule developed in PipeDream [12].”
MegaBlocks: Efficient Sparse Training with Mixture-of-Experts (2022, methodology): “We use gradient accumulation for all models and train with a batch size of 512 sequences and the largest micro batch size that does not run out of memory (Narayanan et al., 2021a).”
GLM-130B: An Open Bilingual Pre-trained Model (2022, methodology): “…the model into sequential stages for each parallel group, and to further minimize bubbles introduced by pipeline, we leverage the PipeDream-Flush (Narayanan et al., 2021) implementation from DeepSpeed (Rasley et al., 2020) to train GLM-130B with a relative big global batch size (4,224) to reduce…”
DAPPLE: a pipelined data parallel approach for training large models (2020, background): “Some researchers have been seeking for the optimal placement strategy to assign operations in a DNN to different devices[21, 35, 37, 46] to further improve system efficiency.”
ERNIE 3.0 Titan: Exploring Larger-scale Knowledge Enhanced Pre-training for Language Understanding and Generation (2021, methodology): “The 4D hybrid parallelism includes data parallelism (DP) [60, 61], intra-layer tensor model parallelism (MP) [62, 57], inter-layer pipeline model parallelism (PP) [32, 63, 64, 65, 33], and the improved sharded data parallelism (Sharded) based on ZeRO [58].”
Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models (2022, methodology): “Pipeline model parallelism [24, 29, 33, 34, 53] is another technique that have been proposed for large deep learning models to address the scaling challenges.”
Ring Attention with Blockwise Transformers for Near-Infinite Context (2023, background): “Another avenue of research explores various parallelism methods, including data parallelism [10], tensor parallelism [35], pipeline parallelism [28, 15, 29], sequence parallelism [22, 19, 17], and FSDP [11, 32].”
MiCS: Near-linear Scaling for Training Gigantic Model on Public Cloud (2022, unknown): “PipeMare [42], PipeDream [17], PipeDream-2BW [43] use asynchronous and bounded-staleness for efficient resouce utilization but affect convergence quality.”
Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training (2021, methodology): “Megatron- LM (Shoeybi et al., 2019) created a tensor-parallel implementation of GPT-2 and T5 based on PyTorch, and added pipeline parallelism in later work (Narayanan et al., 2021b).”
DistIR: An Intermediate Representation for Optimizing Distributed Neural Networks (2021, methodology): “Many recent works propose algorithms for automatic distribution, such as MCMC search [17], integer and dynamic programming [24, 33], reinforcement learning [23, 34], and custom algorithms [16, 25].”