計算所王穎研究員團隊聯合華北電力大學與上海科技大學團隊,在計算機體系結構領域頂級會議ASPLOS 2025上發布并開源了名為COMET的創新量化推理框架。該框架通過系統與算法的緊密協同,在權重(W4)、激活(A4)以及KV緩存(KV4)三個關鍵維度實現全4比特量化,打破了傳統量化性能與精度的瓶頸。在LLaMA-70B等大型模型上,僅帶來了0.32的困惑度細微增加,同時實現了端到端推理速度提升2.02倍,樹立了大模型量化推理領域的新標桿。
技術亮點包括:
1. 全面量化壓縮:將權重、激活和KV緩存全部壓縮至4比特,極大程度降低存儲與計算資源占用,且精度損失控制在極低水平。
2. 充分挖掘GPU潛能:通過創新的通道重排技術和異步流水線設計,顯著提高硬件利用率,GPU算力利用率達76%。
3. 簡易集成開源方案:用戶僅需5行代碼即可將COMET框架集成至TensorRT-LLM,實現對LLaMA-3及Qwen等大模型全系列的支持。
4. 卓越實際效果:在70億參數大型模型的長序列推理場景中,最高加速可達3.27倍,顯著降低端側設備推理運行成本約70%。
當前,全球大型模型的能耗問題日益嚴峻,例如GPT-4單次訓練耗電接近數千戶家庭的年用電量,推理過程的能效瓶頸更是不容忽視。傳統量化技術面臨明顯不足:
? 權重量化(W4A16)雖壓縮權重,但激活仍采用高精度,造成GPU算力嚴重浪費;
? 權重和激活雙量化(W8A8)雖然保障了精度,但加速效果受限且在長序列推理中內存需求仍居高不下。
COMET框架通過硬件與算法的深度融合,從GPU架構底層重新設計量化策略和計算核心,徹底解決了權重、激活及KV緩存量化難題,實現了全維度4比特超低精度推理的新突破,為大規模模型的節能高效推理提供了堅實技術支撐。
1. 離群點分布規律的重新認識
經過對LLaMA等主流大模型的深入研究,團隊首次發現激活中的異常值并非隨機分布,而是集中出現在大約10%的特定通道。這一發現為細粒度混合精度量化算法(FMPQ)的開發提供了重要依據。
在實際應用中,團隊設計了一種兼顧模型精度與計算效率的創新方案:首先,通過將激活張量以128通道為單位進行分塊劃分,使其與現代GPU張量核心的64×64×32計算結構高度匹配,確保量化過程與硬件執行緊密結合。針對激活中非均勻分布的異常值,算法引入動態精度調節機制,賦予異常值密集區域8比特高精度表示,而其他部分則采用4比特量化,從而在保持模型性能的同時,大幅提升計算效率。
此外,FMPQ引入了通道重排技術,通過智能置換將分散的異常值通道聚集到同一計算塊內,進一步降低對高精度計算資源的需求,實現更優的性能表現。
2. 動態調度混合精度計算
在實際應用中,結合W4A4和W4A8的混合精度運算常常導致計算負載分布不均,成為性能提升的瓶頸。為解決這一問題,團隊設計了三層異步流水線結構,旨在充分發揮GPU的計算能力,提高整體資源利用率。
整個系統架構從數據傳輸與計算的重疊優化入手,采用雙緩沖技術,在數據預加載的同時啟動計算過程,有效隱藏了讀取延遲,保證張量核心持續高效運行。團隊創新性地打破了傳統將一個計算任務固定分配給單個流式多處理器(SM)的模式,設計了更靈活的SM協同機制,使得閑置的計算單元能夠動態接管鄰近的任務,大幅度降低了局部計算資源的浪費。
此外,COMET框架在任務調度方面也實現了細致的優化。系統根據實時計算負載,動態調整INT4和INT8任務的分配比例,從而將不同SM之間的負載差異控制在5%以內,幾乎實現了均衡的算力分配。這些設計共同提升了整體的計算效率和資源利用率。
3. 極致挖掘硬件指令集潛能
在4比特量化領域,真正的難點不在于量化算法本身,而是數據格式轉換的效率限制。傳統做法通常需要執行多達10條指令才能完成一次4比特到8比特的轉換,極大地影響了部署性能。針對這一瓶頸,研究團隊從硬件指令集入手,深度挖掘GPU的底層能力,通過重構數據存儲格式和指令執行流程,將轉換操作精簡至僅需兩條指令完成。
這一簡化方案依賴兩大關鍵創新:首先,團隊充分利用GPU的PTX指令特性,實現了同時從16位數據塊中提取4個4比特數值,顯著提升了解包速度;其次,采用了“零擴展”替代常用的“符號擴展”技術,結合縮放參數的補償機制,保證了整體精度穩定不受影響。憑借此設計,轉換效率較傳統方法提升了約5倍。
在4比特極限壓縮的背景下,COMET框架通過細粒度混合精度量化,有效降低了量化對模型精度的影響。實際測試數據顯示,LLaMA-7B模型在WikiText2數據集上的困惑度僅由5.68略微上升至5.95,誤差保持在5%以內,展現出較強的穩定性。更為突出的是,針對LLaMA-3-70B等大型模型,困惑度的提升被嚴格控制在0.36以內,精度表現較傳統方案提升了三倍以上。同時,COMET還創新性地將KV緩存壓縮至4比特,顯著減少了長序列任務中的內存需求,降低幅度達75%,且精度損失控制在極低的0.05左右,基本可忽略不計。整體來看,COMET的精度已接近FP16浮點基準,明顯優于常見的W4A4及W8A8量化方法。
在性能方面,COMET同樣實現了顯著突破。以A100 GPU為例,經過深度優化的COMET-W4Ax計算內核在適配LLaMA、Mistral、Qwen等主流模型時,其運行速度是cuBLAS內核的2.75到2.97倍。更為關鍵的是,這個性能提升不僅體現在單次計算內核中,也得以在端到端推理流程中保持。在處理長序列任務(如1024個token輸入)時,整體吞吐量提升達到了3.27倍;即使在較短序列(128個token)條件下,依然能實現1.63倍的加速。得益于異步流水線和動態調度技術的應用,GPU中流式多處理器(SM)的利用效率由原先的45%提升到了76%,有效挖掘了硬件的計算潛能。
COMET-W4Ax的內核和接口現已公開開源,開發者只需簡單幾行代碼便能將其無縫集成到TensorRT-LLM框架中,實現從模型壓縮到推理加速的全流程支持,極大推動大規模模型的普及與廣泛應用。
相關論文《COMET: Towards Practical W4A4KV4 LLMs Serving》已正式發布,可通過鏈接 https://dl.acm.org/doi/pdf/10.1145/3676641.3716252 獲取全文。論文中詳細闡述了研究工作,并附帶完整的開源代碼,歡迎大家下載閱讀并發表意見交流。
我們相信人工智能為普通人提供了一種“增強工具”,并致力于分享全方位的AI知識。在這里,您可以找到最新的AI科普文章、工具評測、提升效率的秘籍以及行業洞察。 歡迎關注“福大大架構師每日一題”,讓AI助力您的未來發展。
特別聲明:以上內容(如有圖片或視頻亦包括在內)為自媒體平臺“網易號”用戶上傳并發布,本平臺僅提供信息存儲服務。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.