8-bit Optimizers via Block-wise Quantization

The rapid growth of deep learning models, from millions to hundreds of billions of parameters, has driven remarkable performance gains but also introduced significant challenges in resource management. Optimizer states, especially in algorithms like Adam and Momentum, consume vast amounts of memory—up to 75% of total training memory—limiting the maximum model size. Prior efforts to reduce memory footprint focused on quantizing weights and activations, but optimizer states remained high-precision due to stability concerns. Recent advances in mixed-precision training (e.g., 16-bit) have alleviated some issues, yet further compression is needed for extremely large models. The main obstacle is maintaining training stability and accuracy when reducing the bit-width of optimizer statistics, which are highly sensitive to quantization errors. This paper builds on prior work in low-bit quantization, introducing a novel approach tailored for optimizer states, aiming to enable efficient training of models exceeding one billion parameters on commodity hardware.

The core challenge lies in compressing optimizer states—such as the first and second moments in Adam—without degrading training stability or convergence. Existing low-bit methods, like 16-bit or fixed-point quantization, struggle with outliers and large gradient variations, leading to divergence or unstable training. Large models exacerbate these issues, as the impact of quantization errors becomes more pronounced. Additionally, input data distributions in NLP tasks are often highly non-uniform, causing further gradient variance. The key problem is designing a quantization scheme that can handle extreme outliers, adapt to wide dynamic ranges, and ensure stable, efficient training across diverse tasks and model sizes, all without requiring hyperparameter tuning.

This work introduces three main innovations: • Block-wise dynamic quantization divides gradient statistics into small, independently normalized blocks, reducing outlier influence and improving quantization accuracy. • Non-linear dynamic quantization extends dynamic tree quantization to handle unipolar data, adapting to the wide dynamic range of optimizer statistics. • A stable embedding layer normalizes input token distributions, reducing gradient variance and enhancing training stability. These components collectively enable the optimizer to operate effectively at 8 bits, preserving performance while drastically reducing memory usage. The approach is compatible with existing optimizers like Adam, requiring only minimal code changes, and is validated across multiple large-scale tasks.

• �� The optimizer state tensors are partitioned into small blocks of size 2048 elements. Each block's maximum absolute value is computed independently, serving as a normalization constant.
• �� Within each block, the tensor is normalized by dividing by this maximum, then quantized using a non-linear scheme that maps values to 8-bit indices via a binary search.
• �� The quantized indices are stored, and during dequantization, lookup tables and multiplication by the block maximum restore approximate original values.
• �� Dynamic quantization adjusts to the magnitude of gradients, ensuring small and large values are represented accurately.
• �� The stable embedding layer employs Xavier initialization and layer normalization to stabilize input token distributions, reducing gradient variance.
• �� During training, the optimizer states are dequantized to 32-bit, updated, then re-quantized, all within registers to avoid extra memory overhead.
• �� CUDA kernels are customized for efficient in-register quantization/dequantization, maintaining high throughput.

The method is evaluated on tasks including large language modeling (1.5B parameters), GLUE, ImageNet classification, WMT translation, and RoBERTa pretraining. All experiments use the same hyperparameters as baseline 32-bit optimizers, with no tuning. Results show that the 8-bit optimizer achieves comparable perplexity, accuracy, and BLEU scores, with significant memory savings. Ablation studies confirm the importance of each component—block-wise, dynamic quantization, and stable embedding layer—for stability and performance. The experiments demonstrate robustness across different hardware platforms and model sizes, validating the approach's scalability and practicality.

In large-scale language modeling, 8-bit Adam maintains perplexity at 9.0, matching 32-bit results, while reducing memory by 8.5GB and increasing training speed by 5%. On GLUE, accuracy is within 0.1%, and on ImageNet, top-1 accuracy remains at 77.2%. The method scales to models over 1.5 billion parameters, with no divergence or instability. Ablation results show that removing block-wise or dynamic quantization causes performance drops or instability, confirming their necessity. The approach also reduces memory footprint by up to 75%, enabling training on smaller GPUs, thus broadening accessibility.

This technique is immediately applicable to training and fine-tuning large NLP and vision models, especially where hardware memory is limited. It allows researchers and industry practitioners to train bigger models without upgrading hardware, reducing costs and energy consumption. The plug-and-play nature facilitates integration into existing frameworks like PyTorch, making large-scale AI more accessible. Long-term, this approach could evolve into fully low-bit training pipelines, including activations and gradients, further democratizing AI development and deployment.

While effective, the method may face challenges with highly non-uniform data distributions or extreme outliers, which could still cause quantization errors. The optimal block size may vary across tasks, requiring tuning for best results. Extending quantization to all training components (activations, gradients) remains future work, which could introduce additional complexity. Computational overhead from non-linear quantization, though minimized, might impact training speed in some scenarios. Further research is needed to address these issues for broader applicability.