最近在技術(shù)社區(qū)刷到 MiniMax-H3 的討論很多人在聊它生成的豎屏短片效果也在聊模型下載后怎么加速推理。但我更關(guān)心的不是短片本身而是這些視頻模型背后一個很容易被帶過的底層問題多卡訓(xùn)練時all-reduce 到底怎樣讓每張卡上的模型保持一致這不是純理論問題。寫過分布式訓(xùn)練腳本的人都知道每次看 loss 曲線時真正想讓所有 GPU 拿到的不是“各自的結(jié)果”而是一份完全一致的、聚合了所有卡信息之后的結(jié)果。all-reduce 就是負(fù)責(zé)這件事的算子。這篇文章想把它講透從語義到 Ring 算法從 PyTorch/DDP 的落地到排查鏈路最后給一個判斷框架。1. 先搞清楚一個容易被帶偏的問題訓(xùn)練時每張卡為什么會有差異1.1 數(shù)據(jù)并行下每張卡看到的本來就是不同數(shù)據(jù)規(guī)模化訓(xùn)練里最常見的并行方式是數(shù)據(jù)并行。把訓(xùn)練集切成 N 份每張卡拿其中一份各自做前向、算 loss、再反向。這里的起始點就是不一致的每張卡輸入的是不同 batch算出來的 loss 和梯度天然不同。不要覺得這有什么問題。數(shù)據(jù)并行本來就是用“數(shù)據(jù)分片”換取“算力擴(kuò)展”每張卡看到的數(shù)據(jù)不同是設(shè)計的一部分。1.2 不一樣從梯度開始最后會傳導(dǎo)到參數(shù)如果每張卡各自按自己算出來的梯度直接更新參數(shù)那么從第一個 step 開始模型權(quán)重就已經(jīng)分叉了。N 張卡跑一段時間后等于 N 個版本不同的模型在并行訓(xùn)練。最后保存模型的時候該信任哪一張卡的 checkpoint結(jié)果通常完全不可復(fù)現(xiàn)。所以需要一種機(jī)制每張卡算好自己的梯度之后先不要急著更新而是把 N 份梯度匯總成一份統(tǒng)一結(jié)果再廣播回每一張卡。所有卡用同一份梯度去更新各自的參數(shù)。因為初始參數(shù)相同、更新規(guī)則相同、拿到的梯度也相同下一步參數(shù)自然一致。1.3 “每卡相同”具體指哪些東西相同這里需要分辨三層狀態(tài)模型參數(shù)每個 step 結(jié)束之后N 張卡上的權(quán)重應(yīng)當(dāng)完全一致。這是最核心的要求。優(yōu)化器狀態(tài)如果每張卡初始狀態(tài)一致、拿到的梯度一致、超參一致優(yōu)化器狀態(tài)也會保持一致。實際運行時會因為浮點運算順序不同出現(xiàn)極小數(shù)值偏差但一般不會改變訓(xùn)練路徑。隨機(jī)狀態(tài)dropout 這類訓(xùn)練期隨機(jī)性可以不一致也不需要一致。如果要做嚴(yán)格復(fù)現(xiàn)需要統(tǒng)一 seed并控制 CUDA 層面引入的非確定性。一個很容易誤解的點是很多人以為“每卡相同”是指計算過程完全一樣。其實不是。輸入不同、計算路徑里可能有隨機(jī)性這些差異都允許存在真正必須相同的是“聚合后的梯度”和“更新后的參數(shù)”。all-reduce 保證的是結(jié)果側(cè)的收斂。1.4 先把主判斷放在這里all-reduce 不是為了讓每張卡“算得一模一樣”而是讓每張卡在拿到其他人的信息之后變成“最終結(jié)果一模一樣”。它是把分布式訓(xùn)練的結(jié)果拉回等價于單卡訓(xùn)練的橋梁。2. all-reduce 的語義一次操作本質(zhì)上要完成“聚合”與“分發(fā)”2.1 Reduce 和 All-Reduce 的區(qū)別在集合通信里reduce 指把多張卡上的張量聚合成一個結(jié)果這個結(jié)果可能只落在某一張卡上。比如 all-gather 是每個節(jié)點都拿到所有人的數(shù)據(jù)然后本地做 reduce這是一種可以實現(xiàn)目標(biāo)的做法但通信成本高。all-reduce 把“聚合”和“分發(fā)”合并成了一個語義N 張卡各自持有一個張量經(jīng)過一次調(diào)用之后每張卡都拿到完全相同的聚合結(jié)果。常見聚合操作有這么幾種SUM直接把 N 份張量相加這是 DDP 梯度同步的默認(rèn)做法。AVG相加后除以 N。MAX / MIN取最大或最小。PROD相乘。在梯度同步場景里最常用的是 SUM 和 AVG。PyTorch 的torch.distributed.all_reduce支持ReduceOp.SUM、ReduceOp.AVG等DDP 內(nèi)部默認(rèn)是 SUM所以使用時要留意 loss 的歸約方式。2.2 集合語義帶來的兩個工程約束all-reduce 不是 A 到 B 的點對點通信而是要求某個集合內(nèi)所有成員一起參與、結(jié)束時狀態(tài)一致。這帶來兩個工程含義。第一個是調(diào)用必須“齊步走”所有 rank 都要執(zhí)行相同的 all-reduce且參與的張量形態(tài)要匹配。如果一個 rank 沒調(diào)用其他 rank 會一直卡在通信等待里。很多分布式腳本“掛住不動”的問題根源就在這里。第二個是結(jié)果需要確定性不管調(diào)用順序怎樣同一 batch 下各卡拿到的最終張量應(yīng)當(dāng)一致。這取決于底層實現(xiàn)的歸約順序也是后面講數(shù)值精度時的一個重要背景。2.3 邏輯上可以拆成“先匯總再廣播”不管底層怎么實現(xiàn)所有 all-reduce 都可以在邏輯上拆成兩步匯總把所有卡的局部張量聚合成一份全局結(jié)果。廣播把這份全局結(jié)果復(fù)制回每一張卡。Ring 算法之所以經(jīng)典是因為它對這兩步做了高效的流水化實現(xiàn)而不是先把數(shù)據(jù)匯總到某個中心再分發(fā)出去。高效點不在語義創(chuàng)新而在通信模式的重新設(shè)計。3. Ring AllReduce為什么它能高效地做到“讓每卡相同”3.1 樸素做法的問題在哪里一個最直接的做法是每張卡把自己的張量廣播給其他所有卡然后每張卡本地求和。兩卡時很直觀但到大規(guī)模集群就有幾個問題。第一消息數(shù)隨卡數(shù)快速增長變成接近 O(N2) 級別的事務(wù)數(shù)量。第二某一時刻多個來源同時發(fā)給同一目的地接收端網(wǎng)卡成為瓶頸。第三沒有中心節(jié)點來統(tǒng)一匯聚很多實現(xiàn)會退化成“集線器模式”實際帶寬浪費嚴(yán)重。3.2 Ring 的核心reduce-scatter all-gatherRing AllReduce 是當(dāng)前多卡訓(xùn)練里最主流實現(xiàn)之一。它的思路是把 N 張卡看成一個環(huán)每張卡只和前后兩個鄰居通信。假設(shè)每張卡持有大小為 S 的梯度張量總共有 N 張卡。算法分兩個階段。第一階段叫 reduce-scatter歸約分散。把本卡大小為 S 的張量切成 N 份。每輪里每張卡把自己手里的某個 chunk 發(fā)給下一個鄰居同時從上一個鄰居收到一個 chunk并累加到自己對應(yīng)的 chunk 上。重復(fù) N-1 輪之后每張卡上有一份 chunk它是所有卡對應(yīng)位置累加后的結(jié)果。第二階段叫 all-gather全收集。每張卡把手里那份累加后的 chunk 繼續(xù)沿環(huán)發(fā)給下一個鄰居。重復(fù) N-1 輪之后每張卡都收集齊了全部累加 chunk。把它們拼接起來就得到完整的 all-reduce 結(jié)果。這個階段不需要計算只做搬運。3.3 通信量為什么近似“2 倍數(shù)據(jù)量”每張卡在第一階段發(fā)送 N-1 次 chunk第二階段同樣發(fā)送 N-1 次 chunk。每張卡發(fā)送和接收的總數(shù)據(jù)量大致是每卡收發(fā)總量 2 × (N-1) × S / N當(dāng) N 較大時約等于 2S。這個數(shù)很關(guān)鍵無論集群里有幾百張卡單個節(jié)點搬進(jìn)搬出的數(shù)據(jù)量主要只由它自己的張量規(guī)模決定不會因為卡數(shù)變多而線性膨脹。所以在消息足夠大的場景下Ring 是帶寬最優(yōu)的近似實現(xiàn)特別適合大模型、大梯度張量。如果每張卡把整份 S 都廣播出去每卡收發(fā)量會變成 2(N-1)S比 Ring 差了 N 倍。Ring 厲害的地方就在這里同樣一次 all-reduce它把“與卡數(shù)相關(guān)的開銷”壓到了只剩常數(shù)級別。3.4 四卡環(huán)的一個文字演示以四卡 A、B、C、D 組成環(huán)為例。reduce-scatter 階段每輪大家同時做“發(fā)一個 chunk、收一個 chunk、累加”。三輪之后A 手里持有某一份完整累加 chunkB、C、D 各持有另外一份完整累加 chunk誰都不完整但每個人手里那份已經(jīng)匯聚了所有卡的信息。all-gather 階段A 把手里 chunk 發(fā)給 B同時從 D 收到所需 chunk。三輪之后A、B、C、D 都擁有全部四個累加結(jié)果拼接后就是完整張量。注意每個環(huán)節(jié)所有卡都在同時收發(fā)鏈路上沒有閑置這是 Ring 能跑滿帶寬的核心原因。4. 從參數(shù)服務(wù)器到 NCCL為什么工業(yè)界默認(rèn)選擇 all-reduce4.1 參數(shù)服務(wù)器不是銀彈在 all-reduce 流行之前參數(shù)服務(wù)器PS是分布式訓(xùn)練里很常見的一種方案。思路直接worker 卡負(fù)責(zé)計算梯度把梯度發(fā)給一組 serverserver 匯總后更新參數(shù)再把新參數(shù)廣播回 worker。優(yōu)點是實現(xiàn)直觀天然支持異步更新。缺點是當(dāng)模型變大、worker 變多時server 的入口帶寬和出口帶寬會變成瓶頸而且 server 節(jié)點一旦出問題整個訓(xùn)練就被拖住。從工程直覺看PS 更適合“通信不是主要矛盾、需要靈活調(diào)度”的場景數(shù)據(jù)并行 all-reduce 更貼近“算力足夠多、想盡量壓低通信開銷”的場景。4.2 NCCL 把 Ring 做成了生產(chǎn)級NCCLNVIDIA Collective Communications Library把環(huán)、樹等多種集合通信算法做成了生產(chǎn)級實現(xiàn)并針對單機(jī) NVLink、多機(jī) InfiniBand 或 RoCE 做了適配。幾個關(guān)鍵優(yōu)化點根據(jù)卡間拓?fù)溥x擇算法比如單機(jī)內(nèi)部用 NVLink 高帶寬鏈路跨機(jī)場景用樹或環(huán)降低鏈路壓力。梯度分桶把多個參數(shù)的小梯度合并成更大的通信塊減少消息數(shù)量。通信和反向計算重疊bucket 填滿就先發(fā)不必等全部梯度算完。支持 RDMA 和 GPUDirect減少數(shù)據(jù)拷貝。這里要強(qiáng)調(diào)一點NCCL 不是只有 Ring。實際使用中它會根據(jù)設(shè)備數(shù)、節(jié)點數(shù)、數(shù)據(jù)規(guī)模自動或手動選擇算法。使用者一開始不必手動干預(yù)但要知道有這些參數(shù)存在。方案通信方式主要瓶頸適合場景參數(shù)服務(wù)器中心節(jié)點匯聚再廣播server 帶寬、單點故障小規(guī)模、異步容忍度高樸素 AllReduce每卡廣播再本地求和消息數(shù)多、接收端擁塞卡數(shù)非常少Ring AllReduce環(huán)上點對點流轉(zhuǎn)延遲隨環(huán)長度增加大梯度、大規(guī)模數(shù)據(jù)并行Tree AllReduce樹狀聚合再廣播根節(jié)點仍是匯聚瓶頸消息較小、延遲敏感4.3 all-reduce 不只出現(xiàn)在梯度同步很多人把 all-reduce 等同于“梯度同步”。其實在大模型訓(xùn)練里它還頻繁出現(xiàn)在張量并行中。以常見的 Megatron 風(fēng)格模型并行舉例一層 MLP 被切到多張卡上每張卡只算一部分神經(jīng)元。前向過程中的某些位置需要把各卡的部分結(jié)果聚合起來才能繼續(xù)算下一層。這個聚合動作往往就是 all-reduce。理解這一點后再看“讓每卡相同”會更完整無論做數(shù)據(jù)并行還是張量并行最終都要保證關(guān)鍵張量在所有卡上是對齊的。5. 在 PyTorch 里DDP 是用 all-reduce 幫你怎么做的5.1 DDP 的三步固定流程實際項目里我一般不會手動寫 all-reduce 來同步梯度因為DistributedDataParallel已經(jīng)封裝好了。封裝不是魔法它只是做幾件固定的事在進(jìn)程組初始化時建立所有 rank 的通信上下文。前向結(jié)束后在反向傳播過程中給每個參數(shù)的梯度注冊鉤子。梯度算出來后按 bucket 分批執(zhí)行 all-reduce把多卡帶來的梯度匯總回每張卡。對應(yīng)到代碼大致是這樣一種結(jié)構(gòu)import torch import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP def train_worker(rank, world_size): dist.init_process_group(nccl, rankrank, world_sizeworld_size) torch.cuda.set_device(rank) model build_model().cuda(rank) model DDP(model, device_ids[rank], bucket_cap_mb25) dataset build_dataset(rankrank, num_replicasworld_size) loader torch.utils.data.DataLoader(dataset, batch_size8) optimizer torch.optim.AdamW(model.parameters(), lr1e-4) for batch in loader: optimizer.zero_grad() loss compute_loss(model, batch) # DDP 的梯度 all-reduce 默認(rèn)是 SUM。 # 如果想等價于“全局平均梯度”需要在 loss 處除以 world_size。 loss loss / world_size loss.backward() optimizer.step()上面是通用寫法具體模型和數(shù)據(jù)集要換成你自己的。寫的時候注意兩件事數(shù)據(jù)加載用DistributedSampler保證 batch 不重疊訓(xùn)練語義正確。loss 是否除以 world_size取決于你想讓優(yōu)化步等價于單卡哪種語義。這里沒有唯一正確答案但必須明確自己選的是哪一種。5.2 幾個會直接影響結(jié)果的 DDP 參數(shù)參數(shù)影響落地建議bucket_cap_mb控制梯度分桶大小影響 all-reduce 次數(shù)先用默認(rèn) 25MB觀察 GPU 利用率和通信耗時再調(diào)find_unused_parameters處理前向里沒用到的參數(shù)設(shè)置不當(dāng)會卡住有動態(tài)分支時設(shè)為 True會帶來額外開銷gradient_as_bucket_view讓梯度視圖直接指向 bucket省內(nèi)存拷貝顯存緊張時可開啟但升級版本后要看 API 變化timeout控制通信等待超時長訓(xùn)練任務(wù)建議設(shè)置明確超時暴露卡住問題很多人忽略bucket_cap_mb但它是影響通信利用率的關(guān)鍵之一。bucket 太小通信次數(shù)多、消息碎bucket 太大通信占用的顯存多、延遲發(fā)布。我會先跑一個穩(wěn)定小任務(wù)再用torch.profiler看通信占比最后才動這個參數(shù)。5.3 混合精度里最容易忽略的一點在 AMP 或 bf16 訓(xùn)練里損失 scale 和梯度精度會影響 all-reduce 的結(jié)果。fp16 梯度做通信時數(shù)值范圍有限。如果 loss scale 不當(dāng)梯度可能下溢。bf16 指數(shù)范圍更大但有效精度較少。直接用 bf16 梯度做 all-reduce結(jié)果會粗糙一些。更穩(wěn)妥的做法是用 fp32 主權(quán)重和梯度做歸約穩(wěn)定之后再嘗試低精度通信優(yōu)化。這種細(xì)節(jié)放在“每卡相同”的話題里尤其重要。因為浮點歸約順序不同各卡拿到的最終值會有極小差異。多數(shù)情況下差異不影響訓(xùn)練但做嚴(yán)格復(fù)現(xiàn)和調(diào)優(yōu)對齊時它就是需要關(guān)注的那一部分。5.4 如果你是沖著“視頻生成模型加速”來的回到開頭提到的 MiniMax-H3 討論。社區(qū)里聊得多的往往是“能不能生成豎屏短片”“怎么加速推理”但很少有人聊模型背后的并行鏈路。我的體感是如果你只是用現(xiàn)成模型跑一次推理all-reduce 不一定出現(xiàn)在你面前但如果你要自己微調(diào)一個類似規(guī)模的視頻生成模型或者想把推理拆到多卡并行all-reduce 遲早會出現(xiàn)。這類視頻生成模型通常有幾個共同點batch 內(nèi)視頻片段較長顯存和算力需求高需要足夠大的 batch 才能讓訓(xùn)練穩(wěn)定對時間維度的處理會帶來更多激活值占用。用 DDP 或更上層框架跑這類訓(xùn)練時梯度 all-reduce 的通信占比往往很可觀。優(yōu)化方向不是調(diào)一個神秘參數(shù)而是先搞清楚通信發(fā)生的位置以及瓶頸是帶寬還是延遲。6. 排查鏈路如果發(fā)現(xiàn)每卡結(jié)果不一致按什么順序查很多人遇到“多卡訓(xùn)練結(jié)果不對”時第一反應(yīng)是懷疑 all-reduce 壞了。但從工程經(jīng)驗看all-reduce 本身很少是根因絕大多數(shù)問題出現(xiàn)在它前后。6.1 先確認(rèn)是不是真的“不一致”不要急著看通信。把同樣數(shù)據(jù)、同樣 seed 跑一次單卡再跑一次多卡對比 loss 曲線和最終指標(biāo)。如果只是存在微小浮點差異很可能不是故障而是歸約順序、TF32、低精度算子導(dǎo)致的數(shù)值抖動。這種情況不需要消除只需要控制在可接受范圍。如果差異很大比如 loss 發(fā)散或者從一個 step 開始就完全不同進(jìn)入下一步。6.2 再查初始化和數(shù)據(jù)按順序檢查每張卡的 seed 是否一致。checkpoint 是否在初始化進(jìn)程組之后正確加載還是每卡都從隨機(jī)權(quán)重開始。DistributedSampler是否正確使用batch 是否重疊。數(shù)據(jù)預(yù)處理中是否有依賴進(jìn)程號的隨機(jī)操作沒有固定。6.3 再查同步機(jī)制和 DDP 狀態(tài)確認(rèn)模型確實被DDP包住而不是仍然用了nn.DataParallel。檢查模型是否有requires_gradFalse的參數(shù)。這類參數(shù)不參與梯度 all-reduce如果依賴它們變化步調(diào)就會出現(xiàn)偏差。檢查是否有 buffer 依賴訓(xùn)練過程更新。DDP 默認(rèn)在 forward 時從 rank 0 同步 buffer但 buffer 不參與梯度歸約。如果各卡獨立更新 buffer就可能不一致。檢查是否有某個參數(shù)沒有被任何 loss 反向使用尤其是有條件分支、動態(tài)結(jié)構(gòu)的模型。此時需要find_unused_parametersTrue。6.4 再查通信環(huán)境如果是多機(jī)多卡還要看NCCL 版本和 PyTorch 自帶的 NCCL 是否匹配。網(wǎng)卡配置比如 RDMA 是否啟用網(wǎng)卡是否綁定到了正確的 CPU 和 NUMA 節(jié)點。NCCL_P2P_DISABLE、NCCL_IB_DISABLE這類環(huán)境變量是否被誤設(shè)。日志里是否有 timeout、network error 等關(guān)鍵詞。6.5 最后才查算子層面的非確定性如果以上都沒問題再考慮算子層面是否設(shè)置了torch.backends.cudnn.deterministicTrue。是否關(guān)閉了 TF32。是否在低精度下啟用了原子累加。張量并行或序列并行里的 all-reduce 次數(shù)和位置是否正確。這套排查順序的核心邏輯是先排除“數(shù)據(jù)或初始化導(dǎo)致起點不同”再排除“同步機(jī)制失效”最后才進(jìn)入“通信過程中數(shù)值不精確”的細(xì)節(jié)。否則很容易在最后一個環(huán)節(jié)浪費時間。7. 一個判斷框架什么時候值得在 all-reduce 上花時間7.1 三個層次的需求判斷不是所有項目都需要深入 all-reduce。這里給一個簡單框架場景你需要關(guān)心什么建議學(xué)習(xí) / 小實驗跑通 DDP理解語義先用默認(rèn) DDP不手動調(diào)通信參數(shù)單機(jī)多卡生產(chǎn)通信占比、bucket、顯存先用 profiler 分析再調(diào) bucket 和重疊多機(jī)多卡 / 大規(guī)模帶寬、拓?fù)?、RDMA、故障恢復(fù)了解 NCCL 拓?fù)浜筒⑿胁呗员匾獣r上 FSDP 或混合并行7.2 all-reduce 不解決什么寫清楚邊界能省掉很多調(diào)試時間它不解決數(shù)據(jù)不均衡。如果某張卡分到的 batch 明顯更大或更復(fù)雜訓(xùn)練效率會被慢卡拖住。它不解決節(jié)點速度不一致。Ring 假設(shè)各節(jié)點速度接近遇到 straggler 時整體都會變慢。它不解決模型并行里的手動調(diào)度。在這類場景里all-reduce 只是眾多算子之一關(guān)鍵是并行切分策略。它不保證推理階段所有卡輸出一致。推理一致性還需要模型本身確定性、精度設(shè)置和輸入處理方式配合。7.3 長期優(yōu)化方向如果你真的需要在一個長訓(xùn)練任務(wù)里持續(xù)壓通信成本方向通常是這幾條梯度壓縮或量化用低精度表示梯度減少通信量但要在小規(guī)模先驗證收斂。通信與計算重疊利用 bucket 和異步 all-reduce把通信藏進(jìn)反向計算里。減少 all-reduce 次數(shù)把過小的參數(shù)合并進(jìn)同一個 bucket。當(dāng)模型大到單卡放不下要上 FSDP 或混合并行那是另一套權(quán)衡不能只在通信參數(shù)上打轉(zhuǎn)。回到題目那句話all-reduce 怎樣讓每卡相同我的答案是它把“各自不同的局部結(jié)果”先聚合成一份“大家共享的全局結(jié)果”再讓每一份副本回到每一張卡。機(jī)制本身看起來簡單真正難的是背后的算法選擇、工程封裝和故障排查。如果你剛接觸分布式訓(xùn)練不要急著手動實現(xiàn) all-reduce。先用 DDP 把一個小模型跑通打開torch.profiler看一眼 GPU 利用率和通信耗時再嘗試調(diào) bucket 和 batch size。當(dāng)你親眼看到通信開銷占了多少時間之后才會真正理解為什么 ring、bucket、低精度、拓?fù)鋬?yōu)化這些東西值得被反復(fù)討論。大模型訓(xùn)練和多卡推理越來越常見但能讓整個系統(tǒng)穩(wěn)定運行的從來不只是一個模型結(jié)構(gòu)而是這些不起眼的集合操作在每一層正確地對齊。