只需單卡RTX 3090,低位元量化訓練就能實現LLaMA-3 8B全參微調

机器之心發表於2024-05-25

AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年,機器之心AIxiv專欄接收報導了2000多篇內容,覆蓋全球各大高校與企業的頂級實驗室,有效促進了學術交流與傳播。如果您有優秀的工作想要分享,歡迎投稿或者聯絡報導。 投稿郵箱:liyazhou@jiqizhixin.com;zhaoyunfeng@jiqizhixin.com

本文由GreenBit.AI團隊撰寫,團隊的核心成員來自德國哈索·普拉特納計算機系統工程院開源技術小組。我們致力於推動開源社群的發展,倡導可持續的機器學習理念。我們的目標是透過提供更具成本效益的解決方案,使人工智慧技術在環境和社會層面產生積極影響。

自 2010 年起,AI 技術歷經多個重大發展階段,深度學習的崛起和 AlphaGo 的標誌性勝利顯著推動了技術前進。尤其是 2022 年底推出的 ChatGPT,彰顯了大語言模型(LLM)的能力達到了前所未有的水平。自此,生成式 AI 大模型迅速進入高速發展期,並被譽為第四次工業革命的驅動力,尤其在推動智慧化和自動化技術在產業升級中有巨大潛力。這些技術正在改變我們處理資訊、進行決策和相互交流的方式,預示著將對經濟和社會各層面帶來深遠的變革。因此,AI 為我們帶來了重大機遇,然而在 AI 技術不斷進步的同時,其產業落地也面臨諸多挑戰,尤其是高昂的成本問題。例如,在商業化過程中,大模型尤其因成本過高而成為企業的一大負擔。持續的技術突破雖然令人鼓舞,但如果落地階段的成本無法控制,便難以持續資助研發並贏得廣泛信任。然而,開源大模型的興起正逐步改變這一局面,它們不僅技術開放,還透過降低使用門檻促進了技術的平等化和快速發展。例如,普通的消費級 GPU 就能夠支援 7B/8B 規模模型的全引數微調操作,可能比採用高成本閉源模型成本低幾個數量級。在這種去中心化的 AI 正規化下,開源模型的應用在保證質量的前提下,可以顯著降低邊際成本,加速技術的商業化程序。此外,觀察顯示,經過量化壓縮的較大模型在效能上往往優於同等大小的預訓練小模型,說明量化壓縮後的模型仍然保持了優秀的能力,這為採用開源模型而非自行重複預訓練提供了充分的理由。

在 AI 技術的迅猛發展中,雲端大模型不斷探索技術的極限,以實現更廣泛的應用和更強大的計算能力。然而,市場對於能夠快速落地和支撐高速成長的智慧應用有著迫切需求,這使得邊緣計算中的大模型 —— 特別是中小型模型如 7B 和 13B 的模型 —— 因其高價效比和良好的可調性而受到青睞。企業更傾向於自行微調這些模型,以確保應用的穩定執行和資料質量的持續控制。此外,透過迴流機制,從應用中收集到的資料可以用於訓練更高效的模型,這種資料的持續最佳化和使用者反饋的精細化調整成為了企業核心競爭力的一部分。儘管雲端模型在處理複雜任務時精度高,但它們面臨的幾個關鍵挑戰不容忽視:

  • 推理服務的基礎設施成本:支援 AI 推理的高效能硬體,尤其是 GPU,不僅稀缺而且價格昂貴,集中式商業運營帶來的邊際成本遞增問題成為 AI 業務從 1 到 10 必須翻越的障礙。
  • 推理延遲:在生產環境中,模型必須快速響應並返回結果,任何延遲都會直接影響使用者體驗和應用效能,這要求基礎設施必須有足夠的處理能力以滿足高效執行的需求。
  • 隱私和資料保護:特別是在涉及敏感資訊的商業應用場景中,使用第三方雲服務處理敏感資料可能會引發隱私和安全問題,這限制了雲模型的使用範圍。

考慮到這些挑戰,邊緣計算提供了一個有吸引力的替代方案。在邊緣裝置上直接執行中小模型不僅能降低資料傳輸的延遲,提高響應速度,而且有助於在本地處理敏感資料,增強資料安全和隱私保護。結合自有資料的實時反饋和迭代更新,AI 應用將更高效和個性化。

