MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

編輯:LRS

【新智元導讀】一個簡單的量化操作,讓你執行大規模語言模型的GPU需求量直接降低一半!

大型語言模型(LLM)雖然效能強勁,但動輒幾百上千億的引數量,對計算裝置還是記憶體的需求量之大,都不是一般公司能承受得住的。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

量化(Quantization)是常見的壓縮操作,透過降低模型權重的精度(如32bit降為8bit),犧牲一部分模型的效能來換取更快的推理速度,更少的記憶體需求。

但對於超過1000億引數量的LLM來說,現有的壓縮方法都無法保持模型的準確率,也無法在硬體上高效地執行。

最近,麻省理工學院和英偉達的研究人員聯合提出了一個

通用後訓練的量化(GPQ, general-purpose post-training quantization)方案SmoothQuant

,對大型語言模型可以高效實現8-bit權重,8-bit啟用(W8A8)的量化,無需訓練也能保持模型的準確率。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

論文連結:https://arxiv。org/pdf/2211。10438。pdf

程式碼連結:https://github。com/mit-han-lab/smoothquant

由於啟用相比權重更難量化,SmoothQuant透過數學等價變換將較難量化的啟用遷移到權重上,實現了對啟用異常值(activation outliers)的平滑處理。

SmoothQuant能夠對所有LLM的各種層中權重和啟用量化到INT8,包括OPT-175B, BLOOM-176B和GLM-130B。

相比現有方法僅對權重進行靚貨,或者對啟用進行混合精度的量化,

SmoothQuant有更高的硬體效率,實現了1。56倍加速,記憶體需求僅為原始LLM的一半,並且在準確率上幾乎沒有損失

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

SmoothQuant同時具有硬體友好的設計,研究人員將SmoothQuant整合進了LLM服務框架FasterTransformer中,實現了更快的推理速度,相比FP16的精度僅需一半數量的GPU

文章的第一作者肖光烜是MIT EECS的一年級博士生,本科畢業於清華大學計算機科學與技術學院。

導師Song Han是MIT EECS的副教授,博士畢業於斯坦福大學,主要研究方向為高效深度學習,曾提出深度壓縮(deep compression)技術,可以將神經網路的尺寸降低一個數量級,而不損失準確率。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

SmoothQuant

量化(Quantization)就是把高精度的值對映到更低精度的離散值,在這篇論文中研究人員主要關注對硬體更高效的整數均勻量化(integer uniform quantization),尤其是INT8。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

量化操作可以在不同的粒度上執行,如per-tensor量化應用於整個權重矩陣,per-token量化應用於啟用中的每個token,per-channel量化應用於權重的每個輸出通道。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

透過對啟用的量化結果進行觀察,研究人員總結出了幾個模式:

1、量化比權重更難量化。

權重的分佈相對更加均勻和平坦,之前的研究結果已經證明將大型語言模型的權重降低到INT8,甚至到INT4對準確率的影響都不大。

2、異常值是啟用量化中的主要難點。

啟用中的異常值通常比正常值要高出100倍左右,導致沒有異常值通道中的量化bits/levels效率很低。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

3、異常值固定在某一通道中出現。

異常值只在很小一部分的通道中才會出現,但如果一個通道中有一個異常值,那該異常值可能會在所有的token中出現。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

給定一個token中所有通道的方差會很大(一些通道會非常大,但大部分很小),但是給定一個通道在所有token度中的方差會很小(異常值通道會很大)。

由於異常值具有持續出現和每個通道內小方差的特點,那如果對啟用執行per-channel量化,其量化誤差將會遠遠小於per-tensor量化。

透過一個簡單的實驗,其結果再次驗證了研究人員的想法,量化到INT8時,per-channel的準確率遠遠高於per-tensor和per-token量化,和FP16基線準確率相差無幾。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

研究人員透過使用一個per-channel平滑因子s來將輸入啟用進行平滑(smooth)。為了保持線性層的數學等價,還需要反向縮放權重。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

由於輸入X通常是由之前的線性操作生成的(如線性層、層norms等),所以就可以很容易地將平滑因子融合到之前層的引數offline,而且不會產生額外縮放的核心呼叫開銷。對於其他情況,比如當輸入來自殘差add時,可以向殘差分支新增一個額外的縮放。

將量化難度從啟用轉移到權重

Smooth的目標是選擇一個per-channel的平滑因子s,使該逆操作更易於量化。

為了減少量化誤差,應該增加所有通道的有效量化位元。當所有通道的最大magnitude相同時,總的有效量化位數將是最大的。

因此,一個最直接的平滑因子選擇就是輸入中每個通道的最大值,可以保證在劃分之後,所有的啟用通道都有相同的最大值,從而實現更容易的量化。

