1. 訓(xùn)練側(cè)顯存測(cè)量到底在測(cè)什么顯存優(yōu)化這件事很多人一上來(lái)就想著怎么省結(jié)果省了半天發(fā)現(xiàn)根本沒(méi)省到點(diǎn)子上。問(wèn)題出在哪出在沒(méi)搞清楚顯存到底被誰(shuí)吃掉了。訓(xùn)練側(cè)的顯存測(cè)量核心目標(biāo)就一個(gè)把顯存賬單拆開(kāi)看清楚每一筆開(kāi)銷的去向然后才能做預(yù)算決策。我見(jiàn)過(guò)太多人拿著nvidia-smi看一眼顯存占用發(fā)現(xiàn)快滿了就開(kāi)始慌然后盲目上梯度檢查點(diǎn)、盲目降batch size最后訓(xùn)練速度掉了一半顯存也沒(méi)省下多少。這種做法的問(wèn)題在于nvidia-smi看到的只是一個(gè)總數(shù)它不會(huì)告訴你這8GB里有3GB是模型參數(shù)、2GB是優(yōu)化器狀態(tài)、1.5GB是激活值、剩下的是碎片和臨時(shí)緩沖區(qū)。你不知道錢花在哪就沒(méi)辦法做預(yù)算。訓(xùn)練側(cè)顯存測(cè)量要回答的問(wèn)題很具體模型參數(shù)占多少、梯度占多少、優(yōu)化器狀態(tài)占多少、激活值占多少、臨時(shí)緩沖區(qū)占多少。這五塊加起來(lái)才是真實(shí)賬單。而且這五塊的比例關(guān)系會(huì)隨著模型規(guī)模、batch size、序列長(zhǎng)度、精度策略的變化而劇烈變化。一個(gè)7B模型在FP16下參數(shù)占14GB但如果用Adam優(yōu)化器優(yōu)化器狀態(tài)就要占28GBFP32的一階動(dòng)量和二階動(dòng)量各14GB再加上梯度14GB光這三項(xiàng)就56GB了還沒(méi)算激活值。所以為什么大家說(shuō)7B模型全量微調(diào)至少要80GB顯存賬就是這么算出來(lái)的。測(cè)量方法上最直接的是用PyTorch的顯存分析工具。torch.cuda.memory_allocated()能拿到當(dāng)前分配的顯存torch.cuda.max_memory_allocated()能拿到峰值。但這兩個(gè)數(shù)字只反映PyTorch分配器層面的情況不包括CUDA上下文本身占用的那幾百M(fèi)B。更細(xì)的拆解需要用torch.cuda.memory_summary()它會(huì)按分配塊大小分類列出。不過(guò)這個(gè)輸出比較原始我一般會(huì)自己寫個(gè)hook在模型forward前后、backward前后分別打點(diǎn)算出每個(gè)階段的增量。還有一個(gè)容易被忽略的點(diǎn)顯存碎片。PyTorch的緩存分配器會(huì)預(yù)留一些顯存不還給系統(tǒng)導(dǎo)致nvidia-smi看到的占用比memory_allocated()高不少。這個(gè)差值在長(zhǎng)時(shí)間訓(xùn)練中會(huì)逐漸增大尤其是當(dāng)你有動(dòng)態(tài)shape的輸入時(shí)。測(cè)量的時(shí)候要把這個(gè)差值也記錄下來(lái)否則你按memory_allocated()做的預(yù)算到了實(shí)際訓(xùn)練時(shí)就會(huì)OOM。注意測(cè)量一定要在真實(shí)訓(xùn)練循環(huán)里做不能只跑一個(gè)forward就完事。因?yàn)榧せ钪档姆逯党霈F(xiàn)在backward階段只測(cè)forward會(huì)嚴(yán)重低估。2. 顯存賬單的五大部分與計(jì)算邏輯2.1 模型參數(shù)與梯度的顯存占用模型參數(shù)的顯存占用最好算參數(shù)量乘以每個(gè)參數(shù)的字節(jié)數(shù)。FP32是4字節(jié)FP16和BF16是2字節(jié)INT8是1字節(jié)。一個(gè)7B模型在FP16下就是7B乘2等于14GB。梯度占用的字節(jié)數(shù)和參數(shù)一致因?yàn)樘荻刃枰蛥?shù)同樣的精度來(lái)保證更新時(shí)不丟信息。所以FP16訓(xùn)練時(shí)參數(shù)加梯度就是28GB。但這里有個(gè)坑很多人以為用FP16訓(xùn)練參數(shù)就是FP16。實(shí)際上PyTorch的AMP自動(dòng)混合精度會(huì)保留一份FP32的master weight。也就是說(shuō)參數(shù)實(shí)際上占了兩份一份FP16用于前向和反向計(jì)算一份FP32用于優(yōu)化器更新。這樣算下來(lái)7B模型的參數(shù)占用是14GBFP16加28GBFP32 master總共42GB。梯度也是FP16一份但優(yōu)化器更新時(shí)需要FP32的梯度所以梯度實(shí)際占用也是14GBFP16加28GBFP32又是42GB。這就是為什么AMP能省顯存但省不了太多——它省的是激活值和部分計(jì)算緩沖區(qū)參數(shù)和優(yōu)化器狀態(tài)的大頭省不掉。2.2 優(yōu)化器狀態(tài)的顯存黑洞優(yōu)化器狀態(tài)是顯存占用里最容易被低估的部分。以Adam為例它需要為每個(gè)參數(shù)維護(hù)一階動(dòng)量m和二階動(dòng)量v都是FP32。所以優(yōu)化器狀態(tài)的顯存等于參數(shù)量乘以4字節(jié)乘以2也就是參數(shù)量乘以8字節(jié)。7B模型就是56GB。加上前面的參數(shù)和梯度已經(jīng)98GB了。這就是為什么全量微調(diào)7B模型至少需要8張80GB的卡來(lái)做數(shù)據(jù)并行——單卡根本放不下。AdamW稍微好一點(diǎn)它把權(quán)重衰減和梯度更新解耦了但動(dòng)量部分還是一樣的。Adafactor通過(guò)分解二階動(dòng)量矩陣來(lái)省顯存能把優(yōu)化器狀態(tài)降到參數(shù)量乘以4字節(jié)左右但收斂性會(huì)受一些影響。Sophia優(yōu)化器用對(duì)角Hessian估計(jì)來(lái)替代二階動(dòng)量也能省不少但實(shí)現(xiàn)復(fù)雜度和調(diào)參難度都上去了。實(shí)際做預(yù)算的時(shí)候我一般按這個(gè)公式估算總顯存等于參數(shù)量乘以2加2加8加激活值加緩沖區(qū)。前面的2是FP16參數(shù)第二個(gè)2是FP16梯度8是Adam的優(yōu)化器狀態(tài)。這個(gè)公式在AMP加AdamW的場(chǎng)景下比較準(zhǔn)誤差在10%以內(nèi)。2.3 激活值的動(dòng)態(tài)波動(dòng)激活值是顯存占用里最動(dòng)態(tài)的部分它跟batch size、序列長(zhǎng)度、模型層數(shù)、隱藏維度都相關(guān)。粗略估算的話激活值顯存約等于batch size乘以序列長(zhǎng)度乘以隱藏維度乘以層數(shù)乘以一個(gè)系數(shù)。這個(gè)系數(shù)取決于具體的網(wǎng)絡(luò)結(jié)構(gòu)Transformer里主要是注意力矩陣和FFN的中間激活。注意力矩陣的顯存是序列長(zhǎng)度的平方乘以batch size乘以頭數(shù)乘以2字節(jié)。序列長(zhǎng)度2048時(shí)這個(gè)平方項(xiàng)就是4M乘以batch size和頭數(shù)后很容易上GB。序列長(zhǎng)度拉到8192平方項(xiàng)變成64M直接爆炸。這就是長(zhǎng)上下文訓(xùn)練顯存吃緊的核心原因。FFN的中間激活是batch size乘以序列長(zhǎng)度乘以4倍隱藏維度乘以2字節(jié)。這個(gè)和序列長(zhǎng)度是線性關(guān)系比注意力矩陣溫和一些。但層數(shù)一多累積起來(lái)也很可觀。梯度檢查點(diǎn)gradient checkpointing就是針對(duì)激活值的優(yōu)化手段。它不保存中間激活而是在backward時(shí)重新計(jì)算。代價(jià)是計(jì)算量增加約30%但激活值顯存能降到原來(lái)的平方根級(jí)別。對(duì)于層數(shù)很深的模型這個(gè) trade-off 非常劃算。2.4 臨時(shí)緩沖區(qū)與碎片臨時(shí)緩沖區(qū)包括CUDA kernel執(zhí)行時(shí)的workspace、通信操作的緩沖區(qū)、以及PyTorch分配器的預(yù)留空間。這部分很難精確測(cè)量但可以通過(guò)對(duì)比nvidia-smi和memory_allocated()的差值來(lái)估算。一般來(lái)說(shuō)這個(gè)差值在500MB到2GB之間模型越大、并行策略越復(fù)雜差值越大。碎片問(wèn)題在動(dòng)態(tài)shape場(chǎng)景下特別嚴(yán)重。比如你訓(xùn)練時(shí)序列長(zhǎng)度不固定PyTorch的緩存分配器會(huì)按最大shape預(yù)留塊導(dǎo)致實(shí)際占用遠(yuǎn)高于理論值。解決辦法是設(shè)置PYTORCH_CUDA_ALLOC_CONF環(huán)境變量啟用expandable_segments讓分配器能合并碎片。這個(gè)設(shè)置在新版PyTorch里效果很明顯我實(shí)測(cè)能把碎片率從20%降到5%以下。3. 預(yù)算決策從測(cè)量結(jié)果到優(yōu)化方案3.1 顯存預(yù)算的分配原則測(cè)完顯存賬單后下一步是做預(yù)算決策。預(yù)算的核心原則是先保訓(xùn)練穩(wěn)定性再保吞吐量最后才考慮省顯存。很多人搞反了順序?yàn)榱耸★@存把batch size降到1結(jié)果訓(xùn)練不穩(wěn)定收斂慢得要命省下來(lái)的顯存也沒(méi)換來(lái)什么好處。我的預(yù)算分配一般是這樣的參數(shù)和優(yōu)化器狀態(tài)是剛性支出沒(méi)法省必須留足。激活值是彈性支出可以通過(guò)梯度檢查點(diǎn)、序列并行、FlashAttention等手段壓縮。臨時(shí)緩沖區(qū)留10%到15%的余量防止峰值OOM。具體到數(shù)字上假設(shè)你有80GB顯存7B模型AMP加AdamW參數(shù)加梯度加優(yōu)化器狀態(tài)大約98GB單卡放不下。這時(shí)候你有幾個(gè)選擇一是上ZeRO Stage 2把優(yōu)化器狀態(tài)和梯度分片到多卡每卡只存一部分二是上LoRA只訓(xùn)練低秩適配器參數(shù)量降到原來(lái)的百分之一三是上QLoRA把基座模型量化到4bit進(jìn)一步壓縮。3.2 ZeRO與FSDP的顯存賬ZeRO Stage 1只分片優(yōu)化器狀態(tài)每卡顯存等于參數(shù)量乘以2加梯度乘以2加優(yōu)化器狀態(tài)除以N。7B模型8卡每卡優(yōu)化器狀態(tài)7GB加上參數(shù)14GB和梯度14GB總共35GB80GB卡能放下。Stage 2再分片梯度每卡梯度降到1.75GB總共約23GB。Stage 3分片參數(shù)每卡參數(shù)1.75GB總共約10GB但通信開(kāi)銷大幅增加。FSDP本質(zhì)上是ZeRO Stage 3的PyTorch原生實(shí)現(xiàn)它把參數(shù)、梯度、優(yōu)化器狀態(tài)都分片每卡只存1/N。但FSDP在forward和backward時(shí)需要all-gather參數(shù)通信量很大。實(shí)際用下來(lái)FSDP在8卡A100上訓(xùn)練7B模型每卡顯存約12GB但吞吐量比DDP低20%左右。這個(gè) trade-off 要看你的瓶頸是顯存還是算力。3.3 AMP的顯存收益與代價(jià)AMP自動(dòng)混合精度是顯存優(yōu)化的第一板斧但它的收益經(jīng)常被高估。AMP省的主要是激活值和部分計(jì)算緩沖區(qū)參數(shù)和優(yōu)化器狀態(tài)的大頭省不掉。實(shí)測(cè)下來(lái)AMP能把激活值顯存降到FP32的60%左右總體顯存節(jié)省約20%到30%。AMP的代價(jià)是數(shù)值穩(wěn)定性。FP16的動(dòng)態(tài)范圍窄梯度容易下溢或上溢。PyTorch的GradScaler通過(guò)動(dòng)態(tài)調(diào)整loss scale來(lái)緩解這個(gè)問(wèn)題但在某些模型結(jié)構(gòu)上仍然會(huì)出現(xiàn)NaN。BF16的動(dòng)態(tài)范圍和FP32一樣不需要loss scaling但精度比FP16低。A100及以上支持BF16V100只支持FP16。選哪個(gè)取決于你的硬件和模型對(duì)精度的敏感度。實(shí)操心得用AMP時(shí)把LayerNorm和softmax強(qiáng)制保留在FP32能顯著提升訓(xùn)練穩(wěn)定性。PyTorch的torch.cuda.amp.autocast默認(rèn)已經(jīng)這么做了但如果你自己寫kernel要注意手動(dòng)指定。4. 實(shí)操測(cè)量流程與工具鏈4.1 測(cè)量腳本的編寫要點(diǎn)我一般會(huì)寫一個(gè)獨(dú)立的測(cè)量腳本不摻在訓(xùn)練代碼里這樣干凈、可復(fù)現(xiàn)。腳本的核心結(jié)構(gòu)是初始化模型和優(yōu)化器構(gòu)造一個(gè)真實(shí)batch的輸入然后分階段打點(diǎn)。import torch from torch.cuda import memory_allocated, max_memory_allocated, reset_peak_memory_stats def measure(model, optimizer, input_ids, labels): reset_peak_memory_stats() base memory_allocated() # forward outputs model(input_ids, labelslabels) loss outputs.loss after_forward memory_allocated() # backward loss.backward() after_backward memory_allocated() # optimizer step optimizer.step() optimizer.zero_grad() after_step memory_allocated() return { base: base, forward_delta: after_forward - base, backward_delta: after_backward - after_forward, step_delta: after_step - after_backward, peak: max_memory_allocated() }這個(gè)腳本跑一次就能拿到四個(gè)關(guān)鍵數(shù)字。forward_delta主要是激活值backward_delta是梯度加激活值的峰值step_delta是優(yōu)化器狀態(tài)的增量。peak是整個(gè)過(guò)程的最大值用來(lái)做OOM判斷。4.2 不同配置的對(duì)比測(cè)量測(cè)量不能只測(cè)一個(gè)配置要測(cè)一組配置做對(duì)比。我一般會(huì)測(cè)這幾組FP32基線、AMP、AMP加梯度檢查點(diǎn)、AMP加梯度檢查點(diǎn)加ZeRO。每組跑三次取平均排除冷啟動(dòng)的影響。對(duì)比的時(shí)候重點(diǎn)看兩個(gè)指標(biāo)峰值顯存和吞吐量。峰值顯存決定你能不能跑起來(lái)吞吐量決定你跑得多快。有時(shí)候省顯存的方案會(huì)把吞吐量砍半這時(shí)候就要算一筆賬省下來(lái)的顯存能不能換來(lái)更大的batch size如果能吞吐量可能反而更高。舉個(gè)例子7B模型FP32訓(xùn)練峰值顯存60GB吞吐量100 samples/s。AMP后峰值45GB吞吐量130 samples/s。AMP加梯度檢查點(diǎn)后峰值30GB吞吐量90 samples/s。雖然梯度檢查點(diǎn)讓吞吐量降了但省下的15GB顯存可以讓你把batch size翻倍實(shí)際吞吐量變成180 samples/s。這就是預(yù)算決策的價(jià)值。4.3 顯存碎片的手動(dòng)清理測(cè)量過(guò)程中如果發(fā)現(xiàn)nvidia-smi和memory_allocated()差值越來(lái)越大說(shuō)明碎片在累積。這時(shí)候可以手動(dòng)調(diào)torch.cuda.empty_cache()但注意這個(gè)操作會(huì)釋放緩存分配器預(yù)留的塊可能導(dǎo)致后續(xù)分配變慢。更好的辦法是設(shè)置PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True讓分配器自己管理碎片。還有一個(gè)技巧是在訓(xùn)練循環(huán)里定期調(diào)用torch.cuda.reset_peak_memory_stats()把峰值統(tǒng)計(jì)重置這樣能更準(zhǔn)確地看到每個(gè)step的顯存波動(dòng)。如果不重置峰值會(huì)一直保留歷史最大值掩蓋了后期的顯存增長(zhǎng)。5. 常見(jiàn)問(wèn)題與排查技巧實(shí)錄5.1 測(cè)量結(jié)果與nvidia-smi對(duì)不上這是最常見(jiàn)的問(wèn)題。memory_allocated()只統(tǒng)計(jì)PyTorch分配器管理的顯存不包括CUDA上下文、cuDNN workspace、NCCL通信緩沖區(qū)。這些加起來(lái)可能有1到2GB。所以nvidia-smi看到的數(shù)字總是比memory_allocated()大。解決辦法是用torch.cuda.memory_snapshot()拿到完整的內(nèi)存快照它會(huì)列出所有分配塊包括那些不在PyTorch管理范圍內(nèi)的。不過(guò)這個(gè)輸出很大一般只在排查問(wèn)題時(shí)用。5.2 激活值峰值出現(xiàn)在意想不到的地方有時(shí)候你會(huì)發(fā)現(xiàn)backward階段的顯存峰值比f(wàn)orward高很多甚至高出一倍。這通常是因?yàn)槟承┎僮鞯腷ackward需要保存forward的中間結(jié)果而這些結(jié)果在forward結(jié)束后并沒(méi)有釋放。比如attention的softmax輸出forward時(shí)算完就丟了但backward需要它來(lái)計(jì)算梯度所以PyTorch會(huì)把它保留到backward。排查方法是給每個(gè)模塊注冊(cè)forward hook和backward hook記錄每個(gè)模塊的輸入輸出顯存。這樣能精確定位到哪個(gè)模塊的backward顯存開(kāi)銷異常。5.3 OOM發(fā)生在optimizer.step()optimizer.step()本身不分配大塊顯存但它會(huì)觸發(fā)參數(shù)更新而參數(shù)更新需要讀取梯度、寫入?yún)?shù)。如果梯度是FP16而參數(shù)是FP32這里會(huì)有一個(gè)隱式的類型轉(zhuǎn)換產(chǎn)生臨時(shí)緩沖區(qū)。Adam的動(dòng)量更新還會(huì)產(chǎn)生中間變量。這些加起來(lái)可能幾百M(fèi)B在顯存已經(jīng)接近上限時(shí)就是壓死駱駝的最后一根稻草。解決辦法是在step之前手動(dòng)torch.cuda.empty_cache()或者把優(yōu)化器狀態(tài)分片到多卡。另一個(gè)辦法是用foreach實(shí)現(xiàn)的優(yōu)化器它把多個(gè)參數(shù)的更新合并成一個(gè)kernel減少臨時(shí)緩沖區(qū)的分配次數(shù)。5.4 梯度檢查點(diǎn)與AMP的兼容問(wèn)題梯度檢查點(diǎn)在backward時(shí)會(huì)重新計(jì)算forward如果和AMP一起用重計(jì)算時(shí)的精度策略要和原forward一致否則梯度會(huì)對(duì)不上。PyTorch的torch.utils.checkpoint默認(rèn)會(huì)保留AMP的autocast狀態(tài)但如果你自己實(shí)現(xiàn)了checkpoint邏輯要注意手動(dòng)傳遞autocast上下文。還有一個(gè)坑是梯度檢查點(diǎn)不能和某些自定義autograd函數(shù)一起用因?yàn)橹赜?jì)算時(shí)這些函數(shù)的forward可能不是確定性的。遇到這種情況要么把自定義函數(shù)排除在checkpoint范圍外要么確保它是確定性的。5.5 多卡訓(xùn)練時(shí)的顯存不均衡數(shù)據(jù)并行時(shí)每卡的顯存占用應(yīng)該基本一致。如果發(fā)現(xiàn)某張卡顯存特別高通常是數(shù)據(jù)加載不均衡或者通信操作卡住了。檢查DataLoader的num_workers和pin_memory設(shè)置確保每個(gè)進(jìn)程拿到的batch大小一致。通信方面NCCL的all-reduce是同步操作如果某張卡算得慢其他卡會(huì)在通信點(diǎn)等待顯存占用會(huì)暫時(shí)升高。排查方法是打印每張卡的memory_allocated()看差異是否超過(guò)5%。如果超過(guò)檢查數(shù)據(jù)分片邏輯和通信配置。問(wèn)題現(xiàn)象可能原因排查方法解決手段nvidia-smi比memory_allocated高2GBCUDA上下文和通信緩沖區(qū)memory_snapshot預(yù)留2GB余量backward峰值遠(yuǎn)高于forward中間激活未釋放模塊級(jí)hook梯度檢查點(diǎn)step時(shí)OOM類型轉(zhuǎn)換臨時(shí)緩沖區(qū)逐步打點(diǎn)foreach優(yōu)化器多卡顯存不均衡數(shù)據(jù)或通信不均衡逐卡打印調(diào)整DataLoader碎片率持續(xù)增長(zhǎng)動(dòng)態(tài)shape監(jiān)控差值expandable_segments6. 預(yù)算決策的實(shí)戰(zhàn)案例6.1 單卡24GB訓(xùn)練7B模型的可行性分析有人問(wèn)單卡24GB能不能訓(xùn)7B模型。按前面的公式算7B模型AMP加AdamW參數(shù)14GBFP16加梯度14GBFP16加優(yōu)化器狀態(tài)56GBFP32總共84GB遠(yuǎn)超24GB。所以全量微調(diào)不可能。但LoRA可以。LoRA只訓(xùn)練低秩矩陣參數(shù)量通常是原模型的0.1%到1%。7B模型的LoRA參數(shù)量約7M到70MFP16下占14MB到140MB。優(yōu)化器狀態(tài)按Adam算是參數(shù)量乘以8也就是112MB到1.1GB。加上基座模型的14GBFP16總共約15GB到16GB。24GB卡能放下還能留出8GB給激活值和緩沖區(qū)。QLoRA更進(jìn)一步把基座模型量化到4bit7B模型只占3.5GB。加上LoRA參數(shù)和優(yōu)化器狀態(tài)總共約5GB。24GB卡能跑得很寬裕甚至能上更大的batch size。6.2 多卡場(chǎng)景下的并行策略選擇如果你有4張24GB卡總共96GB顯存想訓(xùn)7B模型全量微調(diào)。DDP每卡都要存完整的參數(shù)、梯度、優(yōu)化器狀態(tài)每卡84GB放不下。ZeRO Stage 2把優(yōu)化器狀態(tài)和梯度分片到4卡每卡參數(shù)14GB加梯度3.5GB加優(yōu)化器狀態(tài)14GB總共31.5GB還是放不下。ZeRO Stage 3把參數(shù)也分片每卡參數(shù)3.5GB加梯度3.5GB加優(yōu)化器狀態(tài)14GB總共21GB勉強(qiáng)能放下但激活值還沒(méi)算。這時(shí)候要么上梯度檢查點(diǎn)把激活值壓到2GB以內(nèi)要么上CPU offload把優(yōu)化器狀態(tài)放到內(nèi)存。CPU offload的代價(jià)是step速度慢3到5倍但能省下14GB顯存。實(shí)測(cè)下來(lái)4卡24GB加ZeRO Stage 3加梯度檢查點(diǎn)加CPU offload能訓(xùn)7B模型但吞吐量只有DDP的十分之一。所以如果追求速度還是建議上80GB卡。6.3 預(yù)算決策的決策樹(shù)我把預(yù)算決策整理成一個(gè)簡(jiǎn)單的決策樹(shù)方便快速判斷第一步算剛性支出參數(shù)量乘以2加2加8除以并行度。如果這個(gè)數(shù)字超過(guò)單卡顯存的80%考慮LoRA或QLoRA。第二步算激活值batch size乘以序列長(zhǎng)度乘以隱藏維度乘以層數(shù)乘以系數(shù)。如果超過(guò)剩余顯存的50%上梯度檢查點(diǎn)。第三步算碎片余量留10%到15%的顯存給臨時(shí)緩沖區(qū)和碎片。如果不夠調(diào)expandable_segments或減小batch size。第四步測(cè)吞吐量在滿足顯存約束的前提下找吞吐量最大的配置。有時(shí)候大batch加梯度檢查點(diǎn)比小batch不加檢查點(diǎn)更快。這個(gè)決策樹(shù)不是絕對(duì)的但能幫你快速縮小選擇范圍。實(shí)際做的時(shí)候還是要跑測(cè)量腳本驗(yàn)證因?yàn)椴煌Ｐ徒Y(jié)構(gòu)、不同框架版本的顯存行為差異很大。最后分享一個(gè)小技巧在訓(xùn)練腳本里加一個(gè)顯存監(jiān)控回調(diào)每個(gè)step記錄峰值顯存和吞吐量輸出到TensorBoard。這樣你能看到顯存隨訓(xùn)練進(jìn)程的變化趨勢(shì)及時(shí)發(fā)現(xiàn)碎片累積或激活值增長(zhǎng)的問(wèn)題。我靠這個(gè)回調(diào)抓到過(guò)好幾次隱蔽的顯存泄漏都是自定義層里緩存了不該緩存的東西。