在當前的開源模型和工具生態中,儘管存在眾多創新和進步,仍面臨一系列不足之處。首先,這些模型和工具往往並未針對本地部署場景進行最佳化,導致在本地運用時常常受限於算力和記憶體資源。例如,即便是相對較小的 7B 規模模型,也可能需要高達 60GB 的 GPU 視訊記憶體 (需要價格昂貴的 H100/A100 GPU) 來進行全引數微調。此外,市場上可選的預訓練小型模型數量和規模相對有限,大模型的開發團隊往往更專注於追求模型規模的擴充套件而非最佳化較小的模型。另一方面,現有的量化技術雖然在模型推理部署中表現良好,但其主要用途是減少模型部署時的記憶體佔用。量化後的模型權重在微調過程中無法進行最佳化,這限制了開發者在資源有限的情況下使用較大模型的能力。開發者往往希望在微調過程中也能透過量化技術節省記憶體,這一需求尚未得到有效解決。

我們的出發點在於解決上述痛點,併為開源社群貢獻實質性的技術進步。基於 Neural Architecture Search (NAS) 以及相匹配的 Post-Training Quantization (PTQ) 量化方案,我們提供了超過 200 個從不同規模開源大模型序列壓縮而來的低位元量化小模型,這些模型涵蓋了從 110B 到 0.5B 的規模跨度,並優先保證精度和質量,再次重新整理了低位元量化的 SOTA 精度。同時,我們的 NAS 演算法深入考量了模型引數量化排布的硬體友好性,使得這些模型能輕易的在主流計算硬體 (如 Nvidia GPU 和 Apple silicon 晶片硬體平臺) 進行適配,極大地方便了開發者的使用。此外,我們推出了 Bitorch Engine 開源框架以及專為低位元模型訓練設計的 DiodeMix 最佳化器,開發者可以直接對低位元量化模型在量化空間進行全引數監督微調與繼續訓練,實現了訓練與推理表徵的對齊,大幅度壓縮模型開發與部署的中間環節。更短的工程鏈條將大幅度提升工程效率,加快模型與產品迭代。透過結合低位元權重訓練技術和低秩梯度技術,我們就能實現在單卡 RTX 3090 GPU 上對 LLaMA-3 8B 模型進行全引數微調(圖 1)。上述解決方案簡潔有效,不僅節省資源,而且有效地解決了量化模型精度損失的問題。我們將在下文對更多技術細節進行詳細解讀。

圖片

圖 1. 單卡 3090 實現 LLaMA-3 8B 全參微調

模型量化

大模型時代的顯著特徵之一便是模型對計算資源需求的大幅度攀升。GPTQ 與 AWQ 等權重 PTQ 壓縮方案以可擴充套件的方式驗證了大語言模型在 4-bit 表徵上的可靠性,在相比於 FP16 表徵實現 4 倍的權重空間壓縮的同時實現了較小的效能丟失,大幅度降低模型推理所需的硬體資源。與此同時,QLoRA 巧妙地將 4-bit LLM 表徵與 LoRA 技術相結合,將低位元表徵推廣至監督微調階段,並在微調結束後將 LoRA 模組與原始 FP16 模型融合,實現了低資源需求下的模型 Parameter-Efficient-Finetuning (PEFT)。這些前沿的高效工程探索為社群提供了便利的研究工具,大幅度降低了模型研究與產業應用的資源門檻,也進一步激發了學術與產業界對更低位元表徵的想象空間。

相比於 INT4,更低位元的 Round-To-Nearest (RTN) 量化如 INT2 表徵通常要求原始模型有著更平滑的連續引數空間才能保持較低的量化損失,例如,超大規模模型往往存在容量冗餘並有著更優的量化容忍度。而透過 LLM.int8 () 等工作對當前基於 transformer 架構的大語言模型分析研究,我們已經觀察到了模型推理中廣泛存在著系統性啟用湧現現象,少數通道對最終推理結果起著決定性作用。近期 layer-Importance 等多項工作則進一步觀測到了不同深度的 transformer 模組對模型容量的參與度也表現出非均勻分佈的特性。這些不同於小引數模型的分佈特性引申出了大量 Model Pruning 相關研究的同時也為低位元壓縮技術提供了研究啟示。基於這些工作的啟發,我們探索了一種搜尋與校準結合的 Two-stage LLM 低位元量化方案。

