位元組開源大模型量化新思路，2-bit量化模型精度齊平fp16

AIxiv專欄是本站發布學術、技術內容的欄位。過去數年，本站AIxiv專欄接收通報了2,000多篇內容，涵蓋全球各大專院校與企業的頂尖實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或聯絡報道。投稿信箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

隨著深度學習大語言模型的越來越火爆，大語言模型越做越大，使得其推理成本也水漲船高。模型量化，成為熱門的研究主題。

最近，位元組跳動推出了一個全新的量化思路，拋棄傳統的量化範式，從數學最佳化的角度來對量化任務建模。文章放在了 arXiv，程式碼已經開源，可以一鍵複現文中的所有結果。此量化想法是基於數學最佳化，從數學最佳化的角度來對量化任務建模，透過最大化目標函數或最小化損失函數來尋找最優解。這種想法在實驗中獲得了很好的效果，並且取得了令人滿意的結果。

• 論文連結：https://arxiv.org/abs/2404.12759

• 計畫連結： https://github.com/bytedance/decoupleQ

• W2 算符：https://github.com/NVIDIA/TensorRT-LLM/pull/1568

#1. 背景

大規模的快速發展，使得推理成本越來越高。模型量化，作為一個降低推理成本的技術方案，得到了越來越多的關注與研究。然而，在傳統的量化範式下，模型的精確度在極低位元下會迅速下降。基於此，作者們提出了一種新的量化思路，將模型參數解耦為整數部分和浮點部分，從數學優化的角度來對量化任務建模，使得在極低比特下，模型依然能保持較高的精度。這樣做的優勢是明顯的，我們不再需要關注量化特有的問題，例如如何處理敏感通道，如何處理outlier等，而是只需要將量化問題進行數學建模，找到一個合適的最佳化目標函數，然後去求解該函數。

2. 傳統量化

傳統上，我們對一個模型的量化想法是：

其中，是量化前模型的浮點型weights；s 和z 是一個線性變換係數，表示scale 和zero point；α 和β 是整數表示範圍的上下界，例如對於int4 量化，可以取α = - 8, β = 7;表示取整函數，一般是四捨五入到臨近的整數。

關於s 和z 的取值，一般來說，對於非對稱量化，我們可以取：

這樣就將一個分佈在的浮點的weights，線性對應到的區間範圍。

在反量化的時候，一般使用以下公式：

#在這種傳統的量化方案中，我們需要關注許多量化特有的細枝末節問題，例如針對敏感通道，我們有敏感通道的處理方式；針對outlier，我們有outlier 的處理方式。這種頭痛醫頭腳痛醫腳的處理範式，難以應付複雜多變的業務場景。字節跳動的研究者試圖把這些問題都抽像出來，從一個宏觀的角度來看量化問題。我們只需要建立一個抽象的最佳化目標函數，再去求解這個目標函數即可。

3. decoupleQ

#

觀察 (1)~(3) 式在量化中的作用，如果我們換個思路，就會發現其實我們不需要知道 (1) 式和 (2) 式。我們將一個大模型量化完畢以後，交付給下游引擎同學的時候，只需要知道 (3) 式中的和 (s,z) 即可。也就是說，(3) 式中的 (s,z) 完全可以視為一個普通的仿射變換的係數，不需要保留其在 (2) 式中的意義。此仿射變換係數，可以透過數學優化的方法求出來。

進一步挖掘 (3) 式中，我們可以將一個大模型的參數，解耦 (decouple) 成整數部分和浮點部分 (s,z)。這樣解耦以後，模型量化的過程，就可以視為一個求解模型整數部分和浮點部分 (s,z) 的過程。我們可以交替優化求解。為此，要確定最佳化目標函數和其約束條件。

對於一個linear 層，我們可以建立如下優化目標函數：

其中，是該層的輸入，是一個對稱矩陣（若X 的列都不全為零，則H 是正定對稱矩陣）。

一般來說，為了提高量化精度，我們可以對模型的權重用 per-channel 量化。在 per-channel 量化中，在最佳化 (4) 式的時候，的各個欄位是獨立進行最佳化的。所以我們只需要關注其中一列。

至此，優化目標可以寫如下：（為了符號簡潔，文章中重新定義符號）：

其中最佳化目標函數為

其中，w 是的某一列，b 是中對應的列。其他符號的定義與前文相同。

事實上，最佳化目標函數 (6) 與 (4) 是完全一致的，就是反量化過程。

將一個量化問題轉化為形如(5)式的數學最佳化問題，正是 decoupleQ 區別於傳統量化論文的關鍵。 這種轉化，使得我們只需要專注於求解 (5) 式即可，不再需要處理關於量化本身的細枝末節的問題，例如 outlier，等等。

式子 (5) 的解法並不容易，因為有關於的限制條件的存在，特別是這個非凸限制。在論文中，作者給出一種交替求解的思路，即獲得關於 (s,z) 和 w 的良好的初始化以後，交替迭代求解 (s,z) 和 w。在解 (s,z) 的時候，注意到式子 (5) 關於 (s,z) 是一個無限制的二次型，可以直接將目標函數求導，令導數為零，從而得到解析解。在求解 w 的時候，作者採取了兩個層面的近似，第一層近似收斂性更高，但是求解慢；第二層近似採樣 GPTQ【1】的思想，收斂性稍差，但是求解更快。

為了進一步提升量化模型精度，作者指出，除了在layer 層面做mse 最小化以外，還可以在block 層面做mse 最小化，即：