但需要注意的是,啟用的範圍是動態的,對於不同的輸入樣本是不同的。所以研究人員使用預訓練資料集中的校準樣本來估計啟用通道的規模。

由於這個公式將所有的量化困難遷移給了權重,可以發現在這種情況下,權重的量化誤差會很大,導致準確性下降很多。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

另一方面,也可以透過選擇sj = 1/ max(|Wj |),將所有的量化難度從權重推到啟用上。同樣,由於啟用量化誤差過大,模型的效能也不好。因此需要在權重和啟用之間分割量化難度,使它們都易於量化。

研究人員引入一個超引數遷移強度α,來控制要從啟用遷移到權重的難度。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

可以發現,對於大多數模型,例如OPT和BLOOM模型,α=0。5是一個很好的平衡點,可以平均分配量化難度,特別是使用相同的量化器進行權重和啟用。

該公式保證了相應通道的權重和啟用具有相似的最大值,從而共享相同的量化難度。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

對於其他一些啟用異常值比較大的模型,例如GLM-130B有30%的異常值,這對啟用量化來說比較困難,可以選擇一個較大的α(如0。75),將更多的量化難度遷移到權重上。

SmoothQuant應用於Transformer塊

線性層佔據了LLM模型的大部分引數和計算。在預設情況下,SmoothQuant對Transformer中所有線性層的輸入啟用進行比例平滑,並用W8A8對線性層進行量化,在注意力計算中啟用了BMM運算子的量化。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

在流程中,首先用INT8對線性層和注意力層中的BMM等計算量大的運算子的輸入和權重進行量化,而對其他輕量級元素的運算,如Softmax和LayerNorm,保持啟用為FP16,這樣的設計有助於平衡準確性和推理效率。

實驗部分

研究人員選擇了三個大型語言模型用來評估SmoothQuant,包括OPT, BLOOM和GLM-130B;並使用七個zero-shot任務,包括LAMBADA, HellaSwag, PIQA, WinoGrande, OpenBookQA, RTE, COPA等。

實驗結果顯示SmoothQuant可以處理非常大的LLM的量化問題,其啟用更難量化。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

SmoothQuant可以在所有評估資料集上匹配FP16的準確性,而W8A8、ZeroQuant和Outlier Suppression基線產生的結果幾乎是隨機的。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

並且SmoothQuant可以無損地量化所有超過100B引數的開放式LLMs

SmoothQuant的O1和O2級成功地保持了浮點精度,而O3級(per-tensor static)使平均精度下降了0。8%,可能是因為靜態收集的統計資料與真實評估樣本的啟用統計資料之間的差異。

儘管如此,SmoothQuant-O1可以與FP16的準確性相匹配,而SmoothQuant-O3只降低了1%的準確性,明顯優於基線。

SmoothQuant不僅對超過100B引數的非常大的LLM有效,而且對較小的LLM也有穩定的效果,SmoothQuant可以在所有規模的OPT模型上工作,並與INT8量化的FP16精度相匹配。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

為了展示整合到PyTorch和FasterTransformer中的SmoothQuant-O3的速度提升和記憶體節省,研究人員我們測量了一次生成一批4個句子的所有隱藏狀態的端到端延遲,也就是context階段的延遲,並記錄了這個過程中GPU記憶體使用的峰值。

由於Huggingface缺乏對模型並行的支援,所以研究人員只測量了SmoothQuant在單個GPU上的PyTorch實現的效能,因此選擇了OPT-6。7B、OPT-13B和OPT-30B進行評估。

在FasterTransformer庫中,SmoothQuant可以與Tensor Parallelism演算法無縫對接,因此研究人員在OPT-13B、OPT-30B、OPT-66B和OPT-175B上測試SmoothQuant的單GPU和多GPU基準。

在NVIDIA A100 80GB GPU伺服器上進行的實驗結果顯示,基於PyTorch實現的推理延遲和峰值記憶體使用上,SmoothQuant始終比FP16基線快,當序列長度為256時,在OPT-30B上獲得了1。51倍的速度提升。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

還可以看到一個趨勢,即模型越大,加速越明顯,但LLM。int8()幾乎總是比FP16基線慢,這也是由於混合精度啟用表示的巨大開銷造成的。

在記憶體方面,SmoothQuant和LLM。int8()都可以將FP16模型的記憶體用量幾乎減半,而SmoothQuant由於完全使用INT8 GEMM,所以節省的記憶體稍多。

MIT華人提出SmoothQuant量化,記憶體需求降一半,速度提升1.56倍

與FasterTransformer對OPT的FP16實現相比,SmoothQuant-O3在使用單個GPU時可以進一步降低OPT-13B和OPT-30B的執行延遲,最高可達1。56倍。

參考資料:

https://www。reddit。com/r/MachineLearning/comments/z1b2rp/r_smoothquant_accurate_and_efficient_posttraining/