8-bit Optimizers via Block-wise Quantization

By Tim Dettmers, Michael Lewis, Sam Shleifer, Luke S. Zettlemoyer. 2021

509 citations recorded by the source corpus.

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-sourceour 8-bit optimizers as a drop-in replacement that only requires a two-line code change. Increasing model size is an effective way to achieve better performance for given resources (Kaplan et al., 2020; Henighan et al., 2020; Raffel et al., 2019; Lewis et al., 2021). However, training such large models requires storing the model, gradient, and state of the optimizer (e.g., exponentially smoothed sum and squared sum of previous gradients for Adam), all in a fixed amount of available memory. Although significant research has focused on enabling larger model training by reducing or efficiently distributing the memory required for the model parameters (Shoeybi et al., 2019; Lepikhin et al., 2020; Fedus et al., 2021; Brown et al., 2020; Rajbhandari et al., 2020), reducing the memory footprint of optimizer gradient statistics is much less studied. This is a significant missed opportunity since these optimizer states use 33-75% of the total memory footprint during training. For example, the Adam optimizer states for the largest GPT-2 (Radford et al., 2019) and T5 (Raffel et al., 2019) models are 11 GB and 41 GB in size. In this paper, we develop a fast, high-precision non-linear quantization method – block-wise dynamic quantization – that enables stable 8-bit optimizers (e.g., Adam, AdamW, and Momentum) which maintain 32-bit performance at a fraction of the memory footprint and without any changes to the original hyperparameters.1 While most current work uses 32-bit optimizer states, recent high-profile efforts to use 16-bit optimizers report difficultly for large models with more than 1B parameters (Ramesh et al., 2021). Going from 16-bit optimizers to 8-bit optimizers reduces the range of possible values from 2 = 65536 values to just 2 = 256. To our knowledge, this has not been attempted before. Effectively using this very limited range is challenging for three reasons: quantization accuracy, computational efficiency, and large-scale stability. To maintain accuracy, it is critical to introduce some form of non-linear quantization to reduce errors for both common small magnitude values We study 8-bit optimization with current best practice model and gradient representations (typically 16-bit mixed precision), to isolate optimization challenges. Future work could explore further compressing all three. 1 ar X iv :2 11 0. 02 86 1v 2 [ cs .L G ] 2 0 Ju n 20 22 Published as a conference paper at ICLR 2022

Subfields: Artificial Intelligence, Computational Mathematics, Computer Vision and Pattern Recognition, Efficiency & Compression, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Mathematics

In Computer Vision and Pattern Recognition

In Efficiency & Compression

In Optimization & Training

Citations in the corpus

QLoRA: Efficient Finetuning of Quantized LLMs (2023, methodology): “Therefore fast quantile approximation algorithms, such as SRAM quantiles [15], are used to estimate them.”
Stable and low-precision training for large-scale vision-language models (2023, methodology): “While Algorithm 1 is the most straightforward version of SwitchBack, we also present two alternative versions—SwitchBackM and SwitchBackQ— and release triton [55] implementations for all three via the bitsandbytes library [17].”
The case for 4-bit precision: k-bit Inference Scaling Laws (2022, methodology): “Dynamic exponent data types (Dettmers, 2016; Dettmers et al., 2022b) use one bit for the sign, and the number of following zero bits represents the exponent with base 10.”
GLM-130B: An Open Bilingual Pre-trained Model (2022, methodology): “OPT-175B attempted to fix by manually skipping data and adjusting hyper-parameters; BLOOM176B did so via the embedding norm technique (Dettmers et al., 2021).”
LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (2022, neutral): “As such, it is desirable to have multiple scaling constants per tensor, such as block-wise constants (Dettmers et al., 2022), so that the effect of that outliers is confined to each block.”
MetaICL: Learning to Learn In Context (2021, methodology): “To save memory during metatraining, we use an 8-bit approximation (Dettmers et al., 2021) of an Adam optimizer (Kingma and Ba, 2015) and mixed precision (Micikevicius et al.”
What Language Model to Train if You Have One Million GPU Hours? (2022, background): “Many promising additions and tweaks have been proposed to the Transformer architecture (Press et al., 2022; Shazeer, 2020; Dettmers et al., 2021) but there is a lack of systematic investigation to identify which will best transfer to increased scales, and which are idiosyncrasies of specific setups.”
Designing Effective Sparse Expert Models (2022, methodology): “In this work, we use the Adafactor optimizer due to its memory efﬁciency (though recently introduced 8-bit optimizers (Dettmers et al., 2021) may offer better trade-offs).”
MedAlpaca -- An Open-Source Collection of Medical Conversational AI Models and Training Data (2023, unknown): “When combined with LoRA, this strategy further reduces the memory needed for training [13] [14].”
Leveraging large language models for predictive chemistry (2024, unknown): “We also performed some of our experiments by fine-tuning the GPT-J-6B model 52,53 (which has been trained on the Pile dataset 54 ) on consumer hardware using 8-bit quantization 55 and 8-bit optimizers 56 in addition to the low-rank adaptation (LoRA) technique 57 .”
GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection (2024, methodology): “In our work, we mainly consider applying GaLore with 8-bit optimizers (Dettmers et al., 2021) and per-layer weight updates (Lv et al., 2023).”
ST-MoE: Designing Stable and Transferable Sparse Expert Models (2022, methodology): “In this work, we use the Adafactor optimizer due to its memory efﬁciency (though recently introduced 8-bit optimizers (Dettmers et al., 2021) may offer better trade-offs).”
Separable Self-attention for Mobile Vision Transformers (2022, methodology): “Other methods Transformer-based models performance can be improved using different methods, including mixed-precision training [21], efficient optimizers [22, 23], and knowledge distillation [2].”
A Survey of Large Language Models for Healthcare: from Data, Technology, and Applications to Accountability and Ethics (2023, unknown): “Besides, they employed 8-bit matrix multiplication for the feed-forward and attention projection layers [217], along with an 8-bit optimizer [218] to”
Fine-Tuning Language Models with Just Forward Passes (2023, background): “, gradient checkpointing [17], FlashAttention [22], and quantization training [26].”
LlamaFactory: Unified Efficient Fine-Tuning of 100+ Language Models (2024, unknown): “We fine-tune these models using the 8-bit AdamW optimizer (Dettmers et al., 2022b) in bfloat16 precision with activation checkpointing.”
Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads (2024, background): “We target to depict tree attention’s motivation and its performance.”
BloombergGPT: A Large Language Model for Finance (2023, methodology): “Following Le Scao et al. (2022) and first used in Dettmers et al. (2022), the model has an additional layer normalization after token embeddings, formally: h̄1 = LN em(h0) + SA(LN(LN em(h0))), where h0 is the initial token embedding and LN em is the new component of embedding layernormalization.”