1. 分布式訓練到底在解決什么問題先把場景擺出來。你手里有一塊24GB顯存的卡想訓一個7B參數(shù)的模型。光是把權重、梯度、優(yōu)化器狀態(tài)這三樣東西用fp16存下來粗算一下權重14GB梯度14GBAdam的優(yōu)化器狀態(tài)一階動量加二階動量又是28GB加起來56GB打底還沒算激活值。單卡直接爆。這時候你有兩條路要么換更大的卡要么把這份“家當”拆開分給多塊卡一起扛。分布式訓練干的就是后面這件事。但“拆”這個字說起來輕巧拆法不同代價天差地別。最樸素的拆法是數(shù)據(jù)并行每塊卡都存一份完整模型只是喂不同的數(shù)據(jù)算完梯度再同步。這個方案簡單可它有個硬傷模型本身必須塞得進單卡。當模型大到單卡放不下時數(shù)據(jù)并行就失效了于是才有了后面ZeRO、張量并行、流水線并行、上下文并行這些更細粒度的切分手段。我見過不少剛?cè)腴T的朋友一上來就糾結(jié)“該用哪種并行”其實這個問題問反了。正確的順序是先問“我的瓶頸在哪”——是顯存不夠還是通信太慢還是單卡算力吃不滿。瓶頸不同選型完全不同。這篇文章我就按這個思路把DDP、ZeRO、張量并行、流水線并行、上下文并行這幾條路線挨個拆開講講清楚它們各自解決什么、代價是什么、什么場景該上哪個最后再聊聊這幾種方案怎么組合起來用。適合誰看如果你已經(jīng)能跑通單卡訓練正準備往多卡擴或者你已經(jīng)在用DDP但發(fā)現(xiàn)顯存還是不夠、想搞清楚下一步該往哪走那這篇就是寫給你的。純小白也能看我會盡量用生活化的類比把原理講透但涉及具體配置的地方你得有點PyTorch基礎才跟得上。2. 數(shù)據(jù)并行DDP為什么是分布式訓練的起點2.1 DDP的核心機制與它解決的邊界DDPDistributedDataParallel是絕大多數(shù)人接觸分布式訓練的第一站也是理解后面所有方案的基石。它的邏輯非常直白假設你有N塊卡就把同一個模型完整復制N份每塊卡拿一份然后整個訓練集切成N份每塊卡只吃自己那份數(shù)據(jù)。前向傳播各算各的互不干擾反向傳播時每塊卡算出自己那份數(shù)據(jù)產(chǎn)生的梯度然后所有卡把梯度做一次全局平均保證大家更新后的權重完全一致。這里的關鍵動作是梯度同步。DDP不是等所有梯度都算完才開始通信而是把梯度按參數(shù)分桶bucket某個桶的梯度算完就立刻發(fā)起AllReduce讓通信和計算重疊起來。這個設計很聰明因為反向傳播本身就是從后往前逐層算梯度的算完一層就能同步一層不用干等。實測下來只要桶大小調(diào)得合適通信開銷能被計算掩蓋掉大半。那DDP的邊界在哪就一句話模型必須完整地放進單卡。因為每塊卡都存了一份完整模型包括權重、梯度、優(yōu)化器狀態(tài)。前面算過7B模型用Adam訓練光這三樣就56GB24GB的卡根本放不下。所以DDP適合的是中小模型或者你用了梯度檢查點、混合精度這些省顯存的手段之后單卡勉強能塞下的場景。一旦模型大到單卡放不下DDP就無能為力了必須換思路。2.2 DDP實操配置與幾個容易踩的坑啟動DDP最省事的方式是用torchrun不用自己寫多進程啟動邏輯。一個典型的命令長這樣torchrun --nproc_per_node4 --nnodes1 --node_rank0 \ --master_addr127.0.0.1 --master_port29500 \ train.py代碼里初始化進程組的寫法import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group(backendnccl) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) model MyModel().to(local_rank) model DDP(model, device_ids[local_rank])幾個我踩過的坑值得單獨拎出來說。第一DataLoader必須用DistributedSampler否則每塊卡都會讀到全量數(shù)據(jù)等于白切。第二BatchNorm在DDP下要換成SyncBatchNorm因為普通BN只統(tǒng)計單卡的均值和方差各卡統(tǒng)計量不一致會導致訓練不穩(wěn)定。第三隨機種子要按rank設置不然所有卡的數(shù)據(jù)增強、dropout掩碼都一樣等于變相減小了batch size。第四保存checkpoint只在rank 0上做其他卡寫了會互相覆蓋。提示DDP的通信后端在GPU上一定用nccl別用gloo。gloo是給CPU用的在GPU上跑會慢到讓你懷疑人生。還有一個隱蔽的坑如果你的模型里有不參與梯度計算的參數(shù)比如凍結(jié)的embedding記得在DDP里設置find_unused_parametersTrue否則會報錯說某些參數(shù)沒收到梯度。但這個開關會拖慢速度因為它每次都要遍歷計算圖找未使用參數(shù)。更好的做法是提前把這些參數(shù)從優(yōu)化器里排除而不是靠這個開關兜底。3. ZeRO把冗余徹底干掉的三級切分3.1 ZeRO的三級切分邏輯DDP最大的浪費在于冗余。N塊卡每塊都存一份完整的優(yōu)化器狀態(tài)、梯度和權重可這些東西在同步之后是完全一樣的。ZeROZero Redundancy Optimizer的思路就是既然大家存的東西一樣那就別各存各的了切開分給不同的卡需要的時候再通信湊齊。ZeRO分三個階段切得越來越狠ZeRO-1只切優(yōu)化器狀態(tài)。這是最劃算的一刀因為優(yōu)化器狀態(tài)占大頭Adam是權重的兩倍。切完之后每塊卡只存1/N的優(yōu)化器狀態(tài)顯存直接省下一大截通信量增加很少。ZeRO-2在ZeRO-1基礎上再切梯度。每塊卡只保留自己負責那部分參數(shù)對應的梯度其余梯度算完就扔掉。顯存進一步下降通信量略有增加。ZeRO-3連權重也切。每塊卡只存1/N的權重前向和反向傳播時需要哪層參數(shù)就臨時從其他卡把參數(shù)收集過來用完再釋放。這個階段顯存省得最狠但通信量也最大因為每次計算都要AllGather參數(shù)。用生活化的類比DDP是N個人每人背一整套工具箱ZeRO-1是大家把不常用的工具分開放誰用誰去拿ZeRO-2是連常用工具也分開放ZeRO-3是連手上正在用的工具都只拿零件用的時候現(xiàn)組裝。3.2 ZeRO-3的通信代價與適用判斷ZeRO-3省顯存是真的猛但代價也真實。因為每次前向傳播都要把參數(shù)從其他卡收集過來反向傳播又要再收集一次或者用預取來重疊通信量比DDP高出好幾倍。在帶寬有限的集群上ZeRO-3可能慢得讓你想砸機器。那什么時候該上ZeRO-3我的經(jīng)驗判斷是當你用ZeRO-2還是放不下模型時才考慮ZeRO-3。因為ZeRO-2已經(jīng)切了優(yōu)化器狀態(tài)和梯度通常能省下60%以上的顯存很多場景到這一步就夠了。只有當權重本身都成了瓶頸比如百B級模型才需要動ZeRO-3。配置上DeepSpeed和PyTorch FSDP都實現(xiàn)了ZeRO。DeepSpeed的配置文件里zero_optimization的stage字段就是選級別{ zero_optimization: { stage: 2, offload_optimizer: { device: cpu, pin_memory: true }, allgather_partitions: true, overlap_comm: true, reduce_scatter: true } }這里有個很實用的選項offload_optimizer。它把優(yōu)化器狀態(tài)卸載到CPU內(nèi)存進一步省顯存代價是CPU和GPU之間的數(shù)據(jù)傳輸。如果你的機器CPU內(nèi)存夠大、PCIe帶寬還行這個選項能讓你在有限的顯存下訓更大的模型。我實測過offload之后顯存能再降30%左右但每步耗時增加15%到25%值不值取決于你的顯存有多緊張。注意overlap_comm和reduce_scatter這兩個開關建議都打開它們能讓通信和計算重疊、用更高效的通信原語實測能提速10%以上。3.3 ZeRO與DDP的取舍對照維度DDPZeRO-1ZeRO-2ZeRO-3優(yōu)化器狀態(tài)每卡全量切分切分切分梯度每卡全量每卡全量切分切分權重每卡全量每卡全量每卡全量切分顯存節(jié)省基準約4倍約8倍約N倍通信開銷基準略增增加顯著增加適用場景單卡放得下顯存略緊顯存緊張單卡放不下這張表我建議存下來選型時對著看。核心判斷邏輯是能用DDP就用DDPDDP放不下就上ZeRO-1還放不下上ZeRO-2再放不下才上ZeRO-3。別一上來就ZeRO-3那是殺雞用牛刀還慢。4. 張量并行把單層矩陣乘法切開4.1 張量并行的切分原理ZeRO解決的是“存不下”的問題但它沒有解決“單層算得太慢”或者“單層參數(shù)太大”的問題。張量并行Tensor Parallelism切的是模型內(nèi)部的矩陣運算把一層里的權重矩陣按行或按列切開分到不同卡上算算完再拼起來。拿一個最簡單的全連接層舉例輸出Y XAA是權重矩陣。列并行就是把A按列切成[A1, A2]每塊卡算X和Ai的乘積得到Y(jié)的一部分最后拼起來。行并行是把A按行切每塊卡算一部分輸入和對應權重的乘積最后做AllReduce求和。Transformer里最典型的是把多頭注意力按頭切分每個卡負責幾個注意力頭FFN層則按中間維度切分。Megatron-LM是這方面的標桿實現(xiàn)它的切分策略至今被廣泛借鑒。張量并行的優(yōu)點是通信量相對可控因為每層算完只需要一次AllReduce或AllGather而且通信量正比于激活值大小不像ZeRO-3那樣要傳參數(shù)。缺點是它要求把模型結(jié)構拆得很細實現(xiàn)復雜度高而且對卡間帶寬要求極高因為每層都要通信。所以張量并行通常只在單機內(nèi)NVLink互聯(lián)用跨機用會被網(wǎng)絡帶寬拖死。4.2 張量并行的實操要點與通信瓶頸用Megatron-LM或者它的衍生框架時張量并行度tensor parallel size一般設成2、4、8這種且必須整除注意力頭數(shù)和FFN中間維度。比如模型有32個注意力頭你設tp8那每塊卡分4個頭沒問題設tp5就除不盡直接報錯。通信瓶頸是張量并行的命門。我做過一個對比測試同樣8卡tp8單機NVLink跑和tp2跨機跑前者吞吐是后者的2.5倍以上。原因就是張量并行每層都要AllReduce跨機走網(wǎng)絡延遲太高。所以我的建議很明確張量并行盡量限制在單機內(nèi)跨機交給流水線并行。還有一個細節(jié)張量并行下LayerNorm和Dropout這些操作要在切分后的局部做但它們的統(tǒng)計量需要跨卡一致。Megatron的做法是在這些地方插入額外的通信或者干脆把LayerNorm復制到每塊卡上算。這個取舍會影響精度和速度調(diào)的時候要留意。提示張量并行度不是越大越好。tp增大單卡計算量減小但通信占比上升。存在一個最優(yōu)值通常在4到8之間。超過8之后通信開銷往往吃掉計算收益吞吐不升反降。5. 流水線并行按層切分的時間游戲5.1 流水線并行的基本模型與氣泡問題流水線并行Pipeline Parallelism切的是模型的深度把不同的層分到不同的卡上。比如一個32層的模型4塊卡每塊卡放8層。數(shù)據(jù)從第一塊卡流入逐層往后傳最后一塊卡輸出loss梯度再反向傳回來。這個方案聽起來很美但有個致命問題氣泡bubble。假設4塊卡同一時刻只有一塊卡在干活其他三塊在等——第一塊卡算完第一批數(shù)據(jù)傳給第二塊第二塊才能開始這時候第一塊卡閑著。為了讓所有卡都忙起來就得把一個大batch切成多個micro-batch讓它們像流水線一樣錯峰流動。但即使這樣流水線的頭尾仍然會有填充和排空階段這段時間卡是閑置的這就是氣泡。氣泡的比例大致是(P-1)/(MP-1)P是流水線級數(shù)M是micro-batch數(shù)量。想讓氣泡小要么減少P要么增大M。但M增大會增加顯存占用因為要同時保存多個micro-batch的激活值所以這是個權衡。5.2 1F1B調(diào)度與micro-batch調(diào)優(yōu)為了壓氣泡業(yè)界提出了各種調(diào)度策略最經(jīng)典的是1F1BOne Forward One Backward。它的做法是預熱階段先做P-1次前向填滿流水線然后進入穩(wěn)態(tài)每做一次前向就做一次反向保持流水線滿載最后排空階段把剩下的反向做完。相比樸素的“全前向再全反向”1F1B能把激活值顯存占用降到最低因為同一時刻只需要保存P個micro-batch的激活。調(diào)micro-batch數(shù)量是個經(jīng)驗活。太少氣泡大太多顯存爆。我的做法是先按顯存上限估一個M然后逐步減小看吞吐變化找到拐點。通常M在8到32之間比較常見。流水線并行的優(yōu)點是通信量小因為只在層與層之間傳激活值而且可以跨機用對帶寬要求比張量并行低得多。缺點是氣泡帶來的效率損失以及實現(xiàn)上的調(diào)度復雜度。PyTorch的Pipe、DeepSpeed的Pipeline、Megatron的流水線都是成熟實現(xiàn)選一個用就行別自己造輪子。并行方式切分對象通信頻率通信量跨機友好度數(shù)據(jù)并行數(shù)據(jù)每步一次梯度大小好ZeRO狀態(tài)/梯度/權重每步多次較大好張量并行層內(nèi)矩陣每層多次激活值差流水線并行層間每micro-batch激活值好這張表能幫你快速判斷跨機部署優(yōu)先考慮數(shù)據(jù)并行、ZeRO、流水線并行單機內(nèi)追求極致可以上張量并行。6. 上下文并行長序列訓練的新解法6.1 上下文并行要解決的長序列顯存問題前面幾種并行方式切的是模型參數(shù)或數(shù)據(jù)但有一類場景它們都搞不定超長序列。當序列長度拉到32K、128K甚至更長時注意力機制的顯存占用是隨序列長度平方增長的。這時候哪怕模型不大光是激活值就能把顯存撐爆。上下文并行Context Parallelism切的就是序列維度。它把一條長序列切成若干段每塊卡負責一段注意力計算時通過環(huán)形通信ring attention在卡間傳遞KV讓每塊卡都能看到完整的上下文。這樣顯存占用隨卡數(shù)線性下降序列就能拉得更長。這個思路和前面幾種并行是正交的可以疊加使用。比如你可以同時用張量并行切模型、上下文并行切序列、數(shù)據(jù)并行切batch三者互不沖突。這也是為什么現(xiàn)在訓長上下文模型基本都是多種并行組合著上。6.2 環(huán)形注意力的通信模式與實現(xiàn)選擇環(huán)形注意力的核心是每塊卡持有自己那段序列的Q、K、V然后KV沿著環(huán)形在卡間流轉(zhuǎn)。第i步時每塊卡用當前的KV塊算一部分注意力算完把KV傳給下一塊卡同時接收上一塊卡傳來的KV。轉(zhuǎn)一圈下來每塊卡都見過了所有KV注意力就完整了。這個通信模式的好處是通信量正比于序列長度和卡數(shù)而不是平方。但實現(xiàn)起來對通信和計算的overlap要求很高因為KV傳遞和注意力計算要盡量重疊否則通信就會成為瓶頸。目前支持上下文并行的框架不算多Megatron-LM、DeepSpeed、以及一些專門的長上下文訓練庫有實現(xiàn)。選型時要注意框架對注意力實現(xiàn)的兼容性比如FlashAttention和環(huán)形注意力能不能配合這個得看具體版本。注意上下文并行對序列切分的位置有講究。如果序列里有跨段的依賴比如某些位置編碼切分時要保證這些依賴不被破壞否則精度會掉。實踐中通常按固定長度切并在邊界處做特殊處理。7. 組合并行3D并行的實戰(zhàn)配置思路7.1 3D并行的維度分配邏輯真實的大模型訓練幾乎不會只用一種并行。常見的是3D并行數(shù)據(jù)并行乘以張量并行乘以流水線并行。比如一個千億模型可能配置成tp8、pp8、dp16總共1024塊卡。維度怎么分配我的經(jīng)驗是按這個優(yōu)先級來先看單卡顯存能放下多少用張量并行和流水線并行把模型切開直到單卡放得下剩下的卡全部拿去做數(shù)據(jù)并行。因為數(shù)據(jù)并行的通信效率最高、實現(xiàn)最簡單能用數(shù)據(jù)并行解決的就別用更復雜的。具體到數(shù)字張量并行度受限于單機卡數(shù)和NVLink拓撲一般不超過8流水線并行度受限于層數(shù)和氣泡容忍度一般不超過16數(shù)據(jù)并行度就是剩下的卡數(shù)。三者相乘等于總卡數(shù)。7.2 通信組劃分與配置落地組合并行最容易出錯的地方是通信組的劃分。不同的并行維度需要不同的通信組張量并行組內(nèi)做AllReduce流水線并行組間傳激活數(shù)據(jù)并行組內(nèi)同步梯度。這些組必須正確初始化否則要么通信錯亂要么死鎖。以Megatron-LM為例它的初始化邏輯是先按張量并行分組再按流水線并行分組最后按數(shù)據(jù)并行分組。每塊卡根據(jù)自己的rank計算出自己在各個維度上的坐標然后加入對應的通信組。這個邏輯框架幫你封裝好了你只需要傳對參數(shù)torchrun --nproc_per_node8 --nnodes16 \ pretrain_gpt.py \ --tensor-model-parallel-size 8 \ --pipeline-model-parallel-size 8 \ --num-layers 64 \ --hidden-size 8192 \ --num-attention-heads 64這里tp8、pp8單機8卡正好一個張量并行組16臺機器組成8個流水線級數(shù)據(jù)并行度就是16/82。配置的時候要保證num-layers能被pp整除num-attention-heads和hidden-size能被tp整除否則直接報錯。提示組合并行下全局batch size等于micro-batch乘以數(shù)據(jù)并行度乘以梯度累積步數(shù)。調(diào)參時這三個量要一起看別只盯著一個改。8. 常見問題與排查技巧實錄8.1 顯存與通信類問題速查分布式訓練報錯八成逃不出顯存和通信這兩類。我整理了一張速查表遇到問題先對著查現(xiàn)象可能原因排查方向OOM但單卡能跑并行配置不當檢查ZeRO級別、offload是否開啟訓練卡住不動通信死鎖檢查各卡是否都進入了同一個集合通信吞吐遠低于預期通信瓶頸用profiler看通信占比調(diào)整并行度loss不收斂梯度同步錯誤檢查DDP的find_unused_parameters各卡loss不一致隨機種子未按rank設置檢查seed和sampler配置NCCL超時網(wǎng)絡或拓撲問題檢查NCCL環(huán)境變量和網(wǎng)卡綁定8.2 幾個我踩過的深坑第一個坑NCCL超時。默認NCCL的超時是30分鐘但有些集合通信在特定拓撲下會卡很久導致超時中斷。解決辦法是設NCCL_TIMEOUT環(huán)境變量調(diào)大或者用TORCH_NCCL_ASYNC_ERROR_HANDLING讓錯誤及時拋出而不是干等。第二個坑顯存碎片。長時間訓練后顯存碎片會導致明明還有空閑顯存卻OOM。設PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True能緩解這個問題它讓分配器用可擴展的段來管理顯存減少碎片。第三個坑checkpoint不兼容。用不同并行配置保存和加載checkpoint時參數(shù)切分方式不一樣直接加載會錯位。這時候需要用框架提供的轉(zhuǎn)換工具把checkpoint從一種并行布局轉(zhuǎn)成另一種。Megatron和DeepSpeed都有這類工具別自己手寫轉(zhuǎn)換腳本容易出錯。第四個坑數(shù)據(jù)加載成為瓶頸。當計算很快時DataLoader的讀取速度可能跟不上GPU在等數(shù)據(jù)。解決辦法是增大num_workers、用pin_memory、把數(shù)據(jù)預處理做成流式。我見過一個案例光是把num_workers從4調(diào)到16吞吐就漲了20%。8.3 性能調(diào)優(yōu)的實操順序調(diào)優(yōu)別瞎調(diào)按這個順序來先確認單卡能跑滿再上多卡看擴展效率然后逐步加并行維度。每一步都用profiler記錄找到真正的瓶頸再動手。我常用的profiler是PyTorch自帶的torch.profiler能看到每個算子的耗時和通信占比。擴展效率的計算很簡單N卡吞吐除以單卡吞吐再除以N。如果這個值低于0.7說明擴展性有問題得查通信。低于0.5基本就是通信瓶頸要么調(diào)并行度要么換更快的互聯(lián)。最后分享一個我個人的判斷準則任何并行方案先在小規(guī)模上驗證正確性和擴展效率再往大規(guī)模推。我見過太多人一上來就上幾百卡結(jié)果配置錯了跑了一天發(fā)現(xiàn)loss是nan白白燒錢。小規(guī)模跑通、擴展效率達標再放大這是最省時間的路徑。 分布式訓練并行策略全解析：DDP、ZeRO、張量并行與流水線并行選型指南
1. 分布式訓練到底在解決什么問題先把場景擺出來。你手里有一塊24GB顯存的卡想訓一個7B參數(shù)的模型。光是把權重、梯度、優(yōu)化器狀態(tài)這三樣東西用fp16存下來粗算一下權重14GB梯度14GBAdam的優(yōu)化器狀態(tài)一階動量加二階動量又是28GB加起來56GB打底還沒算激活值。單卡直接爆。這時候你有兩條路要么換更大的卡要么把這份“家當”拆開分給多塊卡一起扛。分布式訓練干的就是后面這件事。但“拆”這個字說起來輕巧拆法不同代價天差地別。最樸素的拆法是數(shù)據(jù)并行每塊卡都存一份完整模型只是喂不同的數(shù)據(jù)算完梯度再同步。這個方案簡單可它有個硬傷模型本身必須塞得進單卡。當模型大到單卡放不下時數(shù)據(jù)并行就失效了于是才有了后面ZeRO、張量并行、流水線并行、上下文并行這些更細粒度的切分手段。我見過不少剛?cè)腴T的朋友一上來就糾結(jié)“該用哪種并行”其實這個問題問反了。正確的順序是先問“我的瓶頸在哪”——是顯存不夠還是通信太慢還是單卡算力吃不滿。瓶頸不同選型完全不同。這篇文章我就按這個思路把DDP、ZeRO、張量并行、流水線并行、上下文并行這幾條路線挨個拆開講講清楚它們各自解決什么、代價是什么、什么場景該上哪個最后再聊聊這幾種方案怎么組合起來用。適合誰看如果你已經(jīng)能跑通單卡訓練正準備往多卡擴或者你已經(jīng)在用DDP但發(fā)現(xiàn)顯存還是不夠、想搞清楚下一步該往哪走那這篇就是寫給你的。純小白也能看我會盡量用生活化的類比把原理講透但涉及具體配置的地方你得有點PyTorch基礎才跟得上。2. 數(shù)據(jù)并行DDP為什么是分布式訓練的起點2.1 DDP的核心機制與它解決的邊界DDPDistributedDataParallel是絕大多數(shù)人接觸分布式訓練的第一站也是理解后面所有方案的基石。它的邏輯非常直白假設你有N塊卡就把同一個模型完整復制N份每塊卡拿一份然后整個訓練集切成N份每塊卡只吃自己那份數(shù)據(jù)。前向傳播各算各的互不干擾反向傳播時每塊卡算出自己那份數(shù)據(jù)產(chǎn)生的梯度然后所有卡把梯度做一次全局平均保證大家更新后的權重完全一致。這里的關鍵動作是梯度同步。DDP不是等所有梯度都算完才開始通信而是把梯度按參數(shù)分桶bucket某個桶的梯度算完就立刻發(fā)起AllReduce讓通信和計算重疊起來。這個設計很聰明因為反向傳播本身就是從后往前逐層算梯度的算完一層就能同步一層不用干等。實測下來只要桶大小調(diào)得合適通信開銷能被計算掩蓋掉大半。那DDP的邊界在哪就一句話模型必須完整地放進單卡。因為每塊卡都存了一份完整模型包括權重、梯度、優(yōu)化器狀態(tài)。前面算過7B模型用Adam訓練光這三樣就56GB24GB的卡根本放不下。所以DDP適合的是中小模型或者你用了梯度檢查點、混合精度這些省顯存的手段之后單卡勉強能塞下的場景。一旦模型大到單卡放不下DDP就無能為力了必須換思路。2.2 DDP實操配置與幾個容易踩的坑啟動DDP最省事的方式是用torchrun不用自己寫多進程啟動邏輯。一個典型的命令長這樣torchrun --nproc_per_node4 --nnodes1 --node_rank0 \ --master_addr127.0.0.1 --master_port29500 \ train.py代碼里初始化進程組的寫法import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group(backendnccl) local_rank int(os.environ[LOCAL_RANK]) torch.cuda.set_device(local_rank) model MyModel().to(local_rank) model DDP(model, device_ids[local_rank])幾個我踩過的坑值得單獨拎出來說。第一DataLoader必須用DistributedSampler否則每塊卡都會讀到全量數(shù)據(jù)等于白切。第二BatchNorm在DDP下要換成SyncBatchNorm因為普通BN只統(tǒng)計單卡的均值和方差各卡統(tǒng)計量不一致會導致訓練不穩(wěn)定。第三隨機種子要按rank設置不然所有卡的數(shù)據(jù)增強、dropout掩碼都一樣等于變相減小了batch size。第四保存checkpoint只在rank 0上做其他卡寫了會互相覆蓋。提示DDP的通信后端在GPU上一定用nccl別用gloo。gloo是給CPU用的在GPU上跑會慢到讓你懷疑人生。還有一個隱蔽的坑如果你的模型里有不參與梯度計算的參數(shù)比如凍結(jié)的embedding記得在DDP里設置find_unused_parametersTrue否則會報錯說某些參數(shù)沒收到梯度。但這個開關會拖慢速度因為它每次都要遍歷計算圖找未使用參數(shù)。更好的做法是提前把這些參數(shù)從優(yōu)化器里排除而不是靠這個開關兜底。3. ZeRO把冗余徹底干掉的三級切分3.1 ZeRO的三級切分邏輯DDP最大的浪費在于冗余。N塊卡每塊都存一份完整的優(yōu)化器狀態(tài)、梯度和權重可這些東西在同步之后是完全一樣的。ZeROZero Redundancy Optimizer的思路就是既然大家存的東西一樣那就別各存各的了切開分給不同的卡需要的時候再通信湊齊。ZeRO分三個階段切得越來越狠ZeRO-1只切優(yōu)化器狀態(tài)。這是最劃算的一刀因為優(yōu)化器狀態(tài)占大頭Adam是權重的兩倍。切完之后每塊卡只存1/N的優(yōu)化器狀態(tài)顯存直接省下一大截通信量增加很少。ZeRO-2在ZeRO-1基礎上再切梯度。每塊卡只保留自己負責那部分參數(shù)對應的梯度其余梯度算完就扔掉。顯存進一步下降通信量略有增加。ZeRO-3連權重也切。每塊卡只存1/N的權重前向和反向傳播時需要哪層參數(shù)就臨時從其他卡把參數(shù)收集過來用完再釋放。這個階段顯存省得最狠但通信量也最大因為每次計算都要AllGather參數(shù)。用生活化的類比DDP是N個人每人背一整套工具箱ZeRO-1是大家把不常用的工具分開放誰用誰去拿ZeRO-2是連常用工具也分開放ZeRO-3是連手上正在用的工具都只拿零件用的時候現(xiàn)組裝。3.2 ZeRO-3的通信代價與適用判斷ZeRO-3省顯存是真的猛但代價也真實。因為每次前向傳播都要把參數(shù)從其他卡收集過來反向傳播又要再收集一次或者用預取來重疊通信量比DDP高出好幾倍。在帶寬有限的集群上ZeRO-3可能慢得讓你想砸機器。那什么時候該上ZeRO-3我的經(jīng)驗判斷是當你用ZeRO-2還是放不下模型時才考慮ZeRO-3。因為ZeRO-2已經(jīng)切了優(yōu)化器狀態(tài)和梯度通常能省下60%以上的顯存很多場景到這一步就夠了。只有當權重本身都成了瓶頸比如百B級模型才需要動ZeRO-3。配置上DeepSpeed和PyTorch FSDP都實現(xiàn)了ZeRO。DeepSpeed的配置文件里zero_optimization的stage字段就是選級別{ zero_optimization: { stage: 2, offload_optimizer: { device: cpu, pin_memory: true }, allgather_partitions: true, overlap_comm: true, reduce_scatter: true } }這里有個很實用的選項offload_optimizer。它把優(yōu)化器狀態(tài)卸載到CPU內(nèi)存進一步省顯存代價是CPU和GPU之間的數(shù)據(jù)傳輸。如果你的機器CPU內(nèi)存夠大、PCIe帶寬還行這個選項能讓你在有限的顯存下訓更大的模型。我實測過offload之后顯存能再降30%左右但每步耗時增加15%到25%值不值取決于你的顯存有多緊張。注意overlap_comm和reduce_scatter這兩個開關建議都打開它們能讓通信和計算重疊、用更高效的通信原語實測能提速10%以上。3.3 ZeRO與DDP的取舍對照維度DDPZeRO-1ZeRO-2ZeRO-3優(yōu)化器狀態(tài)每卡全量切分切分切分梯度每卡全量每卡全量切分切分權重每卡全量每卡全量每卡全量切分顯存節(jié)省基準約4倍約8倍約N倍通信開銷基準略增增加顯著增加適用場景單卡放得下顯存略緊顯存緊張單卡放不下這張表我建議存下來選型時對著看。核心判斷邏輯是能用DDP就用DDPDDP放不下就上ZeRO-1還放不下上ZeRO-2再放不下才上ZeRO-3。別一上來就ZeRO-3那是殺雞用牛刀還慢。4. 張量并行把單層矩陣乘法切開4.1 張量并行的切分原理ZeRO解決的是“存不下”的問題但它沒有解決“單層算得太慢”或者“單層參數(shù)太大”的問題。張量并行Tensor Parallelism切的是模型內(nèi)部的矩陣運算把一層里的權重矩陣按行或按列切開分到不同卡上算算完再拼起來。拿一個最簡單的全連接層舉例輸出Y XAA是權重矩陣。列并行就是把A按列切成[A1, A2]每塊卡算X和Ai的乘積得到Y(jié)的一部分最后拼起來。行并行是把A按行切每塊卡算一部分輸入和對應權重的乘積最后做AllReduce求和。Transformer里最典型的是把多頭注意力按頭切分每個卡負責幾個注意力頭FFN層則按中間維度切分。Megatron-LM是這方面的標桿實現(xiàn)它的切分策略至今被廣泛借鑒。張量并行的優(yōu)點是通信量相對可控因為每層算完只需要一次AllReduce或AllGather而且通信量正比于激活值大小不像ZeRO-3那樣要傳參數(shù)。缺點是它要求把模型結(jié)構拆得很細實現(xiàn)復雜度高而且對卡間帶寬要求極高因為每層都要通信。所以張量并行通常只在單機內(nèi)NVLink互聯(lián)用跨機用會被網(wǎng)絡帶寬拖死。4.2 張量并行的實操要點與通信瓶頸用Megatron-LM或者它的衍生框架時張量并行度tensor parallel size一般設成2、4、8這種且必須整除注意力頭數(shù)和FFN中間維度。比如模型有32個注意力頭你設tp8那每塊卡分4個頭沒問題設tp5就除不盡直接報錯。通信瓶頸是張量并行的命門。我做過一個對比測試同樣8卡tp8單機NVLink跑和tp2跨機跑前者吞吐是后者的2.5倍以上。原因就是張量并行每層都要AllReduce跨機走網(wǎng)絡延遲太高。所以我的建議很明確張量并行盡量限制在單機內(nèi)跨機交給流水線并行。還有一個細節(jié)張量并行下LayerNorm和Dropout這些操作要在切分后的局部做但它們的統(tǒng)計量需要跨卡一致。Megatron的做法是在這些地方插入額外的通信或者干脆把LayerNorm復制到每塊卡上算。這個取舍會影響精度和速度調(diào)的時候要留意。提示張量并行度不是越大越好。tp增大單卡計算量減小但通信占比上升。存在一個最優(yōu)值通常在4到8之間。超過8之后通信開銷往往吃掉計算收益吞吐不升反降。5. 流水線并行按層切分的時間游戲5.1 流水線并行的基本模型與氣泡問題流水線并行Pipeline Parallelism切的是模型的深度把不同的層分到不同的卡上。比如一個32層的模型4塊卡每塊卡放8層。數(shù)據(jù)從第一塊卡流入逐層往后傳最后一塊卡輸出loss梯度再反向傳回來。這個方案聽起來很美但有個致命問題氣泡bubble。假設4塊卡同一時刻只有一塊卡在干活其他三塊在等——第一塊卡算完第一批數(shù)據(jù)傳給第二塊第二塊才能開始這時候第一塊卡閑著。為了讓所有卡都忙起來就得把一個大batch切成多個micro-batch讓它們像流水線一樣錯峰流動。但即使這樣流水線的頭尾仍然會有填充和排空階段這段時間卡是閑置的這就是氣泡。氣泡的比例大致是(P-1)/(MP-1)P是流水線級數(shù)M是micro-batch數(shù)量。想讓氣泡小要么減少P要么增大M。但M增大會增加顯存占用因為要同時保存多個micro-batch的激活值所以這是個權衡。5.2 1F1B調(diào)度與micro-batch調(diào)優(yōu)為了壓氣泡業(yè)界提出了各種調(diào)度策略最經(jīng)典的是1F1BOne Forward One Backward。它的做法是預熱階段先做P-1次前向填滿流水線然后進入穩(wěn)態(tài)每做一次前向就做一次反向保持流水線滿載最后排空階段把剩下的反向做完。相比樸素的“全前向再全反向”1F1B能把激活值顯存占用降到最低因為同一時刻只需要保存P個micro-batch的激活。調(diào)micro-batch數(shù)量是個經(jīng)驗活。太少氣泡大太多顯存爆。我的做法是先按顯存上限估一個M然后逐步減小看吞吐變化找到拐點。通常M在8到32之間比較常見。流水線并行的優(yōu)點是通信量小因為只在層與層之間傳激活值而且可以跨機用對帶寬要求比張量并行低得多。缺點是氣泡帶來的效率損失以及實現(xiàn)上的調(diào)度復雜度。PyTorch的Pipe、DeepSpeed的Pipeline、Megatron的流水線都是成熟實現(xiàn)選一個用就行別自己造輪子。并行方式切分對象通信頻率通信量跨機友好度數(shù)據(jù)并行數(shù)據(jù)每步一次梯度大小好ZeRO狀態(tài)/梯度/權重每步多次較大好張量并行層內(nèi)矩陣每層多次激活值差流水線并行層間每micro-batch激活值好這張表能幫你快速判斷跨機部署優(yōu)先考慮數(shù)據(jù)并行、ZeRO、流水線并行單機內(nèi)追求極致可以上張量并行。6. 上下文并行長序列訓練的新解法6.1 上下文并行要解決的長序列顯存問題前面幾種并行方式切的是模型參數(shù)或數(shù)據(jù)但有一類場景它們都搞不定超長序列。當序列長度拉到32K、128K甚至更長時注意力機制的顯存占用是隨序列長度平方增長的。這時候哪怕模型不大光是激活值就能把顯存撐爆。上下文并行Context Parallelism切的就是序列維度。它把一條長序列切成若干段每塊卡負責一段注意力計算時通過環(huán)形通信ring attention在卡間傳遞KV讓每塊卡都能看到完整的上下文。這樣顯存占用隨卡數(shù)線性下降序列就能拉得更長。這個思路和前面幾種并行是正交的可以疊加使用。比如你可以同時用張量并行切模型、上下文并行切序列、數(shù)據(jù)并行切batch三者互不沖突。這也是為什么現(xiàn)在訓長上下文模型基本都是多種并行組合著上。6.2 環(huán)形注意力的通信模式與實現(xiàn)選擇環(huán)形注意力的核心是每塊卡持有自己那段序列的Q、K、V然后KV沿著環(huán)形在卡間流轉(zhuǎn)。第i步時每塊卡用當前的KV塊算一部分注意力算完把KV傳給下一塊卡同時接收上一塊卡傳來的KV。轉(zhuǎn)一圈下來每塊卡都見過了所有KV注意力就完整了。這個通信模式的好處是通信量正比于序列長度和卡數(shù)而不是平方。但實現(xiàn)起來對通信和計算的overlap要求很高因為KV傳遞和注意力計算要盡量重疊否則通信就會成為瓶頸。目前支持上下文并行的框架不算多Megatron-LM、DeepSpeed、以及一些專門的長上下文訓練庫有實現(xiàn)。選型時要注意框架對注意力實現(xiàn)的兼容性比如FlashAttention和環(huán)形注意力能不能配合這個得看具體版本。注意上下文并行對序列切分的位置有講究。如果序列里有跨段的依賴比如某些位置編碼切分時要保證這些依賴不被破壞否則精度會掉。實踐中通常按固定長度切并在邊界處做特殊處理。7. 組合并行3D并行的實戰(zhàn)配置思路7.1 3D并行的維度分配邏輯真實的大模型訓練幾乎不會只用一種并行。常見的是3D并行數(shù)據(jù)并行乘以張量并行乘以流水線并行。比如一個千億模型可能配置成tp8、pp8、dp16總共1024塊卡。維度怎么分配我的經(jīng)驗是按這個優(yōu)先級來先看單卡顯存能放下多少用張量并行和流水線并行把模型切開直到單卡放得下剩下的卡全部拿去做數(shù)據(jù)并行。因為數(shù)據(jù)并行的通信效率最高、實現(xiàn)最簡單能用數(shù)據(jù)并行解決的就別用更復雜的。具體到數(shù)字張量并行度受限于單機卡數(shù)和NVLink拓撲一般不超過8流水線并行度受限于層數(shù)和氣泡容忍度一般不超過16數(shù)據(jù)并行度就是剩下的卡數(shù)。三者相乘等于總卡數(shù)。7.2 通信組劃分與配置落地組合并行最容易出錯的地方是通信組的劃分。不同的并行維度需要不同的通信組張量并行組內(nèi)做AllReduce流水線并行組間傳激活數(shù)據(jù)并行組內(nèi)同步梯度。這些組必須正確初始化否則要么通信錯亂要么死鎖。以Megatron-LM為例它的初始化邏輯是先按張量并行分組再按流水線并行分組最后按數(shù)據(jù)并行分組。每塊卡根據(jù)自己的rank計算出自己在各個維度上的坐標然后加入對應的通信組。這個邏輯框架幫你封裝好了你只需要傳對參數(shù)torchrun --nproc_per_node8 --nnodes16 \ pretrain_gpt.py \ --tensor-model-parallel-size 8 \ --pipeline-model-parallel-size 8 \ --num-layers 64 \ --hidden-size 8192 \ --num-attention-heads 64這里tp8、pp8單機8卡正好一個張量并行組16臺機器組成8個流水線級數(shù)據(jù)并行度就是16/82。配置的時候要保證num-layers能被pp整除num-attention-heads和hidden-size能被tp整除否則直接報錯。提示組合并行下全局batch size等于micro-batch乘以數(shù)據(jù)并行度乘以梯度累積步數(shù)。調(diào)參時這三個量要一起看別只盯著一個改。8. 常見問題與排查技巧實錄8.1 顯存與通信類問題速查分布式訓練報錯八成逃不出顯存和通信這兩類。我整理了一張速查表遇到問題先對著查現(xiàn)象可能原因排查方向OOM但單卡能跑并行配置不當檢查ZeRO級別、offload是否開啟訓練卡住不動通信死鎖檢查各卡是否都進入了同一個集合通信吞吐遠低于預期通信瓶頸用profiler看通信占比調(diào)整并行度loss不收斂梯度同步錯誤檢查DDP的find_unused_parameters各卡loss不一致隨機種子未按rank設置檢查seed和sampler配置NCCL超時網(wǎng)絡或拓撲問題檢查NCCL環(huán)境變量和網(wǎng)卡綁定8.2 幾個我踩過的深坑第一個坑NCCL超時。默認NCCL的超時是30分鐘但有些集合通信在特定拓撲下會卡很久導致超時中斷。解決辦法是設NCCL_TIMEOUT環(huán)境變量調(diào)大或者用TORCH_NCCL_ASYNC_ERROR_HANDLING讓錯誤及時拋出而不是干等。第二個坑顯存碎片。長時間訓練后顯存碎片會導致明明還有空閑顯存卻OOM。設PYTORCH_CUDA_ALLOC_CONFexpandable_segments:True能緩解這個問題它讓分配器用可擴展的段來管理顯存減少碎片。第三個坑checkpoint不兼容。用不同并行配置保存和加載checkpoint時參數(shù)切分方式不一樣直接加載會錯位。這時候需要用框架提供的轉(zhuǎn)換工具把checkpoint從一種并行布局轉(zhuǎn)成另一種。Megatron和DeepSpeed都有這類工具別自己手寫轉(zhuǎn)換腳本容易出錯。第四個坑數(shù)據(jù)加載成為瓶頸。當計算很快時DataLoader的讀取速度可能跟不上GPU在等數(shù)據(jù)。解決辦法是增大num_workers、用pin_memory、把數(shù)據(jù)預處理做成流式。我見過一個案例光是把num_workers從4調(diào)到16吞吐就漲了20%。8.3 性能調(diào)優(yōu)的實操順序調(diào)優(yōu)別瞎調(diào)按這個順序來先確認單卡能跑滿再上多卡看擴展效率然后逐步加并行維度。每一步都用profiler記錄找到真正的瓶頸再動手。我常用的profiler是PyTorch自帶的torch.profiler能看到每個算子的耗時和通信占比。擴展效率的計算很簡單N卡吞吐除以單卡吞吐再除以N。如果這個值低于0.7說明擴展性有問題得查通信。低于0.5基本就是通信瓶頸要么調(diào)并行度要么換更快的互聯(lián)。最后分享一個我個人的判斷準則任何并行方案先在小規(guī)模上驗證正確性和擴展效率再往大規(guī)模推。我見過太多人一上來就上幾百卡結(jié)果配置錯了跑了一天發(fā)現(xiàn)loss是nan白白燒錢。小規(guī)模跑通、擴展效率達標再放大這是最省時間的路徑。