大語言模型權重、激活的全方位低bit可微量化，已集成進商用APP

機器之心專欄

作者：邵文琪

模型量化是模型壓縮與加速中的一項關鍵技術，其將模型權重與激活值量化至低 bit，以允許模型佔用更少的內存開銷並加快推理速度。對於具有海量參數的大語言模型而言，模型量化顯得更加重要。例如，GPT-3 模型的 175B 參數當使用 FP16 格式加載時，需消耗 350GB 的內存，需要至少 5 張 80GB 的 A100 GPU。

但若是可以將 GPT-3 模型的權重壓縮至 3bit，則可以實現單張 A100-80GB 完成所有模型權重的加載。

現有的大語言模型後訓練量化算法依賴於手工制定量化參數，優於缺乏相應的優化過程，導致面對低 bit 量化時，現有的方法都表現出顯著的性能下降。儘管量化感知訓練在確定最佳量化配置方面是有效的，但它需要引入大量額外的訓練開銷和訓練數據。尤其是大語言模型本身的計算量進一步阻礙了量化感知訓練在大預言模型量化上的應用。

這引出一個問題：我們能否在保持後訓練量化的時間和數據效率的同時，達到量化感知訓練的性能？

爲了解決大語言模型後訓練量化中的量化參數優化問題，來自上海人工智能實驗室、香港大學、香港中文大學的研究者們提出了《OmniQuant: Omnidirectionally Calibrated Quantization for Large Language Models》。該算法同時支持大語言模型中的權重與激活值的量化，且覆蓋多種量化 bit 位設置。

arXiv 論文地址：https://arxiv.org/abs/2308.13137

OpenReview 論文地址：https://openreview.net/forum?id=8Wuvhh0LYW

代碼地址：https://github.com/OpenGVLab/OmniQuant

框架方法

如上圖所示，OmniQuant 是一種針對大語言模型（LLM）的可微分量化技術，同時支持僅權重量化和權重激活值同時量化。並且，其在實現高性能量化模型的同時，保持了後訓練量化的訓練時間高效性和數據高效性。例如，OmniQuant 可在單卡 A100-40GB 上，在 1-16 小時內完成對 LLaMA-7B ~ LLaMA70B 模型量化參數的更新。

爲了達到這個目標，OmniQuant 採用了一個 Block-wise 量化誤差最小化框架。同時，OmniQuant 設計了兩種新穎的策略來增加可學習的量化參數，包括可學習的權重裁剪（Learnable Weight Clipping，LWC），以減輕量化權重的難度，以及一個可學習的等價轉換（Learnable Equivalent Transformation, LET），進一步將量化的挑戰從激活值轉移到權重。

此外，OmniQuant 引入的所有可學習參數在量化完成後可以被融合消除，量化模型可以基於現有工具完成在多平臺的部署，包括 GPU、Android、IOS 等等。

Block-wise 量化誤差最小化

OmniQuant 提出了一個新的優化流程，該流程採用 Block-wise 量化誤差最小化，並且以可微分的方式優化額外的量化參數。其中，優化目標公式化如下：

其中 F 代表 LLM 中一個變換器塊的映射函數，W 和 X 分別是全精度權重和激活，和分別代表權重和激活量化器，和分別是可學習的權重裁剪（LWC）和可學習的等價轉換（LET）中的量化參數。OmniQuant 安裝 Block-wise 量化按順序量化一個 Transformer Block 中的參數，然後再移動到下一個。

可學習的權重裁剪 (LWC)

等價轉換在模型權重和激活值之間進行量級遷移。OmniQuant 採用的可學習等價轉換使得在參數優化過程中會使得模型權重的分佈隨着訓練不斷地發生改變。此前直接學習權重裁剪閾值的方法 [1,2] 只適用於權重分佈不發生劇烈改變的情況，否則會難以收斂。基於此問題，與以往方法直接學習權重裁剪閾值不同，LWC 通過以下方式優化裁剪強度：

其中 ⌊⋅⌉ 表示取整操作。N 是目標位數。和 W 分別表示量化後和全精度的權重。h 是權重的歸一化因子，z 是零點值。裁剪（clamp）操作限制量化值在 N 位整數的範圍內，即。在上式中，和分別是權重上界和下界的可學習裁剪強度。因此，在優化目標函數中。

可學習的等價轉換 (LET)

除了通過優化裁剪閾值來實現更適合量化的權重的 LWC 之外，OmniQuant 通過 LET 進一步降低激活值的量化難度。考慮到 LLM 激活值中的異常值是存在於特定通道，以前的方法如 SmoothQuant [3], Outlier Supression+[4] 通過數學上的等價轉換將量化的難度從激活值轉移到權重。

然而，手工選擇或者貪心搜索得到的等價轉換參數會限制量化模型的性能。得益於 Block-wise 量化誤差最小化的引入，OmniQuant 的 LET 可以以一種可微分的方式確定最優的等價轉換參數。受 Outlier Suppression+~\citep {outlier-plus} 的啓發，採用了通道級的縮放和通道級的移位來操縱激活分佈，爲激活值中的異常值問題提供了一個有效的解決方案。具體來說，OmniQuant 探索了線性層和注意力操作中的等價轉換。

