只需單卡RTX 3090,低比特量化訓(xùn)練就能實(shí)現(xiàn)LLaMA-3 8B全參微調(diào)
自 2010 年起,AI 技術(shù)歷經(jīng)多個(gè)重大發(fā)展階段,深度學(xué)習(xí)的崛起和 AlphaGo 的標(biāo)志性勝利顯著推動(dòng)了技術(shù)前進(jìn)。尤其是 2022 年底推出的 ChatGPT,彰顯了大語言模型(LLM)的能力達(dá)到了前所未有的水平。自此,生成式 AI 大模型迅速進(jìn)入高速發(fā)展期,并被譽(yù)為第四次工業(yè)革命的驅(qū)動(dòng)力,尤其在推動(dòng)智能化和自動(dòng)化技術(shù)在產(chǎn)業(yè)升級(jí)中有巨大潛力。這些技術(shù)正在改變我們處理信息、進(jìn)行決策和相互交流的方式,預(yù)示著將對(duì)經(jīng)濟(jì)和社會(huì)各層面帶來深遠(yuǎn)的變革。因此,AI 為我們帶來了重大機(jī)遇,然而在 AI 技術(shù)不斷進(jìn)步的同時(shí),其產(chǎn)業(yè)落地也面臨諸多挑戰(zhàn),尤其是高昂的成本問題。例如,在商業(yè)化過程中,大模型尤其因成本過高而成為企業(yè)的一大負(fù)擔(dān)。持續(xù)的技術(shù)突破雖然令人鼓舞,但如果落地階段的成本無法控制,便難以持續(xù)資助研發(fā)并贏得廣泛信任。然而,開源大模型的興起正逐步改變這一局面,它們不僅技術(shù)開放,還通過降低使用門檻促進(jìn)了技術(shù)的平等化和快速發(fā)展。例如,普通的消費(fèi)級(jí) GPU 就能夠支持 7B/8B 規(guī)模模型的全參數(shù)微調(diào)操作,可能比采用高成本閉源模型成本低幾個(gè)數(shù)量級(jí)。在這種去中心化的 AI 范式下,開源模型的應(yīng)用在保證質(zhì)量的前提下,可以顯著降低邊際成本,加速技術(shù)的商業(yè)化進(jìn)程。此外,觀察顯示,經(jīng)過量化壓縮的較大模型在性能上往往優(yōu)于同等大小的預(yù)訓(xùn)練小模型,說明量化壓縮后的模型仍然保持了優(yōu)秀的能力,這為采用開源模型而非自行重復(fù)預(yù)訓(xùn)練提供了充分的理由。
在 AI 技術(shù)的迅猛發(fā)展中,云端大模型不斷探索技術(shù)的極限,以實(shí)現(xiàn)更廣泛的應(yīng)用和更強(qiáng)大的計(jì)算能力。然而,市場對(duì)于能夠快速落地和支撐高速成長的智能應(yīng)用有著迫切需求,這使得邊緣計(jì)算中的大模型 —— 特別是中小型模型如 7B 和 13B 的模型 —— 因其高性價(jià)比和良好的可調(diào)性而受到青睞。企業(yè)更傾向于自行微調(diào)這些模型,以確保應(yīng)用的穩(wěn)定運(yùn)行和數(shù)據(jù)質(zhì)量的持續(xù)控制。此外,通過回流機(jī)制,從應(yīng)用中收集到的數(shù)據(jù)可以用于訓(xùn)練更高效的模型,這種數(shù)據(jù)的持續(xù)優(yōu)化和用戶反饋的精細(xì)化調(diào)整成為了企業(yè)核心競爭力的一部分。盡管云端模型在處理復(fù)雜任務(wù)時(shí)精度高,但它們面臨的幾個(gè)關(guān)鍵挑戰(zhàn)不容忽視:
- 推理服務(wù)的基礎(chǔ)設(shè)施成本:支持 AI 推理的高性能硬件,尤其是 GPU,不僅稀缺而且價(jià)格昂貴,集中式商業(yè)運(yùn)營帶來的邊際成本遞增問題成為 AI 業(yè)務(wù)從 1 到 10 必須翻越的障礙。
- 推理延遲:在生產(chǎn)環(huán)境中,模型必須快速響應(yīng)并返回結(jié)果,任何延遲都會(huì)直接影響用戶體驗(yàn)和應(yīng)用性能,這要求基礎(chǔ)設(shè)施必須有足夠的處理能力以滿足高效運(yùn)行的需求。
- 隱私和數(shù)據(jù)保護(hù):特別是在涉及敏感信息的商業(yè)應(yīng)用場景中,使用第三方云服務(wù)處理敏感數(shù)據(jù)可能會(huì)引發(fā)隱私和安全問題,這限制了云模型的使用范圍。
考慮到這些挑戰(zhàn),邊緣計(jì)算提供了一個(gè)有吸引力的替代方案。在邊緣設(shè)備上直接運(yùn)行中小模型不僅能降低數(shù)據(jù)傳輸?shù)难舆t,提高響應(yīng)速度,而且有助于在本地處理敏感數(shù)據(jù),增強(qiáng)數(shù)據(jù)安全和隱私保護(hù)。結(jié)合自有數(shù)據(jù)的實(shí)時(shí)反饋和迭代更新,AI 應(yīng)用將更高效和個(gè)性化。
在當(dāng)前的開源模型和工具生態(tài)中,盡管存在眾多創(chuàng)新和進(jìn)步,仍面臨一系列不足之處。首先,這些模型和工具往往并未針對(duì)本地部署場景進(jìn)行優(yōu)化,導(dǎo)致在本地運(yùn)用時(shí)常常受限于算力和內(nèi)存資源。例如,即便是相對(duì)較小的 7B 規(guī)模模型,也可能需要高達(dá) 60GB 的 GPU 顯存 (需要價(jià)格昂貴的 H100/A100 GPU) 來進(jìn)行全參數(shù)微調(diào)。此外,市場上可選的預(yù)訓(xùn)練小型模型數(shù)量和規(guī)模相對(duì)有限,大模型的開發(fā)團(tuán)隊(duì)往往更專注于追求模型規(guī)模的擴(kuò)展而非優(yōu)化較小的模型。另一方面,現(xiàn)有的量化技術(shù)雖然在模型推理部署中表現(xiàn)良好,但其主要用途是減少模型部署時(shí)的內(nèi)存占用。量化后的模型權(quán)重在微調(diào)過程中無法進(jìn)行優(yōu)化,這限制了開發(fā)者在資源有限的情況下使用較大模型的能力。開發(fā)者往往希望在微調(diào)過程中也能通過量化技術(shù)節(jié)省內(nèi)存,這一需求尚未得到有效解決。
我們的出發(fā)點(diǎn)在于解決上述痛點(diǎn),并為開源社區(qū)貢獻(xiàn)實(shí)質(zhì)性的技術(shù)進(jìn)步?;?Neural Architecture Search (NAS) 以及相匹配的 Post-Training Quantization (PTQ) 量化方案,我們提供了超過 200 個(gè)從不同規(guī)模開源大模型序列壓縮而來的低比特量化小模型,這些模型涵蓋了從 110B 到 0.5B 的規(guī)??缍?,并優(yōu)先保證精度和質(zhì)量,再次刷新了低比特量化的 SOTA 精度。同時(shí),我們的 NAS 算法深入考量了模型參數(shù)量化排布的硬件友好性,使得這些模型能輕易的在主流計(jì)算硬件 (如 Nvidia GPU 和 Apple silicon 芯片硬件平臺(tái)) 進(jìn)行適配,極大地方便了開發(fā)者的使用。此外,我們推出了 Bitorch Engine 開源框架以及專為低比特模型訓(xùn)練設(shè)計(jì)的 DiodeMix 優(yōu)化器,開發(fā)者可以直接對(duì)低比特量化模型在量化空間進(jìn)行全參數(shù)監(jiān)督微調(diào)與繼續(xù)訓(xùn)練,實(shí)現(xiàn)了訓(xùn)練與推理表征的對(duì)齊,大幅度壓縮模型開發(fā)與部署的中間環(huán)節(jié)。更短的工程鏈條將大幅度提升工程效率,加快模型與產(chǎn)品迭代。通過結(jié)合低比特權(quán)重訓(xùn)練技術(shù)和低秩梯度技術(shù),我們就能實(shí)現(xiàn)在單卡 RTX 3090 GPU 上對(duì) LLaMA-3 8B 模型進(jìn)行全參數(shù)微調(diào)(圖 1)。上述解決方案簡潔有效,不僅節(jié)省資源,而且有效地解決了量化模型精度損失的問題。我們將在下文對(duì)更多技術(shù)細(xì)節(jié)進(jìn)行詳細(xì)解讀。
圖 1. 單卡 3090 實(shí)現(xiàn) LLaMA-3 8B 全參微調(diào)
模型量化
大模型時(shí)代的顯著特征之一便是模型對(duì)計(jì)算資源需求的大幅度攀升。GPTQ 與 AWQ 等權(quán)重 PTQ 壓縮方案以可擴(kuò)展的方式驗(yàn)證了大語言模型在 4-bit 表征上的可靠性,在相比于 FP16 表征實(shí)現(xiàn) 4 倍的權(quán)重空間壓縮的同時(shí)實(shí)現(xiàn)了較小的性能丟失,大幅度降低模型推理所需的硬件資源。與此同時(shí),QLoRA 巧妙地將 4-bit LLM 表征與 LoRA 技術(shù)相結(jié)合,將低比特表征推廣至監(jiān)督微調(diào)階段,并在微調(diào)結(jié)束后將 LoRA 模塊與原始 FP16 模型融合,實(shí)現(xiàn)了低資源需求下的模型 Parameter-Efficient-Finetuning (PEFT)。這些前沿的高效工程探索為社區(qū)提供了便利的研究工具,大幅度降低了模型研究與產(chǎn)業(yè)應(yīng)用的資源門檻,也進(jìn)一步激發(fā)了學(xué)術(shù)與產(chǎn)業(yè)界對(duì)更低比特表征的想象空間。
相比于 INT4,更低比特的 Round-To-Nearest (RTN) 量化如 INT2 表征通常要求原始模型有著更平滑的連續(xù)參數(shù)空間才能保持較低的量化損失,例如,超大規(guī)模模型往往存在容量冗余并有著更優(yōu)的量化容忍度。而通過 LLM.int8 () 等工作對(duì)當(dāng)前基于 transformer 架構(gòu)的大語言模型分析研究,我們已經(jīng)觀察到了模型推理中廣泛存在著系統(tǒng)性激活涌現(xiàn)現(xiàn)象,少數(shù)通道對(duì)最終推理結(jié)果起著決定性作用。近期 layer-Importance 等多項(xiàng)工作則進(jìn)一步觀測到了不同深度的 transformer 模塊對(duì)模型容量的參與度也表現(xiàn)出非均勻分布的特性。這些不同于小參數(shù)模型的分布特性引申出了大量 Model Pruning 相關(guān)研究的同時(shí)也為低比特壓縮技術(shù)提供了研究啟示?;谶@些工作的啟發(fā),我們探索了一種搜索與校準(zhǔn)結(jié)合的 Two-stage LLM 低比特量化方案。
(1) 首先,我們利用 NAS 相關(guān)方法對(duì)大語言模型參數(shù)空間的量化敏感性進(jìn)行搜索與排序,利用經(jīng)典的混合精度表征實(shí)現(xiàn)模型參數(shù)中的最優(yōu)比特位分配。為降低模型量化后的大規(guī)模硬件部署難度,我們放棄了復(fù)雜的矢量量化與 INT3 表征等設(shè)計(jì),采用經(jīng)典 Group-wise MinMax Quantizer,同時(shí)僅選擇 INT4 (group size 128) 與 INT2 (group size 64) 作為基礎(chǔ)量化表征。相對(duì)簡單的 Quantizer 設(shè)計(jì)一方面降低了計(jì)算加速內(nèi)核的設(shè)計(jì)復(fù)雜度與跨平臺(tái)部署難度,另一方面也對(duì)優(yōu)化方案提出了更高的要求。為此,我們探索了 Layer-mix 與 Channel-mix 兩種排布下的混合精度搜索空間。其中,Channel-mix 量化由于能更好適配 transformer 架構(gòu)的系統(tǒng)性激活涌現(xiàn)現(xiàn)象,往往能達(dá)到更低的量化損失,而 Layer-mix 量化在具備更優(yōu)的硬件友好度的同時(shí)仍然保持了極佳的模型容量。利用高效的混合精度 NAS 算法,我們能在數(shù)小時(shí)內(nèi)基于低端 GPU 如 RTX 3090 上完成對(duì) Qwen1.5 110B 大模型的量化排布統(tǒng)計(jì),并基于統(tǒng)計(jì)特性在數(shù)十秒內(nèi)完成任意低比特量級(jí)模型的最優(yōu)架構(gòu)搜索。我們觀察到,僅僅基于搜索與重要性排序,已經(jīng)可以快速構(gòu)造極強(qiáng)的低比特模型。
(2) 搜索得到模型的量化排布后,我們引入了一種基于離線知識(shí)蒸餾的可擴(kuò)展 PTQ 校準(zhǔn)算法,以應(yīng)對(duì)超低比特量化 (如 2 到 3-bit) 帶來的累積分布漂移問題,僅需使用不超過 512 個(gè)樣本的多源校準(zhǔn)數(shù)據(jù)集,即可在數(shù)小時(shí)內(nèi)使用單張 A100 GPU 完成 0.5B-110B 大語言模型的 PTQ 校準(zhǔn)。盡管額外的校準(zhǔn)步驟引入了更長的壓縮時(shí)間,從經(jīng)典低比特以及量化感知訓(xùn)練(QAT)相關(guān)研究的經(jīng)驗(yàn)中我們可以了解到,這是構(gòu)建低量化損失模型的必要條件。而隨著當(dāng)前開源社區(qū) 100B + 大模型的持續(xù)涌現(xiàn) (如 Command R plus、Qwen1.5 110B, LLama3 400B),如何構(gòu)建高效且可擴(kuò)展的量化壓縮方案將是 LLM 社區(qū)系統(tǒng)工程研究的重要組成部分,也是我們持續(xù)關(guān)注的方向。我們經(jīng)驗(yàn)性的證明,搜索與校準(zhǔn)相結(jié)合的低比特量化方案在推進(jìn)低量化損耗表征模型的同時(shí),在開源社區(qū)模型架構(gòu)適配、硬件預(yù)期管理等方面都有著顯著優(yōu)勢(shì)。
性能分析
基于 Two-stage 量化壓縮方案,我們提供了超過 200 個(gè)從不同規(guī)模開源大模型序列壓縮而來的低比特量化小模型,涵蓋了最新的 Llama3、Phi-3、Qwen1.5 以及 Mistral 等系列。我們利用 EleutherAI 的 lm-evaluation-harness 庫等對(duì)低比特量化模型的真實(shí)性能與產(chǎn)業(yè)場景定位進(jìn)行了探索。其中我們的 4-bit 量化校準(zhǔn)方案基本實(shí)現(xiàn)了相對(duì)于 FP16 的 lossless 壓縮。基于混合 INT4 與 INT2 表征實(shí)現(xiàn)的 sub-4 bit 量化校準(zhǔn)方案在多項(xiàng) zero-shot 評(píng)測結(jié)果表明,搜索與少量數(shù)據(jù)校準(zhǔn)的經(jīng)典 INT2 量化表征已經(jīng)足夠維持 LLM 在語言模型閱讀理解 (BoolQ,RACE,ARC-E/C)、常識(shí)推理 (Winogr, Hellaswag, PIQA) 以及自然語言推理 (WIC, ANLI-R1, ANLI-R2, ANLI-R3) 方面的核心能力。
表 1. 低比特量化模型 zero-shot 評(píng)測示例
同時(shí)我們?cè)噲D利用進(jìn)一步的 few-shot 消融對(duì)比實(shí)驗(yàn)探索超低比特在產(chǎn)業(yè)應(yīng)用中的定位,有趣的現(xiàn)象是,INT2 表征為主體的超低比特 (bpw: 2.2/2.5) 模型在 5-shot 幫助下即可實(shí)現(xiàn)推理能力的大幅度提升。這種對(duì)少量示例樣本的利用能力表明,低比特壓縮技術(shù)在構(gòu)造容量有限但足夠 “聰明” 的語言模型方面已經(jīng)接近價(jià)值兌現(xiàn)期,配合檢索增強(qiáng) (RAG) 等技術(shù)適合構(gòu)造更具備成本效益的模型服務(wù)。
表 2. 低比特量化模型 5-shot 評(píng)測示例
考慮到當(dāng)前少樣本 PTQ 校準(zhǔn)僅僅引入了有限的計(jì)算資源 (校準(zhǔn)數(shù)據(jù)集 < 512),以我們開源的低比特模型作為初始化進(jìn)行更充分的全參數(shù)量化訓(xùn)練將進(jìn)一步提升低比特模型在實(shí)際任務(wù)中的表現(xiàn),我們已經(jīng)為這一需求的高效實(shí)現(xiàn)提供了定制化開源工具。
開源工具
我們推出了三款實(shí)用的工具來輔助這些模型的使用,并計(jì)劃未來持續(xù)優(yōu)化和擴(kuò)展。
Bitorch Engine (BIE) 是一款前沿的神經(jīng)網(wǎng)絡(luò)計(jì)算庫,其設(shè)計(jì)理念旨在為現(xiàn)代 AI 研究與開發(fā)找到靈活性與效率的最佳平衡。BIE 基于 PyTorch,為低位量化的神經(jīng)網(wǎng)絡(luò)操作定制了一整套優(yōu)化的網(wǎng)絡(luò)組件,這些組件不僅能保持深度學(xué)習(xí)模型的高精度和準(zhǔn)確性,同時(shí)也大幅降低了計(jì)算資源的消耗。它是實(shí)現(xiàn)低比特量化 LLM 全參數(shù)微調(diào)的基礎(chǔ)。此外,BIE 還提供了基于 CUTLASS 和 CUDA 的 kernel,支持 1-8 bit 量化感知訓(xùn)練。我們還開發(fā)了專為低比特組件設(shè)計(jì)的優(yōu)化器 DiodeMix,有效解決了量化訓(xùn)練與推理表征的對(duì)齊問題。在開發(fā)過程中,我們發(fā)現(xiàn) PyTorch 原生不支持低比特張量的梯度計(jì)算,為此我們對(duì) PyTorch 進(jìn)行了少量調(diào)整,提供了支持低比特梯度計(jì)算的修改版,以方便社區(qū)利用這一功能。目前我們?yōu)?BIE 提供了基于 Conda 和 Docker 的兩種安裝方式。而完全基于 Pip 的預(yù)編譯安裝版本也會(huì)在近期提供給社區(qū),方便開發(fā)者可以更便捷的使用。
green-bit-llm 是為 GreenBitAI low-bit LLM 專門開發(fā)的工具包。該工具包支持云端和消費(fèi)級(jí) GPU 上的高性能推理,并與 Bitorch Engine 配合,完全兼容 transformers、PEFT 和 TRL 等主流訓(xùn)練 / 微調(diào)框架,支持直接使用量化 LLM 進(jìn)行全參數(shù)微調(diào)和 PEFT。目前,它已經(jīng)兼容了多個(gè)低比特模型序列,詳見表 3。
表 3. 已支持低比特模型序列信息
以目前最新的開源大模型 Llama-3 8b base 模型為例,我們選擇它的 2.2/2.5/3.0 bit 作為全參數(shù)量化監(jiān)督微調(diào) (Q-SFT) 對(duì)象,使用 huggingface 庫托管的 tatsu-lab/alpaca 數(shù)據(jù)集 (包含 52000 個(gè)指令微調(diào)樣本) 進(jìn)行 1 個(gè) epoch 的最小指令微調(diào)對(duì)齊訓(xùn)練測試,模型完全在量化權(quán)重空間進(jìn)行學(xué)習(xí)而不涉及常規(guī)的 LoRA 參數(shù)優(yōu)化與融合等后處理步驟,訓(xùn)練結(jié)束后即可直接實(shí)現(xiàn)高性能的量化推理部署。在這一例子中,我們選擇不更新包括 Embedding,、LayerNorm 以及 lm.head 在內(nèi)的任何其他 FP16 參數(shù),以驗(yàn)證量化空間學(xué)習(xí)的有效性。傳統(tǒng) LoRA 微調(diào)、Q-SFT 配合 Galore 優(yōu)化器微調(diào)以及單純使用 Q-SFT 微調(diào)對(duì)低比特 Llama 3 8B 模型能力的影響如表 4 中所示。
表 4. Q-SFT 對(duì)量化 LLM 的 zero-shot 能力影響
相比于傳統(tǒng)的 LoRA 微調(diào) + 后量化的工程組合,Q-SFT 直接在推理模型量化空間進(jìn)行學(xué)習(xí)的方式大幅度簡化了大模型從開發(fā)到部署之間的工程鏈條,同時(shí)實(shí)現(xiàn)了更好的模型微調(diào)效果,為更具成本效益的模型擴(kuò)展提供了解決方案。此外,推理模型與訓(xùn)練模型的表征對(duì)齊也為更絲滑的端側(cè)學(xué)習(xí)與應(yīng)用提供了可能性,例如,在端側(cè)進(jìn)行非全參數(shù) Q-SFT 將成為更可靠的端側(cè)優(yōu)化管線。
在我們的自研低比特模型以外,green-bit-llm 完全兼容 AutoGPTQ 系列 4-bit 量化壓縮模型,這意味著 huggingface 現(xiàn)存的 2848 個(gè) 4-bit GPTQ 模型都可以基于 green-bit-llm 在量化參數(shù)空間進(jìn)行低資源繼續(xù)學(xué)習(xí) / 微調(diào)。作為 LLM 部署生態(tài)最受歡迎的壓縮格式之一,現(xiàn)有 AutoGPTQ 愛好者可以基于 green-bit-llm 在模型訓(xùn)練和推理之間進(jìn)行無縫切換,不需要引入新的工程節(jié)點(diǎn)。
Q-SFT 和 Bitorch-Engine 能在低資源環(huán)境下穩(wěn)定工作的關(guān)鍵是,我們探索了專用于低比特大模型的 DiodeMix 優(yōu)化器,以緩解 FP16 梯度與量化空間的表征 Mismatch 影響,量化參數(shù)更新過程被巧妙的轉(zhuǎn)換為基于組間累積梯度的相對(duì)大小排序問題。更符合量化參數(shù)空間的高效定制優(yōu)化器將是我們未來持續(xù)探索的重要方向之一。
gbx-lm 工具將 GreenBitAI 的低比特模型適配至蘋果的 MLX 框架,進(jìn)而能在蘋果芯片上高效的運(yùn)行大模型。目前已支持模型的加載和生成等基礎(chǔ)操作。此外,該工具還提供了一個(gè)示例,遵循我們提供的詳盡指南,用戶可以在蘋果電腦上迅速建立一個(gè)本地聊天演示頁面,如下圖所示。
我們非常期待與更多開發(fā)者一同推動(dòng)開源社區(qū)的發(fā)展。