1. 從單卡到多卡為什么數(shù)據(jù)并行是繞不開的第一課很多人第一次接觸多GPU訓(xùn)練腦子里冒出來的第一個(gè)念頭是“模型太大一張卡塞不下”。這個(gè)直覺只對(duì)了一半。真正在工程里跑起來絕大多數(shù)場景其實(shí)是“模型塞得下但訓(xùn)練太慢”。一個(gè)中等規(guī)模的問答模型單卡跑完一個(gè)epoch可能要十幾個(gè)小時(shí)調(diào)一次超參就得等一整天這種節(jié)奏根本沒法迭代。多GPU訓(xùn)練要解決的核心矛盾從來不是“能不能跑”而是“能不能快速試錯(cuò)”。數(shù)據(jù)并行Data Parallelism是這里面最樸素也最實(shí)用的方案。它的邏輯簡單到可以用一句話說清把同一個(gè)模型復(fù)制到每張卡上把一個(gè)大batch切成若干份每張卡各算各的梯度最后把梯度匯總求平均再統(tǒng)一更新。你可以把它想象成一個(gè)小組做同一套卷子每個(gè)人分幾道題做完之后把答案對(duì)一遍取共識(shí)最后所有人用同一份標(biāo)準(zhǔn)答案修正自己的理解。這個(gè)“對(duì)答案”的過程術(shù)語叫梯度同步。為什么數(shù)據(jù)并行能成立關(guān)鍵在于隨機(jī)梯度下降的可加性。假設(shè)一個(gè)batch有256條樣本單卡算出來的梯度是這256條樣本梯度的平均。如果拆成4張卡各算64條每張卡算出的局部梯度再求平均數(shù)學(xué)上等價(jià)于原來那個(gè)大batch的梯度。這不是近似是嚴(yán)格相等。所以數(shù)據(jù)并行不會(huì)改變模型的收斂行為只是把計(jì)算壓力分?jǐn)偭?。這一點(diǎn)很重要它意味著你不需要重新調(diào)學(xué)習(xí)率除非你同時(shí)放大了總batch size。但這里有個(gè)容易被忽略的細(xì)節(jié)batch size的放大效應(yīng)。當(dāng)你從單卡batch 32變成4卡各32、總batch 128時(shí)雖然梯度方向在期望上一致但方差變小了實(shí)際訓(xùn)練動(dòng)態(tài)會(huì)變。經(jīng)驗(yàn)做法是總batch size放大k倍學(xué)習(xí)率也相應(yīng)放大但不要線性放大通常取平方根或者更保守的系數(shù)。我見過太多人直接4倍學(xué)習(xí)率拉滿結(jié)果loss直接炸掉還以為是多卡同步出了問題其實(shí)是優(yōu)化器步長太激進(jìn)。還有一個(gè)認(rèn)知誤區(qū)多GPU訓(xùn)練不是“免費(fèi)的加速”?？ㄔ蕉嗤ㄐ砰_銷占比越大。4卡的時(shí)候可能加速比有3.58卡可能只有616卡可能掉到9甚至更低。因?yàn)槊繌埧ㄋ阃晏荻群蠖家茸盥哪菑埧ㄈ缓笞鲆淮稳滞竭@個(gè)同步的延遲隨著卡數(shù)增加而增加。所以選卡數(shù)不是越多越好要看你的模型計(jì)算量和通信量的比值。計(jì)算密集型的模型比如大矩陣乘法多的適合多卡通信密集型的比如小模型、頻繁同步的多卡反而可能更慢。2. 梯度同步的兩種實(shí)現(xiàn)路徑參數(shù)服務(wù)器與全規(guī)約理解了數(shù)據(jù)并行的原理接下來要面對(duì)一個(gè)工程問題梯度到底怎么匯總這不是一個(gè)“隨便選一個(gè)就行”的問題不同實(shí)現(xiàn)路徑的性能差異可能達(dá)到數(shù)倍。2.1 參數(shù)服務(wù)器模式的工作機(jī)制參數(shù)服務(wù)器Parameter Server是最早被廣泛采用的方案。它的結(jié)構(gòu)很直白有一組服務(wù)器節(jié)點(diǎn)負(fù)責(zé)保存全局參數(shù)一組工作節(jié)點(diǎn)負(fù)責(zé)計(jì)算梯度。工作節(jié)點(diǎn)算完梯度后推送給服務(wù)器服務(wù)器聚合后再把新參數(shù)廣播回去。這個(gè)模式的好處是擴(kuò)展性好理論上可以掛很多工作節(jié)點(diǎn)服務(wù)器只負(fù)責(zé)聚合不參與計(jì)算。但它的缺點(diǎn)也很明顯。所有工作節(jié)點(diǎn)都要和服務(wù)器通信服務(wù)器很容易成為帶寬瓶頸。而且服務(wù)器節(jié)點(diǎn)本身如果掛了整個(gè)訓(xùn)練就崩了。在實(shí)際部署中參數(shù)服務(wù)器模式通常需要精心設(shè)計(jì)分片策略把不同的參數(shù)切分到不同的服務(wù)器上避免單點(diǎn)壓力過大。這套東西維護(hù)起來相當(dāng)復(fù)雜對(duì)于中小規(guī)模訓(xùn)練來說性價(jià)比不高。2.2 全規(guī)約模式為什么成為主流全規(guī)約All-Reduce是現(xiàn)在更流行的方案。它的思路是沒有中心服務(wù)器所有卡地位平等通過一種叫環(huán)形規(guī)約的算法讓每張卡只和左右鄰居通信經(jīng)過兩輪傳遞就能讓所有卡都拿到全局梯度的平均值。環(huán)形規(guī)約的精妙之處在于它把通信量均攤了。假設(shè)有N張卡每張卡的數(shù)據(jù)量是D傳統(tǒng)做法是每張卡都發(fā)給中心節(jié)點(diǎn)中心節(jié)點(diǎn)壓力是N×D。環(huán)形規(guī)約則是每張卡只發(fā)D/N給鄰居經(jīng)過2×(N-1)步傳遞所有卡都拿到完整結(jié)果。通信量從O(N)降到了O(1)級(jí)別相對(duì)于單卡數(shù)據(jù)量。這就是為什么8卡、16卡還能保持不錯(cuò)加速比的原因。注意全規(guī)約要求所有卡的參數(shù)量完全一致且通信是同步阻塞的。如果某張卡算得特別慢整個(gè)環(huán)都會(huì)被拖住。所以多卡訓(xùn)練里負(fù)載均衡比單卡時(shí)更重要要確保每張卡分到的數(shù)據(jù)量差不多避免某張卡成為短板。2.3 兩種模式的選型建議對(duì)比維度參數(shù)服務(wù)器全規(guī)約通信模式星型中心聚合環(huán)形去中心化擴(kuò)展性理論好實(shí)際受服務(wù)器帶寬限制中等規(guī)模優(yōu)秀超大規(guī)模需分層容錯(cuò)性服務(wù)器故障即全局故障單卡故障影響環(huán)但可重建實(shí)現(xiàn)復(fù)雜度高需維護(hù)服務(wù)器集群低框架內(nèi)置支持適用場景超大規(guī)模稀疏模型中小規(guī)模稠密模型對(duì)于問答類模型這種參數(shù)量中等、計(jì)算稠密的場景全規(guī)約幾乎是默認(rèn)選擇。主流框架的分布式模塊都內(nèi)置了全規(guī)約實(shí)現(xiàn)你只需要改幾行代碼就能啟用不需要自己搭服務(wù)器集群。3. 動(dòng)手實(shí)現(xiàn)從單卡代碼到多卡并行的改造清單理論說再多不如把代碼跑起來。這一節(jié)我以一個(gè)典型的問答模型訓(xùn)練腳本為例拆解從單卡到多卡需要改哪些地方。假設(shè)你已經(jīng)有了一份能跑的單卡訓(xùn)練代碼結(jié)構(gòu)大概是數(shù)據(jù)加載、模型定義、優(yōu)化器、訓(xùn)練循環(huán)。3.1 進(jìn)程初始化與設(shè)備綁定多卡訓(xùn)練的第一步是讓每張卡跑一個(gè)獨(dú)立的進(jìn)程。這里最容易踩的坑是設(shè)備綁定順序。如果你不顯式指定每張卡用哪個(gè)GPU框架可能會(huì)讓所有進(jìn)程都去搶同一張卡結(jié)果顯存爆掉。import torch.distributed as dist import torch def setup(rank, world_size): dist.init_process_group( backendnccl, init_methodenv://, world_sizeworld_size, rankrank ) torch.cuda.set_device(rank)這段代碼里rank是當(dāng)前進(jìn)程的編號(hào)world_size是總進(jìn)程數(shù)。torch.cuda.set_device(rank)這行是關(guān)鍵它確保第rank號(hào)進(jìn)程只用第rank號(hào)GPU。backendnccl是NVIDIA顯卡的通信后端如果是CPU集群就用gloo但GPU訓(xùn)練基本都用nccl。提示init_methodenv://表示從環(huán)境變量讀取主節(jié)點(diǎn)地址和端口。啟動(dòng)腳本時(shí)需要設(shè)置MASTER_ADDR和MASTER_PORT否則進(jìn)程之間找不到彼此。3.2 數(shù)據(jù)加載器的分布式采樣單卡時(shí)DataLoader按順序或隨機(jī)打亂取數(shù)據(jù)。多卡時(shí)如果每張卡都獨(dú)立打亂可能會(huì)取到重復(fù)樣本導(dǎo)致一個(gè)epoch實(shí)際看到的樣本數(shù)變少。正確的做法是用DistributedSampler它保證每張卡取到的數(shù)據(jù)互不重疊且合起來正好覆蓋整個(gè)數(shù)據(jù)集。from torch.utils.data.distributed import DistributedSampler sampler DistributedSampler( dataset, num_replicasworld_size, rankrank, shuffleTrue ) dataloader DataLoader( dataset, batch_sizebatch_size_per_gpu, samplersampler, num_workers4 )注意batch_size_per_gpu是單卡batch size總batch size是它乘以卡數(shù)。每個(gè)epoch開始前要調(diào)用sampler.set_epoch(epoch)否則每次epoch的打亂順序都一樣影響訓(xùn)練效果。這個(gè)細(xì)節(jié)很多人會(huì)漏掉導(dǎo)致模型在第二個(gè)epoch開始就失去了隨機(jī)性。3.3 模型包裝與梯度同步模型需要用DistributedDataParallelDDP包一層。DDP會(huì)在反向傳播時(shí)自動(dòng)觸發(fā)梯度同步你不需要手動(dòng)寫all_reduce。from torch.nn.parallel import DistributedDataParallel as DDP model MyQAModel().to(rank) model DDP(model, device_ids[rank])DDP的同步發(fā)生在loss.backward()的時(shí)候。它會(huì)把每張卡的梯度做一次全規(guī)約平均然后每張卡用相同的梯度更新參數(shù)。因?yàn)槌跏紖?shù)相同、梯度相同、優(yōu)化器相同所以更新后的參數(shù)也保持相同。這就是為什么DDP不需要在每次前向傳播前同步參數(shù)只需要同步梯度。有一個(gè)性能優(yōu)化點(diǎn)DDP支持梯度分桶。它不會(huì)等所有梯度算完才通信而是把梯度分成若干桶算完一個(gè)桶就通信一個(gè)桶實(shí)現(xiàn)計(jì)算和通信的重疊。這個(gè)特性默認(rèn)開啟但桶的大小可以調(diào)。如果模型層數(shù)很多、每層參數(shù)量小適當(dāng)增大桶的大小能減少通信次數(shù)。3.4 損失函數(shù)的歸一化陷阱多卡訓(xùn)練時(shí)損失函數(shù)的計(jì)算方式會(huì)直接影響梯度同步的正確性。如果你用的是CrossEntropyLoss的默認(rèn)reductionmean每張卡算的是自己那部分樣本的平均損失。但不同卡上的樣本數(shù)可能不一樣尤其是最后一個(gè)batch直接平均會(huì)導(dǎo)致梯度權(quán)重偏差。正確的做法是每張卡的損失乘以該卡的樣本數(shù)同步時(shí)求和再除以總樣本數(shù)?；蛘吒唵未_保每個(gè)batch的樣本數(shù)能被卡數(shù)整除避免最后一個(gè)不完整的batch。我在實(shí)際項(xiàng)目里傾向于后者因?yàn)閷?shí)現(xiàn)簡單代價(jià)只是丟棄少量樣本對(duì)訓(xùn)練影響微乎其微。# 假設(shè)每張卡batch size相同 loss criterion(outputs, labels) loss.backward() # DDP自動(dòng)處理梯度平均如果確實(shí)無法整除可以用reductionsum然后手動(dòng)除以全局batch size但這樣每張卡的loss尺度不一致需要額外小心。4. 實(shí)測中的性能瓶頸與調(diào)優(yōu)手段代碼改完能跑起來只是第一步。真正讓多卡訓(xùn)練發(fā)揮出應(yīng)有的加速比需要處理一系列性能問題。這一節(jié)分享幾個(gè)我在實(shí)際項(xiàng)目中反復(fù)遇到的瓶頸和對(duì)應(yīng)的調(diào)優(yōu)手段。4.1 通信開銷的量化與判斷怎么判斷你的訓(xùn)練是計(jì)算瓶頸還是通信瓶頸一個(gè)簡單的辦法是把卡數(shù)翻倍看訓(xùn)練速度是否也接近翻倍。如果4卡比2卡只快了1.5倍說明通信開銷已經(jīng)很明顯了。更精確的做法是用profiler工具看時(shí)間分布。如果all_reduce或broadcast占用的時(shí)間超過總時(shí)間的20%就說明通信是瓶頸。這時(shí)候可以考慮幾個(gè)方向增大單卡batch size提高計(jì)算量攤薄通信占比、減少同步頻率比如梯度累積、或者換更快的互聯(lián)比如NVLink代替PCIe。注意梯度累積是一種用時(shí)間換通信的策略。比如每4個(gè)batch才同步一次梯度相當(dāng)于把有效batch size放大了4倍但通信次數(shù)減少了4倍。代價(jià)是訓(xùn)練動(dòng)態(tài)會(huì)變化需要重新調(diào)學(xué)習(xí)率。4.2 數(shù)據(jù)加載的隱藏成本多卡訓(xùn)練時(shí)數(shù)據(jù)加載很容易成為瓶頸尤其是當(dāng)數(shù)據(jù)預(yù)處理比較復(fù)雜時(shí)。每張卡都要獨(dú)立加載和預(yù)處理數(shù)據(jù)如果CPU核心不夠數(shù)據(jù)加載速度跟不上GPU計(jì)算速度GPU就會(huì)空轉(zhuǎn)。解決辦法有幾個(gè)增加num_workers但不要超過CPU核心數(shù)、把預(yù)處理結(jié)果緩存到內(nèi)存或磁盤、用更高效的數(shù)據(jù)格式比如把原始文本預(yù)先tokenize好存成二進(jìn)制。我見過一個(gè)案例把在線tokenize改成離線預(yù)處理后訓(xùn)練速度直接提升了40%因?yàn)閠okenize的CPU開銷被完全消除了。還有一個(gè)容易忽略的點(diǎn)共享內(nèi)存的限制。PyTorch的DataLoader用共享內(nèi)存?zhèn)鬟f數(shù)據(jù)如果num_workers太多或者數(shù)據(jù)太大可能會(huì)觸發(fā)共享內(nèi)存不足的錯(cuò)誤。這時(shí)候需要調(diào)整系統(tǒng)的共享內(nèi)存大小或者減少worker數(shù)量。4.3 顯存優(yōu)化與batch size的平衡多卡訓(xùn)練時(shí)每張卡的顯存壓力并沒有減少因?yàn)槟Ｐ褪峭暾麖?fù)制的。所以單卡能跑多大batch多卡還是那么大。但多卡的好處是你可以用更小的單卡batch通過增加卡數(shù)來達(dá)到相同的總batch size從而降低單卡顯存壓力。如果單卡顯存還是不夠可以考慮梯度檢查點(diǎn)Gradient Checkpointing用計(jì)算時(shí)間換顯存空間。它的原理是不保存中間激活值反向傳播時(shí)重新計(jì)算。這會(huì)增加約30%的計(jì)算時(shí)間但能省下大量顯存讓你能跑更大的模型或batch。另一個(gè)技巧是混合精度訓(xùn)練。用float16做前向和反向float32做參數(shù)更新能省一半顯存還能利用Tensor Core加速。但要注意loss scaling避免梯度下溢。主流框架都有自動(dòng)混合精度的接口幾行代碼就能啟用。5. 那些文檔里不會(huì)寫的踩坑記錄這一節(jié)記錄幾個(gè)我在多卡訓(xùn)練中真實(shí)踩過的坑每一個(gè)都花了半天以上才定位到原因。希望你看完能少走彎路。5.1 進(jìn)程啟動(dòng)方式的坑多卡訓(xùn)練的啟動(dòng)方式有好幾種torchrun、mp.spawn、手動(dòng)起多個(gè)進(jìn)程。我強(qiáng)烈推薦用torchrun它是官方維護(hù)的啟動(dòng)器能自動(dòng)處理進(jìn)程管理和故障恢復(fù)。手動(dòng)起進(jìn)程的方式看起來靈活但很容易出現(xiàn)端口沖突、進(jìn)程殘留、日志混亂等問題。用torchrun時(shí)一個(gè)常見的錯(cuò)誤是MASTER_PORT被占用。如果同一臺(tái)機(jī)器上跑多個(gè)訓(xùn)練任務(wù)端口會(huì)沖突。解決辦法是每次啟動(dòng)時(shí)隨機(jī)選一個(gè)端口或者用--rdzv_backendc10d讓啟動(dòng)器自動(dòng)協(xié)商端口。還有一個(gè)坑是進(jìn)程數(shù)不等于卡數(shù)。如果你有8張卡但只啟動(dòng)了4個(gè)進(jìn)程那只有4張卡在工作另外4張閑置。反過來如果啟動(dòng)了16個(gè)進(jìn)程但只有8張卡會(huì)有兩個(gè)進(jìn)程搶同一張卡直接顯存溢出。啟動(dòng)前一定要確認(rèn)nproc_per_node等于實(shí)際卡數(shù)。5.2 隨機(jī)種子的同步問題單卡訓(xùn)練時(shí)你只需要設(shè)一個(gè)隨機(jī)種子。多卡訓(xùn)練時(shí)如果每張卡用相同的種子數(shù)據(jù)增強(qiáng)和dropout會(huì)產(chǎn)生相同的隨機(jī)模式相當(dāng)于降低了隨機(jī)性。正確的做法是每張卡的種子設(shè)為base_seed rank保證各卡隨機(jī)性獨(dú)立同時(shí)整體可復(fù)現(xiàn)。但這里有個(gè)矛盾如果你想讓訓(xùn)練完全可復(fù)現(xiàn)又需要所有卡的隨機(jī)性可控。我的經(jīng)驗(yàn)是在調(diào)試階段用固定種子確?？蓮?fù)現(xiàn)在正式訓(xùn)練時(shí)用動(dòng)態(tài)種子提高泛化。切換的時(shí)候記得把相關(guān)配置都改掉不要只改一處。5.3 評(píng)估階段的單卡陷阱訓(xùn)練用多卡評(píng)估時(shí)很多人會(huì)忘記切換回單卡模式。如果評(píng)估也用多卡每個(gè)進(jìn)程都會(huì)算一遍完整的驗(yàn)證集然后你可能會(huì)把結(jié)果重復(fù)統(tǒng)計(jì)。更隱蔽的問題是DistributedSampler在評(píng)估時(shí)會(huì)打亂數(shù)據(jù)順序?qū)е履銦o法對(duì)應(yīng)預(yù)測結(jié)果和原始樣本。正確的做法是評(píng)估時(shí)用DistributedSampler但設(shè)置shuffleFalse并且只在rank 0上做指標(biāo)計(jì)算和日志輸出。其他rank算完結(jié)果后通過gather操作匯總到rank 0。如果嫌麻煩評(píng)估時(shí)直接切回單卡用torch.no_grad()跑一遍雖然慢一點(diǎn)但省心。5.4 檢查點(diǎn)保存的競爭條件多卡訓(xùn)練時(shí)如果每個(gè)進(jìn)程都保存檢查點(diǎn)會(huì)互相覆蓋最后存下來的可能是某個(gè)中間狀態(tài)。正確做法是只在rank 0上保存其他rank跳過。但要注意保存前需要確保所有rank的模型狀態(tài)一致通常DDP會(huì)自動(dòng)保證這一點(diǎn)但如果你在保存前做了額外的參數(shù)操作就需要手動(dòng)同步。加載檢查點(diǎn)時(shí)也有坑。如果檢查點(diǎn)是在單卡上保存的加載到多卡時(shí)需要先加載到CPU再分發(fā)到各卡。直接load_state_dict可能會(huì)因?yàn)樵O(shè)備不匹配報(bào)錯(cuò)。用map_location參數(shù)指定設(shè)備能解決大部分問題。6. 從能跑到跑得好進(jìn)階優(yōu)化思路當(dāng)你的多卡訓(xùn)練能穩(wěn)定跑起來之后下一步是讓它跑得更快、更穩(wěn)、更省資源。這一節(jié)分享幾個(gè)進(jìn)階方向適合已經(jīng)跨過入門階段的讀者。6.1 梯度壓縮與通信優(yōu)化如果通信確實(shí)是瓶頸可以考慮梯度壓縮。核心思想是只傳梯度的部分信息比如量化到低精度、只傳大于閾值的稀疏梯度、或者用誤差補(bǔ)償保留被丟棄的信息。這些方法能顯著減少通信量但實(shí)現(xiàn)復(fù)雜度高且可能影響收斂。對(duì)于問答模型這種稠密梯度場景量化到float16通常能省一半帶寬且對(duì)精度影響很小。另一個(gè)方向是通信與計(jì)算重疊。DDP已經(jīng)做了梯度分桶的重疊但你可以進(jìn)一步把前向傳播也納入重疊范圍。比如在等待梯度同步的時(shí)候提前開始下一個(gè)batch的前向計(jì)算。這需要更精細(xì)的流水線設(shè)計(jì)框架層面支持有限通常需要自己寫調(diào)度邏輯。6.2 彈性訓(xùn)練與容錯(cuò)大規(guī)模訓(xùn)練時(shí)硬件故障是常態(tài)。一張卡過熱降頻、一個(gè)進(jìn)程意外退出都可能導(dǎo)致整個(gè)訓(xùn)練中斷。彈性訓(xùn)練的思路是允許進(jìn)程動(dòng)態(tài)加入和退出訓(xùn)練不中斷。這需要框架支持動(dòng)態(tài)進(jìn)程組和檢查點(diǎn)恢復(fù)。實(shí)現(xiàn)彈性訓(xùn)練的關(guān)鍵是定期保存檢查點(diǎn)并且檢查點(diǎn)要包含優(yōu)化器狀態(tài)、學(xué)習(xí)率調(diào)度器狀態(tài)、當(dāng)前epoch和step?；謴?fù)時(shí)從最近的檢查點(diǎn)繼續(xù)而不是從頭開始。檢查點(diǎn)的保存頻率需要權(quán)衡太頻繁影響訓(xùn)練速度太稀疏則故障時(shí)損失太多進(jìn)度。我的經(jīng)驗(yàn)是每30分鐘保存一次同時(shí)保留最近3個(gè)檢查點(diǎn)防止檢查點(diǎn)本身損壞。6.3 混合并行數(shù)據(jù)并行之外的選擇當(dāng)模型大到單卡放不下時(shí)數(shù)據(jù)并行就不夠了。這時(shí)候需要模型并行把模型的不同層放到不同卡上。更常見的是混合并行數(shù)據(jù)并行和模型并行結(jié)合比如8張卡分成2組每組4張卡做數(shù)據(jù)并行組內(nèi)做模型并行。混合并行的配置復(fù)雜度比純數(shù)據(jù)并行高一個(gè)數(shù)量級(jí)。你需要手動(dòng)指定哪些層放在哪些卡上還要處理跨卡的前向和反向傳播。對(duì)于問答模型來說除非參數(shù)量超過單卡顯存極限否則不建議上混合并行維護(hù)成本太高。6.4 監(jiān)控與日志讓訓(xùn)練過程可見多卡訓(xùn)練時(shí)日志是最重要的調(diào)試工具。你需要記錄每張卡的loss、梯度范數(shù)、學(xué)習(xí)率、吞吐量。如果某張卡的loss異常偏高可能是數(shù)據(jù)分片有問題如果梯度范數(shù)突然變大可能是遇到了臟數(shù)據(jù)。我習(xí)慣用TensorBoard或類似的工具做可視化把每張卡的指標(biāo)畫在同一張圖上。這樣一眼就能看出哪張卡拖后腿。另外記錄每個(gè)step的耗時(shí)分布也很有用能幫你判斷瓶頸在計(jì)算、通信還是數(shù)據(jù)加載。提示日志文件要按rank分開寫否則多個(gè)進(jìn)程同時(shí)寫同一個(gè)文件會(huì)亂序?？梢杂胠ogging模塊配置每個(gè)rank寫自己的文件最后再合并分析。7. 一個(gè)完整的訓(xùn)練腳本骨架把前面的內(nèi)容串起來這里給一個(gè)可以直接參考的訓(xùn)練腳本骨架。省略了模型定義和數(shù)據(jù)預(yù)處理的具體細(xì)節(jié)重點(diǎn)展示多卡相關(guān)的部分。import os import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data.distributed import DistributedSampler from torch.utils.data import DataLoader def main(): rank int(os.environ[LOCAL_RANK]) world_size int(os.environ[WORLD_SIZE]) dist.init_process_group(backendnccl) torch.cuda.set_device(rank) dataset MyQADataset() sampler DistributedSampler(dataset, num_replicasworld_size, rankrank) dataloader DataLoader(dataset, batch_size32, samplersampler, num_workers4) model MyQAModel().to(rank) model DDP(model, device_ids[rank]) optimizer torch.optim.AdamW(model.parameters(), lr1e-4) criterion torch.nn.CrossEntropyLoss() for epoch in range(num_epochs): sampler.set_epoch(epoch) model.train() for step, batch in enumerate(dataloader): inputs batch[input_ids].to(rank) labels batch[labels].to(rank) outputs model(inputs) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() if rank 0 and step % 100 0: print(fEpoch {epoch}, Step {step}, Loss {loss.item():.4f}) if rank 0: torch.save(model.module.state_dict(), fcheckpoint_epoch{epoch}.pt) dist.destroy_process_group() if __name__ __main__: main()啟動(dòng)命令torchrun --nproc_per_node4 --master_port29500 train.py這個(gè)骨架覆蓋了多卡訓(xùn)練的核心要素進(jìn)程初始化、分布式采樣、DDP包裝、梯度同步、檢查點(diǎn)保存。你可以在此基礎(chǔ)上添加混合精度、梯度累積、學(xué)習(xí)率調(diào)度等功能。8. 一些個(gè)人體會(huì)多卡訓(xùn)練這件事入門門檻其實(shí)不高框架已經(jīng)把大部分臟活累活封裝好了。真正難的是調(diào)優(yōu)和排錯(cuò)。我見過太多人卡在“能跑但跑不快”的階段最后放棄多卡回到單卡。其實(shí)只要理解了通信和計(jì)算的關(guān)系大部分性能問題都有明確的解決方向。另一個(gè)體會(huì)是不要迷信卡數(shù)。4卡能解決的問題不要上8卡?？ㄔ蕉嗾{(diào)試成本越高故障概率越大。我現(xiàn)在的習(xí)慣是先用2卡驗(yàn)證代碼正確性再用4卡跑正式訓(xùn)練除非確實(shí)需要更大規(guī)模否則不輕易增加卡數(shù)。最后多卡訓(xùn)練的日志和監(jiān)控比單卡重要十倍。單卡時(shí)你盯著一個(gè)loss曲線就行多卡時(shí)你要同時(shí)關(guān)注多張卡的指標(biāo)、通信延遲、數(shù)據(jù)加載速度。前期多花點(diǎn)時(shí)間搭好監(jiān)控后期排錯(cuò)能省下大量時(shí)間。