線性層中的等價轉換：線性層接受輸入的令牌序列，其中 T 是令牌長度，是權重矩陣和偏置向量的乘積。數學上等效的線性層表達爲：

其中 Y 代表輸出，和分別是通道級的縮放和移位參數，和分別是等效激活、權重和偏置，⊘ 和 ⊙ 代表元素級的除法和乘法。通過上式等價轉換，激活值被轉換爲更易於量化的形式，代價是增加了權重的量化難度。從這個意義上說，LWC 可以提高由 LET 實現的模型量化性能，因爲它使權重更易於量化。最後，OmniQuant 對轉換後的激活和權重進行量化，如下所示

其中 Q_a 是普通的 MinMax 量化器，Q_w 是帶有可學習權重裁剪（即所提出的 LWC）的 MinMax 量化器。

注意力操作中的等價轉換：除了線性層之外，注意力操作也佔據了 LLM 的大部分計算。此外，LLM 的自迴歸推理模式需要爲每個 token 存儲鍵值（KV）緩存，這對於長序列來說導致了巨大的內存需求。因此，OmniQuant 也考慮將自主力計算中的 Q/K/V 矩陣量化爲低位。具體來說，自注意力矩陣中的可學習等效變換可以寫爲：

其中縮放因子。自注意力計算中的量化計算表達爲。在這裡 OmniQuant 也使用 MinMax 量化方案作爲來量化矩陣。所以，最終優化目標函數中的。

僞代碼

OmniQuant 的僞算法如上圖所示。注意，LWC 與 LET 引入的額外參數在模型量化完後都可以被消除，即 OmniQuant 不會給量化模型引入任何額外開銷，因此其可直接適配於現有的量化部署工具。

實驗性能

上圖顯示了 OmniQuant 在 LLaMA 模型上僅權重量化結果的實驗結果，更多 OPT 模型結果詳見原文。可以看出，OmniQuant 在各種 LLM 模型（OPT、LLaMA-1、LLaMA-2）以及多樣化的量化配置（包括 W2A16、W2A16g128、W2A16g64、W3A16、W3A16g128、W4A16 和 W4A16g128）中，始終優於以前的 LLM 僅權重量化方法。同時，這些實驗表明了 OmniQuant 的通用性，能夠適應多種量化配置。例如，儘管 AWQ [5] 在分組量化方面特別有效，但 OmniQuant 在通道級和分組級量化中均顯示出更優的性能。此外，隨着量化比特位數的減少，OmniQuant 的性能優勢變得更加明顯。

在權重和激活值都量化的設置中中，實驗主要關注點在於 W6A6 和 W4A4 量化。實驗設置中排除了 W8A8 量化，因爲與全精度模型相比，此前的 SmoothQuant 幾乎可以實現無損的 W8A8 模型量化。上圖顯示了 OmniQuant 在 LLaMA 模型上權重激活值都量化結果的實驗結果。值得注意的是，在 W4A4 量化的不同模型中，OmniQuant 顯著提高了平均準確率，增幅在 + 4.99% ∼ +11.80% 之間。特別是在 LLaMA-7B 模型中，OmniQuant 甚至以 + 6.22% 的顯著差距超越了最近的量化感知訓練方法 LLM-QAT [6]。這一改進證明了引入額外可學習參數的有效性，這比量化感知訓練所採用的全局權重調整更爲有益。

同時，使用 OmniQuant 量化的模型可以在 MLC-LLM [7] 上實現無縫部署。上圖展示了 LLaMA 系列量化模型在 NVIDIA A100-80G 上的內存需求和推理速度。

Weights Memory (WM) 代表量化權重存儲，而 Running Memory (RM) 表示推理過程中的內存，後者更高是因爲保留了某些激活值。推理速度是通過生成 512 個令牌來衡量的。顯而易見，與 16 位全精度模型相比，量化模型顯著減少了內存使用。而且，W4A16g128 和 W2A16g128 量化幾乎使推理速度翻倍。

值得注意的是，MLC-LLM [7] 也支持 OmniQuant 量化模型在其餘平臺的部署，包括 Android 手機和 IOS 手機。如上圖所示，近期的應用 Private LLM 即是利用 OmniQuant 算法來完成 LLM 在 iPhone、iPad，macOS 等多平臺的內存高效部署。

總結

OmniQuant 是一種將量化推進到到低比特格式的先進大語言模型量化算法。OmniQuant 的核心原則是保留原始的全精度權重的同時添加可學習的量化參數。它利用可學習的權重才接和等價變換來優化權重和激活值的量化兼容性。在融合梯度更新的同時，OmniQuant 保持了與現有的 PTQ 方法相當的訓練時間效率和數據效率。此外，OmniQuant 還確保了硬件兼容性，因爲其添加的可訓練參數可以被融合到原模型中不帶來任何額外開銷。

Reference

[1] Pact: Parameterized clipping activation for quantized neural networks.

[2] LSQ: Learned step size quantization.

[3] Smoothquant: Accurate and efficient post-training quantization for large language models.

[4] Outlier suppression+: Accurate quantization of large language models by equivalent and optimal shifting and scaling.

[5] Awq: Activation-aware weight quantization for llm compression and acceleration.

[6] Llm-qat: Data-free quantization aware training for large language models.

[7] MLC-LLM：https://github.com/mlc-ai/mlc-llm