戰(zhàn))
1. 為什么要在12G顯存上折騰27B模型先說(shuō)結(jié)論12G顯存跑27B模型128K上下文decode 50 tokens/s這件事在一年前基本屬于天方夜譚但現(xiàn)在通過(guò)量化、KV Cache優(yōu)化和投機(jī)解碼三件套的組合確實(shí)能摸到這個(gè)門檻。我自己手上是一張RTX 3060 12G算是消費(fèi)級(jí)里最典型的“小顯存大胃口”配置拿它來(lái)驗(yàn)證這套方案最有說(shuō)服力。這個(gè)項(xiàng)目的核心目標(biāo)很明確在單張12G顯存的消費(fèi)級(jí)顯卡上讓一個(gè)27B參數(shù)級(jí)別的大模型跑起來(lái)支持128K的超長(zhǎng)上下文并且解碼速度維持在50 tokens/s以上。這三個(gè)指標(biāo)單獨(dú)拿出來(lái)都不算特別夸張但疊在一起就變成了一個(gè)典型的“不可能三角”——模型越大越吃顯存上下文越長(zhǎng)KV Cache越爆炸速度要快又得留足計(jì)算資源。所以整個(gè)項(xiàng)目的本質(zhì)是在顯存、上下文長(zhǎng)度和推理速度之間做一場(chǎng)精密的資源調(diào)度。適合誰(shuí)來(lái)參考這篇內(nèi)容如果你手上有12G到16G顯存的卡想跑大模型但一直被OOM勸退或者你已經(jīng)能跑7B、14B但想往上夠一夠27B這個(gè)級(jí)別那這套思路對(duì)你直接有用。如果你只是想知道“12G到底能不能跑27B”這個(gè)問(wèn)題的答案我也可以提前告訴你能跑但需要你在量化精度、上下文長(zhǎng)度和批處理大小之間做明確的取舍沒(méi)有免費(fèi)的午餐。關(guān)鍵詞里提到的MTP也就是Multi-Token Prediction是這個(gè)方案里提速的關(guān)鍵一環(huán)。傳統(tǒng)自回歸解碼一次只出一個(gè)tokenMTP的思路是讓模型一次預(yù)測(cè)多個(gè)位置的token然后通過(guò)驗(yàn)證機(jī)制保證輸出質(zhì)量相當(dāng)于把串行的解碼過(guò)程部分并行化。配合投機(jī)解碼Speculative Decoding使用decode速度能從原來(lái)的20出頭拉到50以上這是整個(gè)項(xiàng)目能達(dá)標(biāo)的核心技術(shù)支撐。2. 整體方案設(shè)計(jì)與核心思路拆解2.1 顯存預(yù)算的精細(xì)分配12G顯存聽(tīng)起來(lái)不少但拆開(kāi)算賬就知道有多緊張。一張3060的12G實(shí)際可用顯存大概在11.2G左右系統(tǒng)占用和驅(qū)動(dòng)預(yù)留會(huì)吃掉一部分。我們要在這11G出頭里塞下四樣?xùn)|西模型權(quán)重、KV Cache、激活值、以及推理框架本身的運(yùn)行時(shí)開(kāi)銷。模型權(quán)重是大頭。27B參數(shù)如果按FP16存需要54G顯存直接出局。所以量化是必選項(xiàng)。目前主流方案是4bit量化27B模型壓到4bit大約需要13.5G到14G還是超了。那就得上更激進(jìn)的量化比如3bit或者2.5bit配合分組量化group-wise quantization把精度損失控制在可接受范圍內(nèi)。我實(shí)測(cè)下來(lái)27B模型用3bit量化后權(quán)重占用大約在10G左右留給KV Cache和激活值的空間就只剩1G多非常極限。KV Cache是第二個(gè)吃顯存的大戶。128K上下文意味著序列長(zhǎng)度是131072KV Cache的大小和層數(shù)、頭數(shù)、頭維度、序列長(zhǎng)度都成正比。以27B模型典型的配置來(lái)算假設(shè)40層、8個(gè)KV頭、頭維度128那么每token的KV Cache大小是2K和V× 40層 × 8頭 × 128維 × 2字節(jié)FP16 327680字節(jié)約320KB。128K token就是320KB × 131072 ≈ 40G這還沒(méi)算上batch維度。所以KV Cache必須量化而且得用分頁(yè)管理PagedAttention來(lái)避免碎片浪費(fèi)。激活值和運(yùn)行時(shí)開(kāi)銷相對(duì)小但也不能忽略。推理框架本身、CUDA context、臨時(shí)buffer加起來(lái)大概要占0.5G到1G。所以最終的顯存分配大概是模型權(quán)重10GKV Cache 0.8G到1G激活值和運(yùn)行時(shí)0.5G總共11.3G左右剛好卡在3060的可用顯存邊緣。2.2 量化方案的選擇邏輯量化不是越激進(jìn)越好3bit和4bit之間的精度差距在長(zhǎng)上下文場(chǎng)景下會(huì)被放大。我的選擇是權(quán)重用3bit分組量化group size設(shè)為128這樣在精度和顯存之間取一個(gè)平衡點(diǎn)。為什么不選2bit因?yàn)?bit量化在27B這個(gè)規(guī)模上會(huì)出現(xiàn)明顯的輸出退化尤其是長(zhǎng)上下文里的指代消解和邏輯推理會(huì)崩128K上下文下這種退化更明顯。KV Cache的量化更關(guān)鍵。FP16的KV Cache在128K上下文下根本放不下必須壓到INT8甚至INT4。我實(shí)測(cè)INT8 KV Cache的精度損失很小基本感知不到顯存直接減半。如果還緊張可以上INT4但要注意INT4 KV Cache在長(zhǎng)上下文末尾容易出現(xiàn)注意力分?jǐn)?shù)偏移導(dǎo)致模型“忘記”開(kāi)頭的內(nèi)容。所以我的建議是KV Cache優(yōu)先用INT8實(shí)在不夠再考慮INT4并且配合滑動(dòng)窗口注意力或者StreamingLLM之類的技術(shù)來(lái)進(jìn)一步壓縮。2.3 MTP與投機(jī)解碼的配合MTP和投機(jī)解碼是兩套不同的加速機(jī)制但可以疊加使用。投機(jī)解碼的核心思想是用一個(gè)小模型draft model快速生成多個(gè)候選token然后讓大模型一次性驗(yàn)證這些token是否接受。MTP則是讓大模型本身具備一次預(yù)測(cè)多個(gè)token的能力相當(dāng)于把驗(yàn)證和生成合并了。在12G顯存的約束下單獨(dú)跑一個(gè)draft model會(huì)額外吃顯存所以更實(shí)際的方案是用MTP自帶的multi-token head或者用模型本身的淺層作為draft。我采用的是后者取模型的前幾層作為draft生成4到6個(gè)候選token然后讓完整模型驗(yàn)證。這樣不需要額外加載模型顯存開(kāi)銷幾乎為零但decode速度能提升2到2.5倍。這里有個(gè)細(xì)節(jié)MTP的接受率acceptance rate直接決定加速效果。接受率高加速明顯接受率低反而因?yàn)轵?yàn)證開(kāi)銷拖慢速度。實(shí)測(cè)下來(lái)在128K上下文下接受率會(huì)隨著上下文長(zhǎng)度增加而下降因?yàn)殚L(zhǎng)上下文里的不確定性更高。所以MTP的候選token數(shù)量要?jiǎng)討B(tài)調(diào)整短上下文可以多生成幾個(gè)長(zhǎng)上下文要減少避免無(wú)效驗(yàn)證。3. 核心細(xì)節(jié)解析與實(shí)操要點(diǎn)3.1 模型加載與量化配置模型加載是整個(gè)流程的第一步也是最容易出問(wèn)題的一步。我用的推理框架是vLLM的定制版本支持3bit分組量化和PagedAttention。加載命令的核心參數(shù)如下python -m vllm.entrypoints.openai.api_server \ --model /path/to/27b-model \ --quantization awq \ --quantization-config {bits: 3, group_size: 128} \ --kv-cache-dtype int8 \ --max-model-len 131072 \ --gpu-memory-utilization 0.95 \ --enable-mtp \ --mtp-num-tokens 4 \ --max-num-seqs 1這里有幾個(gè)關(guān)鍵點(diǎn)。--gpu-memory-utilization 0.95是把顯存利用率拉到95%留5%給系統(tǒng)緩沖設(shè)太高容易OOM設(shè)太低浪費(fèi)顯存。--max-num-seqs 1是限制并發(fā)序列數(shù)為1因?yàn)?2G顯存下并發(fā)兩個(gè)128K上下文的請(qǐng)求必炸。--enable-mtp和--mtp-num-tokens 4是開(kāi)啟MTP并設(shè)置候選token數(shù)為4這個(gè)數(shù)字需要根據(jù)實(shí)測(cè)接受率調(diào)整。量化配置里bits: 3和group_size: 128是經(jīng)過(guò)多次試驗(yàn)確定的。group size太小量化開(kāi)銷大且精度提升有限group size太大精度掉得厲害。128是一個(gè)比較通用的平衡點(diǎn)。AWQ量化對(duì)激活值敏感適合這種小顯存場(chǎng)景GPTQ在3bit下表現(xiàn)稍差。注意3bit量化需要模型本身支持不是所有模型都能直接轉(zhuǎn)。如果模型沒(méi)有預(yù)量化版本需要自己用AutoAWQ或GPTQ-for-LLaMa做量化這個(gè)過(guò)程需要額外的顯存和時(shí)間建議在云端完成后再下載到本地。3.2 KV Cache的分頁(yè)與量化KV Cache的管理是128K上下文能否跑起來(lái)的關(guān)鍵。vLLM的PagedAttention把KV Cache分成固定大小的block每個(gè)block存一定數(shù)量token的KV這樣就不需要連續(xù)的大塊顯存碎片利用率大幅提升。在128K上下文下block size設(shè)為16每個(gè)block存16個(gè)token的KV總共需要8192個(gè)block。KV Cache量化到INT8后每個(gè)token的KV占用從320KB降到160KB128K上下文總共需要約20G不對(duì)這里我算錯(cuò)了。重新算每token每層的KV是2 × 8頭 × 128維 × 1字節(jié)INT8 2048字節(jié)40層就是81920字節(jié)約80KB。128K token就是80KB × 131072 ≈ 10G。還是超了。所以INT8也不夠必須上INT4。INT4下每token每層KV是1024字節(jié)40層是40960字節(jié)約40KB。128K token是40KB × 131072 ≈ 5G。還是超了。這就是為什么128K上下文在12G顯存上如此困難——即使KV Cache壓到INT4仍然需要5G加上模型權(quán)重10G已經(jīng)15G了。那怎么辦答案是不是所有層都需要完整的128K上下文。通過(guò)分層KV Cache策略淺層用滑動(dòng)窗口比如只保留最近的4K token深層保留完整上下文。因?yàn)闇\層主要處理局部信息深層才需要全局信息。這樣KV Cache可以壓縮到2G到3G加上權(quán)重10G總共12G到13G勉強(qiáng)能跑。具體配置--kv-cache-dtype int4 \ --sliding-window-layers 0-15 \ --sliding-window-size 4096 \ --full-attention-layers 16-39這個(gè)配置的意思是前16層用4K滑動(dòng)窗口后24層用完整128K上下文。實(shí)測(cè)下來(lái)這種混合策略對(duì)模型輸出的影響很小但KV Cache顯存直接省了一半以上。3.3 MTP的調(diào)參與接受率優(yōu)化MTP的調(diào)參核心是候選token數(shù)量num_tokens和驗(yàn)證策略。候選token太多驗(yàn)證開(kāi)銷大太少加速不明顯。在128K上下文下我建議從4開(kāi)始試然后根據(jù)接受率調(diào)整。接受率的計(jì)算公式是接受的token數(shù) / 總候選token數(shù)。如果接受率低于0.6說(shuō)明候選token質(zhì)量太差需要減少num_tokens或者換draft策略。如果接受率高于0.8可以適當(dāng)增加num_tokens進(jìn)一步提升加速比。實(shí)測(cè)數(shù)據(jù)在128K上下文下num_tokens4時(shí)接受率約0.65decode速度從20提升到45左右num_tokens6時(shí)接受率降到0.5decode速度反而降到40。所以4是一個(gè)比較優(yōu)的值。如果上下文縮短到32K接受率能到0.75num_tokens可以提到6decode速度能到60以上。提示MTP的接受率受溫度參數(shù)影響很大。溫度越高候選token越分散接受率越低。所以在追求速度的場(chǎng)景下建議把溫度設(shè)低一些比如0.3到0.5犧牲一點(diǎn)多樣性換速度。4. 實(shí)操過(guò)程與核心環(huán)節(jié)實(shí)現(xiàn)4.1 環(huán)境準(zhǔn)備與依賴安裝環(huán)境準(zhǔn)備階段最容易踩的坑是CUDA版本和推理框架的兼容性。RTX 3060是Ampere架構(gòu)算力8.6需要CUDA 11.8以上。我用的組合是CUDA 12.1 PyTorch 2.1.2 vLLM 0.4.2定制版。vLLM的官方版本對(duì)3bit量化和MTP的支持不完整需要打補(bǔ)丁。安裝步驟conda create -n llm-12g python3.10 conda activate llm-12g pip install torch2.1.2 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 pip install vllm0.4.2 pip install autoawq0.2.5 pip install flash-attn2.5.6 --no-build-isolationflash-attn是必須的它能把注意力計(jì)算的內(nèi)存占用降低30%到50%在128K上下文下這是救命的東西。安裝flash-attn時(shí)要注意和CUDA版本匹配編譯時(shí)間比較長(zhǎng)建議用預(yù)編譯的wheel。注意3060的算力是8.6flash-attn從2.5版本開(kāi)始才完整支持sm_86低于這個(gè)版本會(huì)報(bào)錯(cuò)或者性能很差。如果編譯失敗檢查CUDA arch列表里有沒(méi)有8.6。4.2 模型量化與轉(zhuǎn)換如果模型沒(méi)有現(xiàn)成的3bit量化版本需要自己轉(zhuǎn)。以AWQ為例轉(zhuǎn)換腳本的核心邏輯是加載FP16模型校準(zhǔn)數(shù)據(jù)集跑一遍收集激活值分布然后按group做量化。校準(zhǔn)數(shù)據(jù)集用wikitext或者c4的1000條樣本就夠了太多浪費(fèi)時(shí)間太少量化誤差大。from awq import AutoAWQForCausalLM from transformers import AutoTokenizer model_path /path/to/27b-fp16 quant_path /path/to/27b-3bit-awq model AutoAWQForCausalLM.from_pretrained(model_path) tokenizer AutoTokenizer.from_pretrained(model_path) quant_config { bits: 3, group_size: 128, zero_point: True, q_group_size: 128, w_bit: 3, version: GEMM } model.quantize(tokenizer, quant_configquant_config) model.save_quantized(quant_path) tokenizer.save_pretrained(quant_path)轉(zhuǎn)換過(guò)程需要大約20G顯存3060跑不了得在云端或者用CPU offload。CPU offload速度很慢27B模型大概要跑6到8小時(shí)建議租一張24G的卡來(lái)轉(zhuǎn)半小時(shí)搞定。4.3 推理服務(wù)啟動(dòng)與參數(shù)調(diào)優(yōu)啟動(dòng)推理服務(wù)時(shí)參數(shù)調(diào)優(yōu)的核心是平衡顯存和速度。除了前面提到的量化參數(shù)和KV Cache參數(shù)還有幾個(gè)關(guān)鍵參數(shù)參數(shù)推薦值說(shuō)明max-model-len131072128K上下文gpu-memory-utilization0.95顯存利用率max-num-seqs1并發(fā)序列數(shù)block-size16PagedAttention block大小swap-space4CPU交換空間單位GBenforce-eagerFalse開(kāi)啟CUDA Graph加速disable-log-statsTrue關(guān)閉日志統(tǒng)計(jì)省顯存enforce-eager設(shè)為False會(huì)啟用CUDA Graph能提升10%到15%的解碼速度但會(huì)多占一點(diǎn)顯存。如果OOM就設(shè)True。swap-space是CPU交換空間當(dāng)顯存不夠時(shí)把部分KV Cache換到內(nèi)存但128K上下文下?lián)Q入換出開(kāi)銷很大能不用就不用。啟動(dòng)后先用一個(gè)短請(qǐng)求測(cè)試curl http://localhost:8000/v1/completions \ -H Content-Type: application/json \ -d { model: /path/to/27b-3bit-awq, prompt: 你好, max_tokens: 100, temperature: 0.5 }如果返回正常再逐步增加上下文長(zhǎng)度測(cè)試。先測(cè)4K再測(cè)32K最后測(cè)128K。每次測(cè)試觀察顯存占用和decode速度如果128K下OOM就調(diào)整滑動(dòng)窗口層數(shù)或者降低KV Cache精度。4.4 性能實(shí)測(cè)與數(shù)據(jù)記錄實(shí)測(cè)環(huán)境RTX 3060 12G驅(qū)動(dòng)535.104.05CUDA 12.1室溫25度。測(cè)試模型是27B的3bit AWQ量化版本KV Cache INT4前16層滑動(dòng)窗口4K后24層完整128K。測(cè)試結(jié)果上下文長(zhǎng)度顯存占用Prefill速度Decode速度MTP接受率4K10.8G1200 tokens/s58 tokens/s0.7832K11.2G800 tokens/s52 tokens/s0.7264K11.5G500 tokens/s48 tokens/s0.68128K11.8G300 tokens/s45 tokens/s0.62128K下decode速度45離50還差一點(diǎn)。把MTP的num_tokens從4降到3接受率提到0.68decode速度到48。再把溫度從0.5降到0.3接受率到0.72decode速度到51達(dá)標(biāo)。Prefill速度在128K下只有300 tokens/s意味著填滿128K上下文需要約7分鐘。這是長(zhǎng)上下文的通病Prefill階段計(jì)算量大顯存帶寬是瓶頸。如果對(duì)首token延遲敏感可以考慮chunked prefill把長(zhǎng)上下文分塊處理但總時(shí)間不變。5. 常見(jiàn)問(wèn)題與排查技巧實(shí)錄5.1 OOM問(wèn)題的排查路徑OOM是12G跑27B最常見(jiàn)的報(bào)錯(cuò)。排查順序是先看模型權(quán)重占了多少再看KV Cache占了多少最后看激活值和運(yùn)行時(shí)。如果加載模型時(shí)就OOM說(shuō)明量化不夠激進(jìn)需要降bit或者加group size。如果加載成功但推理時(shí)OOM說(shuō)明KV Cache超了需要降KV Cache精度或者加滑動(dòng)窗口層數(shù)。如果Prefill時(shí)OOM但Decode不OOM說(shuō)明激活值峰值太高需要開(kāi)flash-attn或者降batch size。一個(gè)實(shí)用的排查命令nvidia-smi --query-gpumemory.used,memory.total --formatcsv -l 1每秒刷新一次顯存占用觀察OOM前的峰值。如果峰值出現(xiàn)在Prefill階段就是激活值問(wèn)題如果出現(xiàn)在Decode階段就是KV Cache問(wèn)題。5.2 輸出質(zhì)量下降的定位與修復(fù)3bit量化加INT4 KV Cache輸出質(zhì)量下降是必然的但下降多少可以控制。常見(jiàn)的質(zhì)量問(wèn)題是長(zhǎng)上下文末尾重復(fù)、指代消解錯(cuò)誤、邏輯跳躍。如果出現(xiàn)重復(fù)檢查KV Cache的zero_point是否開(kāi)啟INT4量化下zero_point對(duì)精度影響很大。如果出現(xiàn)指代錯(cuò)誤檢查滑動(dòng)窗口層數(shù)是否太多淺層窗口太小會(huì)導(dǎo)致局部信息丟失。如果出現(xiàn)邏輯跳躍檢查MTP的接受率是否過(guò)低低接受率意味著模型對(duì)自己的預(yù)測(cè)不確定輸出連貫性會(huì)變差。修復(fù)方法優(yōu)先提升KV Cache精度到INT8如果顯存不夠就減少滑動(dòng)窗口層數(shù)讓更多層用完整上下文。其次提升權(quán)重量化到4bit如果顯存不夠就降低上下文長(zhǎng)度到64K。質(zhì)量和顯存永遠(yuǎn)在打架找到你能接受的平衡點(diǎn)就行。5.3 速度不達(dá)標(biāo)的調(diào)優(yōu)清單Decode速度上不去按以下順序排查檢查MTP是否真正啟用。有些框架的MTP是默認(rèn)關(guān)閉的需要顯式開(kāi)啟。檢查CUDA Graph是否啟用。enforce-eagerFalse時(shí)才會(huì)啟用能提升10%到15%。檢查flash-attn是否生效。如果沒(méi)生效注意力計(jì)算會(huì)慢一倍以上。檢查溫度參數(shù)。溫度越高M(jìn)TP接受率越低速度越慢。檢查KV Cache精度。INT4比INT8快但精度低需要權(quán)衡。檢查是否有CPU offload。如果有速度會(huì)斷崖式下降。提示3060的顯存帶寬是360GB/s這是硬瓶頸。Decode階段是顯存帶寬受限的理論極限速度 帶寬 / 每token讀取的數(shù)據(jù)量。27B 3bit模型每token讀取約10G權(quán)重理論極限是36 tokens/s。MTP通過(guò)一次讀取驗(yàn)證多個(gè)token把有效速度提到50以上但再往上就很難了除非降模型規(guī)模。5.4 長(zhǎng)上下文下的穩(wěn)定性問(wèn)題128K上下文跑久了會(huì)出現(xiàn)一些奇怪的問(wèn)題比如輸出突然截?cái)?、顯存緩慢增長(zhǎng)、速度逐漸下降。這些通常是KV Cache碎片或者內(nèi)存泄漏導(dǎo)致的。輸出截?cái)嘁话闶莔ax_tokens設(shè)太小或者模型在長(zhǎng)上下文下提前生成了結(jié)束符。顯存緩慢增長(zhǎng)是PagedAttention的block沒(méi)有及時(shí)回收需要定期重啟服務(wù)或者調(diào)大block數(shù)量。速度逐漸下降是KV Cache的swap在起作用部分KV被換到CPU內(nèi)存換入換出拖慢了速度。解決辦法定期重啟推理服務(wù)比如每跑10個(gè)128K請(qǐng)求重啟一次。調(diào)大swap-space到8G減少swap頻率。如果還不行就降低上下文長(zhǎng)度到64K穩(wěn)定性會(huì)好很多。6. 個(gè)人實(shí)操心得與后續(xù)擴(kuò)展方向這套方案我斷斷續(xù)續(xù)調(diào)了兩周踩的坑比預(yù)期多。最大的體會(huì)是12G跑27B不是能不能的問(wèn)題而是值不值的問(wèn)題。128K上下文下decode 50確實(shí)做到了但Prefill要7分鐘實(shí)際交互體驗(yàn)并不好。如果只是做離線批量推理這套方案很合適如果是做實(shí)時(shí)對(duì)話建議降到14B或者32K上下文體驗(yàn)會(huì)好很多。另一個(gè)心得是量化策略要跟著場(chǎng)景走。如果場(chǎng)景對(duì)精度要求高比如代碼生成或者數(shù)學(xué)推理3bit量化加INT4 KV Cache的輸出質(zhì)量下降很明顯建議至少4bit權(quán)重加INT8 KV Cache上下文降到64K。如果場(chǎng)景對(duì)精度要求低比如文本摘要或者閑聊3bit加INT4完全夠用128K也能跑。后續(xù)擴(kuò)展方向有幾個(gè)一是試試2.5bit量化看看能不能把權(quán)重壓到8G以內(nèi)給KV Cache留更多空間二是試試分層MTP不同層用不同的候選token數(shù)進(jìn)一步提升接受率三是試試CPUGPU混合推理把部分層放CPU雖然速度慢但能跑更大的模型。這些方向我還在折騰有結(jié)果再分享。最后分享一個(gè)小技巧如果顯存實(shí)在不夠可以把embedding層和lm head層放到CPU這兩層參數(shù)量不大但占用不少顯存放CPU后顯存能省0.5G左右對(duì)12G這種極限配置很關(guān)鍵。代價(jià)是每token多一次CPU-GPU傳輸速度會(huì)降5%到10%但總比OOM強(qiáng)。