【DeepSeek】開源三箭齊發,梁文鋒親自上陣!雙向平行LLM訓練飆升


DeepSeek開源第四天,連更三個項目。DualPipe、EPLB、以及計算與通訊重疊機制的最佳化平行策略,讓大模型訓練更快,成本更低,還能保持頂尖性能。

開源周第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。

該層次負載平衡策略產生的結果,展示了如下專家複製與分配方案。

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文件。

參考資料:

https://x.com/deepseek_ai/status/1894931931554558199 (新智元)