[Paper Review] 8-bit Optimizers via Block-wise Quantization

Tim Dettmers, Mike Lewis+2 more|arXiv (Cornell University)|Oct 6, 2021

Advanced Neural Network Applications58 references4 citations

TL;DR

This paper introduces 8-bit optimizers using block-wise dynamic quantization to maintain 32-bit optimizer performance with drastically reduced memory usage. By quantizing optimizer states (e.g., Adam, AdamW) in independent blocks with dynamic quantization and a stable embedding layer, the method achieves full training stability and speed on models up to 1.5B parameters without hyperparameter changes.

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.

Motivation & Objective

• Address the high memory footprint of 32-bit optimizer states (33–75% of total memory) in large model training.
• Overcome challenges in 8-bit optimization: quantization accuracy, computational efficiency, and training stability for models >1B parameters.
• Enable practical, high-performance 8-bit optimizers that are drop-in replacements for 32-bit counterparts.
• Isolate the impact of quantization on optimizer states by using 16-bit mixed precision for model and gradients.
• Ensure stability and performance across diverse NLP and vision tasks without modifying original hyperparameters.

Proposed method

• Apply block-wise quantization: split optimizer state tensors into independent blocks for parallel processing and localized normalization.
• Use dynamic quantization—non-linear quantization with adaptive exponent per block—to preserve precision for both small and large magnitude values.
• Introduce a stable embedding layer that normalizes highly non-uniform input token distributions to reduce gradient variance.
• Implement custom CUDA kernels for efficient block-wise normalization and quantization/dequantization operations.
• Maintain 32-bit performance by performing optimizer updates in 32-bit, then quantizing the state to 8-bit for storage.
• Use a lookup table for dequantization during backward passes, enabling fast and accurate state recovery.

Experimental results

Research questions

• RQ1 Can 8-bit optimizer states maintain 32-bit training performance and stability in large models (>1B parameters) without hyperparameter tuning?
• RQ2 How does block-wise quantization improve stability and efficiency compared to tensor-wide quantization?
• RQ3 What role does dynamic quantization play in reducing quantization error for extreme value ranges in optimizer states?
• RQ4 To what extent does the stable embedding layer mitigate gradient variance caused by non-uniform token distributions?
• RQ5 Can 8-bit optimizers be deployed as a drop-in replacement for 32-bit optimizers across diverse NLP and vision tasks?

Key findings

• The proposed 8-bit optimizers maintain 32-bit performance across 1.5B and 355M parameter language modeling, GLUE fine-tuning, ImageNet classification, WMT’14/16 translation, MoCo v2 pretraining+finetuning, and RoBERTa pretraining.
• Memory footprint of optimizer states is reduced by up to 80% compared to 32-bit counterparts, with no change to original hyperparameters.
• Block-wise quantization improves training stability by isolating outliers within blocks, reducing their impact on global quantization.
• Dynamic quantization enables high-precision representation across a 5-order-of-magnitude range of optimizer state values, critical for large-scale training.
• The stable embedding layer is essential for NLP tasks, as it reduces gradient variance from non-uniform token inputs, enabling stable 8-bit training.
• Ablation studies confirm that all components—block-wise quantization, dynamic quantization, and stable embedding layer—are necessary for full performance and stability.