[論文レビュー] 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）を独立したブロックに分割して動的量子化し、安定した埋め込み層を併用することで、ハイパーパramータの変更なしに最大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％）を低減すること。
• 10億パラメータを超えるモデルにおける8ビット最適化の課題を克服すること：量子化の精度、計算効率、学習安定性。
• 32ビットの代替として即座に利用可能な実用的で高性能な8ビット最適化手法を実現すること。
• モデルと勾配には16ビットミックスドプレシジョンを用いることで、最適化状態への量子化の影響を明確に分離すること。
• 元のハイパーパramータを変更せず、多様な自然言語処理（NLP）およびビジョンタスクで安定性とパフォーマンスを確保すること。

提案手法

• ブロック単位の量子化を適用：最適化状態テンソルを独立したブロックに分割し、並列処理と局所的正規化を可能にする。
• 動的量子化を用いる：各ブロックごとに適応的な指数を有する非線形量子化により、小さな値と大きな値の両方の精度を保持する。
• 非一様な入力トークン分布に起因する勾配分散を低減するため、安定した埋め込み層を導入する。
• 効率的なブロック単位の正規化および量子化／逆量子化演算を実装するためのカスタムCUDAカーネルを採用する。
• 最適化状態を32ビットで更新し、その後8ビットに量子化して保存することで、32ビット性能を維持する。
• 逆量子化のためのルックアップテーブルを用いることで、バックワードパス中に高速かつ正確な状態回復を実現する。

実験結果

リサーチクエスチョン

• RQ1 10億パラメータを超える大規模モデルにおいて、ハイパーパramータのチューニングなしに8ビット最適化状態が32ビット学習性能と安定性を維持できるか？
• RQ2 テンソル全体の量子化と比較して、ブロック単位の量子化が安定性と効率性をどのように向上させるか？
• RQ3 最適化状態の極端な値の範囲において、動的量子化が量子化誤差をどのように低減するか？
• RQ4 非一様なトークン分布に起因する勾配分散を、安定した埋め込み層がどの程度低減できるか？
• RQ5 8ビット最適化手法が、多様なNLPおよびビジョンタスクにおいて32ビット最適化手法の即座の置き換えとして利用可能か？

主な発見

• 提案手法の8ビット最適化手法は、15億および3億5500万パラメータの言語モデリング、GLUE微調整、ImageNet分類、WMT’14/16翻訳、MoCo v2事前学習＋微調整、RoBERTa事前学習において、32ビット性能を維持した。
• 32ビットの対応する最適化状態と比較して、最適化状態のメモリ使用量を最大80％まで削減でき、元のハイパーパramータに変更がなかった。
• ブロック単位の量子化により、外れ値をブロック内に局在化させ、グローバルな量子化への影響を低減することで、学習安定性が向上した。
• 動的量子化により、最適化状態の値が5桁のオーダーにわたり高精度な表現が可能となり、大規模学習において不可欠である。
• NLPタスクにおいて、非一様なトークン入力に起因する勾配分散を低減するため、安定した埋め込み層は不可欠である。
• アブレーションスタディの結果、ブロック単位の量子化、動的量子化、安定した埋め込み層のすべてのコンponentsが、完全なパフォーマンスと安定性を達成するために必要であることが確認された。