背景

在 MoE 模型的訓(xùn)練過程中，EP rank 之間的 A2A 通信在端到端時(shí)間中占據(jù)了相當(dāng)大比重，對(duì)訓(xùn)練效率影響很大，特別是對(duì)于 Fine-grained MoE model，EP size 會(huì)比較大，跨機(jī)通信基本無法避免。那么要如何減少 EP A2A 對(duì)整體訓(xùn)練效率的影響呢?

首先是盡可能提升 EP A2A 通信本身的效率。DeepSeek 開源的 DeepEP 通信庫(kù)通過將通信拆分為機(jī)內(nèi)通信和機(jī)間通信，并將二者進(jìn)行 overlap，實(shí)現(xiàn)了很高的通信效率。這一工作已經(jīng)集成到 NVIDIA Megatron-LM 中;

其次是讓 EP A2A 通信盡可能和計(jì)算 overlap，用計(jì)算來掩蓋通信的時(shí)間。本文從這個(gè)主題出發(fā)，介紹我們?nèi)绾位?1F1B 的流水線機(jī)制來實(shí)現(xiàn) EP A2A 與計(jì)算的 overlap。

用計(jì)算掩蓋 EP A2A 通信的相關(guān)研究

關(guān)于如何用計(jì)算掩蓋 EP A2A 通信，我們注意到目前有以下這些研究(不完全統(tǒng)計(jì))：

1 - micro batch 內(nèi) overlap

EP A2A 和 expert mlp overlap (minimax-01 方案、faster moe 方案、tutel 方案);

將 Attention + MLP 整體分成多個(gè) block，block 間的 EP A2A 和計(jì)算 overlap (AMPipe)。

2 - micro batch 間 overlap

DeepSeek V3 論文提出了一種新的流水線并行算法：DualPipe。實(shí)現(xiàn)了在兩個(gè) micro batch forward 和 backward 之間的 EP A2A 通信與計(jì)算 overlap。

以上方案有各自的不足之處：

EP A2A 和 expert mlp overlap：其效果比較好的前提是計(jì)算時(shí)間要大于通信時(shí)間，但對(duì)于跨機(jī) A2A 的情況，A2A 的通信可能會(huì)大于 expert mlp 計(jì)算時(shí)間，DeepSeek V3 論文中提到訓(xùn)練過程中 EP A2A 的整體耗時(shí)占比大約是 50%，此時(shí)這一方案的 overlap 效果就大打折扣了;

Attention + MLP 整體分成多個(gè) block 方案：需要分塊 attention，支持起來比較麻煩，而且分了 block 可能會(huì)降低計(jì)算效率，另外這個(gè)方案可能和其他分布式并行策略有沖突，比如 CP;

DeepSeek DualPipe：毋庸置疑是很優(yōu)秀的工作，但是調(diào)度邏輯比較復(fù)雜，難以理解與 debug，另外需要對(duì)現(xiàn)有的分布式訓(xùn)練框架 (如 Megatron-LM) 進(jìn)行大范圍的重構(gòu)才能接入。

我們分析認(rèn)為，Deepseek DualPipe 調(diào)度核心亮點(diǎn)是實(shí)現(xiàn)了 batch 之間 EP A2A 通信和 attention 、mlp 計(jì)算 overlap 的精細(xì)化調(diào)度，從而使得跨機(jī) EP 成為可行的分布式策略。 Dualpipe 起到的主要作用是創(chuàng)造了 overlap 的機(jī)會(huì)，從這個(gè)角度而言，基于 1F1B 的調(diào)度也能做到 batch 之間 forward 和 backward 之間的 EP A2A 通信與計(jì)算的 overlap。

基于 1F1B 的 MoE A2A 通信計(jì)算 overlap

Conventional 1F1B

可以很自然地想到，通過調(diào)整 1F1B 調(diào)度來實(shí)現(xiàn) 1F1B 調(diào)度中穩(wěn)態(tài)階段 f 和 b 的 EP A2A 通信與計(jì)算 overlap。

原始的 1F1B 調(diào)度，最后一個(gè) pp stage 穩(wěn)態(tài)階段的 f 和 b 之間存在數(shù)據(jù)依賴，無法實(shí)現(xiàn) f 和 b 之間的 EP A2A 通信與計(jì)算 overlap。為了解決這個(gè)問題，我們可以把最后一個(gè) pp stage 的 warmup step 加 1。同時(shí)，考慮到 pp stage 間前反向的數(shù)據(jù)依賴，其他 pp stage 也需要調(diào)整 warmup step。 具體的調(diào)度方式如下：

圖 1: 1F1B 調(diào)度下 A2A 基于傳統(tǒng) 1F1B 方案實(shí)現(xiàn)相鄰 micro batch 間計(jì)算和通信的 overlap

