字節開源大模型量化新思路，2-bit量化模型精度齊平fp16
隨著深度學習大語言模型的越來越火爆，大語言模型越做越大，使得其推理成本也水漲船高。模型量化，成為一個熱門的研究課題。
近日，字節跳動語音團隊推出一個全新的量化思路，拋棄傳統的量化范式，從數學優化的角度來對量化任務建模。文章放在了 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)，以進一步提升模型的精度。這樣既保證了模型的泛化性能（源于固定整數部分
），又能在特定子任務上發揮能力（源于微調浮點部分）。在字節跳動的實際業務中，在上一版的模型量化完畢上線以后，下一版更新的時候，也可以只訓練模型中的浮點部分。
本文轉自機器之心 ，作者：機器之心
