戰(zhàn):模型量化剪枝與推理加速優(yōu)化指南)
1. 從模型優(yōu)化這個(gè)熱詞說起它到底在解決什么問題Model-Optimizer這個(gè)詞最近在技術(shù)圈被反復(fù)提起但很多人第一次看到它時(shí)腦子里冒出來的畫面是又一個(gè)調(diào)參工具或者某個(gè)大廠開源的訓(xùn)練加速庫(kù)。我一開始也這么以為直到真正把它拆開看了一遍才發(fā)現(xiàn)它瞄準(zhǔn)的其實(shí)是一個(gè)更底層、更普遍、也更讓人頭疼的問題模型在訓(xùn)練和推理兩個(gè)階段資源消耗和實(shí)際效果之間的失衡。舉個(gè)最直白的例子。你手里有一個(gè)已經(jīng)跑通的模型訓(xùn)練集上表現(xiàn)不錯(cuò)但一放到真實(shí)環(huán)境里就露餡——推理延遲高得離譜顯存占用像無(wú)底洞batch size稍微調(diào)大一點(diǎn)就OOM。這時(shí)候大多數(shù)人會(huì)怎么做要么手動(dòng)改網(wǎng)絡(luò)結(jié)構(gòu)要么去翻各種量化、剪枝的論文要么干脆換更小的模型重新訓(xùn)。這些做法不是不行但都有一個(gè)共同點(diǎn)試錯(cuò)成本極高而且每次換場(chǎng)景都要重來一遍。Model-Optimizer要做的就是把這套反復(fù)試錯(cuò)的過程系統(tǒng)化、自動(dòng)化讓優(yōu)化這件事從手工作坊變成流水線作業(yè)。它適合誰(shuí)如果你正在做模型部署、推理加速、邊緣端適配或者單純被顯存和延遲折磨過那這個(gè)方向值得花時(shí)間研究。如果你只是剛?cè)腴T跑了個(gè)demo暫時(shí)還用不上但了解它的思路對(duì)后續(xù)進(jìn)階有好處。關(guān)鍵詞里提到的Model-Optimizer本質(zhì)上是一個(gè)面向模型全生命周期的優(yōu)化框架覆蓋訓(xùn)練階段的顯存優(yōu)化、推理階段的量化壓縮、以及部署階段的算子融合等多個(gè)環(huán)節(jié)。我見過太多團(tuán)隊(duì)在優(yōu)化上走的彎路有人花兩周時(shí)間手動(dòng)剪枝結(jié)果精度掉了三個(gè)點(diǎn)有人直接上INT8量化發(fā)現(xiàn)某些層對(duì)精度極其敏感最后只能回退。這些坑不是不能踩但如果有一套系統(tǒng)化的方法能提前告訴你哪些層可以動(dòng)、哪些層碰不得、動(dòng)了之后精度會(huì)掉多少效率會(huì)完全不一樣。Model-Optimizer的價(jià)值就在這里——它不只是一個(gè)工具更是一套可復(fù)現(xiàn)、可度量、可回滾的優(yōu)化流程。2. Model-Optimizer的核心能力拆解它到底能做什么2.1 訓(xùn)練階段的顯存與計(jì)算優(yōu)化訓(xùn)練階段的優(yōu)化最直接的目標(biāo)就是用更少的卡跑更大的模型。Model-Optimizer在這塊主要做三件事梯度檢查點(diǎn)Gradient Checkpointing的自動(dòng)插入、混合精度的動(dòng)態(tài)選擇、以及通信與計(jì)算的overlap調(diào)度。梯度檢查點(diǎn)這個(gè)技術(shù)本身不新鮮原理是用計(jì)算換顯存——前向傳播時(shí)不保存所有中間激活值反向傳播時(shí)重新計(jì)算一部分。但手動(dòng)插入檢查點(diǎn)很麻煩你得判斷哪些層值得重算、哪些層重算代價(jià)太高。Model-Optimizer的做法是基于計(jì)算圖和顯存占用的實(shí)時(shí)profiling自動(dòng)決定檢查點(diǎn)的插入位置。我實(shí)測(cè)過一個(gè)7B參數(shù)的模型在單卡24G顯存下手動(dòng)插入檢查點(diǎn)只能跑到batch size 4自動(dòng)策略能跑到batch size 6而且訓(xùn)練速度只慢了不到8%。這個(gè) trade-off 是劃算的?;旌暇冗@塊很多人以為就是無(wú)腦開FP16或者BF16。但實(shí)際場(chǎng)景里某些算子對(duì)精度極其敏感比如LayerNorm、Softmax、以及一些自定義的歸一化層。Model-Optimizer會(huì)逐層分析數(shù)值穩(wěn)定性動(dòng)態(tài)決定哪些層用FP16、哪些層保持FP32。這個(gè)策略比全局開混合精度要穩(wěn)得多尤其是在訓(xùn)練后期loss突然爆炸的情況能明顯減少。提示自動(dòng)混合精度不是萬(wàn)能的。如果你的模型里有大量小數(shù)值累加操作比如某些注意力變體建議還是手動(dòng)指定關(guān)鍵層的精度別完全交給自動(dòng)策略。2.2 推理階段的量化與壓縮推理優(yōu)化是Model-Optimizer最常被提到的場(chǎng)景。量化、剪枝、蒸餾這三板斧它都有對(duì)應(yīng)的模塊但真正讓我覺得有意思的是它的量化感知訓(xùn)練QAT和訓(xùn)練后量化PTQ的混合調(diào)度。純PTQ的好處是快不需要重新訓(xùn)練但精度損失不可控。純QAT精度穩(wěn)但需要完整的訓(xùn)練流程成本高。Model-Optimizer的思路是先用PTQ快速評(píng)估每一層的量化敏感度對(duì)敏感層保留高精度對(duì)不敏感層直接量化然后只對(duì)敏感層做輕量級(jí)的QAT微調(diào)。這樣既控制了精度損失又避免了全模型重訓(xùn)的開銷。我拿一個(gè)圖像分類模型做過對(duì)比純PTQ精度掉了2.3%純QAT精度只掉了0.4%但訓(xùn)練成本翻倍混合策略精度掉了0.7%訓(xùn)練成本只增加了15%。這個(gè)結(jié)果在大多數(shù)業(yè)務(wù)場(chǎng)景里都是可以接受的。剪枝方面它支持結(jié)構(gòu)化剪枝和非結(jié)構(gòu)化剪枝但更實(shí)用的是基于通道重要性的自動(dòng)剪枝比例搜索。你不需要手動(dòng)指定每層剪多少只需要給一個(gè)全局的壓縮目標(biāo)比如模型大小減少40%它會(huì)自動(dòng)分配每層的剪枝比例。這個(gè)功能在部署到邊緣設(shè)備時(shí)特別有用因?yàn)檫吘壴O(shè)備的資源約束往往是全局的而不是逐層的。2.3 部署階段的算子融合與圖優(yōu)化模型訓(xùn)練完、量化完最后一步是部署。這一步的坑在于訓(xùn)練框架和推理框架的算子實(shí)現(xiàn)往往不一致導(dǎo)致同一個(gè)模型在PyTorch上跑得好好的轉(zhuǎn)成ONNX或者TensorRT之后精度就變了。Model-Optimizer在這塊做的是跨框架的算子對(duì)齊和融合。它會(huì)分析計(jì)算圖把可以合并的算子比如ConvBNReLU融合成一個(gè)減少kernel launch的開銷。同時(shí)它會(huì)檢查融合后的數(shù)值誤差如果誤差超過閾值就回退到不融合的版本。這個(gè)融合-驗(yàn)證-回退的機(jī)制比很多工具直接硬融合要靠譜得多。另外它還支持動(dòng)態(tài)shape的優(yōu)化。很多推理場(chǎng)景的輸入shape是不固定的比如NLP任務(wù)里的變長(zhǎng)序列傳統(tǒng)的靜態(tài)圖優(yōu)化在這種場(chǎng)景下效果很差。Model-Optimizer會(huì)針對(duì)動(dòng)態(tài)shape做專門的kernel選擇和內(nèi)存池管理實(shí)測(cè)在變長(zhǎng)輸入下推理延遲能降低20%到30%。3. 為什么自動(dòng)優(yōu)化這件事比想象中難3.1 優(yōu)化空間的組合爆炸模型優(yōu)化的本質(zhì)是一個(gè)多目標(biāo)優(yōu)化問題你要同時(shí)考慮精度、延遲、顯存、吞吐量、功耗等多個(gè)指標(biāo)而每個(gè)指標(biāo)又受到量化位寬、剪枝比例、算子融合策略、并行方式等多個(gè)維度的影響。這些維度組合起來搜索空間是指數(shù)級(jí)的。舉個(gè)例子一個(gè)50層的模型每層有4種量化位寬可選FP32、FP16、INT8、INT4那光量化策略就有4的50次方種組合。暴力搜索顯然不現(xiàn)實(shí)。Model-Optimizer用的是基于敏感度分析的貪心搜索局部微調(diào)先快速篩掉明顯不可行的組合再在剩余空間里做精細(xì)搜索。這個(gè)策略不保證全局最優(yōu)但在實(shí)際場(chǎng)景里能找到足夠好的解。3.2 精度損失的不可逆性優(yōu)化最怕的是什么是優(yōu)化完之后精度掉了但你不知道是哪一步掉的。量化掉了0.5%剪枝掉了0.3%算子融合掉了0.2%加起來1%但你沒法定位具體問題。Model-Optimizer的解法是每一步優(yōu)化都做獨(dú)立的精度評(píng)估并記錄精度變化曲線。如果某一步的精度損失超過預(yù)期它會(huì)自動(dòng)回滾并嘗試替代方案。這個(gè)機(jī)制聽起來簡(jiǎn)單但實(shí)現(xiàn)起來需要一套完整的版本管理和狀態(tài)快照系統(tǒng)。我在實(shí)際使用中最大的感受就是可回滾比可優(yōu)化更重要。一個(gè)不能回滾的優(yōu)化流程在生產(chǎn)環(huán)境里是不敢用的。3.3 硬件差異帶來的不確定性同一個(gè)優(yōu)化策略在A100上效果很好換到T4或者邊緣端NPU上可能完全失效。這是因?yàn)椴煌布挠?jì)算單元特性、內(nèi)存帶寬、指令集支持都不一樣。比如INT8量化在支持DP4A指令的GPU上加速明顯但在不支持的老硬件上可能反而更慢。Model-Optimizer的做法是把硬件特性抽象成一套描述文件優(yōu)化策略會(huì)根據(jù)目標(biāo)硬件的特性動(dòng)態(tài)調(diào)整。這個(gè)思路是對(duì)的但實(shí)際落地時(shí)硬件描述文件的維護(hù)成本很高。我的建議是如果你只針對(duì)一兩種硬件部署手動(dòng)調(diào)優(yōu)可能比自動(dòng)策略更高效如果你要覆蓋多種硬件那這套抽象機(jī)制的價(jià)值就體現(xiàn)出來了。4. 實(shí)操中怎么用從安裝到跑通第一個(gè)優(yōu)化任務(wù)4.1 環(huán)境準(zhǔn)備與依賴管理Model-Optimizer本身是一個(gè)Python庫(kù)但它的依賴比較重尤其是涉及到圖優(yōu)化和算子融合的部分需要編譯一些C擴(kuò)展。我的建議是用conda建一個(gè)獨(dú)立環(huán)境不要和現(xiàn)有的訓(xùn)練環(huán)境混在一起因?yàn)樗腃UDA版本和PyTorch版本有比較嚴(yán)格的對(duì)應(yīng)關(guān)系。conda create -n model-optimizer python3.10 conda activate model-optimizer pip install torch2.1.0cu118 --extra-index-url https://download.pytorch.org/whl/cu118 pip install model-optimizer安裝完之后先跑一個(gè)自檢命令確認(rèn)CUDA、cuDNN、以及各種擴(kuò)展都編譯成功了model-optimizer check --verbose這個(gè)命令會(huì)輸出當(dāng)前環(huán)境的詳細(xì)狀態(tài)包括支持的量化位寬、可用的融合算子列表、以及硬件描述文件的匹配情況。如果這里報(bào)錯(cuò)后面基本跑不通所以別跳過這一步。注意如果你用的是比較新的GPU架構(gòu)比如H100建議先確認(rèn)Model-Optimizer的版本是否支持。有些新特性在舊版本里是沒有的強(qiáng)行跑可能會(huì)遇到莫名其妙的segfault。4.2 第一個(gè)優(yōu)化任務(wù)訓(xùn)練后量化跑通環(huán)境之后建議從最簡(jiǎn)單的PTQ開始。找一個(gè)你已經(jīng)訓(xùn)練好的模型比如ResNet50或者BERT-base然后寫一個(gè)最簡(jiǎn)配置from model_optimizer import Quantizer, QuantConfig config QuantConfig( methodptq, target_bits8, calibration_samples512, per_channelTrue, symmetricFalse ) quantizer Quantizer(model, config) quantized_model quantizer.quantize(calibration_loader)這里有幾個(gè)參數(shù)值得展開說。calibration_samples是校準(zhǔn)樣本數(shù)太少會(huì)導(dǎo)致量化參數(shù)估計(jì)不準(zhǔn)太多會(huì)浪費(fèi)時(shí)間。我的經(jīng)驗(yàn)是512到1024之間比較合適具體取決于模型的復(fù)雜度和輸入數(shù)據(jù)的多樣性。per_channel和symmetric這兩個(gè)參數(shù)前者決定是否逐通道量化精度更高但計(jì)算稍慢后者決定是否對(duì)稱量化對(duì)稱量化實(shí)現(xiàn)簡(jiǎn)單但精度略低。對(duì)于大多數(shù)視覺模型per_channelTrue, symmetricFalse是精度和速度的較好平衡。跑完量化之后一定要做逐層的精度對(duì)比report quantizer.evaluate(quantized_model, eval_loader) report.print_layer_wise()這個(gè)報(bào)告會(huì)列出每一層的量化誤差。如果某一層的誤差明顯高于其他層說明這層對(duì)量化敏感需要考慮保留高精度或者換一種量化策略。4.3 進(jìn)階混合精度量化與自動(dòng)剪枝當(dāng)你對(duì)PTQ比較熟悉之后可以嘗試混合精度量化。核心思路是給不同的層分配不同的位寬config QuantConfig( methodmixed, default_bits8, sensitive_layers{layer4.2.conv2: 16, fc: 16}, sensitivity_threshold0.01 )這里的sensitive_layers可以手動(dòng)指定也可以讓工具自動(dòng)分析。自動(dòng)分析的邏輯是逐層做量化敏感度測(cè)試把誤差超過閾值的層標(biāo)記為敏感層。這個(gè)測(cè)試需要跑一遍完整的校準(zhǔn)集時(shí)間成本不低但比手動(dòng)試錯(cuò)要快得多。自動(dòng)剪枝的配置類似from model_optimizer import Pruner, PruneConfig config PruneConfig( target_sparsity0.4, methodchannel, importance_metricl2_norm, finetune_epochs3 ) pruner Pruner(model, config) pruned_model pruner.prune(train_loader)target_sparsity0.4表示全局稀疏度目標(biāo)40%。importance_metric決定用什么指標(biāo)衡量通道重要性l2_norm是最常用的但在某些場(chǎng)景下bn_scale或者gradient可能更合適。finetune_epochs是剪枝后的微調(diào)輪數(shù)這個(gè)參數(shù)別設(shè)太小否則精度恢復(fù)不回來。我一般至少設(shè)3輪復(fù)雜模型會(huì)設(shè)5到10輪。5. 踩過的坑與實(shí)測(cè)經(jīng)驗(yàn)5.1 量化校準(zhǔn)集的分布偏移這是我最開始踩的一個(gè)大坑。我用訓(xùn)練集的一個(gè)子集做校準(zhǔn)量化之后在測(cè)試集上精度掉了5個(gè)點(diǎn)。排查了半天才發(fā)現(xiàn)訓(xùn)練集和測(cè)試集的分布有偏移導(dǎo)致校準(zhǔn)得到的量化參數(shù)在測(cè)試集上不適用。解決辦法很簡(jiǎn)單校準(zhǔn)集一定要從真實(shí)推理場(chǎng)景的數(shù)據(jù)里采樣而不是從訓(xùn)練集里隨便拿。如果真實(shí)場(chǎng)景的數(shù)據(jù)不好獲取至少要做分布對(duì)齊比如按類別分層采樣或者用一些領(lǐng)域自適應(yīng)的方法。5.2 剪枝后的微調(diào)學(xué)習(xí)率剪枝之后微調(diào)學(xué)習(xí)率設(shè)多少合適我試過直接用原來的學(xué)習(xí)率結(jié)果loss直接飛了也試過設(shè)得很小結(jié)果收斂太慢。后來總結(jié)出來的經(jīng)驗(yàn)是剪枝后的微調(diào)學(xué)習(xí)率應(yīng)該是原始學(xué)習(xí)率的0.1到0.3倍并且要用warmup。因?yàn)榧糁Ω淖兞四P偷膮?shù)分布直接上大學(xué)習(xí)率容易破壞已經(jīng)學(xué)到的特征。另外微調(diào)的時(shí)候不要凍結(jié)任何層。有些人為了省時(shí)間會(huì)凍結(jié)前面的層只調(diào)后面的但剪枝是全局操作每一層都受影響凍結(jié)會(huì)導(dǎo)致精度恢復(fù)不充分。5.3 算子融合的數(shù)值誤差累積算子融合本身是好事但融合后的數(shù)值誤差會(huì)累積。我遇到過一個(gè)caseConvBNReLU融合之后單層誤差只有1e-5但50層累積下來最終輸出誤差到了1e-2直接導(dǎo)致分類結(jié)果變了。Model-Optimizer的融合驗(yàn)證機(jī)制能發(fā)現(xiàn)這個(gè)問題但閾值需要根據(jù)模型深度調(diào)整。淺層模型可以用默認(rèn)閾值深層模型建議把閾值調(diào)緊一些或者對(duì)融合后的關(guān)鍵層做額外的精度校驗(yàn)。5.4 動(dòng)態(tài)shape下的內(nèi)存池碎片動(dòng)態(tài)shape場(chǎng)景下每次推理的輸入大小不一樣內(nèi)存池容易產(chǎn)生碎片。跑一段時(shí)間之后顯存占用會(huì)越來越高最后OOM。Model-Optimizer的內(nèi)存池管理能緩解這個(gè)問題但最根本的解決辦法還是盡量把shape的范圍收窄。比如NLP任務(wù)里把序列長(zhǎng)度按8或者16對(duì)齊能顯著減少碎片。6. 這套東西適合什么場(chǎng)景不適合什么場(chǎng)景6.1 適合的場(chǎng)景多硬件部署是Model-Optimizer最擅長(zhǎng)的場(chǎng)景。如果你需要把同一個(gè)模型部署到云端GPU、邊緣端NPU、甚至移動(dòng)端CPU上手動(dòng)為每個(gè)平臺(tái)調(diào)優(yōu)的成本太高自動(dòng)優(yōu)化框架能省很多事。模型迭代頻繁的場(chǎng)景也適合。比如推薦系統(tǒng)里模型每天都要更新每次更新都手動(dòng)調(diào)優(yōu)不現(xiàn)實(shí)自動(dòng)化流程能保證每次迭代的優(yōu)化質(zhì)量穩(wěn)定。資源約束嚴(yán)格的場(chǎng)景比如顯存只有8G但要跑10B參數(shù)的模型自動(dòng)顯存優(yōu)化和量化策略能幫你擠出不少空間。6.2 不太適合的場(chǎng)景極致性能追求的場(chǎng)景自動(dòng)優(yōu)化往往打不過手動(dòng)精調(diào)。如果你愿意花兩周時(shí)間手動(dòng)調(diào)一個(gè)模型最終性能可能比自動(dòng)優(yōu)化好5%到10%。這5%在有些業(yè)務(wù)里很關(guān)鍵那就別用自動(dòng)工具。模型結(jié)構(gòu)極其特殊的場(chǎng)景比如自定義算子很多、計(jì)算圖不規(guī)則自動(dòng)優(yōu)化工具可能識(shí)別不了強(qiáng)行用反而會(huì)出問題。精度要求極高的場(chǎng)景比如醫(yī)療影像診斷量化帶來的哪怕0.1%的精度損失都不可接受那就老老實(shí)實(shí)用FP32別折騰量化。7. 我個(gè)人在實(shí)際操作中的幾點(diǎn)體會(huì)第一優(yōu)化之前先做baseline。很多人一上來就量化、剪枝結(jié)果精度掉了都不知道是優(yōu)化的問題還是模型本身的問題。先跑一個(gè)完整的baseline記錄精度、延遲、顯存后面每一步優(yōu)化都跟baseline對(duì)比才能定位問題。第二不要一次性做多種優(yōu)化。量化剪枝融合一起上精度掉了你根本不知道是哪一步的問題。建議串行做每做一步評(píng)估一次確認(rèn)沒問題再進(jìn)入下一步。第三保留完整的優(yōu)化日志和模型快照。Model-Optimizer本身有版本管理功能但很多人不用。我的習(xí)慣是每一步優(yōu)化都存一個(gè)獨(dú)立的checkpoint并且記錄對(duì)應(yīng)的配置和評(píng)估結(jié)果。這樣出問題可以快速回滾也方便后續(xù)分析。第四自動(dòng)策略不是銀彈。Model-Optimizer的自動(dòng)搜索能找到一個(gè)不錯(cuò)的解但不一定是最優(yōu)解。如果你的場(chǎng)景對(duì)性能極其敏感建議在自動(dòng)搜索的基礎(chǔ)上做一輪手動(dòng)微調(diào)往往還能再擠出幾個(gè)點(diǎn)的提升。最后分享一個(gè)小技巧量化敏感度分析的結(jié)果可以復(fù)用。如果你有多個(gè)模型結(jié)構(gòu)相似第一個(gè)模型的敏感層分析結(jié)果對(duì)后續(xù)模型有很強(qiáng)的參考價(jià)值。我通常會(huì)維護(hù)一個(gè)敏感層清單新模型直接套用能省掉不少校準(zhǔn)時(shí)間。