該圖片來源于小紅書 Agi Infra 團(tuán)隊(duì)，若您有任何疑問或需要使用該圖片，請(qǐng)聯(lián)系小紅書

我們得到了一個(gè) 1F1B 通信計(jì)算 overlap 的初步方案，但是這個(gè)方案有很大的缺陷：因?yàn)?pipeline parallel rank 之間存在數(shù)據(jù)依賴，warmup step 過大，activation 顯存占用過大。

Interleaved 1F1B

我們轉(zhuǎn)向考慮 interleaved 1F1B，沿用上面調(diào)整 warmup step 的邏輯，驚喜地發(fā)現(xiàn)，通過將穩(wěn)態(tài)的 1F1B stage 中第一個(gè) micro batch 的 fprop 提前到 warmup stage，即 warmup step + 1，就可以在 interleaved 1F1B 實(shí)現(xiàn) 1F1B 穩(wěn)態(tài)階段不同 micro batch 前反向之間的 EP A2A 與計(jì)算的 overlap。

圖 2: 基于 interleaved 1F1B 方案實(shí)現(xiàn)相鄰 micro batch 間計(jì)算和通信的 overlap

該圖片來源于小紅書 Agi Infra 團(tuán)隊(duì)，若您有任何疑問或需要使用該圖片，請(qǐng)聯(lián)系小紅書

相比于原始的 interleaved 1F1B，我們的這個(gè)方案雖然在 warmup stage 增加了一個(gè) micro batch 的前向計(jì)算，但是 1F1B 階段的調(diào)度是前反向細(xì)粒度地交疊調(diào)度，F(xiàn) 在增加 activation 顯存的同時(shí)，B 會(huì)釋放 activation 顯存，使得峰值顯存基本沒有變化;Bubble rate 和原始的 interleaved 1F1B 完全一致。

因此我們最終選擇使用 interleaved 1F1B with A2A overlap 方案。

Fine-grained schedule of 1F1B overlap

具體來說，對(duì)于兩個(gè) micro batch 之間 forward 與 backward 的 EP A2A 通信/計(jì)算 overlap，我們采用了類似 DualPipe 的方案， 拆分了 dw。我們認(rèn)為這是很有必要的。如果計(jì)算與通信 1:1, 那就意味著前向的計(jì)算比通信少，反向的計(jì)算比通信多，前向的計(jì)算不足以 overlap 反向的通信, 需要用 dw 來 overlap 部分反向通信。

我們通過下面一組流程圖來對(duì)比說明 dw 拆分的必要性。Baseline 版本將 forward 和 backward 的計(jì)算部分放到了同一個(gè) stream 中，通信部分放到了另一個(gè) stream 中，通過將 forward 和 backward 的各個(gè)模塊交錯(cuò)排布，消除了同一時(shí)刻計(jì)算流和通信流上 kernel 的數(shù)據(jù)依賴，從而實(shí)現(xiàn)了 batch 間的 EP A2A 和計(jì)算的 overlap。

但如果 EP A2A 的時(shí)間比較長(zhǎng)，超過了計(jì)算時(shí)間，比如下圖里的 F/attn 和 F/mlp，此時(shí) EP A2A 通信還是會(huì)暴露出來。這時(shí)我們可以把 mlp 和 attention 的 dw 從 backward 計(jì)算中拆分出來，那么 B/dispatch 和 B/combine 的執(zhí)行就可以提前，因?yàn)?B/dispatch 和 B/combine 只依賴 dx 的結(jié)果。這樣用來 overlap B/dispatch 的計(jì)算除了 F/mlp 之外還有 W/mlp 的一部分，從下面兩種方案的對(duì)比我們也可以看出，拆分 dw 之后，計(jì)算流上 kernel 的排布更加緊密，整體的執(zhí)行時(shí)間也要比 Baseline 快。

圖 3: 拆分 dw 前后，計(jì)算和通信 overlap 示意圖 該圖片來源于 NVIDIA，若您有任何疑問或需要使用該圖片，請(qǐng)聯(lián)系 NVIDIA

與 Deepseek DualPipe 對(duì)比

注：F - forward chunk time, B - full backward chunk time, W - “backward for weight” chunk, F&B - two mutually overlapped forward and backward chunks.

為了更好地對(duì)比兩種方案的 bubble overhead，接下來我們對(duì) F/B/W 的時(shí)間做一個(gè)合理的假設(shè)(參考圖三)：B=2F=2W，F(xiàn)&B 的耗時(shí)比較難以給出較準(zhǔn)確的估計(jì)，我們合理地認(rèn)為 F&B=B，即前向傳播的時(shí)間完全被反向傳播掩蓋住，這種估計(jì)的合理性在于：

