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.

本文开发了首个使用 8 位统计量进行训练的优化器，采用了分块动态量化进行加速，同时结合了非线性优化和嵌入层技术以提高精度和稳定性，并在一系列任务中展现了较高的性能和较小的内存占用。