8-bit Optimizers via Block-wise Quantization

Abstract

Stateful optimizers maintain gradient statistics over time, e.g., the exponentially smoothed sum (SGD with momentum) or squared sum (Adam) of past gradient values. This state can be used to accelerate optimization compared to plain stochastic gradient descent but uses memory that might otherwise be allocated to model parameters, thereby limiting the maximum size of models trained in practice. In this paper, we develop the first optimizers that use 8-bit statistics while maintaining the performance levels of using 32-bit optimizer states. To overcome the resulting computational, quantization, and stability challenges, we develop block-wise dynamic quantization. Block-wise quantization divides input tensors into smaller blocks that are independently quantized. Each block is processed in parallel across cores, yielding faster optimization and high precision quantization. To maintain stability and performance, we combine block-wise quantization with two additional changes: (1) dynamic quantization, a form of non-linear optimization that is precise for both large and small magnitude values, and (2) a stable embedding layer to reduce gradient variance that comes from the highly non-uniform distribution of input tokens in language models. As a result, our 8-bit optimizers maintain 32-bit performance with a small fraction of the memory footprint on a range of tasks, including 1.5B parameter language modeling, GLUE finetuning, ImageNet classification, WMT'14 machine translation, MoCo v2 contrastive ImageNet pretraining+finetuning, and RoBERTa pretraining, without changes to the original optimizer hyperparameters. We open-source our 8-bit optimizers as a drop-in replacement that only requires a two-line code change.

ICLR2022 spotlight version

References in corpus (28)

• Adam: A Method for Stochastic Optimization
• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• Language Models are Few-Shot Learners
• Improved Baselines with Momentum Contrastive Learning
• Scaling Laws for Neural Language Models
• Zero-Shot Text-to-Image Generation
• Mixed Precision Training
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Trained Ternary Quantization
• Binary Neural Networks: A Survey
• Training Deep Nets with Sublinear Memory Cost
• Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
• GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• A Simple Method for Commonsense Reasoning
• Training deep neural networks with low precision multiplications
• Training Deep Neural Networks with 8-bit Floating Point Numbers
• The Reversible Residual Network: Backpropagation Without Storing Activations
• fairseq: A Fast, Extensible Toolkit for Sequence Modeling
• Scaling Laws for Autoregressive Generative Modeling
• 8-Bit Approximations for Parallelism in Deep Learning
• PipeDream: Fast and Efficient Pipeline Parallel DNN Training
• BASE Layers: Simplifying Training of Large, Sparse Models
• ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
• Computing Extremely Accurate Quantiles Using t-Digests
• Mixed Precision Training With 8-bit Floating Point
• Training DNNs with Hybrid Block Floating Point
• Training Large Neural Networks with Constant Memory using a New Execution Algorithm
• Shifted and Squeezed 8-bit Floating Point format for Low-Precision Training of Deep Neural Networks