AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：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
w2 cuda 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 Chen, Zhaoyang Zhang, 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, 2023
【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, 2023.