Training Deep Nets with Sublinear Memory Cost

By Tianqi Chen, Bing Xu, Chiyuan Zhang, Carlos Guestrin. 2016

1,506 citations recorded by the source corpus.

We propose a systematic approach to reduce the memory consumption of deep neural network training. Specifically, we design an algorithm that costs O( √ n) memory to train a n layer network, with only the computational cost of an extra forward pass per mini-batch. As many of the state-of-the-art models hit the upper bound of the GPU memory, our algorithm allows deeper and more complex models to be explored, and helps advance the innovations in deep learning research. We focus on reducing the memory cost to store the intermediate feature maps and gradients during training. Computation graph analysis is used for automatic in-place operation and memory sharing optimizations. We show that it is possible to trade computation for memory giving a more memory efficient training algorithm with a little extra computation cost. In the extreme case, our analysis also shows that the memory consumption can be reduced to O(logn) with as little as O(n logn) extra cost for forward computation. Our experiments show that we can reduce the memory cost of a 1,000-layer deep residual network from 48G to 7G on ImageNet problems. Similarly, significant memory cost reduction is observed in training complex recurrent neural networks on very long sequences.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, 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 Computer Vision and Pattern Recognition

In Efficiency & Compression

Citations in the corpus

Learning Transferable Visual Models From Natural Language Supervision (2021, methodology): “To save additional memory, gradient checkpointing (Griewank & Walther, 2000; Chen et al., 2016), half-precision Adam statistics (Dhariwal et al.”
ALBERT: A Lite BERT for Self-supervised Learning of Language Representations (2019, methodology): “Existing solutions to the aforementioned problems include model parallelization (Shoeybi et al., 2019) and clever memory management (Chen et al., 2016; Gomez et al., 2017).”
Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (2019, methodology): “As these models become larger, they exceed the memory limit of modern processors, and require additional memory management techniques such as activation checkpointing (Chen et al., 2016).”
Generating Long Sequences with Sparse Transformers (2019, background): “Gradient checkpointing has been shown to be effective in reducing the memory requirements of training deep neural networks (Chen et al., 2016), (Gruslys et al., 2016).”
Efficient Large Scale Language Modeling with Mixtures of Experts (2021, methodology): “We adopt a popular technique called activation checkpointing, which saves memory during training by discarding a subset of activations in the forward pass and recomputing them in the backward pass (Chen et al., 2016).”
ZeRO: Memory optimizations Toward Training Trillion Parameter Models (2019, background): “Activation checkpointing (or activation recomputation) is a common approach to reduce the activation memory by approximately the square root of the total activations at the expense of 33% re-computation overhead [8].”
GSPMD: General and Scalable Parallelization for ML Computation Graphs (2021, methodology): “We follow GPipe’s approach [15], where a key difference from non-pipelined configurations is that we recompute forward pass intermediate results during the backward pass, a technique known as rematerialization [10] to lower peak”
ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep learning (2021, background): “On the positive side, the memory required for activation can be significantly reduced via activation checkpointing [29], which trades off activation memory at the expense of 0.”
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (2021, background): “Activation recomputation [12, 18, 20, 21] is an optional technique that trades off an increase in the number of compute operations performed for additional memory footprint, by running the forward pass a second time just before the backward pass (and stashing only the input activations for a given p”
ZeRO-Offload: Democratizing Billion-Scale Model Training (2021, background): “The first approach trades computation for memory saving from activations (residual memory) by recomputing from checkpoints [4].”
Big Bird: Transformers for Longer Sequences (2020, background): “Use of gradient checkpointing [15] could potentially allow for handling even longer sequences.”
Memory-Efficient Pipeline-Parallel DNN Training (2020, methodology): “PipeDream-2BW’s planner determines the size of each model partition, per-GPU microbatch size, and whether to use memory-saving optimizations like activation recomputation [5,9] while taking into account both the throughput and memory footprint ramifications of these decisions.”
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, background): “[9] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin.”
Reducing Activation Recomputation in Large Transformer Models (2022, background): “Activation recomputation [5] overcomes this memory limitation by storing (or ”checkpointing”) the input activations to a group of layers and recomputing other required activations using an extra forward pass during back-propagation (this is referred to in this paper as full activation…”
Self-attention Does Not Need $O(n^2)$ Memory (2021, methodology): “So we apply checkpointing (Chen et al., 2016) in line 11 to the function that summarizes the individual chunks.”
Efficient Memory Management for Large Language Model Serving with PagedAttention (2023, methodology): “Swapping [23, 42, 55], recomputation [7, 24] and their combination [40] have been utilized to reduce the peak memory of training.”
Longformer: The Long-Document Transformer (2020, methodology): “5 We used gradient checkpointing (Chen et al., 2016) to reduce memory usage, and ran our experiments on 48GB RTX8000 GPUs.”
R2-D2: A Modular Baseline for Open-Domain Question Answering (2021, background): “However, such approach requires tremendous amount of GPU memory, and thus requires employing speed-memory trade-offs such as gradient checkpointing (Chen et al., 2016).”
VirTex: Learning Visual Representations from Textual Annotations (2020, background): “We compensate this by reducing memory usage through gradient checkpointing [102, 103].”