PaperlayerView exact revision on arXiv

Source-linked AI summary

8-bit Optimizers via Block-wise Quantization

Tim Dettmers, Mike Lewis, Sam Shleifer, Luke Zettlemoyer

arXiv:2110.02861v2cs.LG

TL;DR

Stateful optimizer statistics improve optimization but consume memory that could otherwise support larger models. The paper introduces 8-bit optimizers using block-wise dynamic quantization and a stable embedding layer, reporting 32-bit performance across a broad task range with lower memory use. Natural-language tasks still require a stable embedding layer trained with 32-bit optimizer states.

  • Problem

    Stateful optimizer statistics consume memory that could otherwise be allocated to model parameters, limiting practical model size.

  • Method

    The paper combines block-wise quantization, dynamic quantization, and a stable embedding layer to store optimizer statistics in 8 bits.

  • Results

    8-bit optimizers match replicated 32-bit performance across a broad range of tasks while saving up to 8.5 GB for a 1.5B-parameter language model and 2.0 GB for RoBERTa.

  • Takeaways & Limitations

    The optimizers provide a drop-in replacement for 32-bit optimizers without changing hyperparameters and improve access to large-model finetuning on smaller GPUs.

  • Takeaways & Limitations

    For natural-language tasks, 8-bit optimizers require a stable embedding layer with 32-bit optimizer states to reach 32-bit performance.

Abstract

from arXiv · show

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.

1 BACKGROUND

Stateful optimizers accelerate training by tracking gradient statistics, but these states consume substantial memory. Quantization maps tensors to compact representations, with dynamic tree quantization targeting precision across value magnitudes.

  • Stateful optimizers: Stateful optimizers maintain per-parameter gradient statistics over time, including momentum-like and squared-gradient states used for accelerated optimization.Adam and Momentum update these statistics recursively before updating parameters.
  • Memory costs: 4 GB and 8 GB are the 32-bit state costs for Momentum and Adam in a 1B-parameter model, compared with 1 GB and 2 GB using 8-bit quantization.The corresponding costs are 4 and 8 bytes per parameter for 32-bit states.
  • Quantization: Quantization maps k-bit integer indices to real values in a target domain, storing indices after normalization and nearest-value assignment.Dequantization retrieves the mapped value and reverses normalization.
  • Dynamic quantization: Dynamic quantization normalizes by the tensor’s absolute maximum and uses binary search to find nearby quantized values.The normalization constant is N = max(|T|), and values are first mapped into [-1, 1].
  • Dynamic tree quantization: Dynamic tree quantization uses a variable exponent and fraction structure to achieve low quantization error for both small and large magnitudes.Its encoding includes a sign bit, exponent information, and linear quantization of following values.

2 8-BIT OPTIMIZERS

The proposed 8-bit optimizers combine block-wise quantization, dynamic quantization, and a stable embedding layer. These components target outliers, quantization precision, computational throughput, and instability in models with word embeddings.

  • Components: The optimizers combine block-wise quantization, dynamic quantization, and a stable embedding layer to improve precision, efficiency, and optimization stability.Block-wise quantization isolates outliers, dynamic quantization handles small and large values, and stable embeddings address word-embedding instability.
  • Implementation: 8-bit optimizer states are dequantized to 32-bit for updates and requantized for storage, using register-level conversion without temporary memory.The implementation avoids slow GPU-memory copies and is reported to be faster than regular 32-bit optimizers on GPUs.
  • Block-wise quantization: Blocks of size B = 2048 reduce synchronization costs by allowing independent normalization and quantization across GPU cores.For each block, normalization uses its own maximum absolute value.
  • Block-wise quantization: Block-wise quantization isolates outliers, making the largest optimizer states quantized without error and improving robustness and precision.This property is described as essential for good training performance in practice.
  • Dynamic quantization: Dynamic quantization extends dynamic tree quantization for positive Adam states by reusing the sign bit for a fixed fraction bit.The extension targets second-state values spanning roughly 3–5 orders of magnitude.
  • Stable embedding layer: The stable embedding layer normalizes highly non-uniform inputs to reduce extreme gradient variation and support more aggressive quantization.It uses Xavier uniform initialization and layer normalization before adding position embeddings.

3 8-BIT VS 32-BIT OPTIMIZER PERFORMANCE FOR COMMON BENCHMARKS

Across diverse NLP and computer-vision benchmarks, 8-bit optimizers match or exceed replicated 32-bit performance while reducing memory use without changing hyperparameters.

  • The experiments replace 32-bit optimizers with 8-bit optimizers without changing hyperparameters or weight, gradient, and activation precision.Benchmarks cover Adam, AdamW, or Momentum settings and follow public baseline configurations.
  • 8-bit optimizers match or exceed replicated 32-bit performance across GLUE, classification, machine translation, and language-modeling tasks.Table 1 reports median performance across diverse benchmarks and no observed instabilities.
  • 8-bit optimizers save up to 8.5 GB of GPU memory for the largest 1.5B-parameter language model and 2.0 GB for RoBERTa.These savings improve finetuning accessibility for users without large-memory GPUs.
  • The broad benchmark results indicate that 8-bit optimizers are robust replacements that save significant memory while slightly speeding training.The paper reports no additional hyperparameter changes.
  • Table 2 shows that larger models can be finetuned with the same GPU memory as standard 32-bit optimizer training.The comparison uses a batch size of one.

