模型量化是模型壓縮與加速中的一項關鍵技術，其將模型權重與啟用值量化至低 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