Memory Optimization for Deep Networks
arXiv:2010.14501
Abstract
Deep learning is slowly, but steadily, hitting a memory bottleneck. While the tensor computation in top-of-the-line GPUs increased by 32x over the last five years, the total available memory only grew by 2.5x. This prevents researchers from exploring larger architectures, as training large networks requires more memory for storing intermediate outputs. In this paper, we present MONeT, an automatic framework that minimizes both the memory footprint and computational overhead of deep networks. MONeT jointly optimizes the checkpointing schedule and the implementation of various operators. MONeT is able to outperform all prior hand-tuned operations as well as automated checkpointing. MONeT reduces the overall memory requirement by 3x for various PyTorch models, with a 9-16% overhead in computation. For the same computation cost, MONeT requires 1.2-1.8x less memory than current state-of-the-art automated checkpointing frameworks. Our code is available at https://github.com/utsaslab/MONeT.
arXiv comment18 pages, ICLR'21
References in corpus (12)
- PyTorch: An Imperative Style, High-Performance Deep Learning Library
- TensorFlow: A system for large-scale machine learning
- MobileNetV2: Inverted Residuals and Linear Bottlenecks
- cuDNN: Efficient Primitives for Deep Learning
- Mixed Precision Training
- Training Deep Nets with Sublinear Memory Cost
- Generating Long Sequences with Sparse Transformers
- Scalable Methods for 8-bit Training of Neural Networks
- The Reversible Residual Network: Backpropagation Without Storing Activations
- Memory-Efficient Implementation of DenseNets
- Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization
- Backprop with Approximate Activations for Memory-efficient Network Training