解決LLaMA、BERT等部署難題：首個4-bit浮點量化LLM來了
大語言模型 (LLM) 壓縮一直備受關(guān)注，后訓練量化（Post-training Quantization) 是其中一種常用算法，但是現(xiàn)有 PTQ 方法大多數(shù)都是 integer 量化，且當比特數(shù)低于 8 時，量化后模型的準確率會下降非常多。想較于 Integer (INT) 量化，F(xiàn)loating Point (FP) 量化能更好的表示長尾分布，因而越來越多的硬件平臺開始支持 FP 量化。而這篇文章給出了大模型 FP 量化的解決方案。文章發(fā)表在 EMNLP 2023 上。
- 論文地址：https://arxiv.org/abs/2310.16836
- 代碼地址：https://github.com/nbasyl/LLM-FP4
要了解本文，必須要先具備基本的有關(guān) Floating Point Format 以及 Floating Point Quantization 的知識，首先 Floating Point Number 可以用以下公式表示:
s 代表正負符號位 (sign bit)，m 代表尾數(shù)位 (mantissa bits)，e 代表指數(shù)位 (exponent bits)。p 是一個介于 0 到 2^e - 1 之間的值，用來表示當前數(shù)字該被劃分到哪一個指數(shù)區(qū)間，d 取 0 或 1 的值，用來表示第 i 個 mantissa bit。b 是 bias，一個用來調(diào)整 exponent 區(qū)間的整數(shù)值。
接下來介紹 Floating Point Quantization 是怎么運作的，首先輸入值必須經(jīng)過一個 scale and clip 的步驟，先把 input clip 到 Floating Point 能表示的最大區(qū)間 (±Qmax)，如以下公式所示：
可以看到類似于 integer 量化，F(xiàn)P 量化也會加入一個 full-precision 的縮放因子 (scaling factor) 來縮放 input 到合適的區(qū)間。而縮放因子在運算矩陣乘法的時候，和低比特的矩陣乘法分開計算，所以并不會造成很大的 overhead。融入了這個 full-precision 的縮放因子之后，不同的 quantized tensor 能夠被相應地 clip 到不同的最大最小值區(qū)間。在實際使用過程中，會根據(jù)輸入 tensor 的值域確定需要的量化區(qū)間，然后利用公式 (4) 推導出相對應的 bias。注意公式 (4) 里的 bias 可以被用作實數(shù)值的縮放因子，見公式 (2)(3)。
Floating-Point Quantization 的下一個步驟是在決定好量化區(qū)間后把區(qū)間內(nèi)的值分配到相對應的量化區(qū)間內(nèi)，這個步驟稱為 Compare and Quantize：
上圖直觀說明了量化的過程，當前的輸入值，在用公式 5 比較過后，量化到不同的量化區(qū)間中。
在得到量化過的 activation 和 weight 后，這里的 scaling factor 提到前面先計算，而達到如下的 efficient matrix multiplication，完成矩陣乘法的加速：
接著本文指出 FP 量化的準確度，和 exponent bits 的設(shè)定以及量化的區(qū)間息息相關(guān)。
如下圖所示，不同的 FP format (浮點數(shù)的指數(shù)位 / 尾數(shù)位設(shè)定) 之間存在巨大的量化誤差差異，只有當選取合適的 FP format 時，F(xiàn)P Quantization 比 INT Quantization 能更好的表示長尾分布。這個現(xiàn)象也在之前的論文中得到驗證 [1]。
而這篇文章提出了對應的解決方案，用一個 search-based 浮點量化算法，統(tǒng)籌搜索出最適合的浮點數(shù)的指數(shù)位 / 尾數(shù)位設(shè)定以及對應的量化區(qū)間。
除此之外，另一個同時出現(xiàn)在各種不同類別 Transformer 模型 (Bert,LLaMA,ViT) 中的現(xiàn)象也會嚴重影響量化的難度：那就是模型的 activation 中不同 channel 之間的數(shù)量級會有很高的差異，而同 channel 之間的量級十分一致。之前 LLM.int8 [2] 和 SmoothQuant [3] 也有類似的發(fā)現(xiàn)，不過這篇文章指出這個現(xiàn)象不僅僅存在于 LLM 中，并且在其他 Transformer 模型里也有類似現(xiàn)象 如下如所示，LLaMA 與 BERT 以及 DeIT-S 中的 activation 的分布都發(fā)現(xiàn)了類似的現(xiàn)象：
從圖中可以看到，那些異常大的 channel 都比剩余的 channel 大很多，所以在量化 activation tensor 的過程中，量化的精度很大程度會被這些異常值決定，從而抑制其他 channel 值的量化區(qū)間，最終降低整體影響量化精度。這會導致量化的最終結(jié)果崩壞，尤其當比特數(shù)降到一定程度的時候。值得注意的是，只有 tensor-wise 和 token-wise 量化可以在 efficient matrix multipilication 的時候?qū)?scaling factor 提取出來，而 channel-wise 量化是不支持 efficient matrix multipilication 的，見下圖。
為了解決這個問題，同時維持高效率矩陣乘法 (Efficient Matrix Multiplication)，本文利用少量的校正資料集，預先算出 activation 的每個 channel 的最大值，從而計算縮放因子。然后將這個縮放因子一拆為二，拆解成一個 per-tensor 的實數(shù)乘以 per-channel 的 2 的冪。而這個 2 的整數(shù)次方即用 FP 里的 exponent bias 表示。完整的過程可以用以下公式表示：
進一步地，在 calibration 完成之后，這個 per-channel exponent bias 就不再變化，因此可以和 weight quantization 一起進行預計算 (pre-compute)，將這個 per-channel exponent bias 整合進量化后的 weights 中，提高量化精度。完整的過程如以下公式：
可以看到在 pre-shifted 后，原本 activation 中的 full-precision per-channel biases 的位置變成了一個 tensor-wise 的實數(shù) scaling factor ，而被拆解出來的整數(shù) per-channel biases 被移到了 weight 中原本 integer bias 的位置，如公式 4。
從而這個方法 (pre-shifted exponent bias) 能在維持 efficient matrix multiplication 的原則下，更好得提高量化精度，方法的直觀展示如下圖所示：
最后本文展示 Floating Point Quantization (FPQ) 方法，在 LLaMA, BERT 以及 ViTs 模型上，4-bit 量化皆取得了遠超 SOTA 的結(jié)果。特別是，這篇文章展示了 4-bit 量化的 LLaMA-13B 模型，在零樣本推理任務上達到平均 63.1 的分數(shù)，只比完整精度模型低了 5.8 分，且比之前的 SOTA 方法平滑量高出了 12.7，這是目前少數(shù)已知可行的 4-bit 量化方案了。