4 ANALYSIS

The analysis finds that dynamic and block-wise quantization, together with stable embeddings, are important for performance and stability. Sensitivity results indicate that 8-bit Adam can replace 32-bit Adam without further hyperparameter tuning.

  • Sensitivity Analysis: 8-bit Adam with block-wise dynamic quantization is a reliable replacement for 32-bit Adam without further hyperparameter tuning.The comparison varies learning rate, betas, and ϵ individually from baseline settings.
  • Ablation Analysis: The ablation compares small- and large-scale language-model perplexity and training stability against a 32-bit Adam baseline.The study includes component ablations and combinations to expose their interactions.
  • Ablation Analysis: Dynamic quantization, block-wise quantization, and the stable embedding layer are critical for performance or stability.Block-wise quantization is especially critical for large-scale language-model stability.
  • Sensitivity Analysis: The sensitivity analysis changes learning rate, betas, and ϵ individually and evaluates 8-bit Adam with stable embeddings against 32-bit Adam.The baseline hyperparameters are β1=0.9, β2=0.995, ϵ=1e-7, and lr=0.0163.
  • Ablation Analysis: The large-model ablations were not conducted fully because each training run takes one GPU year.This limits the extent of the large-scale ablation analysis.

5 RELATED WORK

Related work reduces training memory through optimizer-state factorization, sharding, parallelism, checkpointing, offloading, and quantization. The paper distinguishes its block-wise dynamic quantization from these approaches and emphasizes compatibility with different parallelism strategies.

  • Optimizer-State Reduction: Adafactor saves memory through row-column factorization of the second Adam state, while an even leaner variant omits the first moment and requires hyperparameter retuning.The omitted-first-moment variant has memory efficiency comparable to 8-bit Adam.
  • Optimizer-State Reduction: AdaGrad and extensions such as SM3 can be more memory-efficient than 8-bit Adam, but the paper reports only an initial AdaGrad comparison.These methods use aggregate training statistics or reduced optimizer-state structures.
  • Distributed Approaches: Optimizer sharding distributes states across accelerators but requires multiple accelerators and data parallelism, whereas 8-bit optimizers work with all kinds of parallelism.The paper states that 8-bit optimizers can reduce sharding communication overhead by 75%.
  • General Memory Reduction Techniques: Model, pipeline, and operator parallelism distribute memory across accelerators, while checkpointing, reversible connections, and offloading operate locally with added computational or communication costs.The paper contrasts these costs with 8-bit optimizers maintaining 32-bit training speed.
  • Quantization Methods and Data Types: The paper is the first to apply 8-bit quantization to optimizer statistics, unlike established quantization work focused on model compression, training, and inference.Common 8-bit formats use static sign, exponent, and fraction bits, which can have large errors for large magnitudes.
  • Related Quantization Methods: HBFP uses a separate exponent per matrix-multiplication tile, whereas block-wise dynamic quantization combines block-wise normalization with a dynamic exponent for each number.The latter provides a broader range of important values according to the comparison.

6 DISCUSSION & LIMITATIONS

The discussion identifies where 8-bit optimizers are most useful and where their guarantees are limited. They reduce optimizer memory and can enable larger-model training, but language tasks require stable embeddings and some instability remains poorly understood.

  • Limitations: Natural-language tasks require a stable embedding layer to reach 32-bit performance with 8-bit optimizers.The authors also report that 32-bit optimizers benefit from the stable embedding layer.
  • Scope and Practical Benefits: 8-bit optimizers are most beneficial for training or finetuning many-parameter models on highly memory-constrained GPUs.Models dominated by activation memory, such as convolutional networks, receive few benefits because optimizer savings scale with parameters.
  • Limitations: Sources of instability in models beyond 1B parameters remain not well understood, including systemic divergence and rare cascading instability from extreme outliers.The authors identify block-wise quantization and stable embeddings as key to stability but frame the instability explanation as a hypothesis.
  • Broader Impact: The optimizers can enable models that previously could not be trained on various GPUs and reduce optimizer memory significantly on single devices without degrading performance.The paper presents this as a way to improve access to larger models, especially for users with fewer resources.

B GLUE SCORE BREAKDOWN

This section breaks down GLUE scores and explains why embedding layers can destabilize low-precision optimization. The stable embedding layer combines enhanced precision, layer normalization, and Xavier initialization to address this issue.

  • Table 4 reports median GLUE scores for each task across 10 random seeds, with the mean computed over task medians.
  • C STABILITY OF EMBEDDING LAYERS: Embedding gradients can vary widely because token inputs are sparse and non-uniform, producing maximum magnitudes 100x larger than other layers.
  • C STABILITY OF EMBEDDING LAYERS: Gradient normalization by total mini-batch tokens makes embedding-gradient magnitudes depend on token frequency and batch composition.
  • C STABILITY OF EMBEDDING LAYERS: The stable embedding layer addresses 8-bit optimizer instability by combining enhanced precision, layer normalization, and Xavier initialization.

