戰(zhàn)：用transformers跑通下載、加載與生成全流程)
說起來(lái)你可能不信最近我把一批開源大模型拉到本地做推理驗(yàn)證真正讓我熬夜的不是注意力機(jī)制怎么計(jì)算也不是某個(gè)層歸一化的數(shù)學(xué)細(xì)節(jié)而是大模型推理--transformers模型下載實(shí)際運(yùn)用這條鏈路里最不起眼的三件事把模型權(quán)重順利下載下來(lái)、用transformers正確加載到設(shè)備上、再讓輸出真正接進(jìn)業(yè)務(wù)邏輯。模型推理的原理在各大博客里已經(jīng)講得很透但實(shí)操中那些版本匹配、顯存翻車、tokenizer截?cái)?、輸出格式失控的問題才是決定一個(gè)項(xiàng)目能不能跑起來(lái)的勝負(fù)手。這篇文章不是原理科普也不打算講訓(xùn)練就是一段很實(shí)在的推理落地記錄。圍繞transformers庫(kù)從環(huán)境搭建、模型下載到生成參數(shù)調(diào)優(yōu)再到實(shí)際接口封裝和對(duì)一批常見坑的排查把整個(gè)流程完整走一遍。適合那種已經(jīng)看過不少大模型資料、正準(zhǔn)備動(dòng)手把手上的開源模型跑起來(lái)、或者跑起來(lái)之后發(fā)現(xiàn)結(jié)果不符合預(yù)期的同學(xué)參考。1. 為什么選 transformers 庫(kù)作為推理底座一次選型復(fù)盤1.1 推理任務(wù)的幾種形態(tài)決定了你該用哪套工具先說我遇到的場(chǎng)景。我需要把幾個(gè)開源模型接進(jìn)一個(gè)相對(duì)完整的業(yè)務(wù)驗(yàn)證系統(tǒng)里任務(wù)包括給一段文本做摘要、生成一段類似風(fēng)格的標(biāo)題、以及做一個(gè)多輪對(duì)話的測(cè)試入口。這些任務(wù)在模型層面都屬于自回歸生成也就是給定一串token讓模型一個(gè)token一個(gè)token地往后預(yù)測(cè)。這類任務(wù)可選的工具其實(shí)不少有的框架專門做高性能推理服務(wù)能把吞吐壓得很高有的框架則更偏研究原型寫起來(lái)靈活但性能一般。我最終選擇transformers作為主力不是因?yàn)樗跇O致性能上有多強(qiáng)而是它有幾個(gè)很現(xiàn)實(shí)的好處生態(tài)覆蓋面廣絕大多數(shù)新開源的模型都能找到對(duì)應(yīng)的權(quán)重格式和加載方式從加載、分詞、生成到后處理一個(gè)庫(kù)全包了不需要在多個(gè)工具之間來(lái)回切換單機(jī)單卡環(huán)境下稍微調(diào)一下參數(shù)就能獲得很可用的生成效果調(diào)試成本低后續(xù)如果要把模型切成量化版本或者換高性能推理框架transformers也能作為權(quán)重轉(zhuǎn)換的通路1.2 環(huán)境配置里最容易被忽略的版本匹配問題很多人在第一步環(huán)境配置就翻車了不是transformers裝不上而是裝好之后加載模型報(bào)各種莫名其妙的錯(cuò)。我自己的經(jīng)驗(yàn)是最先要把三樣?xùn)|西的版本關(guān)系理清楚深度學(xué)習(xí)框架版本、CUDA版本、transformers版本。這里給一個(gè)我實(shí)際用下來(lái)比較穩(wěn)的組合思路# 先確認(rèn)本機(jī)顯卡驅(qū)動(dòng)支持的CUDA版本 nvidia-smi # 再安裝對(duì)應(yīng)版本的深度學(xué)習(xí)框架注意不要用 pip 默認(rèn)裝最新版盲踩 pip install torch --index-url 對(duì)應(yīng)CUDA版本的安裝源踩過一次的坑是驅(qū)動(dòng)已經(jīng)很新但深度學(xué)習(xí)框架默認(rèn)裝成了CPU版本結(jié)果模型加載成功推理慢到離譜后來(lái)一查是安裝時(shí)沒指定正確的index-url。缺了這一行模型照樣能跑但GPU根本沒用上。transformers版本的選擇上我的建議是別追最新也別用太老。太老的版本對(duì)新模型格式的支持不夠比如很多新權(quán)重都開始用Safetensors格式存儲(chǔ)老版本加載就會(huì)出問題。通常我選一個(gè)相對(duì)穩(wěn)定的大版本然后鎖定小版本確保后續(xù)復(fù)現(xiàn)結(jié)果不會(huì)因?yàn)閹?kù)升級(jí)出現(xiàn)變化。在實(shí)際項(xiàng)目中版本鎖定比什么都重要。1.3 從零加載模型的三種姿勢(shì)對(duì)應(yīng)三種不同需求transformers里加載模型的基礎(chǔ)API就是from_pretrained但圍繞它有不同的用法我梳理成三種姿勢(shì)直接加載權(quán)重并自動(dòng)分流到多個(gè)設(shè)備比如一張卡裝不下的時(shí)候讓模型跨卡存放這種方式適合單機(jī)多卡但顯存單張不夠的情況。指定低精度加載比如用4位量化加載模型文件并不會(huì)真的變小但加載進(jìn)顯存后的占用會(huì)顯著降低適合消費(fèi)級(jí)顯卡。先手動(dòng)下載再指定本地目錄加載適合離線環(huán)境或者反復(fù)調(diào)試同一套權(quán)重避免每次都走網(wǎng)絡(luò)。這三種方式我后面都會(huì)談到但先記住一個(gè)原則你選擇的加載方式直接影響后續(xù)推理速度和輸出質(zhì)量這不是一句模型越大越好能概括的。2. 模型下載這件事比你想象的更值得認(rèn)真對(duì)待2.1 不要直接 git clone 整個(gè)倉(cāng)庫(kù)那是給自己找麻煩一開始我犯了一個(gè)特別典型的錯(cuò)誤看到模型倉(cāng)庫(kù)頁(yè)面上寫著git clone想都沒想就復(fù)制了命令開始拉取。小模型還好幾十GB的模型倉(cāng)庫(kù)一旦斷網(wǎng)本地目錄里留下一堆LFS指針文件害得我花了很長(zhǎng)時(shí)間區(qū)分哪些是真正的權(quán)重哪些是下載失敗的占位文件。后來(lái)我徹底轉(zhuǎn)成通過transformers自帶的方式或者直接下載原始文件。實(shí)際項(xiàng)目里我建議按這個(gè)順序處理文件# 用 huggingface_hub 里的 snapshot_download 函數(shù)拉取整個(gè)倉(cāng)庫(kù)快照 # 這個(gè)函數(shù)會(huì)直接處理正確的文件結(jié)構(gòu)并且支持?jǐn)帱c(diǎn)續(xù)傳 from huggingface_hub import snapshot_download model_dir snapshot_download( repo_id某個(gè)開源模型, local_dir./local_model ) print(model_dir)用這個(gè)方式下載產(chǎn)出的目錄結(jié)構(gòu)基本可以直接喂給transformers的from_pretrained不需要額外整理。如果你只需要推理跑起來(lái)常駐文件就那幾類權(quán)重文件、配置文件、分詞器文件、以及說明文件。2.2 緩存目錄會(huì)讓你磁盤悄悄不夠用這里有一個(gè)很隱蔽的問題一直用默認(rèn)方式下載模型權(quán)重會(huì)緩存在系統(tǒng)用戶目錄下的隱藏文件夾里。你以為你把模型文件下載到了項(xiàng)目目錄其實(shí)磁盤上存了兩份一份在項(xiàng)目目錄一份在緩存目錄項(xiàng)目多幾個(gè)之后幾百GB的空間說沒就沒。我的做法是設(shè)置環(huán)境變量指定緩存路徑統(tǒng)一管理export HF_HOME/data/model_cache這樣一來(lái)所有下載的模型都會(huì)集中到一個(gè)大磁盤目錄下不怕系統(tǒng)盤被塞爆也方便統(tǒng)一清理。關(guān)于斷點(diǎn)續(xù)傳transformers的下載機(jī)制已經(jīng)支持得不差網(wǎng)速不穩(wěn)的時(shí)候反復(fù)重試很少出現(xiàn)不得不重新下載整個(gè)文件的情況。2.3 網(wǎng)絡(luò)環(huán)境不理想時(shí)的替換方案不同網(wǎng)絡(luò)環(huán)境下直接從海外倉(cāng)庫(kù)拉權(quán)重經(jīng)常慢到讓人懷疑人生。我的經(jīng)驗(yàn)是別硬等直接用國(guó)內(nèi)可訪問的鏡像源下載速度能快出好幾個(gè)數(shù)量級(jí)。具體做法是在下載命令里把endpoint換成一個(gè)國(guó)內(nèi)鏡像域名其他代碼完全不動(dòng)。export HF_ENDPOINThttps://hf-mirror.com這個(gè)環(huán)境變量對(duì)huggingface_hub的下載操作是全局生效的。改了之后我再執(zhí)行snapshot_download速度明顯上來(lái)。還有一個(gè)小技巧如果某個(gè)模型太大也可以只下載必要的幾個(gè)文件不需要整個(gè)倉(cāng)庫(kù)都拉下來(lái)。比如有些倉(cāng)庫(kù)里塞了不同精度的權(quán)重我只需要7B這個(gè)規(guī)格的safetensors文件和tokenizer文件就沒必要把幾十GB全搬回來(lái)。2.4 下載完整性的驗(yàn)證別等推理時(shí)報(bào)錯(cuò)才發(fā)現(xiàn)有一次我下載完模型直接跑推理結(jié)果加載到一半報(bào)錯(cuò)說文件損壞一查是某一個(gè)大權(quán)重文件下載不完整但沒報(bào)錯(cuò)。從那以后我養(yǎng)成了一個(gè)習(xí)慣下載完先看目錄里是否包含一個(gè).safetensors對(duì)應(yīng)的索引文件并且用transformers提供的方式做一次完整加載測(cè)試。from transformers import AutoConfig # 只加載config做一次快速校驗(yàn)看文件結(jié)構(gòu)是否完整 config AutoConfig.from_pretrained(本地目錄) print(config.model_type)這一步只讀配置文件跑得很快能提早暴露目錄結(jié)構(gòu)層面的問題。完事之后再嘗試做一次最小的模型加載確認(rèn)權(quán)重文件沒被截?cái)唷＿@種先驗(yàn)證再使用的習(xí)慣能省出大量排錯(cuò)時(shí)間。3. 生成參數(shù)里的門道質(zhì)量、速度與穩(wěn)定的三角平衡3.1 采樣方式和基礎(chǔ)參數(shù)的相互作用先明確一個(gè)核心概念transformers里的generate方法默認(rèn)策略和大多數(shù)人理解的不太一樣。如果直接不傳任何參數(shù)跑generate它默認(rèn)采用的是貪心解碼每次都挑概率最高的token結(jié)果穩(wěn)定但容易重復(fù)。想要多樣性就得打開采樣開關(guān)也就是設(shè)置do_sampleTrue再配合temperature和top_p這些參數(shù)。我經(jīng)常被問到temperature設(shè)置成多少最好。這個(gè)問題沒有標(biāo)準(zhǔn)答案但我可以給出一個(gè)不太會(huì)出錯(cuò)的區(qū)間和判斷邏輯摘要、標(biāo)題、關(guān)鍵信息提取這類任務(wù)要求輸出可靠推薦temperature低一點(diǎn)比如0.1到0.3開放對(duì)話、創(chuàng)意寫作、頭腦風(fēng)暴希望輸出有變化可以適當(dāng)提高到0.7以上超過1.0之后文本會(huì)明顯開始胡言亂語(yǔ)一般只用于惡搞場(chǎng)景top_p也很關(guān)鍵。它做的是從累積概率超過某個(gè)閾值的token集合里采樣和temperature的作用維度不同。實(shí)際調(diào)參的時(shí)候我通常把temperature和top_p同時(shí)配合調(diào)整而不是只動(dòng)其中一個(gè)temperature調(diào)低之后top_p可以適度放寬避免文本太死板反過來(lái)temperature調(diào)高后top_p收緊一點(diǎn)能遏制輸出飄得太遠(yuǎn)。3.2 生成長(zhǎng)度與截?cái)嗟牟黄胶鈓ax_new_tokens這個(gè)參數(shù)看起來(lái)就是限制輸出長(zhǎng)度但實(shí)際上和很多細(xì)節(jié)糾纏在一起。模型是邊生成邊輸出的如果你的業(yè)務(wù)需要對(duì)輸出做流式展示卻把max_new_tokens設(shè)置得過大那么一次請(qǐng)求可能要等很久才能返回第一個(gè)token設(shè)置太小又會(huì)截?cái)嗪蟀攵蝺?nèi)容摘要都不完整。我自己處理長(zhǎng)文本摘要時(shí)一般會(huì)做一個(gè)長(zhǎng)度映射邏輯先估算輸入長(zhǎng)度再動(dòng)態(tài)決定max_new_tokens而不是每次用死值。還有一個(gè)更容易被忽略的問題是停止符。一個(gè)模型輸出的開頭和結(jié)尾往往含有特殊token比如有的模型用特定的字符表示一輪對(duì)話結(jié)束有的用另一個(gè)token表示生成結(jié)束。如果你不去處理這些特殊token直接把輸出展示給用戶結(jié)果就是看起來(lái)多了一堆莫名其妙的符號(hào)或者模型已經(jīng)把話說完了還在輸出重復(fù)內(nèi)容。比較好的做法是拿到生成結(jié)果后做一次清洗把特殊token對(duì)應(yīng)的文本替換為空字符串再根據(jù)實(shí)際業(yè)務(wù)需要決定是否保留完整會(huì)話上下文。3.3 實(shí)際推理中的KV Cache與性能暗示不深入講注意力機(jī)制的公式但必須知道一個(gè)東西推理時(shí)模型會(huì)把已經(jīng)算過的歷史狀態(tài)緩存下來(lái)避免每生成一個(gè)token都重新算一遍之前所有的內(nèi)容。這個(gè)緩存就叫KV Cache。在transformers里generate方法默認(rèn)會(huì)啟用這個(gè)機(jī)制所以你連續(xù)生成長(zhǎng)文本的時(shí)候前半段的速度會(huì)比較快越往后越慢因?yàn)槊恳徊揭x取的歷史信息越來(lái)越多。這直接影響實(shí)際應(yīng)用設(shè)計(jì)。如果你的場(chǎng)景是多輪對(duì)話那么把歷史對(duì)話一直累積下去推理會(huì)越來(lái)越慢內(nèi)存和顯存的占用也會(huì)越來(lái)越高。我的處理方式很直接對(duì)超過一定輪數(shù)的會(huì)話做歷史壓縮或者截?cái)嘀槐Ａ糇罱鼛纵啽匾獣r(shí)把更早的內(nèi)容用摘要替代。這樣犧牲了一部分語(yǔ)義連貫性但換來(lái)了響應(yīng)時(shí)間和顯存開銷的穩(wěn)定。4. 顯存不夠時(shí)的幾條活路量化、流式輸出與批量策略4.1 大模型推理顯存翻車現(xiàn)場(chǎng)第一次加載一個(gè)10B級(jí)別模型的時(shí)候我的顯存直接爆了報(bào)錯(cuò)信息里寫著CUDA out of memory而我看訓(xùn)練日志時(shí)明明記得模型參數(shù)量不算夸張。后來(lái)才搞明白推理的時(shí)候除了權(quán)重本身占顯存KV Cache也在占顯存而且它會(huì)隨著生成長(zhǎng)度動(dòng)態(tài)增長(zhǎng)。權(quán)重占了20%KV Cache可能吃掉另外一大塊如果你再開了beam search或者輸出了很長(zhǎng)內(nèi)容顯存壓力還會(huì)繼續(xù)往上走。這里的應(yīng)對(duì)思路不是一句換更大的顯卡能解決的工程上有很多可以做的事情。4.2 低比特量化用少量精度換顯存空間transformers加載模型時(shí)支持直接在加載階段把權(quán)重轉(zhuǎn)為低精度格式。選擇一個(gè)合適的數(shù)值類型可以讓模型在顯存里的占用大幅下降同時(shí)把推理速度提上來(lái)代價(jià)是生成質(zhì)量有一定損耗。我的實(shí)踐經(jīng)驗(yàn)是對(duì)于中文摘要和標(biāo)題生成這類任務(wù)質(zhì)量損耗基本可以接受如果你要處理的是代碼生成或者數(shù)學(xué)推理這類對(duì)精確度要求更高的場(chǎng)景就要謹(jǐn)慎。加載方式很簡(jiǎn)單from transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained( 本地模型目錄, torch_dtypeauto, device_mapauto )通過device_mapauto可以讓transformers自動(dòng)分配哪些層放在哪個(gè)設(shè)備上一張卡放不下就多卡分?jǐn)?。配合低精度加載我現(xiàn)在可以在單卡上跑以前需要雙卡的模型雖然生成速度會(huì)慢一點(diǎn)但至少任務(wù)能起來(lái)。4.3 流式輸出讓用戶先看到第一個(gè)token真實(shí)的業(yè)務(wù)系統(tǒng)里用戶等一個(gè)完整長(zhǎng)回答是很有耐心的考驗(yàn)。與其讓用戶白屏等10秒不如把模型逐步生成的內(nèi)容實(shí)時(shí)流出來(lái)讓用戶看到模型正在逐字打字。這個(gè)做法的技術(shù)本質(zhì)是每次生成一個(gè)token就回調(diào)一次而不是等全部生成完再一次性返回。transformers里直接使用一個(gè)生成選項(xiàng)就能開流式輸出from transformers import TextIteratorStreamer streamer TextIteratorStreamer(tokenizer) generation_kwargs dict( input_idsinput_ids, max_new_tokens512, streamerstreamer )然后在另一個(gè)線程里啟動(dòng)生成主線程里不斷從streamer里取新文本。這樣響應(yīng)延遲感覺上會(huì)小很多而且用戶體驗(yàn)很像真人打字。這里有個(gè)細(xì)節(jié)流式輸出時(shí)最后清洗特殊token的邏輯不能漏否則用戶會(huì)看到結(jié)尾跳出奇怪的字符。4.4 并發(fā)請(qǐng)求的批量處理思路當(dāng)多個(gè)請(qǐng)求同時(shí)進(jìn)來(lái)如果每個(gè)請(qǐng)求都重新加載模型顯存和算力都會(huì)被浪費(fèi)。常規(guī)做法是模型常駐內(nèi)存多個(gè)請(qǐng)求共享同一個(gè)模型實(shí)例但要注意并發(fā)安全。transformers的generate內(nèi)部是線程安全的框架層面不會(huì)崩但如果你在回調(diào)里同時(shí)修改了同一個(gè)模型的狀態(tài)還是可能出問題。面對(duì)高并發(fā)我通常是事先做一個(gè)簡(jiǎn)單的排隊(duì)機(jī)制或者將請(qǐng)求按輸入長(zhǎng)度分組同一批長(zhǎng)度接近的請(qǐng)求一起進(jìn)模型執(zhí)行。這個(gè)叫動(dòng)態(tài)批處理它能讓GPU的利用率顯著提升代價(jià)是代碼邏輯稍微復(fù)雜一些。5. 從驗(yàn)證腳本到可用服務(wù)實(shí)際運(yùn)用里的三個(gè)案例5.1 示例一把單次摘要能力封裝成帶參數(shù)的接口任務(wù)很簡(jiǎn)單給定一篇長(zhǎng)文本返回一句摘要。我寫了一個(gè)函數(shù)把之前講的參數(shù)都暴露出來(lái)方便不同場(chǎng)景微調(diào)。from transformers import AutoTokenizer, AutoModelForCausalLM tokenizer AutoTokenizer.from_pretrained(model_dir) model AutoModelForCausalLM.from_pretrained(model_dir, device_mapauto) def generate_summary(text, max_length200, temperature0.2): prompt f請(qǐng)對(duì)以下內(nèi)容進(jìn)行摘要\n{text}\n摘要 inputs tokenizer(prompt, return_tensorspt).to(model.device) outputs model.generate( inputs.input_ids, max_new_tokensmax_length, temperaturetemperature, do_sampleTrue, top_p0.9 ) result tokenizer.decode(outputs[0], skip_special_tokensTrue) return result這個(gè)封裝足夠應(yīng)對(duì)很多驗(yàn)證場(chǎng)景。注意我用了skip_special_tokensTrue這是解碼時(shí)清洗特殊token最快的方式比自己寫字符串替換穩(wěn)妥得多。5.2 示例二做結(jié)構(gòu)化輸出時(shí)的格式約束有一類任務(wù)不是讓模型自由發(fā)揮而是要它填一個(gè)JSON、提取一份關(guān)鍵詞列表、或者生成一段固定格式的內(nèi)容。模型本身不懂JSON格式它只會(huì)按文本概率生成所以輸出很容易出現(xiàn)格式錯(cuò)誤。我的經(jīng)驗(yàn)是給模型一個(gè)模板引導(dǎo)在prompt里把期望的格式寫清楚甚至給出示例。比如請(qǐng)從下面的文本中提取三個(gè)關(guān)鍵詞以“關(guān)鍵詞”開頭每個(gè)關(guān)鍵詞用“、”分隔。這種方式比單純說請(qǐng)?zhí)崛￡P(guān)鍵詞可靠得多。如果對(duì)格式要求更嚴(yán)格還需要在后處理階段做一次格式校驗(yàn)解析失敗就重試一次或者把輸出剪裁到合法片段。這里沒有銀彈prompt設(shè)計(jì)加后處理校驗(yàn)才是穩(wěn)妥方案。5.3 示例三多輪對(duì)話的上下文管理多輪對(duì)話是很多人想做的落地場(chǎng)景但直接只把當(dāng)前用戶輸入丟給模型模型是沒有記憶的。需要把歷史對(duì)話拼成一個(gè)prompt再交給模型。一個(gè)典型的拼接格式是用戶... 助手... 用戶... 助手當(dāng)歷史過長(zhǎng)超過了模型的上下文窗口就需要做截?cái)?。我通常?huì)留出一個(gè)固定長(zhǎng)度的位置給最新的對(duì)話把最早的歷史丟掉。另外每輪生成完之后的助手回復(fù)要保存起來(lái)作為下一輪的歷史這樣對(duì)話才有連貫性。6. 一路跑下來(lái)最常踩的坑速查清單為了不讓你在同樣的問題上再浪費(fèi)一次時(shí)間我把這段實(shí)操中最常遇到的幾個(gè)坑整理成一張表每個(gè)坑都會(huì)直接告訴你解決辦法?，F(xiàn)象根本原因解決辦法加載模型報(bào)文件缺失下載不完整或目錄結(jié)構(gòu)不對(duì)特別是缺了索引文件重新用snapshot_download下載先加載config做校驗(yàn)推理速度極慢GPU利用率很低深度學(xué)習(xí)框架裝成了CPU版本檢查torch.cuda.is_available()按正確源重裝輸出里混合奇怪符號(hào)或亂碼特殊token沒清洗或者tokenizer沒加載對(duì)應(yīng)文件解碼時(shí)用skip_special_tokensTrue補(bǔ)齊tokenizer相關(guān)文件生成內(nèi)容不斷重復(fù)、繞圈temperature和top_p搭配不當(dāng)或未開啟多樣采樣調(diào)整采樣參數(shù)適當(dāng)提高temperature再配合top_p顯存直接溢出模型權(quán)重加KV Cache超過可用顯存換低精度加載、降低max_new_tokens、限制對(duì)話歷史長(zhǎng)度下載速度像老牛拉車默認(rèn)訪問海外倉(cāng)庫(kù)設(shè)定國(guó)內(nèi)鏡像環(huán)境變量或手動(dòng)下載后離線加載每次生成結(jié)果都一樣沒設(shè)do_sample默認(rèn)貪心解碼設(shè)置do_sampleTrue并調(diào)整temperature和top_p最后再分享一個(gè)小技巧。如果是認(rèn)真要在一個(gè)項(xiàng)目里反復(fù)迭代別把模型目錄散落各處也別讓每次下載都單獨(dú)走一次網(wǎng)絡(luò)。搭一個(gè)本地模型倉(cāng)庫(kù)專門存放下載好的模型每個(gè)模型一個(gè)子目錄用統(tǒng)一的加載函數(shù)去讀取。這樣不僅能省流量還能保證每次跑的權(quán)重完全一致排錯(cuò)的時(shí)候只需要對(duì)比參數(shù)和代碼不需要懷疑模型文件出了問題。大模型推理這條路原理上的門檻反而沒有工程上的門檻高把下載、加載、生成、落地這些環(huán)節(jié)理順你就能真正把模型用起來(lái)了。 大模型推理落地實(shí)戰(zhàn)：用transformers跑通下載、加載與生成全流程
說起來(lái)你可能不信最近我把一批開源大模型拉到本地做推理驗(yàn)證真正讓我熬夜的不是注意力機(jī)制怎么計(jì)算也不是某個(gè)層歸一化的數(shù)學(xué)細(xì)節(jié)而是大模型推理--transformers模型下載實(shí)際運(yùn)用這條鏈路里最不起眼的三件事把模型權(quán)重順利下載下來(lái)、用transformers正確加載到設(shè)備上、再讓輸出真正接進(jìn)業(yè)務(wù)邏輯。模型推理的原理在各大博客里已經(jīng)講得很透但實(shí)操中那些版本匹配、顯存翻車、tokenizer截?cái)?、輸出格式失控的問題才是決定一個(gè)項(xiàng)目能不能跑起來(lái)的勝負(fù)手。這篇文章不是原理科普也不打算講訓(xùn)練就是一段很實(shí)在的推理落地記錄。圍繞transformers庫(kù)從環(huán)境搭建、模型下載到生成參數(shù)調(diào)優(yōu)再到實(shí)際接口封裝和對(duì)一批常見坑的排查把整個(gè)流程完整走一遍。適合那種已經(jīng)看過不少大模型資料、正準(zhǔn)備動(dòng)手把手上的開源模型跑起來(lái)、或者跑起來(lái)之后發(fā)現(xiàn)結(jié)果不符合預(yù)期的同學(xué)參考。1. 為什么選 transformers 庫(kù)作為推理底座一次選型復(fù)盤1.1 推理任務(wù)的幾種形態(tài)決定了你該用哪套工具先說我遇到的場(chǎng)景。我需要把幾個(gè)開源模型接進(jìn)一個(gè)相對(duì)完整的業(yè)務(wù)驗(yàn)證系統(tǒng)里任務(wù)包括給一段文本做摘要、生成一段類似風(fēng)格的標(biāo)題、以及做一個(gè)多輪對(duì)話的測(cè)試入口。這些任務(wù)在模型層面都屬于自回歸生成也就是給定一串token讓模型一個(gè)token一個(gè)token地往后預(yù)測(cè)。這類任務(wù)可選的工具其實(shí)不少有的框架專門做高性能推理服務(wù)能把吞吐壓得很高有的框架則更偏研究原型寫起來(lái)靈活但性能一般。我最終選擇transformers作為主力不是因?yàn)樗跇O致性能上有多強(qiáng)而是它有幾個(gè)很現(xiàn)實(shí)的好處生態(tài)覆蓋面廣絕大多數(shù)新開源的模型都能找到對(duì)應(yīng)的權(quán)重格式和加載方式從加載、分詞、生成到后處理一個(gè)庫(kù)全包了不需要在多個(gè)工具之間來(lái)回切換單機(jī)單卡環(huán)境下稍微調(diào)一下參數(shù)就能獲得很可用的生成效果調(diào)試成本低后續(xù)如果要把模型切成量化版本或者換高性能推理框架transformers也能作為權(quán)重轉(zhuǎn)換的通路1.2 環(huán)境配置里最容易被忽略的版本匹配問題很多人在第一步環(huán)境配置就翻車了不是transformers裝不上而是裝好之后加載模型報(bào)各種莫名其妙的錯(cuò)。我自己的經(jīng)驗(yàn)是最先要把三樣?xùn)|西的版本關(guān)系理清楚深度學(xué)習(xí)框架版本、CUDA版本、transformers版本。這里給一個(gè)我實(shí)際用下來(lái)比較穩(wěn)的組合思路# 先確認(rèn)本機(jī)顯卡驅(qū)動(dòng)支持的CUDA版本 nvidia-smi # 再安裝對(duì)應(yīng)版本的深度學(xué)習(xí)框架注意不要用 pip 默認(rèn)裝最新版盲踩 pip install torch --index-url 對(duì)應(yīng)CUDA版本的安裝源踩過一次的坑是驅(qū)動(dòng)已經(jīng)很新但深度學(xué)習(xí)框架默認(rèn)裝成了CPU版本結(jié)果模型加載成功推理慢到離譜后來(lái)一查是安裝時(shí)沒指定正確的index-url。缺了這一行模型照樣能跑但GPU根本沒用上。transformers版本的選擇上我的建議是別追最新也別用太老。太老的版本對(duì)新模型格式的支持不夠比如很多新權(quán)重都開始用Safetensors格式存儲(chǔ)老版本加載就會(huì)出問題。通常我選一個(gè)相對(duì)穩(wěn)定的大版本然后鎖定小版本確保后續(xù)復(fù)現(xiàn)結(jié)果不會(huì)因?yàn)閹?kù)升級(jí)出現(xiàn)變化。在實(shí)際項(xiàng)目中版本鎖定比什么都重要。1.3 從零加載模型的三種姿勢(shì)對(duì)應(yīng)三種不同需求transformers里加載模型的基礎(chǔ)API就是from_pretrained但圍繞它有不同的用法我梳理成三種姿勢(shì)直接加載權(quán)重并自動(dòng)分流到多個(gè)設(shè)備比如一張卡裝不下的時(shí)候讓模型跨卡存放這種方式適合單機(jī)多卡但顯存單張不夠的情況。指定低精度加載比如用4位量化加載模型文件并不會(huì)真的變小但加載進(jìn)顯存后的占用會(huì)顯著降低適合消費(fèi)級(jí)顯卡。先手動(dòng)下載再指定本地目錄加載適合離線環(huán)境或者反復(fù)調(diào)試同一套權(quán)重避免每次都走網(wǎng)絡(luò)。這三種方式我后面都會(huì)談到但先記住一個(gè)原則你選擇的加載方式直接影響后續(xù)推理速度和輸出質(zhì)量這不是一句模型越大越好能概括的。2. 模型下載這件事比你想象的更值得認(rèn)真對(duì)待2.1 不要直接 git clone 整個(gè)倉(cāng)庫(kù)那是給自己找麻煩一開始我犯了一個(gè)特別典型的錯(cuò)誤看到模型倉(cāng)庫(kù)頁(yè)面上寫著git clone想都沒想就復(fù)制了命令開始拉取。小模型還好幾十GB的模型倉(cāng)庫(kù)一旦斷網(wǎng)本地目錄里留下一堆LFS指針文件害得我花了很長(zhǎng)時(shí)間區(qū)分哪些是真正的權(quán)重哪些是下載失敗的占位文件。后來(lái)我徹底轉(zhuǎn)成通過transformers自帶的方式或者直接下載原始文件。實(shí)際項(xiàng)目里我建議按這個(gè)順序處理文件# 用 huggingface_hub 里的 snapshot_download 函數(shù)拉取整個(gè)倉(cāng)庫(kù)快照 # 這個(gè)函數(shù)會(huì)直接處理正確的文件結(jié)構(gòu)并且支持?jǐn)帱c(diǎn)續(xù)傳 from huggingface_hub import snapshot_download model_dir snapshot_download( repo_id某個(gè)開源模型, local_dir./local_model ) print(model_dir)用這個(gè)方式下載產(chǎn)出的目錄結(jié)構(gòu)基本可以直接喂給transformers的from_pretrained不需要額外整理。如果你只需要推理跑起來(lái)常駐文件就那幾類權(quán)重文件、配置文件、分詞器文件、以及說明文件。2.2 緩存目錄會(huì)讓你磁盤悄悄不夠用這里有一個(gè)很隱蔽的問題一直用默認(rèn)方式下載模型權(quán)重會(huì)緩存在系統(tǒng)用戶目錄下的隱藏文件夾里。你以為你把模型文件下載到了項(xiàng)目目錄其實(shí)磁盤上存了兩份一份在項(xiàng)目目錄一份在緩存目錄項(xiàng)目多幾個(gè)之后幾百GB的空間說沒就沒。我的做法是設(shè)置環(huán)境變量指定緩存路徑統(tǒng)一管理export HF_HOME/data/model_cache這樣一來(lái)所有下載的模型都會(huì)集中到一個(gè)大磁盤目錄下不怕系統(tǒng)盤被塞爆也方便統(tǒng)一清理。關(guān)于斷點(diǎn)續(xù)傳transformers的下載機(jī)制已經(jīng)支持得不差網(wǎng)速不穩(wěn)的時(shí)候反復(fù)重試很少出現(xiàn)不得不重新下載整個(gè)文件的情況。2.3 網(wǎng)絡(luò)環(huán)境不理想時(shí)的替換方案不同網(wǎng)絡(luò)環(huán)境下直接從海外倉(cāng)庫(kù)拉權(quán)重經(jīng)常慢到讓人懷疑人生。我的經(jīng)驗(yàn)是別硬等直接用國(guó)內(nèi)可訪問的鏡像源下載速度能快出好幾個(gè)數(shù)量級(jí)。具體做法是在下載命令里把endpoint換成一個(gè)國(guó)內(nèi)鏡像域名其他代碼完全不動(dòng)。export HF_ENDPOINThttps://hf-mirror.com這個(gè)環(huán)境變量對(duì)huggingface_hub的下載操作是全局生效的。改了之后我再執(zhí)行snapshot_download速度明顯上來(lái)。還有一個(gè)小技巧如果某個(gè)模型太大也可以只下載必要的幾個(gè)文件不需要整個(gè)倉(cāng)庫(kù)都拉下來(lái)。比如有些倉(cāng)庫(kù)里塞了不同精度的權(quán)重我只需要7B這個(gè)規(guī)格的safetensors文件和tokenizer文件就沒必要把幾十GB全搬回來(lái)。2.4 下載完整性的驗(yàn)證別等推理時(shí)報(bào)錯(cuò)才發(fā)現(xiàn)有一次我下載完模型直接跑推理結(jié)果加載到一半報(bào)錯(cuò)說文件損壞一查是某一個(gè)大權(quán)重文件下載不完整但沒報(bào)錯(cuò)。從那以后我養(yǎng)成了一個(gè)習(xí)慣下載完先看目錄里是否包含一個(gè).safetensors對(duì)應(yīng)的索引文件并且用transformers提供的方式做一次完整加載測(cè)試。from transformers import AutoConfig # 只加載config做一次快速校驗(yàn)看文件結(jié)構(gòu)是否完整 config AutoConfig.from_pretrained(本地目錄) print(config.model_type)這一步只讀配置文件跑得很快能提早暴露目錄結(jié)構(gòu)層面的問題。完事之后再嘗試做一次最小的模型加載確認(rèn)權(quán)重文件沒被截?cái)唷＿@種先驗(yàn)證再使用的習(xí)慣能省出大量排錯(cuò)時(shí)間。3. 生成參數(shù)里的門道質(zhì)量、速度與穩(wěn)定的三角平衡3.1 采樣方式和基礎(chǔ)參數(shù)的相互作用先明確一個(gè)核心概念transformers里的generate方法默認(rèn)策略和大多數(shù)人理解的不太一樣。如果直接不傳任何參數(shù)跑generate它默認(rèn)采用的是貪心解碼每次都挑概率最高的token結(jié)果穩(wěn)定但容易重復(fù)。想要多樣性就得打開采樣開關(guān)也就是設(shè)置do_sampleTrue再配合temperature和top_p這些參數(shù)。我經(jīng)常被問到temperature設(shè)置成多少最好。這個(gè)問題沒有標(biāo)準(zhǔn)答案但我可以給出一個(gè)不太會(huì)出錯(cuò)的區(qū)間和判斷邏輯摘要、標(biāo)題、關(guān)鍵信息提取這類任務(wù)要求輸出可靠推薦temperature低一點(diǎn)比如0.1到0.3開放對(duì)話、創(chuàng)意寫作、頭腦風(fēng)暴希望輸出有變化可以適當(dāng)提高到0.7以上超過1.0之后文本會(huì)明顯開始胡言亂語(yǔ)一般只用于惡搞場(chǎng)景top_p也很關(guān)鍵。它做的是從累積概率超過某個(gè)閾值的token集合里采樣和temperature的作用維度不同。實(shí)際調(diào)參的時(shí)候我通常把temperature和top_p同時(shí)配合調(diào)整而不是只動(dòng)其中一個(gè)temperature調(diào)低之后top_p可以適度放寬避免文本太死板反過來(lái)temperature調(diào)高后top_p收緊一點(diǎn)能遏制輸出飄得太遠(yuǎn)。3.2 生成長(zhǎng)度與截?cái)嗟牟黄胶鈓ax_new_tokens這個(gè)參數(shù)看起來(lái)就是限制輸出長(zhǎng)度但實(shí)際上和很多細(xì)節(jié)糾纏在一起。模型是邊生成邊輸出的如果你的業(yè)務(wù)需要對(duì)輸出做流式展示卻把max_new_tokens設(shè)置得過大那么一次請(qǐng)求可能要等很久才能返回第一個(gè)token設(shè)置太小又會(huì)截?cái)嗪蟀攵蝺?nèi)容摘要都不完整。我自己處理長(zhǎng)文本摘要時(shí)一般會(huì)做一個(gè)長(zhǎng)度映射邏輯先估算輸入長(zhǎng)度再動(dòng)態(tài)決定max_new_tokens而不是每次用死值。還有一個(gè)更容易被忽略的問題是停止符。一個(gè)模型輸出的開頭和結(jié)尾往往含有特殊token比如有的模型用特定的字符表示一輪對(duì)話結(jié)束有的用另一個(gè)token表示生成結(jié)束。如果你不去處理這些特殊token直接把輸出展示給用戶結(jié)果就是看起來(lái)多了一堆莫名其妙的符號(hào)或者模型已經(jīng)把話說完了還在輸出重復(fù)內(nèi)容。比較好的做法是拿到生成結(jié)果后做一次清洗把特殊token對(duì)應(yīng)的文本替換為空字符串再根據(jù)實(shí)際業(yè)務(wù)需要決定是否保留完整會(huì)話上下文。3.3 實(shí)際推理中的KV Cache與性能暗示不深入講注意力機(jī)制的公式但必須知道一個(gè)東西推理時(shí)模型會(huì)把已經(jīng)算過的歷史狀態(tài)緩存下來(lái)避免每生成一個(gè)token都重新算一遍之前所有的內(nèi)容。這個(gè)緩存就叫KV Cache。在transformers里generate方法默認(rèn)會(huì)啟用這個(gè)機(jī)制所以你連續(xù)生成長(zhǎng)文本的時(shí)候前半段的速度會(huì)比較快越往后越慢因?yàn)槊恳徊揭x取的歷史信息越來(lái)越多。這直接影響實(shí)際應(yīng)用設(shè)計(jì)。如果你的場(chǎng)景是多輪對(duì)話那么把歷史對(duì)話一直累積下去推理會(huì)越來(lái)越慢內(nèi)存和顯存的占用也會(huì)越來(lái)越高。我的處理方式很直接對(duì)超過一定輪數(shù)的會(huì)話做歷史壓縮或者截?cái)嘀槐Ａ糇罱鼛纵啽匾獣r(shí)把更早的內(nèi)容用摘要替代。這樣犧牲了一部分語(yǔ)義連貫性但換來(lái)了響應(yīng)時(shí)間和顯存開銷的穩(wěn)定。4. 顯存不夠時(shí)的幾條活路量化、流式輸出與批量策略4.1 大模型推理顯存翻車現(xiàn)場(chǎng)第一次加載一個(gè)10B級(jí)別模型的時(shí)候我的顯存直接爆了報(bào)錯(cuò)信息里寫著CUDA out of memory而我看訓(xùn)練日志時(shí)明明記得模型參數(shù)量不算夸張。后來(lái)才搞明白推理的時(shí)候除了權(quán)重本身占顯存KV Cache也在占顯存而且它會(huì)隨著生成長(zhǎng)度動(dòng)態(tài)增長(zhǎng)。權(quán)重占了20%KV Cache可能吃掉另外一大塊如果你再開了beam search或者輸出了很長(zhǎng)內(nèi)容顯存壓力還會(huì)繼續(xù)往上走。這里的應(yīng)對(duì)思路不是一句換更大的顯卡能解決的工程上有很多可以做的事情。4.2 低比特量化用少量精度換顯存空間transformers加載模型時(shí)支持直接在加載階段把權(quán)重轉(zhuǎn)為低精度格式。選擇一個(gè)合適的數(shù)值類型可以讓模型在顯存里的占用大幅下降同時(shí)把推理速度提上來(lái)代價(jià)是生成質(zhì)量有一定損耗。我的實(shí)踐經(jīng)驗(yàn)是對(duì)于中文摘要和標(biāo)題生成這類任務(wù)質(zhì)量損耗基本可以接受如果你要處理的是代碼生成或者數(shù)學(xué)推理這類對(duì)精確度要求更高的場(chǎng)景就要謹(jǐn)慎。加載方式很簡(jiǎn)單from transformers import AutoModelForCausalLM model AutoModelForCausalLM.from_pretrained( 本地模型目錄, torch_dtypeauto, device_mapauto )通過device_mapauto可以讓transformers自動(dòng)分配哪些層放在哪個(gè)設(shè)備上一張卡放不下就多卡分?jǐn)?。配合低精度加載我現(xiàn)在可以在單卡上跑以前需要雙卡的模型雖然生成速度會(huì)慢一點(diǎn)但至少任務(wù)能起來(lái)。4.3 流式輸出讓用戶先看到第一個(gè)token真實(shí)的業(yè)務(wù)系統(tǒng)里用戶等一個(gè)完整長(zhǎng)回答是很有耐心的考驗(yàn)。與其讓用戶白屏等10秒不如把模型逐步生成的內(nèi)容實(shí)時(shí)流出來(lái)讓用戶看到模型正在逐字打字。這個(gè)做法的技術(shù)本質(zhì)是每次生成一個(gè)token就回調(diào)一次而不是等全部生成完再一次性返回。transformers里直接使用一個(gè)生成選項(xiàng)就能開流式輸出from transformers import TextIteratorStreamer streamer TextIteratorStreamer(tokenizer) generation_kwargs dict( input_idsinput_ids, max_new_tokens512, streamerstreamer )然后在另一個(gè)線程里啟動(dòng)生成主線程里不斷從streamer里取新文本。這樣響應(yīng)延遲感覺上會(huì)小很多而且用戶體驗(yàn)很像真人打字。這里有個(gè)細(xì)節(jié)流式輸出時(shí)最后清洗特殊token的邏輯不能漏否則用戶會(huì)看到結(jié)尾跳出奇怪的字符。4.4 并發(fā)請(qǐng)求的批量處理思路當(dāng)多個(gè)請(qǐng)求同時(shí)進(jìn)來(lái)如果每個(gè)請(qǐng)求都重新加載模型顯存和算力都會(huì)被浪費(fèi)。常規(guī)做法是模型常駐內(nèi)存多個(gè)請(qǐng)求共享同一個(gè)模型實(shí)例但要注意并發(fā)安全。transformers的generate內(nèi)部是線程安全的框架層面不會(huì)崩但如果你在回調(diào)里同時(shí)修改了同一個(gè)模型的狀態(tài)還是可能出問題。面對(duì)高并發(fā)我通常是事先做一個(gè)簡(jiǎn)單的排隊(duì)機(jī)制或者將請(qǐng)求按輸入長(zhǎng)度分組同一批長(zhǎng)度接近的請(qǐng)求一起進(jìn)模型執(zhí)行。這個(gè)叫動(dòng)態(tài)批處理它能讓GPU的利用率顯著提升代價(jià)是代碼邏輯稍微復(fù)雜一些。5. 從驗(yàn)證腳本到可用服務(wù)實(shí)際運(yùn)用里的三個(gè)案例5.1 示例一把單次摘要能力封裝成帶參數(shù)的接口任務(wù)很簡(jiǎn)單給定一篇長(zhǎng)文本返回一句摘要。我寫了一個(gè)函數(shù)把之前講的參數(shù)都暴露出來(lái)方便不同場(chǎng)景微調(diào)。from transformers import AutoTokenizer, AutoModelForCausalLM tokenizer AutoTokenizer.from_pretrained(model_dir) model AutoModelForCausalLM.from_pretrained(model_dir, device_mapauto) def generate_summary(text, max_length200, temperature0.2): prompt f請(qǐng)對(duì)以下內(nèi)容進(jìn)行摘要\n{text}\n摘要 inputs tokenizer(prompt, return_tensorspt).to(model.device) outputs model.generate( inputs.input_ids, max_new_tokensmax_length, temperaturetemperature, do_sampleTrue, top_p0.9 ) result tokenizer.decode(outputs[0], skip_special_tokensTrue) return result這個(gè)封裝足夠應(yīng)對(duì)很多驗(yàn)證場(chǎng)景。注意我用了skip_special_tokensTrue這是解碼時(shí)清洗特殊token最快的方式比自己寫字符串替換穩(wěn)妥得多。5.2 示例二做結(jié)構(gòu)化輸出時(shí)的格式約束有一類任務(wù)不是讓模型自由發(fā)揮而是要它填一個(gè)JSON、提取一份關(guān)鍵詞列表、或者生成一段固定格式的內(nèi)容。模型本身不懂JSON格式它只會(huì)按文本概率生成所以輸出很容易出現(xiàn)格式錯(cuò)誤。我的經(jīng)驗(yàn)是給模型一個(gè)模板引導(dǎo)在prompt里把期望的格式寫清楚甚至給出示例。比如請(qǐng)從下面的文本中提取三個(gè)關(guān)鍵詞以“關(guān)鍵詞”開頭每個(gè)關(guān)鍵詞用“、”分隔。這種方式比單純說請(qǐng)?zhí)崛￡P(guān)鍵詞可靠得多。如果對(duì)格式要求更嚴(yán)格還需要在后處理階段做一次格式校驗(yàn)解析失敗就重試一次或者把輸出剪裁到合法片段。這里沒有銀彈prompt設(shè)計(jì)加后處理校驗(yàn)才是穩(wěn)妥方案。5.3 示例三多輪對(duì)話的上下文管理多輪對(duì)話是很多人想做的落地場(chǎng)景但直接只把當(dāng)前用戶輸入丟給模型模型是沒有記憶的。需要把歷史對(duì)話拼成一個(gè)prompt再交給模型。一個(gè)典型的拼接格式是用戶... 助手... 用戶... 助手當(dāng)歷史過長(zhǎng)超過了模型的上下文窗口就需要做截?cái)?。我通常?huì)留出一個(gè)固定長(zhǎng)度的位置給最新的對(duì)話把最早的歷史丟掉。另外每輪生成完之后的助手回復(fù)要保存起來(lái)作為下一輪的歷史這樣對(duì)話才有連貫性。6. 一路跑下來(lái)最常踩的坑速查清單為了不讓你在同樣的問題上再浪費(fèi)一次時(shí)間我把這段實(shí)操中最常遇到的幾個(gè)坑整理成一張表每個(gè)坑都會(huì)直接告訴你解決辦法?，F(xiàn)象根本原因解決辦法加載模型報(bào)文件缺失下載不完整或目錄結(jié)構(gòu)不對(duì)特別是缺了索引文件重新用snapshot_download下載先加載config做校驗(yàn)推理速度極慢GPU利用率很低深度學(xué)習(xí)框架裝成了CPU版本檢查torch.cuda.is_available()按正確源重裝輸出里混合奇怪符號(hào)或亂碼特殊token沒清洗或者tokenizer沒加載對(duì)應(yīng)文件解碼時(shí)用skip_special_tokensTrue補(bǔ)齊tokenizer相關(guān)文件生成內(nèi)容不斷重復(fù)、繞圈temperature和top_p搭配不當(dāng)或未開啟多樣采樣調(diào)整采樣參數(shù)適當(dāng)提高temperature再配合top_p顯存直接溢出模型權(quán)重加KV Cache超過可用顯存換低精度加載、降低max_new_tokens、限制對(duì)話歷史長(zhǎng)度下載速度像老牛拉車默認(rèn)訪問海外倉(cāng)庫(kù)設(shè)定國(guó)內(nèi)鏡像環(huán)境變量或手動(dòng)下載后離線加載每次生成結(jié)果都一樣沒設(shè)do_sample默認(rèn)貪心解碼設(shè)置do_sampleTrue并調(diào)整temperature和top_p最后再分享一個(gè)小技巧。如果是認(rèn)真要在一個(gè)項(xiàng)目里反復(fù)迭代別把模型目錄散落各處也別讓每次下載都單獨(dú)走一次網(wǎng)絡(luò)。搭一個(gè)本地模型倉(cāng)庫(kù)專門存放下載好的模型每個(gè)模型一個(gè)子目錄用統(tǒng)一的加載函數(shù)去讀取。這樣不僅能省流量還能保證每次跑的權(quán)重完全一致排錯(cuò)的時(shí)候只需要對(duì)比參數(shù)和代碼不需要懷疑模型文件出了問題。大模型推理這條路原理上的門檻反而沒有工程上的門檻高把下載、加載、生成、落地這些環(huán)節(jié)理順你就能真正把模型用起來(lái)了。