(1) 首先,我們利用 NAS 相關方法對大語言模型引數空間的量化敏感性進行搜尋與排序,利用經典的混合精度表徵實現模型引數中的最優位元位分配。為降低模型量化後的大規模硬體部署難度,我們放棄了複雜的向量量化與 INT3 表徵等設計,採用經典 Group-wise MinMax Quantizer,同時僅選擇 INT4 (group size 128) 與 INT2 (group size 64) 作為基礎量化表徵。相對簡單的 Quantizer 設計一方面降低了計算加速核心的設計複雜度與跨平臺部署難度,另一方面也對最佳化方案提出了更高的要求。為此,我們探索了 Layer-mix 與 Channel-mix 兩種排佈下的混合精度搜尋空間。其中,Channel-mix 量化由於能更好適配 transformer 架構的系統性啟用湧現現象,往往能達到更低的量化損失,而 Layer-mix 量化在具備更優的硬體友好度的同時仍然保持了極佳的模型容量。利用高效的混合精度 NAS 演算法,我們能在數小時內基於低端 GPU 如 RTX 3090 上完成對 Qwen1.5 110B 大模型的量化排布統計,並基於統計特性在數十秒內完成任意低位元量級模型的最優架構搜尋。我們觀察到,僅僅基於搜尋與重要性排序,已經可以快速構造極強的低位元模型。

(2) 搜尋得到模型的量化排布後,我們引入了一種基於離線知識蒸餾的可擴充套件 PTQ 校準演算法,以應對超低位元量化 (如 2 到 3-bit) 帶來的累積分佈漂移問題,僅需使用不超過 512 個樣本的多源校準資料集,即可在數小時內使用單張 A100 GPU 完成 0.5B-110B 大語言模型的 PTQ 校準。儘管額外的校準步驟引入了更長的壓縮時間,從經典低位元以及量化感知訓練(QAT)相關研究的經驗中我們可以瞭解到,這是構建低量化損失模型的必要條件。而隨著當前開源社群 100B + 大模型的持續湧現 (如 Command R plus、Qwen1.5 110B, LLama3 400B),如何構建高效且可擴充套件的量化壓縮方案將是 LLM 社群系統工程研究的重要組成部分,也是我們持續關注的方向。我們經驗性的證明,搜尋與校準相結合的低位元量化方案在推進低量化損耗表徵模型的同時,在開源社群模型架構適配、硬體預期管理等方面都有著顯著優勢。

效能分析

基於 Two-stage 量化壓縮方案,我們提供了超過 200 個從不同規模開源大模型序列壓縮而來的低位元量化小模型,涵蓋了最新的 Llama3、Phi-3、Qwen1.5 以及 Mistral 等系列。我們利用 EleutherAI 的 lm-evaluation-harness 庫等對低位元量化模型的真實效能與產業場景定位進行了探索。其中我們的 4-bit 量化校準方案基本實現了相對於 FP16 的 lossless 壓縮。基於混合 INT4 與 INT2 表徵實現的 sub-4 bit 量化校準方案在多項 zero-shot 評測結果表明,搜尋與少量資料校準的經典 INT2 量化表徵已經足夠維持 LLM 在語言模型閱讀理解 (BoolQ,RACE,ARC-E/C)、常識推理 (Winogr, Hellaswag, PIQA) 以及自然語言推理 (WIC, ANLI-R1, ANLI-R2, ANLI-R3) 方面的核心能力。

圖片表 1. 低位元量化模型 zero-shot 評測示例

同時我們試圖利用進一步的 few-shot 消融對比實驗探索超低位元在產業應用中的定位,有趣的現象是,INT2 表徵為主體的超低位元 (bpw: 2.2/2.5) 模型在 5-shot 幫助下即可實現推理能力的大幅度提升。這種對少量示例樣本的利用能力表明,低位元壓縮技術在構造容量有限但足夠 “聰明” 的語言模型方面已經接近價值兌現期,配合檢索增強 (RAG) 等技術適合構造更具備成本效益的模型服務。

圖片

表 2. 低位元量化模型 5-shot 評測示例

考慮到當前少樣本 PTQ 校準僅僅引入了有限的計算資源 (校準資料集 < 512),以我們開源的低位元模型作為初始化進行更充分的全引數量化訓練將進一步提升低位元模型在實際任務中的表現,我們已經為這一需求的高效實現提供了定製化開源工具。

開源工具

我們推出了三款實用的工具來輔助這些模型的使用,並計劃未來持續最佳化和擴充套件。

