大模型的訓練和微調對顯存要求很高，優化器狀態是顯存主要開銷之一。近日，清華大學朱軍、陳鍵飛團隊提出了用于神經網絡訓練的 4 比特優化器，節省了模型訓練的內存開銷，同時能達到與全精度優化器相當的準確率。

4 比特優化器在眾多預訓練和微調任務上進行了實驗，在保持準確率無損的情況下可將微調 LLaMA-7B 的顯存開銷降低多達 57%。

論文：https://arxiv.org/abs/2309.01507

代碼：https://github.com/thu-ml/low-bit-optimizers

模型訓練的內存瓶頸

從 GPT-3，Gopher 到 LLaMA，大模型有更好的性能已成為業界的共識。但相比之下，單個 GPU 的顯存大小卻增長緩慢，這讓顯存成為了大模型訓練的主要瓶頸，如何在有限的 GPU 內存下訓練大模型成為了一個重要的難題。

為此，我們首先需要明確消耗顯存的來源有哪些。事實上來源有三類，分別是：

1. 「數據顯存」，包括輸入的數據和神經網絡每層輸出的激活值，它的大小直接受到 batch size 以及圖像分辨率 / 上下文長度的影響；

2. 「模型顯存」，包括模型參數，梯度，以及優化器狀態（optimizer states），它的大小與模型參數數量呈正比；

3. 「臨時顯存」，包括 GPU kernel 計算時用到的臨時內存和其他緩存等。隨著模型規模的增大，模型顯存的占比逐漸增大，成為主要瓶頸。

優化器狀態的大小由使用哪種優化器決定。當前，訓練 Transformer 往往使用 AdamW 優化器，它們在訓練過程中需要存儲并更新兩個優化器狀態，即一階和二階矩（first and second moments）。如果模型參數量為 N，那么 AdamW 中優化器狀態的數量為 2N，這顯然是一筆極大的顯存開銷。

以 LLaMA-7B 為例，該模型含的參數數量大約 7B，如果使用全精度（32 比特）的 AdamW 優化器對它進行微調，那么優化器狀態所占用的顯存大小約為 52.2GB。此外，雖然樸素的 SGD 優化器不需要額外狀態，節省了優化器狀態所占用的內存，但是模型的性能難以保證。因此，本文主要關注如何減少模型內存中的優化器狀態，同時保證優化器的性能不受損。

節省優化器內存的方法

目前在訓練算法方面，節省優化器顯存開銷的方法主要有三類：

1. 通過低秩分解（Factorization）的思路對優化器狀態進行低秩近似（low-rank approximation）；

2. 通過只訓練一小部分參數來避免保存大多數的優化器狀態，例如 LoRA；

3. 基于壓縮 （compression）的方法，使用低精度數值格式來表示優化器狀態。

特別的，Dettmers et al. （ICLR 2022）針對 SGD with momentum 和 AdamW 提出了相應的 8 比特優化器，通過使用分塊量化（block-wise quantization）和動態指數數值格式（dynamic exponential numerical format）的技術，在語言建模、圖像分類、自監督學習、機器翻譯等任務上達到了與原有的全精度優化器相匹配的效果。

本文在基礎上，將優化器狀態的數值精度進一步降低至 4 比特，提出了針對不同優化器狀態的量化方法，最終提出了 4 比特 AdamW 優化器。同時，本文探索了將 壓縮和低秩分解方法結合的可能性，提出了 4 比特 Factor 優化器，這種混合式的優化器同時享有好的性能和更好的內存高效性。本文在眾多經典的任務上對 4 比特優化器進行了評估，包括自然語言理解、圖像分類、機器翻譯和大模型的指令微調。

在所有的任務上，4 比特優化器達到了與全精度優化器可比的效果，同時能夠占用更少的內存。

問題設置

基于壓縮的內存高效優化器的框架

首先，我們需要了解如何將壓縮操作引入到通常使用的優化器中，這由算法 1 給出。其中，A 是一個基于梯度的優化器（例如 SGD 或 AdamW）。該優化器輸入現有的參數 w，梯度 g 和優化器狀態 s，輸出新的參數和優化器狀態。在算法 1 中，全精度的 s_t 是暫時存在的，而低精度的 (s_t ) ? 會持久地保存在 GPU 內存中。這種方式能夠節省顯存的重要原因是：神經網絡的參數往往由每層的參數向量拼接而成。因此，優化器更新也是逐層 / 張量進行，進而在算法 1 下，最多只有一個參數的優化器狀態以全精度的形式留在內存中，其他層對應的優化器狀態都處于被壓縮的狀態。

主要的壓縮方法：量化（quantization）

量化是用低精度數值來表示高精度數據的技術，本文將量化的操作解耦為兩部分：歸一化（normalization）和映射（mapping），從而能夠更加輕量級的設計并實驗新的量化方法。歸一化和映射兩個操作依次以逐元素的形式施加在全精度數據上。歸一化負責將張量中的每個元素投射到單位區間，其中張量歸一化（per-tensor normalization）和分塊歸一化（block-wise normalization）分別如下定義：

不同歸一化方法的粒度不同，處理異常值的能力會有所區別，同時帶來的額外內存開銷也不同。而映射（mapping）操作負責將歸一化的數值映射到低精度能夠表示的整數。正式地講，給定位寬 b（即量化后每個數值使用 b 比特來表示）和預先定義的函數 T

映射操作被定義為：

