實戰(zhàn)：從顯存優(yōu)化到LoRA選型)
簡介面向需要上手大模型微調(diào)的開發(fā)者與研究者這份資源以DeepSpeed多卡訓練為主線完整覆蓋ChatGLM微調(diào)實戰(zhàn)流程并提供可直接運行的模塊化源碼與流程教程。項目包含詳細的環(huán)境搭建、數(shù)據(jù)準備與格式化、LoRA/P-tuning/Freeze等多類微調(diào)方案同時講解ds_config.json等關(guān)鍵配置參數(shù)、訓練啟動腳本寫法和常見踩坑點有助于讀者跨越單卡顯存限制快速搭建多卡并行訓練環(huán)境。壓縮包共17個文件以Python訓練腳本11個、Shell啟動腳本3個、JSON配置文件2個和Markdown說明文檔1個為主大小僅118KB目錄劃分明確便于按模塊對照學習。已有267人瀏覽學習。借助源碼中的模型加載、數(shù)據(jù)加載、訓練循環(huán)、推理腳本和評估測試等模塊讀者不僅能復現(xiàn)ChatGLM多卡微調(diào)還能理解DeepSpeed的內(nèi)存優(yōu)化與梯度累積機制是一份可直接落地的優(yōu)質(zhì)實戰(zhàn)教程。1. 多卡ChatGLM微調(diào)為什么大家都卡在DeepSpeed這一關(guān)把ChatGLM-6B塞進一張A100-80G里做推理是輕松的但一旦開始全量微調(diào)Adam優(yōu)化器的狀態(tài)一上來顯存立刻不夠用。很多人對“多卡微調(diào)”的第一反應是插4張卡等于4倍顯存真正跑起來才發(fā)現(xiàn)權(quán)重、梯度、優(yōu)化器狀態(tài)這堆東西不切開4張卡也只是各自抱著同一份副本在空轉(zhuǎn)。DeepSpeed就是為解決這個問題出現(xiàn)的它把模型狀態(tài)切到各卡上讓基于DeepSpeed的多卡ChatGLM微調(diào)從“跑得動”變成“跑得起”。這篇實戰(zhàn)筆記按環(huán)境搭建、LoRA與全量選型、deepspeed_config.json參數(shù)、多卡啟動和踩坑排查的順序給出一套可以直接復制到你自己機器上的完整流程。適合手里有2到8張GPU、想把ChatGLM微調(diào)成垂直場景模型的算法工程師和獨立開發(fā)者。2. 微調(diào)前的選型全量微調(diào)和LoRA微調(diào)先算一筆賬2.1 全量微調(diào)為什么在多卡上這么貴全量微調(diào)意味著模型每一層參數(shù)都在更新訓練過程中需要同時存放五樣東西權(quán)重本身、梯度、優(yōu)化器的一階動量、二階動量以及混合精度下的FP32主權(quán)重。工程上常用一個近似數(shù)來算賬每1B參數(shù)大約要占16GB的模型狀態(tài)。ChatGLM-6B換算下來光是參數(shù)、梯度和優(yōu)化器狀態(tài)就要接近96GB。一張A100-80G單卡裝不下即便勉強裝下batch只能設(shè)為1訓練效率低到?jīng)]法看。多卡場景里最直接的做法是Data Parallel也就是每張卡復制一份完整模型各自算梯度再做一次allreduce合并。4張卡就是4份96GB顯存不僅沒省反而因為通信開銷變得更慢。DeepSpeed在這里的核心價值是ZeRO它把原本每張卡都保存一份的優(yōu)化器狀態(tài)、梯度甚至參數(shù)切分到不同卡上通信只發(fā)生在真正需要同步的時刻。這也是“多卡微調(diào)為什么繞不開DeepSpeed”的最直接原因?qū)τ谝粋€6B量級的模型光模型狀態(tài)就已經(jīng)超過單卡物理上限不用ZeRO只能靠堆顯存解決而堆顯存的性價比太低。這里還需要區(qū)分一個概念模型狀態(tài)和激活值。模型狀態(tài)指的是參數(shù)、梯度、優(yōu)化器狀態(tài)激活值是前向傳播過程中每一層輸出的中間張量。很多人以為開了DeepSpeed就能解決所有顯存問題結(jié)果發(fā)現(xiàn)batch一大照樣OOM這就是激活值在瓶頸上。后面第4章會給出一個完整的顯存估算方法先把模型狀態(tài)算明白再結(jié)合batch和序列長度去反推激活值的余量。2.2 LoRA微調(diào)是什么意思參數(shù)少到什么程度剛接觸大模型微調(diào)的人常問LoRA微調(diào)是什么意思。簡單說LoRA凍結(jié)原模型的全部權(quán)重在每一層通常是注意力層的線性映射旁邊掛一組低秩矩陣A和B只訓練這兩個小矩陣。最終更新量等于alpha/r乘以AB的結(jié)果再疊加到原始權(quán)重上。由于原始權(quán)重完全不更新可訓練參數(shù)量會從6B級別直接掉到百萬級別。以ChatGLM-6B為例在query_key_value上掛LoRAr設(shè)為8可訓練參數(shù)大約只有三百萬到四百萬比原始參數(shù)小了三個數(shù)量級。這會帶來兩個直接好處一是模型狀態(tài)顯存占用大幅下降因為需要保存的參數(shù)、梯度和優(yōu)化器狀態(tài)都只針對這小部分可訓練參數(shù)二是訓練過程不容易把基座的語言能力帶偏特別適合數(shù)據(jù)量不大、目標明確的垂直微調(diào)場景。from peft import LoraConfig, get_peft_model lora_config LoraConfig( r8, # 低秩矩陣的秩常用 8/16/32 lora_alpha16, # 縮放系數(shù)實際更新量乘 alpha/r target_modules[query_key_value], # ChatGLM 注意力層線性映射的參數(shù)名 lora_dropout0.1, biasnone, ) model get_peft_model(model, lora_config) model.print_trainable_parameters()這里有幾個參數(shù)需要根據(jù)實際任務調(diào)r太小表達能力不夠r太大可訓練參數(shù)變多顯存優(yōu)勢和微調(diào)“安全邊際”都會縮水。lora_alpha和r的比值決定LoRA更新的影響幅度一般alpha取r的1到2倍。target_modules要根據(jù)模型結(jié)構(gòu)改ChatGLM的注意力層里query、key、value是合并在一個名為query_key_value的線性層里的所以要寫這個參數(shù)名如果微調(diào)ChatGLM2或ChatGLM3這個名稱可能不同先打印模型結(jié)構(gòu)確認再填。對比全量微調(diào)LoRA微調(diào)的選擇邏輯可以歸納為數(shù)據(jù)量小、任務垂直度高、希望快速迭代選LoRA數(shù)據(jù)量足夠大、顯卡資源充足、追求極限效果才考慮全量微調(diào)。對于多數(shù)真實項目LoRA都是第一版方案因為它能在一個晚上跑完而且效果差距往往沒有想象中大。對比項全量微調(diào)LoRA微調(diào)更新參數(shù)范圍全部參數(shù)低秩矩陣模型狀態(tài)顯存約16GB/B參數(shù)約0.1GB/B參數(shù)訓練速度慢快3到5倍基座能力保持容易遺忘較好適用場景數(shù)據(jù)量大、資源足垂直微調(diào)、快速迭代2.3 ChatGLM的數(shù)據(jù)拼接和標簽掩碼Prefix-LM結(jié)構(gòu)下的labelsChatGLM是Prefix-LM結(jié)構(gòu)左側(cè)的輸入可以做雙向注意力右側(cè)的生成部分只能看左側(cè)。這和GPT的Causal-LM不一樣訓練數(shù)據(jù)不能直接當普通CausalLM來拼接。社區(qū)里最常見的做法是把指令和回答拼成一條完整文本在回答開始的位置之前把標簽全部設(shè)為-100計算loss時這些位置被忽略。def pack_chat_input(tokenizer, source, answer, max_len2048): # 拼接格式和ChatGLM官方prompt模板保持一致 text 問 source \n答 answer tokenizer.eos_token ids tokenizer.encode(text, max_lengthmax_len, truncationTrue) # 計算prompt部分的長度這部分不參與loss計算 prompt 問 source \n答 prompt_len len(tokenizer.encode(prompt)) labels [-100] * len(ids) # 先把所有位置設(shè)為-100 labels[prompt_len:] ids[prompt_len:] # 回答部分保留真實token return {input_ids: ids, labels: labels}這段代碼的核心邏輯是“回答token要學習指令token不學習”。-100是PyTorch CrossEntropyLoss約定俗成的ignore索引模型在算loss時會自動跳過這些位置。注意這里的prompt_len用的是tokenizer.encode不是字符長度因為中文文本的token切分不是按字來的。訓練循環(huán)里有一個特別容易翻車的細節(jié)用了DeepSpeed之后反向傳播和參數(shù)更新必須交給engine不能直接調(diào)model.backward()或optimizer.step()。原因在于ZeRO把模型狀態(tài)切到了多張卡上只有engine知道當前rank上持有哪部分狀態(tài)。for step, batch in enumerate(dataloader): outputs engine( input_idsbatch[input_ids].cuda(), labelsbatch[labels].cuda(), ) engine.backward(outputs.loss) engine.step()這里的engine就是deepspeed.initialize返回的模型包裝對象。前向接口和原始模型基本一致但backward和step必須走engine。如果自己額外再調(diào)一次optimizer.step()輕則梯度更新錯亂重則訓練過程直接崩掉。3. 環(huán)境搭建與最小復現(xiàn)從安裝deepspeed到多卡啟動3.1 安裝deepspeed包先過CUDA這一關(guān)搜索“安裝deepspeed包一直報錯”你會發(fā)現(xiàn)十個里有八個是同一個原因pip在安裝時嘗試編譯CUDA拓展但系統(tǒng)里找不到CUDA工具鏈。報錯信息里通常會有一大段gcc、nvcc相關(guān)的日志很多人看到就懵了其實解決辦法很簡單先確認PyTorch的CUDA版本再檢查nvcc能不能用。# 先確認torch的CUDA版本再決定裝哪個deepspeed python -c import torch; print(torch.version.cuda) # 安裝前確認nvcc可用 nvcc --version # 常規(guī)安裝方式 pip install deepspeed如果pip直接從源碼編譯導致報錯可以先用預編譯的wheel包緩解。另一個常見做法是設(shè)置CUDA_HOME環(huán)境變量指向CUDA Toolkit的安裝目錄再執(zhí)行pip install deepspeed。安裝完成后用ds_report命令檢查deepspeed是否識別到當前GPU、CUDA版本和NCCL版本。這個命令會打印一份環(huán)境診斷報告看到“OK”字樣說明基礎(chǔ)環(huán)境沒有問題。需要注意deepspeed會在第一次啟動訓練時做JIT編譯編譯一些CUDA算子到緩存目錄所以第一次啟動多等一兩分鐘是正常的。如果多次啟動都卡在同一處編譯多半是系統(tǒng)缺了gcc或者編譯依賴看日志里的具體報錯再補裝對應依賴。3.2 寫一個最小訓練腳本train_chatglm_lora.py環(huán)境就緒后先別急著上全量微調(diào)我一般用LoRA把整個鏈路先跑通。下面這個腳本是完整可運行的最小版本把模型加載、LoRA包裝、DeepSpeed初始化和訓練循環(huán)放在一起總共不到40行。import torch import deepspeed from transformers import AutoTokenizer, AutoModel from peft import LoraConfig, get_peft_model def train(): model_name THUDM/chatglm-6b tokenizer AutoTokenizer.from_pretrained(model_name, trust_remote_codeTrue) # bf16要求Ampere及以上架構(gòu)V100等舊卡請改成fp16 model AutoModel.from_pretrained( model_name, trust_remote_codeTrue, torch_dtypetorch.bfloat16, ) # 用LoRA凍結(jié)大部分參數(shù)只訓練query_key_value上的低秩矩陣 model get_peft_model(model, LoraConfig( r8, lora_alpha16, target_modules[query_key_value], lora_dropout0.1, biasnone, )) # 梯度檢查點用計算換顯存多卡顯存緊張時建議打開 model.gradient_checkpointing_enable() engine, optimizer, _, lr_scheduler deepspeed.initialize( modelmodel, model_parameters[p for p in model.parameters() if p.requires_grad], configds_config.json, ) for step, batch in enumerate(dataloader): outputs engine( input_idsbatch[input_ids].cuda(), labelsbatch[labels].cuda(), ) engine.backward(outputs.loss) engine.step()這段代碼里有兩個地方最容易出錯。一是AutoModel.from_pretrained必須帶trust_remote_codeTrue因為ChatGLM的模型結(jié)構(gòu)通過遠程代碼加載二是deepspeed.initialize的model_parameters參數(shù)只傳可訓練參數(shù)LoRA模式下就是那幾個低秩矩陣不要傳model.parameters()的完整列表否則DeepSpeed會把凍結(jié)參數(shù)也納入優(yōu)化器管理白白浪費顯存。dataloader需要配合第2.3節(jié)的pack_chat_input函數(shù)把每條樣本處理成input_ids和labels。注意collate時要動態(tài)padding到batch內(nèi)最長而不是固定到2048否則顯存大部分都被無效的padding token吃掉了。3.3 啟動命令單機多卡和多機多卡的差異訓練腳本寫好后啟動方式用deepspeed命令不是python命令。單機4卡的最簡命令如下deepspeed --num_gpus 4 --master_port 29500 train_chatglm_lora.pymaster_port是分布式訓練的控制通信端口默認29500。如果同一臺機器上同時跑多個訓練任務端口會沖突啟動時會報端口被占用的錯誤換一個比如29501就行。如果你用的是SLURM之類的調(diào)度系統(tǒng)還要帶上deepspeed_multinode腳本一起用這里不展開。多機多卡稍微復雜一點每臺機器上都要執(zhí)行啟動命令區(qū)別只在node_rank和master_addr# 在第一個節(jié)點執(zhí)行 deepspeed --num_nodes 2 --num_gpus 8 --node_rank 0 \ --master_addr 192.168.1.10 --master_port 29500 train_chatglm_lora.py # 在第二個節(jié)點執(zhí)行 deepspeed --num_nodes 2 --num_gpus 8 --node_rank 1 \ --master_addr 192.168.1.10 --master_port 29500 train_chatglm_lora.py多機場景下master_addr必須填第一臺機器的IP而且兩臺機器要能互相訪問。這里有個容易踩的坑如果機器之間只有以太網(wǎng)沒有InfiniBand通信速度會成為訓練瓶頸需要顯式設(shè)置NCCL_SOCKET_IFNAME指向?qū)嶋H網(wǎng)卡接口名。判斷方法是在訓練日志里看NCCL初始化時用了哪塊網(wǎng)卡如果是lo或docker0訓練速度大概率不正常。4. deepspeed_config.json先設(shè)這幾個參數(shù)再談訓練收斂4.1 ZeRO Stage選1、2還是3zero123的區(qū)別第一次寫deepspeed_config.json卡住最多的問題就是deepspeed zero123的區(qū)別。簡單說Stage 1只切分優(yōu)化器狀態(tài)Stage 2在Stage 1基礎(chǔ)上連梯度也切分Stage 3更進一步把模型參數(shù)本身也切分到各卡上。Stage越高顯存越省但卡間通信量越大訓練速度越慢。ChatGLM-6B微調(diào)我的建議是優(yōu)先Stage 2。6B模型的參數(shù)量不算特別大Stage 2已經(jīng)能把優(yōu)化器狀態(tài)和梯度的重復存儲省掉顯存能被壓到接近“每卡只存一份完整權(quán)重部分狀態(tài)”的水平。Stage 3雖然有額外收益但通信開銷會拖慢迭代速度在4卡或8卡規(guī)模下性價比不高。只有當顯存實在不夠、batch小到?jīng)]法看的時候才考慮Stage 3或者把優(yōu)化器狀態(tài)offload到CPU。DeepSpeed初始化時會打印一份運行配置摘要里面包括當前用的是哪個Stage、每張卡預估的顯存使用量。啟動日志不要直接跳過這些信息比任何第三方工具都準。Stage 3還要注意一個問題參數(shù)被切分后前向和反向過程中需要動態(tài)收集參數(shù)所以eval和保存模型時要額外處理對新手來說這是比顯存更麻煩的坑。4.2 我習慣先寫好的幾個配置項下面這份ds_config.json是我做ChatGLM微調(diào)時最常使用的配置模板每一項都直接影響訓練能否收斂。{ train_batch_size: 32, gradient_accumulation_steps: 4, gradient_clipping: 1.0, bf16: { enabled: true }, zero_optimization: { stage: 2, offload_optimizer: { device: cpu } }, scheduler: { type: WarmupDecayLR, params: { warmup_min_lr: 0, warmup_max_lr: 2e-5, warmup_num_steps: 200 } }, steps_per_print: 20 }第一眼要搞清楚的是train_batch_size、gradient_accumulation_steps和實際每卡batch三者的關(guān)系。train_batch_size是全局batch也就是所有卡累積多少條樣本做一次參數(shù)更新。實際每卡micro batch的計算公式是train_batch_size除以卡數(shù)再除以梯度累積步數(shù)。以這份配置為例4卡、累積4步、全局32那么每卡micro batch是2。很多人在這一步憑感覺填數(shù)字結(jié)果實際batch大小和預期完全不是一回事loss曲線自然不對。配置項取值建議作用train_batch_size16到64全局batch決定每次更新用多少樣本gradient_accumulation_steps2到8梯度累積步數(shù)等效增大batchgradient_clipping1.0防止梯度爆炸微調(diào)大模型必備bf16enabled: true半精度訓練A100及以上建議開offload_optimizerdevice: cpu優(yōu)化器狀態(tài)放內(nèi)存顯存緊張時開warmup_max_lr1e-5到5e-5峰值學習率6B模型別超過5e-5bf16和fp16的區(qū)別也值得多提一句。bf16在A100、H100這些Ampere之后架構(gòu)的卡上效果更好指數(shù)位更多不容易出現(xiàn)loss溢出V100等舊架構(gòu)不支持bf16只能開fp16這時要額外關(guān)注loss scale防止訓練中期loss變成NaN。4.3 用顯存公式反推可用的batch和序列長度顯存占用粗略分為兩塊模型狀態(tài)加激活值。模型狀態(tài)可以用第2.1節(jié)的16GB/B參數(shù)估算激活值則和micro batch大小、序列長度、模型層數(shù)強相關(guān)。ChatGLM-6B在max_seq_len為2048時激活值隨batch增長的斜率非常明顯很多人以為batch設(shè)2沒問題結(jié)果跑到第100步就OOM。我建議的排查順序是先固定序列長度把micro batch調(diào)到1跑一輪用nvidia-smi看顯存占用基線再逐步往上加。# 每秒刷新一次顯存情況觀察訓練過程中的峰值占用 watch -n 1 nvidia-smi看重點不是顯存總量而是進程占用那一欄。如果micro batch為1時已經(jīng)用了接近全部顯存說明瓶頸在激活值這時候調(diào)低max_len或者關(guān)掉冗余輸入更有效。如果micro batch為1時顯存還有余量才把batch往上加。DeepSpeed在啟動時也會給一個顯存預估可以和實際觀測對照幫助判斷配置是否生效。這里有一個容易誤判的地方打開了offload_optimizer之后訓練剛開始時顯存會明顯下降但CPU內(nèi)存占用會漲上去。如果服務器內(nèi)存本身不多offload反而可能引發(fā)內(nèi)存交換訓練速度驟降。觀察指標時不能只看GPU顯存還要帶著看內(nèi)存占用。5. DeepSpeed多卡微調(diào)ChatGLM的四個典型踩坑現(xiàn)場5.1 NCCL超時卡在Waiting for the barrier現(xiàn)象啟動多卡訓練后日志卡在分布式初始化階段幾分鐘后拋NCCLTimeoutError或者某個rank直接掛掉。原因最常見的是master_addr或master_port配置不對多卡之間的控制通信建立不起來其次是機器上有多塊網(wǎng)卡NCCL選錯了網(wǎng)卡導致通信走了一個慢速通道甚至在不通的網(wǎng)口上。還有一種情況是防火墻攔了分布式訓練的端口單機多卡時不太會遇到多機多卡時概率很高。解決先確認所有節(jié)點能互相ping通再檢查deepspeed啟動命令里的master_addr和master_port是否一致。多機場景顯式指定通信網(wǎng)卡export NCCL_SOCKET_IFNAMEeth0 export NCCL_IB_DISABLE1NCCL_IB_DISABLE1是讓NCCL不走InfiniBand只走以太網(wǎng)。如果機器之間的IB網(wǎng)絡(luò)沒有正確配置與其讓它自動探測失敗不如直接關(guān)掉。日志里能看到NCCL初始化時用的網(wǎng)卡名字如果發(fā)現(xiàn)選錯了再用NCCL_SOCKET_IFNAME強制指定。5.2 訓練中途OOM先分清是激活值還是模型狀態(tài)壓爆的現(xiàn)象訓練能啟動但跑到某個step突然報CUDA out of memory而且每次崩的位置不一樣。原因大多數(shù)時候是動態(tài)padding沒有生效batch里某條樣本特別長把激活值峰值瞬間拉高。另一個常見原因是micro batch算錯了以為設(shè)置了梯度累積batch就小了實際上每卡的micro batch仍然是全局batch除以卡數(shù)沒有除以累積步數(shù)。解決先按第4.3節(jié)的方法把micro batch降到1試跑如果不再OOM說明是batch和序列長度的問題如果batch為1仍然OOM說明模型狀態(tài)部分還有優(yōu)化空間。此時從三個方向同時壓打開gradient_checkpointing啟用梯度檢查點、把offload_optimizer設(shè)備設(shè)為cpu、以及確認沒有把凍結(jié)參數(shù)傳進優(yōu)化器。這三招用完ChatGLM-6B的LoRA微調(diào)在24GB顯存卡上都可以跑起來。注意開了gradient_checkpointing之后前向速度會變慢這是正常的它本質(zhì)上是犧牲計算換顯存。5.3 loss不降或亂跳學習率、warmup和數(shù)據(jù)順序一起背鍋現(xiàn)象訓練了上千步loss要么原地不動要么忽高忽低看起來完全沒有收斂特征。原因把之前訓練小模型的經(jīng)驗搬到大模型上學習率用1e-4甚至更高對6B模型來說步長太大loss自然亂跳。另一個是warmup步數(shù)太少或沒設(shè)置調(diào)度器一上來就給滿學習率。還有一個容易被忽略的點dataloader沒有shuffle模型反復看到同一條指令loss曲線會出現(xiàn)周期性波動。解決把warmup_max_lr降到2e-5warmup_num_steps設(shè)為總步數(shù)的5%到10%。檢查dataloader是否設(shè)置了shuffleTrue。如果用的是fp16而不是bf16還要觀察訓練日志里的loss scale出現(xiàn)NaN或不正常跳變說明fp16的精度控制出了問題建議升級到支持bf16的顯卡或者調(diào)整fp16配置。這里要特別說一個誤區(qū)不要因為loss在下降就認定訓練正常。大模型微調(diào)初期loss下降可以很快但真正的考驗是第6章的驗證環(huán)節(jié)生成質(zhì)量和loss下降并不完全掛鉤。5.4 微調(diào)完導出eval時模型像沒訓練過現(xiàn)象訓練日志里loss已經(jīng)降得很低但eval時模型回答和原版幾乎一樣甚至更差。原因最常見的是保存模型時只保存了DeepSpeed的引擎權(quán)重沒有保存LoRA adapter。DeepSpeed的save_checkpoint保存的是引擎狀態(tài)直接拿去推理完全用不上。另一個原因是沒有merge LoRA權(quán)重推理時忘了加載adapter。解決訓練結(jié)束后用PeFT的save_pretrained保存LoRA權(quán)重推理時先加載adapter再合并# 訓練結(jié)束保存LoRA權(quán)重 engine.module.save_pretrained(lora_adapter) # 推理時加載并合并 from transformers import AutoModel model AutoModel.from_pretrained(THUDM/chatglm-6b, trust_remote_codeTrue) model.load_adapter(lora_adapter) model model.merge_and_unload()注意推理時加載基座模型的精度要和訓練時一致否則merge可能出現(xiàn)數(shù)值不匹配的問題生成結(jié)果看起來就像沒訓練過。eval時使用訓練時的prompt模板也很關(guān)鍵微調(diào)時用的格式是“問xxx\n答”eval時如果換成其他格式模型生成質(zhì)量會受到影響。6. 從loss下降走向業(yè)務能用驗證流程和一個顯存技巧6.1 用固定prompt集驗證效果別只看loss訓練loss下降只是“模型記住了訓練數(shù)據(jù)”的必要條件不是充分條件。我的做法是準備一組固定業(yè)務問題訓練過程中每隔幾百步用生成模式跑一遍觀察回答質(zhì)量的變化趨勢。這一步能比loss曲線更早暴露出模型是否過擬合、是否丟失了對話能力。生成測試代碼可以掛在訓練循環(huán)里def eval_generate(engine, tokenizer, questions, max_new_tokens128): engine.module.eval() with torch.no_grad(): for q in questions: prompt 問 q \n答 inputs tokenizer(prompt, return_tensorspt).to(engine.device) out engine.module.generate( **inputs, max_new_tokensmax_new_tokens, do_sampleFalse, ) print(tokenizer.decode(out[0], skip_special_tokensTrue)) engine.module.train()固定prompt集的規(guī)模不用大10條左右就夠了但要覆蓋正常提問、邊界提問和易混淆提問三類。do_sample設(shè)為False保證可對比如果模型訓練中關(guān)了dropout這一步的生成結(jié)果是完全確定的。每500步跑一次保存生成結(jié)果到日志文件后面就可以回看不同訓練階段的質(zhì)量變化。另外把一部分數(shù)據(jù)留作held-out集驗證集loss比訓練loss更早上升就是過擬合信號。6.2 把顯存再壓一檔動態(tài)padding和梯度檢查點組合第5.2節(jié)提到動態(tài)padding能解決一部分OOM問題這里給出一個可用的collate實現(xiàn)。核心思路是把batch內(nèi)的樣本padding到同一長度同時用-100掩蓋無效的標簽位置。def collate_batch(batch): max_len max(len(x[input_ids]) for x in batch) input_ids [] labels [] for x in batch: pad_len max_len - len(x[input_ids]) input_ids.append(x[input_ids] [0] * pad_len) labels.append(x[labels] [-100] * pad_len) return { input_ids: torch.tensor(input_ids), labels: torch.tensor(labels), }這里把padding token的id直接寫成了0是因為ChatGLM的pad_token_id就是0。如果換了其他模型記得用tokenizer.pad_token_id獲取不要寫死。梯度檢查點配合動態(tài)padding是當前單卡24GB或4卡121GB這種中低配環(huán)境里把ChatGLM-6B LoRA微調(diào)跑起來的關(guān)鍵組合。我自己在做這類多卡微調(diào)項目時養(yǎng)成了一個習慣每次改完配置先跑一個100步的smoke test只驗證三件事——loss下降方向?qū)Σ粚?、顯存占用是否有大幅波動、保存下來的checkpoint能不能重新加載。三關(guān)全過再放長訓練。這個習慣幫我在多機多卡場景下避開了大量看似玄學的翻車很多問題其實早在小規(guī)模測試階段就會暴露。希望幫到你。本文還有配套的精品資源點擊獲取 DeepSpeed多卡ChatGLM微調(diào)實戰(zhàn)：從顯存優(yōu)化到LoRA選型
簡介面向需要上手大模型微調(diào)的開發(fā)者與研究者這份資源以DeepSpeed多卡訓練為主線完整覆蓋ChatGLM微調(diào)實戰(zhàn)流程并提供可直接運行的模塊化源碼與流程教程。項目包含詳細的環(huán)境搭建、數(shù)據(jù)準備與格式化、LoRA/P-tuning/Freeze等多類微調(diào)方案同時講解ds_config.json等關(guān)鍵配置參數(shù)、訓練啟動腳本寫法和常見踩坑點有助于讀者跨越單卡顯存限制快速搭建多卡并行訓練環(huán)境。壓縮包共17個文件以Python訓練腳本11個、Shell啟動腳本3個、JSON配置文件2個和Markdown說明文檔1個為主大小僅118KB目錄劃分明確便于按模塊對照學習。已有267人瀏覽學習。借助源碼中的模型加載、數(shù)據(jù)加載、訓練循環(huán)、推理腳本和評估測試等模塊讀者不僅能復現(xiàn)ChatGLM多卡微調(diào)還能理解DeepSpeed的內(nèi)存優(yōu)化與梯度累積機制是一份可直接落地的優(yōu)質(zhì)實戰(zhàn)教程。1. 多卡ChatGLM微調(diào)為什么大家都卡在DeepSpeed這一關(guān)把ChatGLM-6B塞進一張A100-80G里做推理是輕松的但一旦開始全量微調(diào)Adam優(yōu)化器的狀態(tài)一上來顯存立刻不夠用。很多人對“多卡微調(diào)”的第一反應是插4張卡等于4倍顯存真正跑起來才發(fā)現(xiàn)權(quán)重、梯度、優(yōu)化器狀態(tài)這堆東西不切開4張卡也只是各自抱著同一份副本在空轉(zhuǎn)。DeepSpeed就是為解決這個問題出現(xiàn)的它把模型狀態(tài)切到各卡上讓基于DeepSpeed的多卡ChatGLM微調(diào)從“跑得動”變成“跑得起”。這篇實戰(zhàn)筆記按環(huán)境搭建、LoRA與全量選型、deepspeed_config.json參數(shù)、多卡啟動和踩坑排查的順序給出一套可以直接復制到你自己機器上的完整流程。適合手里有2到8張GPU、想把ChatGLM微調(diào)成垂直場景模型的算法工程師和獨立開發(fā)者。2. 微調(diào)前的選型全量微調(diào)和LoRA微調(diào)先算一筆賬2.1 全量微調(diào)為什么在多卡上這么貴全量微調(diào)意味著模型每一層參數(shù)都在更新訓練過程中需要同時存放五樣東西權(quán)重本身、梯度、優(yōu)化器的一階動量、二階動量以及混合精度下的FP32主權(quán)重。工程上常用一個近似數(shù)來算賬每1B參數(shù)大約要占16GB的模型狀態(tài)。ChatGLM-6B換算下來光是參數(shù)、梯度和優(yōu)化器狀態(tài)就要接近96GB。一張A100-80G單卡裝不下即便勉強裝下batch只能設(shè)為1訓練效率低到?jīng)]法看。多卡場景里最直接的做法是Data Parallel也就是每張卡復制一份完整模型各自算梯度再做一次allreduce合并。4張卡就是4份96GB顯存不僅沒省反而因為通信開銷變得更慢。DeepSpeed在這里的核心價值是ZeRO它把原本每張卡都保存一份的優(yōu)化器狀態(tài)、梯度甚至參數(shù)切分到不同卡上通信只發(fā)生在真正需要同步的時刻。這也是“多卡微調(diào)為什么繞不開DeepSpeed”的最直接原因?qū)τ谝粋€6B量級的模型光模型狀態(tài)就已經(jīng)超過單卡物理上限不用ZeRO只能靠堆顯存解決而堆顯存的性價比太低。這里還需要區(qū)分一個概念模型狀態(tài)和激活值。模型狀態(tài)指的是參數(shù)、梯度、優(yōu)化器狀態(tài)激活值是前向傳播過程中每一層輸出的中間張量。很多人以為開了DeepSpeed就能解決所有顯存問題結(jié)果發(fā)現(xiàn)batch一大照樣OOM這就是激活值在瓶頸上。后面第4章會給出一個完整的顯存估算方法先把模型狀態(tài)算明白再結(jié)合batch和序列長度去反推激活值的余量。2.2 LoRA微調(diào)是什么意思參數(shù)少到什么程度剛接觸大模型微調(diào)的人常問LoRA微調(diào)是什么意思。簡單說LoRA凍結(jié)原模型的全部權(quán)重在每一層通常是注意力層的線性映射旁邊掛一組低秩矩陣A和B只訓練這兩個小矩陣。最終更新量等于alpha/r乘以AB的結(jié)果再疊加到原始權(quán)重上。由于原始權(quán)重完全不更新可訓練參數(shù)量會從6B級別直接掉到百萬級別。以ChatGLM-6B為例在query_key_value上掛LoRAr設(shè)為8可訓練參數(shù)大約只有三百萬到四百萬比原始參數(shù)小了三個數(shù)量級。這會帶來兩個直接好處一是模型狀態(tài)顯存占用大幅下降因為需要保存的參數(shù)、梯度和優(yōu)化器狀態(tài)都只針對這小部分可訓練參數(shù)二是訓練過程不容易把基座的語言能力帶偏特別適合數(shù)據(jù)量不大、目標明確的垂直微調(diào)場景。from peft import LoraConfig, get_peft_model lora_config LoraConfig( r8, # 低秩矩陣的秩常用 8/16/32 lora_alpha16, # 縮放系數(shù)實際更新量乘 alpha/r target_modules[query_key_value], # ChatGLM 注意力層線性映射的參數(shù)名 lora_dropout0.1, biasnone, ) model get_peft_model(model, lora_config) model.print_trainable_parameters()這里有幾個參數(shù)需要根據(jù)實際任務調(diào)r太小表達能力不夠r太大可訓練參數(shù)變多顯存優(yōu)勢和微調(diào)“安全邊際”都會縮水。lora_alpha和r的比值決定LoRA更新的影響幅度一般alpha取r的1到2倍。target_modules要根據(jù)模型結(jié)構(gòu)改ChatGLM的注意力層里query、key、value是合并在一個名為query_key_value的線性層里的所以要寫這個參數(shù)名如果微調(diào)ChatGLM2或ChatGLM3這個名稱可能不同先打印模型結(jié)構(gòu)確認再填。對比全量微調(diào)LoRA微調(diào)的選擇邏輯可以歸納為數(shù)據(jù)量小、任務垂直度高、希望快速迭代選LoRA數(shù)據(jù)量足夠大、顯卡資源充足、追求極限效果才考慮全量微調(diào)。對于多數(shù)真實項目LoRA都是第一版方案因為它能在一個晚上跑完而且效果差距往往沒有想象中大。對比項全量微調(diào)LoRA微調(diào)更新參數(shù)范圍全部參數(shù)低秩矩陣模型狀態(tài)顯存約16GB/B參數(shù)約0.1GB/B參數(shù)訓練速度慢快3到5倍基座能力保持容易遺忘較好適用場景數(shù)據(jù)量大、資源足垂直微調(diào)、快速迭代2.3 ChatGLM的數(shù)據(jù)拼接和標簽掩碼Prefix-LM結(jié)構(gòu)下的labelsChatGLM是Prefix-LM結(jié)構(gòu)左側(cè)的輸入可以做雙向注意力右側(cè)的生成部分只能看左側(cè)。這和GPT的Causal-LM不一樣訓練數(shù)據(jù)不能直接當普通CausalLM來拼接。社區(qū)里最常見的做法是把指令和回答拼成一條完整文本在回答開始的位置之前把標簽全部設(shè)為-100計算loss時這些位置被忽略。def pack_chat_input(tokenizer, source, answer, max_len2048): # 拼接格式和ChatGLM官方prompt模板保持一致 text 問 source \n答 answer tokenizer.eos_token ids tokenizer.encode(text, max_lengthmax_len, truncationTrue) # 計算prompt部分的長度這部分不參與loss計算 prompt 問 source \n答 prompt_len len(tokenizer.encode(prompt)) labels [-100] * len(ids) # 先把所有位置設(shè)為-100 labels[prompt_len:] ids[prompt_len:] # 回答部分保留真實token return {input_ids: ids, labels: labels}這段代碼的核心邏輯是“回答token要學習指令token不學習”。-100是PyTorch CrossEntropyLoss約定俗成的ignore索引模型在算loss時會自動跳過這些位置。注意這里的prompt_len用的是tokenizer.encode不是字符長度因為中文文本的token切分不是按字來的。訓練循環(huán)里有一個特別容易翻車的細節(jié)用了DeepSpeed之后反向傳播和參數(shù)更新必須交給engine不能直接調(diào)model.backward()或optimizer.step()。原因在于ZeRO把模型狀態(tài)切到了多張卡上只有engine知道當前rank上持有哪部分狀態(tài)。for step, batch in enumerate(dataloader): outputs engine( input_idsbatch[input_ids].cuda(), labelsbatch[labels].cuda(), ) engine.backward(outputs.loss) engine.step()這里的engine就是deepspeed.initialize返回的模型包裝對象。前向接口和原始模型基本一致但backward和step必須走engine。如果自己額外再調(diào)一次optimizer.step()輕則梯度更新錯亂重則訓練過程直接崩掉。3. 環(huán)境搭建與最小復現(xiàn)從安裝deepspeed到多卡啟動3.1 安裝deepspeed包先過CUDA這一關(guān)搜索“安裝deepspeed包一直報錯”你會發(fā)現(xiàn)十個里有八個是同一個原因pip在安裝時嘗試編譯CUDA拓展但系統(tǒng)里找不到CUDA工具鏈。報錯信息里通常會有一大段gcc、nvcc相關(guān)的日志很多人看到就懵了其實解決辦法很簡單先確認PyTorch的CUDA版本再檢查nvcc能不能用。# 先確認torch的CUDA版本再決定裝哪個deepspeed python -c import torch; print(torch.version.cuda) # 安裝前確認nvcc可用 nvcc --version # 常規(guī)安裝方式 pip install deepspeed如果pip直接從源碼編譯導致報錯可以先用預編譯的wheel包緩解。另一個常見做法是設(shè)置CUDA_HOME環(huán)境變量指向CUDA Toolkit的安裝目錄再執(zhí)行pip install deepspeed。安裝完成后用ds_report命令檢查deepspeed是否識別到當前GPU、CUDA版本和NCCL版本。這個命令會打印一份環(huán)境診斷報告看到“OK”字樣說明基礎(chǔ)環(huán)境沒有問題。需要注意deepspeed會在第一次啟動訓練時做JIT編譯編譯一些CUDA算子到緩存目錄所以第一次啟動多等一兩分鐘是正常的。如果多次啟動都卡在同一處編譯多半是系統(tǒng)缺了gcc或者編譯依賴看日志里的具體報錯再補裝對應依賴。3.2 寫一個最小訓練腳本train_chatglm_lora.py環(huán)境就緒后先別急著上全量微調(diào)我一般用LoRA把整個鏈路先跑通。下面這個腳本是完整可運行的最小版本把模型加載、LoRA包裝、DeepSpeed初始化和訓練循環(huán)放在一起總共不到40行。import torch import deepspeed from transformers import AutoTokenizer, AutoModel from peft import LoraConfig, get_peft_model def train(): model_name THUDM/chatglm-6b tokenizer AutoTokenizer.from_pretrained(model_name, trust_remote_codeTrue) # bf16要求Ampere及以上架構(gòu)V100等舊卡請改成fp16 model AutoModel.from_pretrained( model_name, trust_remote_codeTrue, torch_dtypetorch.bfloat16, ) # 用LoRA凍結(jié)大部分參數(shù)只訓練query_key_value上的低秩矩陣 model get_peft_model(model, LoraConfig( r8, lora_alpha16, target_modules[query_key_value], lora_dropout0.1, biasnone, )) # 梯度檢查點用計算換顯存多卡顯存緊張時建議打開 model.gradient_checkpointing_enable() engine, optimizer, _, lr_scheduler deepspeed.initialize( modelmodel, model_parameters[p for p in model.parameters() if p.requires_grad], configds_config.json, ) for step, batch in enumerate(dataloader): outputs engine( input_idsbatch[input_ids].cuda(), labelsbatch[labels].cuda(), ) engine.backward(outputs.loss) engine.step()這段代碼里有兩個地方最容易出錯。一是AutoModel.from_pretrained必須帶trust_remote_codeTrue因為ChatGLM的模型結(jié)構(gòu)通過遠程代碼加載二是deepspeed.initialize的model_parameters參數(shù)只傳可訓練參數(shù)LoRA模式下就是那幾個低秩矩陣不要傳model.parameters()的完整列表否則DeepSpeed會把凍結(jié)參數(shù)也納入優(yōu)化器管理白白浪費顯存。dataloader需要配合第2.3節(jié)的pack_chat_input函數(shù)把每條樣本處理成input_ids和labels。注意collate時要動態(tài)padding到batch內(nèi)最長而不是固定到2048否則顯存大部分都被無效的padding token吃掉了。3.3 啟動命令單機多卡和多機多卡的差異訓練腳本寫好后啟動方式用deepspeed命令不是python命令。單機4卡的最簡命令如下deepspeed --num_gpus 4 --master_port 29500 train_chatglm_lora.pymaster_port是分布式訓練的控制通信端口默認29500。如果同一臺機器上同時跑多個訓練任務端口會沖突啟動時會報端口被占用的錯誤換一個比如29501就行。如果你用的是SLURM之類的調(diào)度系統(tǒng)還要帶上deepspeed_multinode腳本一起用這里不展開。多機多卡稍微復雜一點每臺機器上都要執(zhí)行啟動命令區(qū)別只在node_rank和master_addr# 在第一個節(jié)點執(zhí)行 deepspeed --num_nodes 2 --num_gpus 8 --node_rank 0 \ --master_addr 192.168.1.10 --master_port 29500 train_chatglm_lora.py # 在第二個節(jié)點執(zhí)行 deepspeed --num_nodes 2 --num_gpus 8 --node_rank 1 \ --master_addr 192.168.1.10 --master_port 29500 train_chatglm_lora.py多機場景下master_addr必須填第一臺機器的IP而且兩臺機器要能互相訪問。這里有個容易踩的坑如果機器之間只有以太網(wǎng)沒有InfiniBand通信速度會成為訓練瓶頸需要顯式設(shè)置NCCL_SOCKET_IFNAME指向?qū)嶋H網(wǎng)卡接口名。判斷方法是在訓練日志里看NCCL初始化時用了哪塊網(wǎng)卡如果是lo或docker0訓練速度大概率不正常。4. deepspeed_config.json先設(shè)這幾個參數(shù)再談訓練收斂4.1 ZeRO Stage選1、2還是3zero123的區(qū)別第一次寫deepspeed_config.json卡住最多的問題就是deepspeed zero123的區(qū)別。簡單說Stage 1只切分優(yōu)化器狀態(tài)Stage 2在Stage 1基礎(chǔ)上連梯度也切分Stage 3更進一步把模型參數(shù)本身也切分到各卡上。Stage越高顯存越省但卡間通信量越大訓練速度越慢。ChatGLM-6B微調(diào)我的建議是優(yōu)先Stage 2。6B模型的參數(shù)量不算特別大Stage 2已經(jīng)能把優(yōu)化器狀態(tài)和梯度的重復存儲省掉顯存能被壓到接近“每卡只存一份完整權(quán)重部分狀態(tài)”的水平。Stage 3雖然有額外收益但通信開銷會拖慢迭代速度在4卡或8卡規(guī)模下性價比不高。只有當顯存實在不夠、batch小到?jīng)]法看的時候才考慮Stage 3或者把優(yōu)化器狀態(tài)offload到CPU。DeepSpeed初始化時會打印一份運行配置摘要里面包括當前用的是哪個Stage、每張卡預估的顯存使用量。啟動日志不要直接跳過這些信息比任何第三方工具都準。Stage 3還要注意一個問題參數(shù)被切分后前向和反向過程中需要動態(tài)收集參數(shù)所以eval和保存模型時要額外處理對新手來說這是比顯存更麻煩的坑。4.2 我習慣先寫好的幾個配置項下面這份ds_config.json是我做ChatGLM微調(diào)時最常使用的配置模板每一項都直接影響訓練能否收斂。{ train_batch_size: 32, gradient_accumulation_steps: 4, gradient_clipping: 1.0, bf16: { enabled: true }, zero_optimization: { stage: 2, offload_optimizer: { device: cpu } }, scheduler: { type: WarmupDecayLR, params: { warmup_min_lr: 0, warmup_max_lr: 2e-5, warmup_num_steps: 200 } }, steps_per_print: 20 }第一眼要搞清楚的是train_batch_size、gradient_accumulation_steps和實際每卡batch三者的關(guān)系。train_batch_size是全局batch也就是所有卡累積多少條樣本做一次參數(shù)更新。實際每卡micro batch的計算公式是train_batch_size除以卡數(shù)再除以梯度累積步數(shù)。以這份配置為例4卡、累積4步、全局32那么每卡micro batch是2。很多人在這一步憑感覺填數(shù)字結(jié)果實際batch大小和預期完全不是一回事loss曲線自然不對。配置項取值建議作用train_batch_size16到64全局batch決定每次更新用多少樣本gradient_accumulation_steps2到8梯度累積步數(shù)等效增大batchgradient_clipping1.0防止梯度爆炸微調(diào)大模型必備bf16enabled: true半精度訓練A100及以上建議開offload_optimizerdevice: cpu優(yōu)化器狀態(tài)放內(nèi)存顯存緊張時開warmup_max_lr1e-5到5e-5峰值學習率6B模型別超過5e-5bf16和fp16的區(qū)別也值得多提一句。bf16在A100、H100這些Ampere之后架構(gòu)的卡上效果更好指數(shù)位更多不容易出現(xiàn)loss溢出V100等舊架構(gòu)不支持bf16只能開fp16這時要額外關(guān)注loss scale防止訓練中期loss變成NaN。4.3 用顯存公式反推可用的batch和序列長度顯存占用粗略分為兩塊模型狀態(tài)加激活值。模型狀態(tài)可以用第2.1節(jié)的16GB/B參數(shù)估算激活值則和micro batch大小、序列長度、模型層數(shù)強相關(guān)。ChatGLM-6B在max_seq_len為2048時激活值隨batch增長的斜率非常明顯很多人以為batch設(shè)2沒問題結(jié)果跑到第100步就OOM。我建議的排查順序是先固定序列長度把micro batch調(diào)到1跑一輪用nvidia-smi看顯存占用基線再逐步往上加。# 每秒刷新一次顯存情況觀察訓練過程中的峰值占用 watch -n 1 nvidia-smi看重點不是顯存總量而是進程占用那一欄。如果micro batch為1時已經(jīng)用了接近全部顯存說明瓶頸在激活值這時候調(diào)低max_len或者關(guān)掉冗余輸入更有效。如果micro batch為1時顯存還有余量才把batch往上加。DeepSpeed在啟動時也會給一個顯存預估可以和實際觀測對照幫助判斷配置是否生效。這里有一個容易誤判的地方打開了offload_optimizer之后訓練剛開始時顯存會明顯下降但CPU內(nèi)存占用會漲上去。如果服務器內(nèi)存本身不多offload反而可能引發(fā)內(nèi)存交換訓練速度驟降。觀察指標時不能只看GPU顯存還要帶著看內(nèi)存占用。5. DeepSpeed多卡微調(diào)ChatGLM的四個典型踩坑現(xiàn)場5.1 NCCL超時卡在Waiting for the barrier現(xiàn)象啟動多卡訓練后日志卡在分布式初始化階段幾分鐘后拋NCCLTimeoutError或者某個rank直接掛掉。原因最常見的是master_addr或master_port配置不對多卡之間的控制通信建立不起來其次是機器上有多塊網(wǎng)卡NCCL選錯了網(wǎng)卡導致通信走了一個慢速通道甚至在不通的網(wǎng)口上。還有一種情況是防火墻攔了分布式訓練的端口單機多卡時不太會遇到多機多卡時概率很高。解決先確認所有節(jié)點能互相ping通再檢查deepspeed啟動命令里的master_addr和master_port是否一致。多機場景顯式指定通信網(wǎng)卡export NCCL_SOCKET_IFNAMEeth0 export NCCL_IB_DISABLE1NCCL_IB_DISABLE1是讓NCCL不走InfiniBand只走以太網(wǎng)。如果機器之間的IB網(wǎng)絡(luò)沒有正確配置與其讓它自動探測失敗不如直接關(guān)掉。日志里能看到NCCL初始化時用的網(wǎng)卡名字如果發(fā)現(xiàn)選錯了再用NCCL_SOCKET_IFNAME強制指定。5.2 訓練中途OOM先分清是激活值還是模型狀態(tài)壓爆的現(xiàn)象訓練能啟動但跑到某個step突然報CUDA out of memory而且每次崩的位置不一樣。原因大多數(shù)時候是動態(tài)padding沒有生效batch里某條樣本特別長把激活值峰值瞬間拉高。另一個常見原因是micro batch算錯了以為設(shè)置了梯度累積batch就小了實際上每卡的micro batch仍然是全局batch除以卡數(shù)沒有除以累積步數(shù)。解決先按第4.3節(jié)的方法把micro batch降到1試跑如果不再OOM說明是batch和序列長度的問題如果batch為1仍然OOM說明模型狀態(tài)部分還有優(yōu)化空間。此時從三個方向同時壓打開gradient_checkpointing啟用梯度檢查點、把offload_optimizer設(shè)備設(shè)為cpu、以及確認沒有把凍結(jié)參數(shù)傳進優(yōu)化器。這三招用完ChatGLM-6B的LoRA微調(diào)在24GB顯存卡上都可以跑起來。注意開了gradient_checkpointing之后前向速度會變慢這是正常的它本質(zhì)上是犧牲計算換顯存。5.3 loss不降或亂跳學習率、warmup和數(shù)據(jù)順序一起背鍋現(xiàn)象訓練了上千步loss要么原地不動要么忽高忽低看起來完全沒有收斂特征。原因把之前訓練小模型的經(jīng)驗搬到大模型上學習率用1e-4甚至更高對6B模型來說步長太大loss自然亂跳。另一個是warmup步數(shù)太少或沒設(shè)置調(diào)度器一上來就給滿學習率。還有一個容易被忽略的點dataloader沒有shuffle模型反復看到同一條指令loss曲線會出現(xiàn)周期性波動。解決把warmup_max_lr降到2e-5warmup_num_steps設(shè)為總步數(shù)的5%到10%。檢查dataloader是否設(shè)置了shuffleTrue。如果用的是fp16而不是bf16還要觀察訓練日志里的loss scale出現(xiàn)NaN或不正常跳變說明fp16的精度控制出了問題建議升級到支持bf16的顯卡或者調(diào)整fp16配置。這里要特別說一個誤區(qū)不要因為loss在下降就認定訓練正常。大模型微調(diào)初期loss下降可以很快但真正的考驗是第6章的驗證環(huán)節(jié)生成質(zhì)量和loss下降并不完全掛鉤。5.4 微調(diào)完導出eval時模型像沒訓練過現(xiàn)象訓練日志里loss已經(jīng)降得很低但eval時模型回答和原版幾乎一樣甚至更差。原因最常見的是保存模型時只保存了DeepSpeed的引擎權(quán)重沒有保存LoRA adapter。DeepSpeed的save_checkpoint保存的是引擎狀態(tài)直接拿去推理完全用不上。另一個原因是沒有merge LoRA權(quán)重推理時忘了加載adapter。解決訓練結(jié)束后用PeFT的save_pretrained保存LoRA權(quán)重推理時先加載adapter再合并# 訓練結(jié)束保存LoRA權(quán)重 engine.module.save_pretrained(lora_adapter) # 推理時加載并合并 from transformers import AutoModel model AutoModel.from_pretrained(THUDM/chatglm-6b, trust_remote_codeTrue) model.load_adapter(lora_adapter) model model.merge_and_unload()注意推理時加載基座模型的精度要和訓練時一致否則merge可能出現(xiàn)數(shù)值不匹配的問題生成結(jié)果看起來就像沒訓練過。eval時使用訓練時的prompt模板也很關(guān)鍵微調(diào)時用的格式是“問xxx\n答”eval時如果換成其他格式模型生成質(zhì)量會受到影響。6. 從loss下降走向業(yè)務能用驗證流程和一個顯存技巧6.1 用固定prompt集驗證效果別只看loss訓練loss下降只是“模型記住了訓練數(shù)據(jù)”的必要條件不是充分條件。我的做法是準備一組固定業(yè)務問題訓練過程中每隔幾百步用生成模式跑一遍觀察回答質(zhì)量的變化趨勢。這一步能比loss曲線更早暴露出模型是否過擬合、是否丟失了對話能力。生成測試代碼可以掛在訓練循環(huán)里def eval_generate(engine, tokenizer, questions, max_new_tokens128): engine.module.eval() with torch.no_grad(): for q in questions: prompt 問 q \n答 inputs tokenizer(prompt, return_tensorspt).to(engine.device) out engine.module.generate( **inputs, max_new_tokensmax_new_tokens, do_sampleFalse, ) print(tokenizer.decode(out[0], skip_special_tokensTrue)) engine.module.train()固定prompt集的規(guī)模不用大10條左右就夠了但要覆蓋正常提問、邊界提問和易混淆提問三類。do_sample設(shè)為False保證可對比如果模型訓練中關(guān)了dropout這一步的生成結(jié)果是完全確定的。每500步跑一次保存生成結(jié)果到日志文件后面就可以回看不同訓練階段的質(zhì)量變化。另外把一部分數(shù)據(jù)留作held-out集驗證集loss比訓練loss更早上升就是過擬合信號。6.2 把顯存再壓一檔動態(tài)padding和梯度檢查點組合第5.2節(jié)提到動態(tài)padding能解決一部分OOM問題這里給出一個可用的collate實現(xiàn)。核心思路是把batch內(nèi)的樣本padding到同一長度同時用-100掩蓋無效的標簽位置。def collate_batch(batch): max_len max(len(x[input_ids]) for x in batch) input_ids [] labels [] for x in batch: pad_len max_len - len(x[input_ids]) input_ids.append(x[input_ids] [0] * pad_len) labels.append(x[labels] [-100] * pad_len) return { input_ids: torch.tensor(input_ids), labels: torch.tensor(labels), }這里把padding token的id直接寫成了0是因為ChatGLM的pad_token_id就是0。如果換了其他模型記得用tokenizer.pad_token_id獲取不要寫死。梯度檢查點配合動態(tài)padding是當前單卡24GB或4卡121GB這種中低配環(huán)境里把ChatGLM-6B LoRA微調(diào)跑起來的關(guān)鍵組合。我自己在做這類多卡微調(diào)項目時養(yǎng)成了一個習慣每次改完配置先跑一個100步的smoke test只驗證三件事——loss下降方向?qū)Σ粚?、顯存占用是否有大幅波動、保存下來的checkpoint能不能重新加載。三關(guān)全過再放長訓練。這個習慣幫我在多機多卡場景下避開了大量看似玄學的翻車很多問題其實早在小規(guī)模測試階段就會暴露。希望幫到你。本文還有配套的精品資源點擊獲取