Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization

Abstract

We formalize the problem of trading-off DNN training time and memory requirements as the tensor rematerialization optimization problem, a generalization of prior checkpointing strategies. We introduce Checkmate, a system that solves for optimal rematerialization schedules in reasonable times (under an hour) using off-the-shelf MILP solvers or near-optimal schedules with an approximation algorithm, then uses these schedules to accelerate millions of training iterations. Our method scales to complex, realistic architectures and is hardware-aware through the use of accelerator-specific, profile-based cost models. In addition to reducing training cost, Checkmate enables real-world networks to be trained with up to 5.1x larger input sizes. Checkmate is an open-source project, available at https://github.com/parasj/checkmate.

In Proceedings of 3rd Conference Machine Learning and Systems 2020 (MLSys 2020)

References in corpus (13)

• Very Deep Convolutional Networks for Large-Scale Image Recognition
• Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems
• Large Scale GAN Training for High Fidelity Natural Image Synthesis
• An Analysis of Deep Neural Network Models for Practical Applications
• DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs
• Training Deep Nets with Sublinear Memory Cost
• Generating Long Sequences with Sparse Transformers
• HDNET: Exploiting HD Maps for 3D Object Detection
• The Reversible Residual Network: Backpropagation Without Storing Activations
• An Empirical Model of Large-Batch Training
• The OoO VLIW JIT Compiler for GPU Inference

Cited by in corpus (14)

• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
• On-device Training: A First Overview on Existing Systems
• PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
• TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models
• Stateful Large Language Model Serving with Pensieve
• ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training
• Memory Optimization for Deep Networks
• Mist: Efficient Distributed Training of Large Language Models via Memory-Parallelism Co-Optimization
• The CoRa Tensor Compiler: Compilation for Ragged Tensors with Minimal Padding
• High-Capacity Expert Binary Networks
• Varuna: Scalable, Low-cost Training of Massive Deep Learning Models
• Hydra: A System for Large Multi-Model Deep Learning
• Sequential Aggregation and Rematerialization: Distributed Full-batch Training of Graph Neural Networks on Large Graphs