因此，如何設計恰當的 T 對于減小量化誤差有很重要的作用。本文主要考慮線性映射（linear）和動態指數映射（dynamic exponent）。最后，去量化的過程就是按順序施加映射（mapping）和歸一化（normalization）的逆算子。

一階矩的壓縮方法

以下主要針對 AdamW 的優化器狀態（一階矩和二階矩）提出不同的量化方法。對于一階矩，本文的量化方法主要基于 Dettmers et al. （ICLR 2022）的方法，使用分塊歸一化（塊大小為 2048）和動態指數映射。

在初步的實驗中，我們直接將位寬從 8 比特降低至 4 比特，發現一階矩對于量化十分魯棒，在很多任務上已經達到匹配的效果，但也在一部分任務上出現性能上的損失。為了進一步提高性能，我們仔細研究了一階矩的模式，發現在單個張量中存在很多異常值。

此前的工作對于參數和激活值的異常值的模式已有一定的研究，參數的分布較為平滑，而激活值則具有按照 channel 分布的特點。本文發現，優化器狀態中異常值的分布較為復雜，其中有些張量的異常值分布在固定的行，而另外一些張量的異常值分布在固定的列。

對于異常值按列分布的張量，以行為優先的分塊歸一化可能會遇到困難。因此，本文提出采用更小的塊，塊大小為 128，這能夠在減小量化誤差的同時使額外的內存開銷保持在可控的范圍內。下圖展示了不同塊大小的量化誤差。

二階矩的壓縮方法

與一階矩相比，二階矩的量化更加困難并且會帶來訓練的不穩定性。本文確定了零點問題是量化二階矩的主要瓶頸，此外針對病態的異常值分布提出了改進的歸一化方法：rank-1 normalization。本文也嘗試了對二階矩的分解方法（factorization）。

零點問題

在參數、激活值、梯度的量化中，零點往往是不可缺少的，并且在也是量化后頻率最高的點。但是，在 Adam 的迭代公式中，更新的大小正比于二階矩的 -1/2 次方，因此在零附近的范圍內改變會極大影響更新的大小，進而造成不穩定。

下圖以直方圖的形式展示了量化前后 Adam 二階矩 -1/2 次方的分布， 即 h (v)=1/(√v+10^(-6) )。如果將零點包括在內（圖 b），那么大多數值都被推到了 10^6, 從而導致極大的近似誤差。一個簡單的辦法是在動態指數映射中將零點移除，在這樣做之后（圖 c），對二階矩的近似變得更加精確。在實際情況中，為了有效利用低精度數值的表達能力，我們提出采用移除零點的線性映射，在實驗中取得了很好的效果。

Rank-1 歸一化

基于一階矩和二階矩復雜的異常值分布，并受 SM3 優化器所啟發，本文提出了一種新的歸一化方法，命名為 rank-1 歸一化。對一個非負的矩陣張量 x∈R^(n×m), 它的一維統計量定義為：

進而 rank-1 歸一化可以被定義為：

rank-1 歸一化以更細粒度的方式利用了張量的一維信息，能夠更聰明且有效地處理按行分布或按列分布的異常值。此外，rank-1 歸一化能夠簡單的推廣到高維張量中，并且隨著張量規模的增大，它所產生的額外內存開銷要小于分塊歸一化。

此外，本文發現 Adafactor 優化器中對于二階矩的低秩分解方法能夠有效的避免零點問題，因此也對低秩分解和量化方法的結合進行了探索。下圖展示了針對二階矩的一系列消融實驗，證實了零點問題是量化二階矩的瓶頸，同時也驗證了 rank-1 歸一化，低秩分解方法的有效性。

實驗結果

研究根據所觀察的現象和使用的方式，最終提出兩種低精度優化器：4 比特 AdamW 和 4 比特 Factor，并與其他優化器進行對比，包括 8 比特 AdamW，Adafactor, SM3。研究選擇在廣泛的任務上進行評估，包括自然語言理解、圖像分類、機器翻譯和大模型的指令微調。下表展示了各優化器在不同任務上的表現。

可以看到，在所有的微調任務上，包括 NLU，QA，NLG，4 比特優化器可以匹配甚至超過 32 比特 AdamW，同時在所有的預訓練任務上，CLS，MT，4 比特優化器達到與全精度可比的水平。從指令微調的任務中可以看到，4 比特 AdamW 并不會破壞預訓練模型的能力，同時能較好地使它們獲得遵守指令的能力。

之后，我們測試了 4 比特優化器的內存和計算效率，結果如下表所示。相比 8 比特優化器，本文提出的 4 比特優化器能夠節省更多內存，在 LLaMA-7B 微調的實驗中最高節省 57.7%。此外，我們提供了 4 比特 AdamW 的融合算子版本，它能夠在節省內存的同時不影響計算效率。對于 LLaMA-7B 的指令微調任務，由于緩存壓力減小，4 比特 AdamW 也為訓練帶來了加速效果。詳細的實驗設置和結果可參考論文鏈接。

替換一行代碼即可在 PyTorch 中使用

importlpmm

optimizer=lpmm.optim.AdamW(model.parameters(),lr=1e-3,betas=(0.9,0.999))

我們提供了開箱即用的 4 比特優化器，只需要將原有的優化器替換為 4 比特優化器即可，目前支持 Adam 和 SGD 的低精度版本。同時，我們也提供了修改量化參數的接口，以支持定制化的使用場景。