在這一步，作者在一個transformer block 層面，將其中的各個linear 層量化完畢以後，固定他們的整數部分，而去微調浮點部分(s,z) 以及layer norm 的相關參數。實驗表面，這一步微調，能夠進一步提升模型精度。

4. W2 算子實現

要對量化後的模型進行推理，需要量化算子的支持，在業界沒有現成的w2a16 的算子可用，作者們基於Tensorrt-LLM 中的w4 算子開發了w2 的Gemm cuda kernel, 實現了w2a16 模型的高效推理。

量化模型本身是以 2bit weight 的形式載入和儲存在顯存中，因此會佔用比較小的顯存。我們的 cuda kernel 在執行時間將 2bit 的 weight 載入到暫存器中，再利用硬體指令高效轉換成 bf16 的形式與 activation 進行 gemm 運算。因為我們的場景受限於 latency, generation 階段的 batchsize 比較小，此時矩陣乘受限於 weight 的訪存，這種實作會大幅減少訪存量，提升模型的效能。在實作過程中，結合了演算法搜尋以及 SpiltK Parallel Reduce，進一步提升模型的效能，實測在 batchsize=1 的情況下，在 L 卡上 w2a16 Gemm 效能相比 w4a16 能提升 1.4x-1.7x 不等。

算術連結：https://github.com/NVIDIA/TensorRT-LLM/pull/1568

           ^{kernel的實作原理}

5. 實驗

#作者在文章中給出了位元組跳動內部的ASR 實驗結果，和開源的實驗比較結果：

其中內部實驗結果為：

在該表格中，作者以 word err rate (WER) 來衡量 ASR 的準確率。作者嘗試使用不同的方法將模型量化為 W2A16g64。量化前的浮點模型的wer 是6.68%，使用GPTQ【1】量化以後是6.83%，帶有block 最小化的decoupleQ 量化以後的wer 是6.70%，該結果與量化前的浮點模型的wer 很接近。同時也 report 了量化所需的耗時。量化高精度的代價，是量化耗時較長。在實際業務中，在使用 decoupleQ 對模型量化完畢以後，固定整數部分，使用有標籤資料集微調 scale 和 zero，模模型精度有進一步的提升。

開源比較實驗結果是：

此表格是 decoupleQ 和其他方法在 Llama-1/2 上的量化結果比較。以 perplexity (PPL) 作為評估指標。可以看出，在同樣的量化配置下，deoucpleQ 的 PPL 在絕大多數時候會低於其他方法。

6. 業務效益

decoupleQ 量化技術在位元組跳動語音部門現在被廣泛使用。已經上線於語音生成模型（Text-to-Speech），語音辨識模型（automic speech recognition）等等，落地於豆包、飛書、抖音等產品。大量上線業務表明，基於decoupleQ 的量化，W4A16 的推理精度已經完全能和fp16/bf16 推理持平；W2A16 的精度只略差於fp16/bf16 精度（對浮點部分sft 以後，精度能和fp16/bf16 持平）。儘管論文中只介紹了 weight-only 的量化，但是在實際業務中，在 weight 獲得良好的量化以後，對 activation 的量化也便能簡單許多。

在硬體加速上相比fp16、w8fp16、w4fp16 獲得了不錯的加速效果，在小batch 下w2 矩陣乘的性能相比fp16 提升5-6 倍，相比w4 提升1.5-1.7 倍。在內部業務模型上，w2fp16 相比fp16 性能有3-5 倍的提升， 相比w4fp16 性能有1.25-1.4 倍的性能提升，同時也會使得模型weight 佔用顯存大幅下降，為runtime 的顯存利用提供更多空間。

7. 總結與討論

在總結與討論部分，作者也指出了decoupleQ 目前存在的兩個風險：

1. decoupleQ 旨在使用數學最佳化的手段，將量化前後的L2 loss 最小化。然而，layer 層面或block 層面的L2 loss 最小化，未必能代表最終模型精度的最優；

2. 在(5) 式和(7) 式的最佳化過程中，在求解和(s,z) 的時候，只針對一小部分校準資料來求解，這使得decoupleQ 容易對校準資料過擬合。

儘管如此，作者也指出，將模型參數解耦為整數部分和浮點部分的思路，是很有意義的。如果有標籤資料集存在，我們可以量化完畢以後，固定整數部分，使用有標籤資料集來專門訓練 (s,z)，以進一步提升模型的精確度。這樣既保證了模型的泛化表現（源自於固定整數部分），又能在特定子任務上發揮能力（源自於微調浮點部分）。在位元組跳動的實際業務中，在上一版的模型量化完畢上線以後，下一版更新的時候，也可以只訓練模型中的浮點部分。

^{參考文獻：}

^{#【1】Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Optq : Accurate quantization for generative pretrained transformers. In The Eleventh International Conference on Learning Representations, 2022.}

^{#【2】Wenqi Shao, Mengzhao Chenen, Zengang, PZhangqi, Peng Xu, Lirui Zhao, Zhiqian Li, Kaipeng Zhang, Peng Gao, Yu Qiao, and Ping Luo. Omniquant: Omnidirectionally calibrated quantization for large language models.  arXiv preprint arXiv:2308.13137,##37,##37,}#【3】Ji Lin, Jiaming Tang, Haotian Tang, Shang Yang, Xingyu Dang, and Song Han. Awq: Activation-aware weight quantization for llm compression and acceleration. arXiv preprint arXiv:2306.00978.

##

以上是位元組開源大模型量化新思路，2-bit量化模型精度齊平fp16的詳細內容。更多資訊請關注PHP中文網其他相關文章！