大模型的訓(xùn)練和微調(diào)對顯存要求很高，優(yōu)化器狀態(tài)是顯存主要開銷之一。近日，清華大學(xué)朱軍、陳鍵飛團(tuán)隊提出了用于神經(jīng)網(wǎng)絡(luò)訓(xùn)練的 4 比特優(yōu)化器，節(jié)省了模型訓(xùn)練的內(nèi)存開銷，同時能達(dá)到與全精度優(yōu)化器相當(dāng)?shù)臏?zhǔn)確率。

4 比特優(yōu)化器在眾多預(yù)訓(xùn)練和微調(diào)任務(wù)上進(jìn)行了實驗，在保持準(zhǔn)確率無損的情況下可將微調(diào) LLaMA-7B 的顯存開銷降低多達(dá) 57%。

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

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

模型訓(xùn)練的內(nèi)存瓶頸

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

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

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

2. 「模型顯存」，包括模型參數(shù)，梯度，以及優(yōu)化器狀態(tài)（optimizer states），它的大小與模型參數(shù)數(shù)量呈正比；

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

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

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

節(jié)省優(yōu)化器內(nèi)存的方法

目前在訓(xùn)練算法方面，節(jié)省優(yōu)化器顯存開銷的方法主要有三類：

1. 通過低秩分解（Factorization）的思路對優(yōu)化器狀態(tài)進(jìn)行低秩近似（low-rank approximation）；

2. 通過只訓(xùn)練一小部分參數(shù)來避免保存大多數(shù)的優(yōu)化器狀態(tài)，例如 LoRA；

3. 基于壓縮 （compression）的方法，使用低精度數(shù)值格式來表示優(yōu)化器狀態(tài)。

特別的，Dettmers et al. （ICLR 2022）針對 SGD with momentum 和 AdamW 提出了相應(yīng)的 8 比特優(yōu)化器，通過使用分塊量化（block-wise quantization）和動態(tài)指數(shù)數(shù)值格式（dynamic exponential numerical format）的技術(shù)，在語言建模、圖像分類、自監(jiān)督學(xué)習(xí)、機(jī)器翻譯等任務(wù)上達(dá)到了與原有的全精度優(yōu)化器相匹配的效果。

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

在所有的任務(wù)上，4 比特優(yōu)化器達(dá)到了與全精度優(yōu)化器可比的效果，同時能夠占用更少的內(nèi)存。

問題設(shè)置

基于壓縮的內(nèi)存高效優(yōu)化器的框架

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

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

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

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

映射操作被定義為：

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

一階矩的壓縮方法

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

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

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

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

二階矩的壓縮方法

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

零點問題

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

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

Rank-1 歸一化

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

進(jìn)而 rank-1 歸一化可以被定義為：

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

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

實驗結(jié)果

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

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

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

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

import lpmm

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

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