D QUANTIZATION ERROR ANALYSIS

The error analysis evaluates how often Adam state-value combinations are used and how accurately they are quantized. Block-wise dynamic quantization has the smallest overlap between frequently used and high-error regions, with absolute error 0.0061.

  • The analysis measures usage and mean absolute or relative error across the 256×256 combinations of Adam’s two 8-bit states.The first state is the exponentially smoothed running sum, and the second is the exponentially smoothed running squared sum.
  • An overlap between high-use and high-error regions indicates frequent large deviations associated with unstable training.
  • 0.0061 absolute Adam quantization error is achieved by block-wise dynamic quantization, versus 0.0067 for dynamic quantization.The analysis also indicates that block-wise dynamic quantization produces rarer large errors, which likely contributes to improved optimization stability.
  • Block-wise dynamic quantization has the smallest overlap between regions of high use and high error.
  • Table 5 benchmarks the average time for 100 optimizer updates per billion parameters using a large normal-distribution sample.

F ADDITIONAL QUANTIZATION DATA TYPES

This section compares additional quantization data types and develops quantile quantization as a distribution-based alternative. Quantile quantization can reduce average error but is too slow for practical training, while SRAM-Quantiles accelerates its estimation.

  • Quantile quantization has twice as low average quantization error as dynamic quantization for normal distributions but sporadic large errors and poor model performance.
  • Quantile quantization is too slow for practical use even with state-of-the-art quantile estimation algorithms.
  • Inverse Dynamic Quantization reallocates dynamic-tree exponent and fraction bits based on the hypothesis that large Adam updates are more important.
  • F.2 QUANTILE QUANTIZATION: A LOSSY MINIMUM ENTROPY ENCODING: Quantile quantization assigns 2^k bins equal probability mass and uses their midpoints as quantization values.
  • G SRAM-QUANTILES: A FAST QUANTILE ESTIMATION ALGORITHM: SRAM-Quantiles estimates sample quantiles using a divide-and-conquer strategy that performs sorting solely in fast SRAM.

F.3 VISUALIZATION: DYNAMIC VS LINEAR QUANTIZATION VS QUANTILE QUANTIZATION

The visualization compares how linear, dynamic, and quantile quantization map values and distribute quantization error. Dynamic quantization allocates precision to small and large magnitudes, while quantile quantization requires expensive quantile estimation.

  • Dynamic quantization allocates most bits to both large and small values, unlike the mappings shown for linear and quantile quantization.
  • Estimating tensor quantiles requires sorting, making the process computationally expensive for training.
  • G SRAM-QUANTILES: A FAST QUANTILE ESTIMATION ALGORITHM: SRAM-Quantiles estimates tensor-wide quantiles by computing subset eCDFs in SRAM and atomically averaging the subset quantiles in DRAM.Each subset contains about 4096 32-bit values.
  • G SRAM-QUANTILES: A FAST QUANTILE ESTIMATION ALGORITHM: 15x faster SRAM than DRAM motivates restructuring quantile estimation around GPU memory hierarchies.
  • G SRAM-QUANTILES: A FAST QUANTILE ESTIMATION ALGORITHM: 0.064 ns per stream element is reported for estimating 256 quantiles, compared with 300 ns and 5 ns for two general algorithms.

H ADAGRAD COMPARISONS

The paper examines how 8-bit optimizers behave with AdaGrad, whose accumulated statistics span a wide range over training. In the reported 209M-parameter language-model experiment, the 8-bit methods underperform Adam in perplexity, and possible improvements are left for future work.

  • Interpretation: AdaGrad’s gradient statistics may be difficult to quantize because they average gradients over the entire course of training.The paper presents this as one hypothesis for the poorer 8-bit performance.
  • Comparison and future work: AdaGrad underperforms Adam in perplexity, and adding momentum is proposed as a possible future improvement.The paper leaves such improvements for future work.

I STABLE EMBEDDING LAYER ABLATIONS

The stable embedding ablations show that Xavier initialization and layer normalization improve performance in the small-model setup. However, the experiments cannot establish whether the layer stabilizes models beyond 1B parameters because large-scale instability testing exceeds the computational budget.

  • Stable Embedding Layer Ablations: Xavier initialization and layer normalization improve performance in the stable embedding layer ablations.The ablations use a 200M language-model setup, three random seeds, and median perplexity.
  • Stable Embedding Layer Ablations: The stabilizing effect of the Stable Embedding layer beyond 1B parameters cannot be studied within the available computational budget.A single model above 1B parameters takes roughly 300 GPU days, and multiple random seeds are needed to study instability.
  • Stable Embedding Layer Ablations: 32-bit optimizer states do not improve perplexity or stability at the tested small scale.The paper notes that these states might still affect stability in large-scale models.
Loading 2110.02861v2…