DeepSeek開源三箭齊發,梁文峰親自上陣!雙向并行LLM訓練飆升
開源周第4天,DeepSeek放出的是——優化并行策略,一共三個項目。
- DualPipe:一種用于V3/R1模型訓練中實現計算與通信重疊的雙向流水線并行算法
- EPLB:一個針對V3/R1的專家并行負載均衡工具
- 深入分析V3/R1模型中的計算與通信重疊機制
值得一提的是,DualPipe是由三個人——Jiashi Li、Chengqi Deng和梁文峰共同研發。
有網友對此表示,這是一個顛覆性的突破。
「DualPipe和EPLB讓大模型的訓練更快、更便宜、更可持續,同時保持頂尖性能」。
開源僅剩最后一天,評論區網友紛紛許愿:明天一定要來個大的。
DualPipe
DualPipe是DeepSeek-V3技術報告中提出的一種創新雙向流水線并行算法。
它能夠實現前向與后向計算和通信階段的完全重疊,同時有效減少流水線氣泡(空閑時間)。
調度方案
上圖展示了在8個流水線并行階段和20個micro-batches情況下,DualPipe在兩個方向上的調度示例。
由于反向方向的微批次與前向方向對稱,為了簡化圖示,這里省略了反向方向的batch ID。圖中由同一個黑色邊框包圍的兩個單元格表示同時進行的計算和通信操作,它們實現了相互重疊。
流水線氣泡與內存使用比較
圖中,??代表前向塊(forward chunk)的執行時間,??代表完整后向塊(full backward chunk)的執行時間,??代表「權重后向」(backward for weights)塊的執行時間,而??&??則表示同時執行且相互重疊的前向和后向塊的執行時間。
快速入門
使用示例如下:
python example.py
注意:在實際生產環境中,需要根據模塊特點來實現一個定制化的overlapped_forward_backward方法。
專家并行負載均衡器(EPLB)
在使用專家并行(EP)時,不同的專家模塊會被分配到不同的GPU上。由于各個專家的計算負載會隨當前任務而變化,因此保持各GPU間負載均衡至關重要。
如DeepSeek-V3論文所述,研究人員采用了冗余專家(redundant experts)策略,對高負載專家進行復制。
隨后,通過啟發式算法將這些復制的專家合理分配到各GPU上,確保計算資源的平衡利用。
此外,由于DeepSeek-V3采用了組內限制專家路由(group-limited expert routing)機制,研究團隊盡可能將同一組的專家放置在同一節點上,以減少節點間的數據傳輸開銷。
為了便于復現和部署,DeepSeek在eplb.py文件中開源了EP負載均衡算法。該算法能夠根據估計的專家負載,計算出均衡的專家復制和放置方案。
需要說明的是,專家負載的具體預測方法不在此代碼庫的討論范圍內,一種常用的方法是采用歷史統計數據的滑動平均值。
算法原理
負載均衡算法提供了兩種策略,適用于不同場景:
· 層次負載均衡(Hierarchical Load Balancing)
當服務器節點數量能夠整除專家組數量時,研究人員采用層次負載均衡策略,來充分利用組內限制專家路由機制。
首先,他們將專家組均勻分配到各節點,確保節點間負載平衡;然后,在每個節點內部復制專家模型;最后,將復制后的專家打包分配到各個GPU上,實現GPU間的負載均衡。
這種層次化策略特別適用于預填充階段(prefilling stage),此時專家并行規模較小。
· 全局負載均衡(Global Load Balancing)
在其他情況下,研究人員采用全局負載均衡策略,不考慮專家組的限制,直接在全局范圍內復制專家并分配到各個GPU上。這種策略更適合解碼階段使用,此時專家并行規模較大。
接口示例
負載均衡器的核心函數是eplb.rebalance_experts。
下面的代碼展示了一個雙層混合專家模型(MoE)的示例,每層包含12個專家。
DeepSeek為每層引入了4個冗余專家,總計16個專家副本被分配到2個計算節點上,每個節點配有4個GPU。
import torch
import eplb
weight = torch.tensor([[ 90, 132, 40, 61, 104, 165, 39, 4, 73, 56, 183, 86],
[ 20, 107, 104, 64, 19, 197, 187, 157, 172, 86, 16, 27]])
num_replicas = 16
num_groups = 4
num_nodes = 2
num_gpus = 8
phy2log, log2phy, logcnt = eplb.rebalance_experts(weight, num_replicas, num_groups, num_nodes, num_gpus)
print(phy2log)
# Output:
# tensor([[ 5, 6, 5, 7, 8, 4, 3, 4, 10, 9, 10, 2, 0, 1, 11, 1],
# [ 7, 10, 6, 8, 6, 11, 8, 9, 2, 4, 5, 1, 5, 0, 3, 1]])
該層次負載均衡策略產生的結果,展示了如下專家復制與分配方案。
DeepSeek基礎設施中的性能剖析數據
在這里,DeepSeek公開分享來自訓練和推理框架的性能剖析數據,旨在幫助社區更深入地理解通信與計算重疊策略以及相關底層實現細節。
這些剖析數據是通過PyTorch Profiler工具獲取的。
你可以下載后在Chrome瀏覽器中訪問chrome://tracing(或在Edge瀏覽器中訪問edge://tracing)直接進行可視化查看。
需要說明的是,為了便于剖析,研究人員模擬了一個完全均衡的MoE路由策略。
訓練過程
訓練剖析數據展示了,研究人員在DualPipe中如何實現單對前向和后向計算塊的重疊策略。每個計算塊包含4個MoE層。
并行配置與DeepSeek-V3預訓練設置保持一致:采用EP64、TP1,序列長度為4K。
為簡化剖析過程,未包含流水線并行(PP)通信部分。
推理過程
· 預填充
在預填充階段,剖析配置采用EP32和TP1的配置(與DeepSeek V3/R1實際在線部署一致),提示長度設為4K,每GPU批處理量為16K個token。
研究人員在預填充階段使用兩個micro-batches來實現計算與all-to-all通信的重疊,同時確保注意力機制的計算負載在兩個micro-batches間保持平衡——這意味著同一條提示信息可能會被分割到不同micro-batches中處理。
· 解碼
解碼階段的剖析配置采用EP128、TP1,提示長度4K(與實際在線部署配置非常接近),每GPU批處理量為128個請求。
與預填充類似,解碼階段也利用兩個micro-batches來重疊計算和all-to-all通信。
然而不同的是,解碼過程中的全聯通通信不占用GPU流處理器(SM):RDMA消息發出后,所有GPU流處理器立即被釋放,系統在完成計算后等待全聯通通信完成。
關于all-to-all通信實現的更多技術細節,請參考DeepEP文檔。