在拆分 dw 之前，反向傳播的計(jì)算和通信時(shí)間都大于與其 overlap 的前向計(jì)算的時(shí)間;拆分 dw 之后，dw 的計(jì)算時(shí)間可以進(jìn)一步被 overlap;

計(jì)算和通信的 overlap 在相當(dāng)程度上會(huì)對(duì)計(jì)算 kernel 產(chǎn)生 overhead;

據(jù)以上公式，我們可以估算出兩種方案 bubble overhead 的比例：

可以看出，vpp=8 大致是兩種方案的 bubble overhead 比值的分界點(diǎn)，我們從一些常見模型的最佳配置里抽離出 vpp 的設(shè)置如下：

注：* 該模型的最佳配置為內(nèi)部測(cè)試結(jié)果，訓(xùn)練腳本尚未開源

因此，多數(shù)情況下兩種方案的 bubble overhead 是比較接近的。另外大多數(shù)情況下 micro batch 的數(shù)量比較多，PP bubble 在端到端的耗時(shí)占比處于比較低的水平，此時(shí) PP bubble 對(duì)整體的性能影響比較有限。

關(guān)于 PP bubble 帶來的 Activation memory 開銷，從上面常見模型的最優(yōu)配置可以看出，PP size 通常小于 VPP，此時(shí) 1F1B with A2A overlap 方案的 Activation memory 開銷要低于 DualPipe。

總結(jié)

針對(duì) MoE 模型訓(xùn)練過程中 EP A2A 通信占比過高的問題，我們分析了現(xiàn)有的優(yōu)化策略，認(rèn)為在 batch 間進(jìn)行通信計(jì)算 overlap 是比較高效的實(shí)現(xiàn)方案。參考 DeepSeek 提出的 DualPipe，我們?cè)O(shè)計(jì)了一種基于現(xiàn)有的 interleaved 1F1B 策略實(shí)現(xiàn) batch 間計(jì)算通信 overlap 的方案，該方案具有以下幾個(gè)特點(diǎn)：

基于現(xiàn)有的 1F1B 方案實(shí)現(xiàn)了和 DualPipe 相似的 batch 間 EP A2A 與計(jì)算 overlap 的效果。

與 DualPipe 相比，Interleaved 1F1B 大多數(shù)情況下展現(xiàn)出較接近的 bubble overhead。另外大多數(shù)訓(xùn)練場(chǎng)景下，micro batch 的數(shù)量都比較多，此時(shí)這兩種方案下 bubble rate 在端到端訓(xùn)練中的耗時(shí)占比都非常低。

不需要存儲(chǔ)一份額外的模型參數(shù)。

能很好地兼容 NVIDIA Megatron-Core 現(xiàn)有的代碼結(jié)構(gòu)和分布式策略。

NVIDIA 技術(shù)團(tuán)隊(duì)正與小紅書技術(shù)團(tuán)隊(duì)一起將該優(yōu)化方案實(shí)現(xiàn)到 Megatron-Core 中，目前，前期的 PoC 驗(yàn)證和 Design 已完成，預(yù)計(jì) 3 月中下旬可提供 EA 版本供試用。

關(guān)于作者

劉振海

小紅書大模型訓(xùn)練框架開發(fā)工程師

目前負(fù)責(zé)小紅書大模型訓(xùn)練框架的開發(fā)與優(yōu)化工作。

李屏天

NVIDIA 加速計(jì)算專家

目前主要從事大模型訓(xùn)練相關(guān)的優(yōu)化工作。

張順康

NVIDIA 加速計(jì)算專家

目前負(fù)責(zé) Megatron-Core (MCore) 訓(xùn)練框架開發(fā)。

劉宏斌

NVIDIA 加速計(jì)算專家

專注于大語(yǔ)言模型訓(xùn)練的功能開發(fā)及性能調(diào)優(yōu)，深度參與 Megatron-LM 及 NeMo 的相關(guān)開發(fā)工作，目前主要負(fù)責(zé) FP8 訓(xùn)練及 MoE 模型性能優(yōu)化工作。

顏?zhàn)咏?/p>

NVIDIA 加速計(jì)算專家

專注于大規(guī)模語(yǔ)言模型 (LLM) 訓(xùn)練系統(tǒng)。目前，他是 Megatron Core MoE 的核心開發(fā)者之一，主要負(fù)責(zé)該系統(tǒng)的工程開發(fā)與性能優(yōu)化工作。

黃雪

NVIDIA 解決方案架構(gòu)師

主要負(fù)責(zé)深度學(xué)習(xí)訓(xùn)練方面工作，在深度學(xué)習(xí)框架、超大規(guī)模模型訓(xùn)練，分布式模型訓(xùn)練加速優(yōu)化等技術(shù)方向有豐富的研究經(jīng)驗(yàn)。