[论文解读] 8-bit Optimizers via Block-wise Quantization

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

Advanced Neural Network Applications参考文献 58被引用 4

一句话总结

本文提出使用逐块动态量化技术的8位优化器，可在大幅降低内存使用的同时保持32位优化器的性能。通过在独立块中对优化器状态（如Adam、AdamW）进行动态量化，并结合稳定的嵌入层，该方法在不改变超参数的前提下，实现了对参数量达15亿的模型的完整训练稳定性与速度。

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.

研究动机与目标

• 解决大规模模型训练中32位优化器状态（占总内存的33–75%）带来的高内存占用问题。
• 克服8位优化中的挑战：量化精度、计算效率以及参数量超过10亿的模型的训练稳定性。
• 实现可即插即用的实用型高性能8位优化器，作为32位优化器的直接替代品。
• 通过使用16位混合精度表示模型和梯度，隔离量化对优化器状态的影响。
• 在不修改原始超参数的前提下，确保在多种自然语言处理和视觉任务中具备稳定性和高性能。

提出的方法

• 应用逐块量化：将优化器状态张量划分为独立块，实现并行处理与局部归一化。
• 使用动态量化——每块自适应指数的非线性量化，以保留小值和大值的精度。
• 引入稳定的嵌入层，对高度非均匀的输入词元分布进行归一化，以降低梯度方差。
• 实现自定义CUDA内核，高效执行逐块归一化及量化/反量化操作。
• 通过在32位下执行优化器更新，再将状态量化为8位进行存储，以保持32位性能。
• 在反向传播中使用查表法进行反量化，实现快速且精确的状态恢复。

实验结果

研究问题

• RQ1 在不进行超参数调优的情况下，8位优化器状态是否能在大模型（>10亿参数）中维持32位训练性能与稳定性？
• RQ2 与张量级量化相比，逐块量化在提升稳定性和效率方面有何优势？
• RQ3 动态量化在减少优化器状态中极端值范围的量化误差方面起到何种作用？
• RQ4 稳定的嵌入层在多大程度上缓解了由非均匀词元分布引起的梯度方差？
• RQ5 8位优化器是否可作为即插即用的替代方案，部署于多种自然语言处理和视觉任务中？

主要发现

• 所提出的8位优化器在15亿和3.55亿参数的语言建模、GLUE微调、ImageNet分类、WMT’14/16翻译、MoCo v2预训练+微调以及RoBERTa预训练任务中，均保持了32位性能。
• 与32位对应物相比，优化器状态的内存占用最高可降低80%，且无需更改原始超参数。
• 逐块量化通过将异常值隔离在各块内，降低了其对全局量化的影响，从而提升了训练稳定性。
• 动态量化能够在5个数量级范围内的优化器状态值中实现高精度表示，这对大规模训练至关重要。
• 在自然语言处理任务中，稳定的嵌入层至关重要，它通过降低非均匀词元输入引起的梯度方差，实现了稳定的8位训练。
• 消融实验表明，所有组件——逐块量化、动态量化和稳定嵌入层——对于实现完整性能与稳定性均不可或缺。