Bitorch Engine (BIE) 是一款前沿的神經網路計算庫,其設計理念旨在為現代 AI 研究與開發找到靈活性與效率的最佳平衡。BIE 基於 PyTorch,為低位量化的神經網路操作定製了一整套最佳化的網路元件,這些元件不僅能保持深度學習模型的高精度和準確性,同時也大幅降低了計算資源的消耗。它是實現低位元量化 LLM 全引數微調的基礎。此外,BIE 還提供了基於 CUTLASS 和 CUDA 的 kernel,支援 1-8 bit 量化感知訓練。我們還開發了專為低位元元件設計的最佳化器 DiodeMix,有效解決了量化訓練與推理表徵的對齊問題。在開發過程中,我們發現 PyTorch 原生不支援低位元張量的梯度計算,為此我們對 PyTorch 進行了少量調整,提供了支援低位元梯度計算的修改版,以方便社群利用這一功能。目前我們為 BIE 提供了基於 Conda 和 Docker 的兩種安裝方式。而完全基於 Pip 的預編譯安裝版本也會在近期提供給社群,方便開發者可以更便捷的使用。

green-bit-llm 是為 GreenBitAI low-bit LLM 專門開發的工具包。該工具包支援雲端和消費級 GPU 上的高效能推理,並與 Bitorch Engine 配合,完全相容 transformers、PEFT 和 TRL 等主流訓練 / 微調框架,支援直接使用量化 LLM 進行全引數微調和 PEFT。目前,它已經相容了多個低位元模型序列,詳見表 3。

圖片

表 3. 已支援低位元模型序列資訊

以目前最新的開源大模型 Llama-3 8b base 模型為例,我們選擇它的 2.2/2.5/3.0 bit 作為全引數量化監督微調 (Q-SFT) 物件,使用 huggingface 庫託管的 tatsu-lab/alpaca 資料集 (包含 52000 個指令微調樣本) 進行 1 個 epoch 的最小指令微調對齊訓練測試,模型完全在量化權重空間進行學習而不涉及常規的 LoRA 引數最佳化與融合等後處理步驟,訓練結束後即可直接實現高效能的量化推理部署。在這一例子中,我們選擇不更新包括 Embedding,、LayerNorm 以及 lm.head 在內的任何其他 FP16 引數,以驗證量化空間學習的有效性。傳統 LoRA 微調、Q-SFT 配合 Galore 最佳化器微調以及單純使用 Q-SFT 微調對低位元 Llama 3 8B 模型能力的影響如表 4 中所示。

圖片

表 4. Q-SFT 對量化 LLM 的 zero-shot 能力影響

相比於傳統的 LoRA 微調 + 後量化的工程組合,Q-SFT 直接在推理模型量化空間進行學習的方式大幅度簡化了大模型從開發到部署之間的工程鏈條,同時實現了更好的模型微調效果,為更具成本效益的模型擴充套件提供瞭解決方案。此外,推理模型與訓練模型的表徵對齊也為更絲滑的端側學習與應用提供了可能性,例如,在端側進行非全引數 Q-SFT 將成為更可靠的端側最佳化管線。

在我們的自研低位元模型以外,green-bit-llm 完全相容 AutoGPTQ 系列 4-bit 量化壓縮模型,這意味著 huggingface 現存的 2848 個 4-bit GPTQ 模型都可以基於 green-bit-llm 在量化引數空間進行低資源繼續學習 / 微調。作為 LLM 部署生態最受歡迎的壓縮格式之一,現有 AutoGPTQ 愛好者可以基於 green-bit-llm 在模型訓練和推理之間進行無縫切換,不需要引入新的工程節點。

Q-SFT 和 Bitorch-Engine 能在低資源環境下穩定工作的關鍵是,我們探索了專用於低位元大模型的 DiodeMix 最佳化器,以緩解 FP16 梯度與量化空間的表徵 Mismatch 影響,量化引數更新過程被巧妙的轉換為基於組間累積梯度的相對大小排序問題。更符合量化引數空間的高效定製最佳化器將是我們未來持續探索的重要方向之一。

gbx-lm 工具將 GreenBitAI 的低位元模型適配至蘋果的 MLX 框架,進而能在蘋果晶片上高效的執行大模型。目前已支援模型的載入和生成等基礎操作。此外,該工具還提供了一個示例,遵循我們提供的詳盡指南,使用者可以在蘋果電腦上迅速建立一個本地聊天演示頁面,如下圖所示。

圖片

我們非常期待與更多開發者一同推動開源社群的發展。如果你也熱衷於此,並希望與志同道合的夥伴們一起前行,請不要猶豫!我們真誠地歡迎你的加入。你可以透過社群平臺或直接傳送郵件至 team@greenbit.ai 與我們聯絡。

資源連結

  • 模型庫: https://huggingface.co/GreenBitAI
  • BIE: https://github.com/GreenBitAI/bitorch-engine
  • green-bit-llm: https://github.com/GreenBitAI/green-bit-llm
  • gbx-lm: https://github.com/GreenBitAI/gbx-lm

相關文章