ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training
arXiv:2104.14129
Abstract
The increasing size of neural network models has been critical for improvements in their accuracy, but device memory is not growing at the same rate. This creates fundamental challenges for training neural networks within limited memory environments. In this work, we propose ActNN, a memory-efficient training framework that stores randomly quantized activations for back propagation. We prove the convergence of ActNN for general network architectures, and we characterize the impact of quantization on the convergence via an exact expression for the gradient variance. Using our theory, we propose novel mixed-precision quantization strategies that exploit the activation's heterogeneity across feature dimensions, samples, and layers. These techniques can be readily applied to existing dynamic graph frameworks, such as PyTorch, simply by substituting the layers. We evaluate ActNN on mainstream computer vision models for classification, detection, and segmentation tasks. On all these tasks, ActNN compresses the activation to 2 bits on average, with negligible accuracy loss. ActNN reduces the memory footprint of the activation by 12x, and it enables training with a 6.6x to 14x larger batch size.
to be published in ICML 2021
References in corpus (21)
- PyTorch: An Imperative Style, High-Performance Deep Learning Library
- Improved Baselines with Momentum Contrastive Learning
- DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients
- Deep Gradient Compression: Reducing the Communication Bandwidth for Distributed Training
- Mixed Precision Training
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
- MMDetection: Open MMLab Detection Toolbox and Benchmark
- PACT: Parameterized Clipping Activation for Quantized Neural Networks
- Training Deep Nets with Sublinear Memory Cost
- Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
- Deep High-Resolution Representation Learning for Visual Recognition
- Training Deep Neural Networks with 8-bit Floating Point Numbers
- Supporting Very Large Models using Automatic Dataflow Graph Partitioning
- Panoptic Feature Pyramid Networks
- ZeRO-Offload: Democratizing Billion-Scale Model Training
- HAWQ-V2: Hessian Aware trace-Weighted Quantization of Neural Networks
- Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization
- Backprop with Approximate Activations for Memory-efficient Network Training
- A Statistical Framework for Low-bitwidth Training of Deep Neural Networks
- Memory Optimization for Deep Networks
- Dynamic Tensor Rematerialization