Training Deep Nets with Sublinear Memory Cost

Abstract

We propose a systematic approach to reduce the memory consumption of deep neural network training. Specifically, we design an algorithm that costs O(sqrt(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(log n) with as little as O(n log n) 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 with only 30 percent additional running time cost on ImageNet problems. Similarly, significant memory cost reduction is observed in training complex recurrent neural networks on very long sequences.

References in corpus (7)

• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
• Training Very Deep Networks
• Theano: new features and speed improvements
• Identity Mappings in Deep Residual Networks
• Highway Long Short-Term Memory RNNs for Distant Speech Recognition
• Joint Training of Deep Boltzmann Machines

Cited by in corpus (191)

• Learning Transferable Visual Models From Natural Language Supervision
• On the Opportunities and Risks of Foundation Models
• Longformer: The Long-Document Transformer
• ALBERT: A Lite BERT for Self-supervised Learning of Language Representations
• Linformer: Self-Attention with Linear Complexity
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• TorchIO: A Python library for efficient loading, preprocessing, augmentation and patch-based sampling of medical images in deep learning
• Generating Long Sequences with Sparse Transformers
• NVAE: A Deep Hierarchical Variational Autoencoder
• GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism
• Neural Image Compression for Gigapixel Histopathology Image Analysis
• Differentiable Programming Tensor Networks
• Demystifying Parallel and Distributed Deep Learning: An In-Depth Concurrency Analysis
• FILIP: Fine-grained Interactive Language-Image Pre-Training
• The Reversible Residual Network: Backpropagation Without Storing Activations
• SuperNeurons: Dynamic GPU Memory Management for Training Deep Neural Networks
• Tiny Machine Learning: Progress and Futures
• A Field Guide to Federated Optimization
• Database Meets Deep Learning: Challenges and Opportunities
• DeepV2D: Video to Depth with Differentiable Structure from Motion
• Supporting Very Large Models using Automatic Dataflow Graph Partitioning
• Self-supervised Pretraining of Visual Features in the Wild
• Enable Deep Learning on Mobile Devices: Methods, Systems, and Applications
• A Practical Survey on Faster and Lighter Transformers
• Swin Transformer V2: Scaling Up Capacity and Resolution
• Reinventing 2D Convolutions for 3D Images
• Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines
• Streaming convolutional neural networks for end-to-end learning with multi-megapixel images
• Review: Deep Learning in Electron Microscopy
• i-flow: High-dimensional Integration and Sampling with Normalizing Flows
• A General Approach to Seismic Inversion with Automatic Differentiation
• Parameter Hub: a Rack-Scale Parameter Server for Distributed Deep Neural Network Training
• DIET-SNN: Direct Input Encoding With Leakage and Threshold Optimization in Deep Spiking Neural Networks
• Learning Longer-term Dependencies in RNNs with Auxiliary Losses
• Memory-Efficient Implementation of DenseNets
• Bayesian Bits: Unifying Quantization and Pruning
• Evaluating Protein Transfer Learning with TAPE
• ANODE: Unconditionally Accurate Memory-Efficient Gradients for Neural ODEs
• RocketQA: An Optimized Training Approach to Dense Passage Retrieval for Open-Domain Question Answering
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• Merak: An Efficient Distributed DNN Training Framework with Automated 3D Parallelism for Giant Foundation Models
• Memory-Efficient Pipeline-Parallel DNN Training
• Supervised Learning with Projected Entangled Pair States
• ZeRO-Offload: Democratizing Billion-Scale Model Training
• ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
• End-to-end optimization of coherent optical communications over the split-step Fourier method guided by the nonlinear Fourier transform theory
• Memory-Efficient Backpropagation Through Time
• Evaluation of LLM Chatbots for OSINT-based Cyber Threat Awareness
• On-Device Machine Learning: An Algorithms and Learning Theory Perspective
• Geometric Constellation Shaping for Fiber-Optic Channels via End-to-End Learning
• TensorOpt: Exploring the Tradeoffs in Distributed DNN Training with Auto-Parallelism
• Achieving Peak Performance for Large Language Models: A Systematic Review
• GSPMD: General and Scalable Parallelization for ML Computation Graphs
• Adversarial Retriever-Ranker for dense text retrieval
• Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization
• Low-Memory Neural Network Training: A Technical Report
• PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
• Efficient Scaling of Dynamic Graph Neural Networks
• OneFlow: Redesign the Distributed Deep Learning Framework from Scratch
• Adaptive Checkpoint Adjoint Method for Gradient Estimation in Neural ODE
• InversionNet3D: Efficient and Scalable Learning for 3D Full Waveform Inversion
• TinyTL: Reduce Activations, Not Trainable Parameters for Efficient On-Device Learning
• TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models
• Hanayo: Harnessing Wave-like Pipeline Parallelism for Enhanced Large Model Training Efficiency
• Tesseract: Parallelize the Tensor Parallelism Efficiently
• torchgpipe: On-the-fly Pipeline Parallelism for Training Giant Models
• DAPPLE: A Pipelined Data Parallel Approach for Training Large Models
• Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention
• TFLMS: Large Model Support in TensorFlow by Graph Rewriting
• MegDet: A Large Mini-Batch Object Detector
• PLATO-2: Towards Building an Open-Domain Chatbot via Curriculum Learning
• Hidet: Task-Mapping Programming Paradigm for Deep Learning Tensor Programs
• PipeMare: Asynchronous Pipeline Parallel DNN Training
• Instance-aware, Context-focused, and Memory-efficient Weakly Supervised Object Detection
• Yuan 1.0: Large-Scale Pre-trained Language Model in Zero-Shot and Few-Shot Learning
• Training Large Neural Networks with Constant Memory using a New Execution Algorithm
• Learning Neural Event Functions for Ordinary Differential Equations
• Revisiting Locally Supervised Learning: an Alternative to End-to-end Training
• Pollux: Co-adaptive Cluster Scheduling for Goodput-Optimized Deep Learning
• Data Movement Is All You Need: A Case Study on Optimizing Transformers
• HongTu: Scalable Full-Graph GNN Training on Multiple GPUs (via communication-optimized CPU data offloading)
• DeepView: View Synthesis with Learned Gradient Descent
• Understanding GNN Computational Graph: A Coordinated Computation, IO, and Memory Perspective
• Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training
• Efficient Memory Management for GPU-based Deep Learning Systems
• LoCo: Local Contrastive Representation Learning
• A Graph Theoretic Framework of Recomputation Algorithms for Memory-Efficient Backpropagation
• Reducing Energy Bloat in Large Model Training
• Full-Resolution Residual Networks for Semantic Segmentation in Street Scenes
• Towards Crowdsourced Training of Large Neural Networks using Decentralized Mixture-of-Experts
• PLATO-XL: Exploring the Large-scale Pre-training of Dialogue Generation
• Backprop with Approximate Activations for Memory-efficient Network Training
• Blockwise Self-Attention for Long Document Understanding
• Efficient Algorithms for Device Placement of DNN Graph Operators
• SimpleDet: A Simple and Versatile Distributed Framework for Object Detection and Instance Recognition
• Large Scale Language Modeling: Converging on 40GB of Text in Four Hours
• Searching for Winograd-aware Quantized Networks
• Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory
• Improved modularity and new features in ipie: Toward even larger AFQMC calculations on CPUs and GPUs at zero and finite temperatures
• p-Meta: Towards On-device Deep Model Adaptation
• FedDCT: Federated Learning of Large Convolutional Neural Networks on Resource Constrained Devices using Divide and Collaborative Training
• GANBERT: Generative Adversarial Networks with Bidirectional Encoder Representations from Transformers for MRI to PET synthesis
• AutoDDL: Automatic Distributed Deep Learning with Near-Optimal Bandwidth Cost
• An Oracle for Guiding Large-Scale Model/Hybrid Parallel Training of Convolutional Neural Networks
• Differentiable Simulation of Soft Multi-body Systems
• ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training
• Memory Optimization for Deep Networks
• MPipeMoE: Memory Efficient MoE for Pre-trained Models with Adaptive Pipeline Parallelism
• Full deep neural network training on a pruned weight budget
• The Case for Strong Scaling in Deep Learning: Training Large 3D CNNs with Hybrid Parallelism
• Convolutional Networks with Dense Connectivity
• Memory Efficient Model Based Deep Learning Reconstructions for High Spatial Resolution 3D Non-Cartesian Acquisitions
• Randomized Automatic Differentiation
• Making EfficientNet More Efficient: Exploring Batch-Independent Normalization, Group Convolutions and Reduced Resolution Training
• SpecNet: Spectral Domain Convolutional Neural Network
• Amazon SageMaker Model Parallelism: A General and Flexible Framework for Large Model Training
• word2ket: Space-efficient Word Embeddings inspired by Quantum Entanglement
• Fast and Accurate 3D Medical Image Segmentation with Data-swapping Method
• Multi-Stage Pathological Image Classification using Semantic Segmentation
• Improving Automated COVID-19 Grading with Convolutional Neural Networks in Computed Tomography Scans: An Ablation Study
• Towards Practical Few-shot Federated NLP
• DUAL-GLOW: Conditional Flow-Based Generative Model for Modality Transfer
• Deeper Image Quality Transfer: Training Low-Memory Neural Networks for 3D Images
• Aggregated Momentum: Stability Through Passive Damping
• MRI Image Reconstruction via Learning Optimization Using Neural ODEs
• LAMP: Large Deep Nets with Automated Model Parallelism for Image Segmentation
• On the Variance of Unbiased Online Recurrent Optimization
• Echo: Compiler-based GPU Memory Footprint Reduction for LSTM RNN Training
• Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup
• Efficient Memory Management for Deep Neural Net Inference
• UFO-BLO: Unbiased First-Order Bilevel Optimization
• Compressing Gradient Optimizers via Count-Sketches
• Sub-Linear Memory: How to Make Performers SLiM
• NoSENSE: Learned unrolled cardiac MRI reconstruction without explicit sensitivity maps
• 8-bit Optimizers via Block-wise Quantization
• Robustness of Graph Neural Networks at Scale
• Distributed Training and Optimization Of Neural Networks
• Invertible Residual Network with Regularization for Effective Medical Image Segmentation
• Adding Recurrence to Pretrained Transformers for Improved Efficiency and Context Size
• Compressed Deep Networks: Goodbye SVD, Hello Robust Low-Rank Approximation
• Training Optimus Prime, M.D.: Generating Medical Certification Items by Fine-Tuning OpenAI's gpt2 Transformer Model
• Automatic Graph Partitioning for Very Large-scale Deep Learning
• Daydream: Accurately Estimating the Efficacy of Optimizations for DNN Training
• COMET: A Novel Memory-Efficient Deep Learning Training Framework by Using Error-Bounded Lossy Compression
• XEngine: Optimal Tensor Rematerialization for Neural Networks in Heterogeneous Environments
• OneAdapt: Fast Configuration Adaptation for Video Analytics Applications via Backpropagation
• Improving the Resolution of CNN Feature Maps Efficiently with Multisampling
• Layered gradient accumulation and modular pipeline parallelism: fast and efficient training of large language models
• A Multi-site Study of a Breast Density Deep Learning Model for Full-field Digital Mammography Images and Synthetic Mammography Images
• VQ-DRAW: A Sequential Discrete VAE
• Optimal Gradient Checkpoint Search for Arbitrary Computation Graphs
• More Information Supervised Probabilistic Deep Face Embedding Learning
• Varuna: Scalable, Low-cost Training of Massive Deep Learning Models
• Jacobian Determinant of Normalizing Flows
• SANM: A Symbolic Asymptotic Numerical Solver with Applications in Mesh Deformation
• Analysis of memory consumption by neural networks based on hyperparameters
• Dynamic Tensor Rematerialization
• SMYRF: Efficient Attention using Asymmetric Clustering
• EDEN: Enabling Energy-Efficient, High-Performance Deep Neural Network Inference Using Approximate DRAM
• Memory Efficient Meta-Learning with Large Images
• Gradients are Not All You Need
• DistIR: An Intermediate Representation and Simulator for Efficient Neural Network Distribution
• Advancing STT for Low-Resource Real-World Speech
• Differentiate Everything with a Reversible Embeded Domain-Specific Language
• Training DNNs in O(1) memory with MEM-DFA using Random Matrices
• Beyond Nyströmformer -- Approximation of self-attention by Spectral Shifting
• FetusMap: Fetal Pose Estimation in 3D Ultrasound
• Hybrid Attention for Automatic Segmentation of Whole Fetal Head in Prenatal Ultrasound Volumes
• Beyond the Memory Wall: A Case for Memory-centric HPC System for Deep Learning
• BPPSA: Scaling Back-propagation by Parallel Scan Algorithm
• Going Deeper with Lean Point Networks
• Scaling Distributed Deep Learning Workloads beyond the Memory Capacity with KARMA
• Learning Energy Based Inpainting for Optical Flow
• Enabling Binary Neural Network Training on the Edge
• Backdrop: Stochastic Backpropagation
• AccUDNN: A GPU Memory Efficient Accelerator for Training Ultra-deep Neural Networks
• SCRAM: Spatially Coherent Randomized Attention Maps
• Training on the Edge: The why and the how
• Learning Directly from Grammar Compressed Text
• VideoClick: Video Object Segmentation with a Single Click
• Doc2Dict: Information Extraction as Text Generation
• Hydra: A System for Large Multi-Model Deep Learning
• Backpropagation for long sequences: beyond memory constraints with constant overheads
• Physics-Aware Downsampling with Deep Learning for Scalable Flood Modeling
• A Survey on Green Deep Learning
• Sequential Aggregation and Rematerialization: Distributed Full-batch Training of Graph Neural Networks on Large Graphs
• Decoupling Hierarchical Recurrent Neural Networks With Locally Computable Losses
• A Dataset of Information-Seeking Questions and Answers Anchored in Research Papers
• Functional Design of Computation Graph
• Ouroboros: On Accelerating Training of Transformer-Based Language Models
• Back-Projection Pipeline