一、背景

在大規(guī)模分布式訓練場景中，計算和通信的重疊（Overlap）一直是一個關(guān)鍵的研究熱點。隨著硬件性能的提升，計算能力和通信帶寬之間的差距日益顯著。如下圖所示，硬件算力每 2 年大約擴大 3x，而通信帶寬每 2 年只提升 1.4x，這種差距帶來的影響在大規(guī)模訓練任務(wù)中愈加明顯。例如，在使用 H100 和 A100 集群進行 LLM 訓練時，H100 的通信開銷占比通常會高于 A100。這種情況下，通信可能成為了系統(tǒng)性能的瓶頸，因此，如何在計算和通信之間實現(xiàn)高效的 Overlap，已成為優(yōu)化分布式訓練性能的關(guān)鍵策略之一。

實際上，大部分計算和通信的 Overlap 可以理解為生產(chǎn)者和消費者的 Overlap，生產(chǎn)者可以是計算也可以是通信，生產(chǎn)者與消費者可以是短程依賴也可以是長程依賴，通常短程依賴帶來的挑戰(zhàn)更大，而長程依賴則較容易實現(xiàn) Overlap。比如：

• 張量并行（Tensor Parallelism，TP）中的 MatMul 計算與 AllReduce 通信就是一個明顯的短程依賴，AllReduce 通信依賴 MatMul 的計算，并且 AllReduce 通信又會阻塞后續(xù)的計算操作。
• Deepspeed Zero 的模型切分通信與計算可以看作是一種長程依賴。在這種情況下，只要保證 Layer N 計算之前拿到權(quán)重即可，完全可以提前獲取權(quán)重，避免等到實際開始計算時再進行通信。通過這種方式，通信可以在更早的階段與之前的計算操作進行重疊，從而更高效地利用計算和通信資源。

本文中我們簡單介紹一系列針對大規(guī)模訓練場景的計算與通信 Overlap 來優(yōu)化訓練性能的工作，包括 Microsoft 的 CoCoNet、Domino，Google 的 Intra-layer Overlapping via Kernel Fusion，AMD 的 T3，北大的 Centauri，字節(jié)的 Flux 以及中科大的 DHelix 等。

訓練相關(guān)可以參考我們之前的文章：

• ???大規(guī)模分布式 AI 模型訓練系列——數(shù)據(jù)并行???
• ???大規(guī)模分布式 AI 模型訓練系列——張量并行???
• ???大規(guī)模分布式 AI 模型訓練系列——流水線并行???
• ???大規(guī)模分布式 AI 模型訓練系列——專家并行???
• ???大規(guī)模分布式 AI 模型訓練系列——序列并行???
• ??超長序列 LLM 訓練：DeepSpeed Zero & Ulysses & FPDT??
• ??DHelix：跨 Micro-Batch 的通信隱藏，SOTA LLM 訓練性能??
• ??萬字綜述：全面梳理 FP8 訓練和推理技術(shù)??

二、引言

2.1 AllReduce

AllReduce 是集合通信中常見的分布式計算操作，用于多個設(shè)備（比如多個 GPU）之間聚合數(shù)據(jù)的場景，可以包含 Sum、Min、Max 等操作。

對于常見的基于 Ring 的 AllReduce 實現(xiàn)中，通?？梢园?AllReduce 操作看成為一個 ReduceScatter 和一個 AllGather 操作，如下圖所示：

具體的 ReduceScatter 操作如下，每個設(shè)備（GPU）發(fā)送一部分數(shù)據(jù)給下一個設(shè)備，同時接收上一個設(shè)備的數(shù)據(jù)并累加。這個過程進行 K-1 步（假設(shè)有 K 個設(shè)備），ReduceScatter 后每個設(shè)備都包含一部分數(shù)據(jù)的 Sum：

具體的 AllGather 操作如下，每個設(shè)備將其持有的部分結(jié)果發(fā)送給下一個設(shè)備，同時接收上一個設(shè)備的部分結(jié)果，逐步匯集完整的結(jié)果，同樣需要 K-1 步。AllGather 后，每個設(shè)備都包含全量的數(shù)據(jù)：

2.2 LLM Tensor Parallelism AllReduce

當前 LLM 推理通常會采用 Tensor Parallelism（TP）模型并行，以便在多個 GPU 上實現(xiàn)較大 LLM 的推理。對于標準的 Transformer Decoder Only 模型，通常會在每個 Transformer Block 中采用如下的 TP 切分方式：

如下圖 （a）所示，MLP 層的兩個 Linear 層采用先列切（A，Column Parallelism），然后行切（B，Row Parallelism）的方案，這樣兩個 Linear 之間不用通信：

如下圖（b）所示，由于每個 Head 的 Attention，Softmax 都可以獨立計算，因此可以按照 Head 的方式切分（等價于列切分），然后對之后的 Linear 采用行切分（B），這樣 Self-Attention 中間也不用通信：

如上所述，采用先列切、再行切的方式，每個 Transformer Block 中都需要兩個 AllReduce 操作，對于一個 40 層的模型則需要至少 80 個 AllReduce 操作。

2.3 ReduceScatter = All2All + Local Reduce

如下圖所示為 Ring ReduceScatter 的優(yōu)化，可以等效為一個 All2All 操作實現(xiàn)數(shù)據(jù)的重排，然后在 Local 進行 Reduce 操作。此過程只有一個 All2All 的整體通信操作，雖然實際上與 Ring 實現(xiàn)的方式的通信量和計算量沒有變化，但可以避免 K-1 個 Ring Step 的同步，進而可以有效降低時延。

2.4 ZeroBubble

 [2401.10241] Zero Bubble Pipeline Parallelism [1]  中作者提出 Zero Bubble，核心思路是將 Backward 分為兩個部分，一部分計算輸入的梯度，另一部分計算參數(shù)的梯度，如下圖 Figure 1 所示。這里計算輸入的梯度有明確的依賴關(guān)系，也是鏈式法則不斷傳遞的基礎(chǔ)；而計算權(quán)重的梯度卻沒有明確的依賴，甚至可以滯后很多。此外，三個紅色部分計算量相當，這也就是為什么常見的流水線并行（Pipeline Parallelism，PP）中 Backward 的長度為 Forward 的 2 倍。

三、Microsoft - CoCoNet

3.1 摘要

對應的 Paper 為：[ASPLOS 22] [2105.05720] Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloads [2]

CoCoNet 可能是最早提出異步張量并行（Async Tensor Parallelism）的工作，作者提供了一種通用 Kernel 融合的范式，能夠自動生成集合通信與常見計算操作（如 GEMM 和卷積）之間的融合 Kernel。具體來說，CoCoNet 包含：

• 一種領(lǐng)域特定語言（DSL），用于以計算和通信操作的形式表示分布式機器學習程序。
• 一組保持語義的轉(zhuǎn)換規(guī)則，以優(yōu)化程序。
• 一個編譯器，可以生成聯(lián)合優(yōu)化通信和計算的 GPU Kernel。

PS：然而，生成的代碼執(zhí)行效率不及直接采用 cuBlas、cutlass 或 cuDNN 中的高度優(yōu)化 Kernel，主要原因在于為實現(xiàn)這種細粒度計算-通信 Overlap 需要在 Kernel 內(nèi)引入額外同步操作。

3.2 方法

如下圖 Figure 2 所示為 CoCoNet 的工作流：

• 首先，使用 DSL 表示用戶的機器學習算法，該語言同時包含計算（如矩陣乘和卷積）與通信操作（如 AllReduce）。
• 然后，Autotuner 應用一系列變換以優(yōu)化程序，同時確保算法邏輯保持不變。例如，將 AllReduce 與 Dropout 融合為 FusedAllReduce，并使其與矩陣乘法 Overlap。
• 最后，生成對應的通信與計算代碼，并且可以通過 PyTorch 執(zhí)行。?

CoCoNet 提供了 4 種能夠保持語義的轉(zhuǎn)換方案，用于優(yōu)化以 DSL 表示的程序。以下圖 Figure 3 的代碼為例，其主要是實現(xiàn)了 ：矩陣乘法 + AllReduce + Dropout + Add。

具體的幾個轉(zhuǎn)換過程如下：

1. AllReduce 拆分為 ReduceScatter（RS）和AllGather（AG）
2. 操作重排并拆分為等價算子（如下圖 Figure 5）：

• scD 和 scOut 均為 Slice 切片操作。
• agOut 用于收集計算的最終結(jié)果。

3. 算子融合：

• FushedAllReduce（FusedAR）：融合了 ReduceScatter（rsSum）、計算操作（scD 和 ScOut），以及 AllGather（agOut）。
• fusedAR.comp(scOut)  則指定了需要與 FusedAllReduce 融合的計算，返回的 out 則是輸出結(jié)果。

4. Overlap：對一系列生產(chǎn)者-消費者操作，可以執(zhí)行 Overlap 變換。比如有多個數(shù)據(jù)要執(zhí)行上述操作（不同樣本，數(shù)據(jù)的不同 Chunk），則可以實現(xiàn)通信與計算的 Overlap。?

CoCoNet 提供了 AutoTunner，能夠自動探索程序的所有調(diào)度方案的空間，并針對特定底層框架和輸入規(guī)模，返回性能最佳的調(diào)度方案。

以下圖 Figure 7 所示為將其應用到 Megatron-LM 中的 TP+PP 的優(yōu)化案例，具體來說，總共 4 個 GPU，分成兩個 PP Stage，每個 Stage 有 2 個 GPU，使用 TP 切分。比如下圖 GPU(0, 1) 表示 PP Stage 0 的 1 號 GPU。

• (a)：兩個 PP Stage 之間會有一個 PP Stage 0 內(nèi)部的AllReduce操作，以及一個 PP Stage 0 與 Stage 1 之間的P2P操作。
• (b)：將數(shù)據(jù)拆分為多個 Chunk，并使用 ReduceScatter + AllGather 代替 AllReduce，即可實現(xiàn)一定的 Overlap，并減少冗余數(shù)據(jù)傳輸。?

3.3 結(jié)果

如下圖 Figure 1 所示，MatMul 與 AllReduce 的細粒度 Overlap 執(zhí)行可掩蓋 80% 的 MatMul 執(zhí)行時間，并帶來 1.36x 的加速效果。

四、Google - Intra-layer Overlapping via Kernel Fusion

4.1 摘要

對應的論文為：[ASPLOS 23] Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models [3]

在大規(guī)模深度學習模型訓練中，層內(nèi)模型并行化產(chǎn)生的通信開銷可能占據(jù)整體訓練時間的顯著部分，而層內(nèi)模型并行化又對支持大型深度學習模型至關(guān)重要。因此作者提出了一種新穎計算，通過計算與通信的 Overlap 來有效降低通信開銷。

該技術(shù)將識別出的原始集合通信與依賴的計算操作分解為一系列更細粒度的操作，通過創(chuàng)造更多 Overlap 機會并并行執(zhí)行新生成的細粒度通信與計算操作，可以有效隱藏數(shù)據(jù)傳輸時延，實現(xiàn)更優(yōu)的系統(tǒng)利用率。

在 TPU v4 Pod 上評估不同規(guī)模的大模型（10B - 1T 參數(shù)量），所提方案可以實現(xiàn) 1.14x 到 1.38x 的吞吐提升。在 1024 TPU 集群中，500B 參數(shù)量語言模型訓練可以實現(xiàn) 72% 的峰值 FLOPs 利用率。

4.2 方法

4.2.1 背景

如下圖 Figure 1 所示，不同規(guī)模模型在 128-2048 TPU 上訓練，通信開銷可達 22%-42%：

4.2.2 方案概覽

如下圖 Figure 4 展示了 AllGather 場景中的執(zhí)行流程。假設(shè)數(shù)據(jù) A（比如模型權(quán)重） 在初始階段已被切分，每個設(shè)備各持有 A 的一個分片，A0 位于設(shè)備 0，A1 位于設(shè)備 1。

• 現(xiàn)有系統(tǒng)中：兩個設(shè)備均需要通過 AllGather 操作得到完整的數(shù)據(jù) A，即 [A0, A1]，然后開始相應的計算。
• 所提系統(tǒng)中：并不用等待全部數(shù)據(jù)準備就緒再啟動計算。

• 每個設(shè)備異步發(fā)送存儲在當前設(shè)備的數(shù)據(jù)到其他設(shè)備（比如設(shè)備 1 異步發(fā)送 A1 到設(shè)備 0），同時利用已有數(shù)據(jù)開始啟動計算，這樣設(shè)備即在計算，也在通信。
• 當之前結(jié)果計算完成，并且從其他設(shè)備接收完成（比如設(shè)備 1 的 [A1, B1] 已經(jīng)計算完，并且已經(jīng)接收完 A0），開始啟動新數(shù)據(jù)的計算（比如設(shè)備 1 上的 [A0, B1]）。
• 為了得到最終結(jié)果，每個部分結(jié)果需要額外執(zhí)行一次 Dynamic Updata Slice 操作。
• 通過執(zhí)行多次上述操作可以獲得最終結(jié)果，確切的步驟次數(shù)取決于 A 的切片數(shù)。?

同樣地，ReduceScatter 操作可與相應的計算過程 Overlap 執(zhí)行，如下圖 Figure 5 所示。在此例中，C0（C00 與 C01 之和）及 C1（C10 與 C11 之和）分別為 C 在設(shè)備 0 和 設(shè)備 1 上進行 ReduceScatter 后的操作分片。由于基于計算結(jié)果進行通信傳輸，在此情形下，各設(shè)備需要異步傳輸累加結(jié)果分片而非操作數(shù)，累加結(jié)果分片在各設(shè)備上初始化為 0：

• 每輪迭代開始時，各設(shè)備異步發(fā)送累加結(jié)果分片到另一個設(shè)備（例如，首輪迭代中設(shè)備 0 發(fā)送切片 C0 到設(shè)備 C1），并與此同時啟動部分 Einsum 計算。
• 計算完成后，部分結(jié)果在迭代末尾被加到接收的累加結(jié)果分片，比如首輪迭代的結(jié)果 C10 與從設(shè)備 1 接收的結(jié)果分片 C1。?

Kernel Fusion 是一種有效減少慢速主內(nèi)存訪問和 Kernel 啟動開銷的方案，作為最重要的優(yōu)化手段之一，Kernel Fusion 在 XLA 中通過啟發(fā)式方法自動執(zhí)行。因此，針對本文的方案作者也會進一步應用 Kernel Fusion。然而，基于默認啟發(fā)式構(gòu)建的某些融合操作可能損害 Overlap 性能。如下圖 Figure 11 所示，11a 展示了一個簡化的圖結(jié)構(gòu)，為默認的融合策略，其中的灰色方框為融合節(jié)點，白色方框表示一個或多個融合的高階算子指令。其中兩個 Einsum，Einsum_1 有一個異步的 CollectivePermuteDone 輸入，由于 Einsum_0 與 CollectivePermuteDone 相互獨立，預期其能與異步數(shù)據(jù)通信并行執(zhí)行，以實現(xiàn) Overlap。然而，與 Einsum_0 融合的加法操作在 Fusion_0 與 CollectivePermuteDone 之間引入了數(shù)據(jù)依賴，導致第三個節(jié)點順序執(zhí)行。為了避免這種不良融合，啟發(fā)式策略調(diào)整為先將 Add 操作與具有異步 CollectivePermuteDone 操作的 Einsum 進行融合，新生成的圖結(jié)構(gòu)如圖 11b 所示，數(shù)據(jù)通信得以成功與 Fusion_0 Overlap。

4.3 結(jié)果

如下圖 Figure 12 所示為不同模型優(yōu)化前后可以達到的峰值 TFLOPS，可以看出，優(yōu)化后有比較明顯的提升：

五、AMD - T3

5.1 摘要

對應的論文為：[2401.16677] T3: Transparent Tracking & Triggering for Fine-grained Overlap of Compute & Collectives [4]

LLM 在訓練與推理過程中，盡管某些分布式技術(shù)能夠通過計算與通信 Overlap 來隱藏通信開銷，但諸如 TP 等技術(shù)，其通信與模型執(zhí)行本質(zhì)上具有序列化特性。為隱藏這種序列化通信，常見做法是以細粒度方式將其數(shù)據(jù)生成操作交錯進行。然而，在軟件層面實現(xiàn)通信與計算的細粒度交錯頗為復雜，此外，并發(fā)執(zhí)行通常都要求計算與通信共享計算和存儲資源，導致資源競爭，從而削弱 Overlap 的效能。

為應對這些挑戰(zhàn)，作者提出 T3 方案，通過硬件與軟件協(xié)同設(shè)計，透明地 Overlap 序列化通信，同時最小化與計算的資源競爭。

• 在軟件層面，T3 通過簡單配置生產(chǎn)者的輸出地址空間，透明地將生產(chǎn)操作與后續(xù)通信融合，僅需少量軟件改動。
• 在硬件層面，T3 引入輕量級的追蹤與觸發(fā)機制，以協(xié)調(diào)生產(chǎn)者的計算與通信。同時，利用計算增強型內(nèi)存處理通信伴隨的計算任務(wù)。由此，T3 減少了資源競爭，并高效地實現(xiàn)了序列化通信與計算的 Overlap。

對于 T-NLG 等重要 Transformer 模型，T3 使通信密集型子層的平均加速比可以達到 30%（最高47%），平均減少 22% 的傳輸量（最高 36%）。此外，隨著模型規(guī)模的擴大，T3 的優(yōu)勢依然顯著：在約 500B 參數(shù)的 PALM 和 MT-NLG 模型中，子層的平均加速比為 29%。

PS：T3 依賴于特定的硬件特性，如計算增強型存儲器（NMC），這可能需要新的硬件支持或者對現(xiàn)有硬件的修改，這也許是其落地應用的最大挑戰(zhàn)。

5.2 方法

5.2.1 背景

如下圖 Figure 2 所示，Transformer 模型通常會采用 TP 來切分模型，以便能訓練更大規(guī)模模型。然而，這一過程要求在層與層之間執(zhí)行 AllReduce 操作。其中 Figure 2b 為未切分的操作，而 Figure 2c 為切分后跨兩個設(shè)備的操作。在 2b 中，每個設(shè)備僅持有部分權(quán)重。連續(xù)兩個矩陣乘可以先按列切，再按行切，之后通過一個 AllReduce 操作獲取完整結(jié)果。

而這些串行的 AllReduce 操作可能成為性能瓶頸。如下圖 Figure 4 展示了多種常見 Transformer 模型中使用 TP 后各操作的執(zhí)行時間占比?？梢钥闯觯?AllReduce（ReduceScatter + AllGather）的通信占比甚至比 GEMM 還長，比如 Megatron-GPT2 和 T-NLG 在訓練與推理（Prefill）中，通信時間分別高達 34% 和 43%。而且往往算力增長比網(wǎng)絡(luò)帶寬增長更快，這也會進一步加大通信的占比。

5.2.2 對比

在 T3 中，作者也提到上述 Microsoft - CoCoNet 和 Google Intra-layer Overlapping via Kernel Fusion 兩篇論文：

• Microsoft：需要昂貴的細粒度同步機制。
• Google：需要對矩陣乘法 Kernel 進行改動，并且可能對 GPU 軟件基礎(chǔ)設(shè)施造成干擾。
• 此外，上述兩個工作中計算與通信的Overlap 會競爭計算與內(nèi)存帶寬，從而削弱 Overlap 的實際效果。

5.2.3 方案概覽

現(xiàn)代 GPU 首先執(zhí)行生產(chǎn)者 GEMM 并將結(jié)果存儲于本地內(nèi)存中。隨后，啟動集合通信操作。相比之下，T3 系統(tǒng)在 GEMM 生成數(shù)據(jù)的同時立即啟動集合通信操作，以實現(xiàn)細粒度的 Overlap 執(zhí)行。如下圖 Figure 1 所示：

• T3 采用追蹤與觸發(fā)機制監(jiān)控 GEMM 與集合通信操作的進展，并協(xié)調(diào)通信流程，無需額外計算單元（CU）參與。
• 此外，T3 利用近內(nèi)存計算（Near Memory Computing, NMC）進行 Reduce 操作，以減少因通信產(chǎn)生的內(nèi)存訪問。
• 最終，這些優(yōu)化可以在幾乎不修改 Kernel 代碼的情況下透明地實現(xiàn)。?

如下圖 Figure 7 展示了 4 個 GPU 下，ReduceScatter（RS） 操作與 GEMM 的 Overlap 執(zhí)行情況。該 GEMM 根據(jù)輸入數(shù)據(jù)和 Kernel 實現(xiàn)分為多個工作組（WG）階段執(zhí)行，而 RS 則依據(jù)參與設(shè)備數(shù)量分為多個 Step 進行。為了簡化圖示，圖中 GEMM 階段數(shù)比所需的 Ring Step 數(shù)多 1。在每一個 Step 中，GEMM 某一階段的執(zhí)行和其輸出的 Reduce 與前一 Step 輸出的通信并行進行。在第一個 Step 中，GEMM 直接將輸出傳輸?shù)竭h程設(shè)備（remote_update）。后續(xù)的穩(wěn)態(tài) Step 則需通過 DMA（dma_update）完成。對于 N 臺設(shè)備，穩(wěn)態(tài) Step 需執(zhí)行 N-2 次，以處理不同的數(shù)據(jù)塊。

以穩(wěn)態(tài)下的 GPU 0 為例，對于其在 Step 2 的操作，如下圖 Figure 7 所示，GPU 0 在執(zhí)行并生成GEMM  Stage 3 輸出的同時，通過 DMA 接收來自鄰近設(shè)備 GPU 1 的 Stage 3 輸出副本（藍色）。這一過程與 GPU 0 將 GEMM 第二階段數(shù)據(jù)（黃色）Reduce 副本通過 DMA 傳輸至 GPU 3 的操作并行進行，從而實現(xiàn)了通信重疊。在 local_update 和 dma_update 中，T3 利用近內(nèi)存計算（NMC）進行原子內(nèi)存位置更新。為生成 Stage 3 數(shù)據(jù)塊的部分 Reduce 副本，無需額外讀取或占用 GPU 計算單元。一旦操作完成，GPU 0 即啟動對數(shù)據(jù)塊的 dma_update，將其傳輸?shù)脚R近設(shè)備 GPU 3 的內(nèi)存中，如下圖的 Step 3 所示。

5.2.4 硬件設(shè)計

通過一款輕量級且可編程的硬件追蹤單一，可以實現(xiàn)上述自動更新追蹤及 DMA 觸發(fā)機制，從而進一步降低對 GPU 計算單元（CU）的依賴。遠程/DMA 更新操作則通過配置 GEMM 輸出地址映射，輔以微小的應用程序和 Kernel 修改，得以透明執(zhí)行。

為了提升 T3 的性能，作者還對運行時和硬件進行了細微調(diào)整。為實現(xiàn)如圖 Figure 7 中 GEMM 與 RS 的完美 Overlap，作者還對跨 GPU 的 GEMM 工作組調(diào)度進行了錯位安排。此外，還通過引入一種簡單而有效的內(nèi)存控制仲裁（Memory Controller Arbitration，MCA）策略來增強內(nèi)存系統(tǒng)，以管理計算與通信之間的內(nèi)存競爭問題。

如下圖 Figure 8 展示了搭載 T3 增強功能（以橙色標注）的 GPU 執(zhí)行上述穩(wěn)態(tài)步驟的情況。GPU 執(zhí)行 GEMM 運算，為某一 Stage 生成 local 更新（L1）。同時，GPU 接收針對同一 Stage 的 DMA 更新（D1a），并向上一階段發(fā)送 DMA 更新（D1b）。在內(nèi)存控制器處，經(jīng)過改進的MCA 策略對 local 與 DMA 流量進行仲裁，以避免爭用。隨后，更新數(shù)據(jù)被傳輸至經(jīng)過 NMC 增強的 DRAM（L2a，D2a），同時 Tracker 記錄其進度（L2b，D2b）。一旦 Tracker 檢測到內(nèi)存區(qū)域所需的 local 和 DMA 更新，便會觸發(fā)這些更新通過 DMA 傳輸至相鄰 GPU（L3）。

5.3 結(jié)果

由于需要新的硬件支持，因此作者只是使用 Accel-Sim 來模擬評估 T3 的性能，當然，作者也對 Accel-Sim 進行了擴展，以支持多 GPU 系統(tǒng)。

如下圖 Figure 16 所示，作者通過模擬評估了本文方案的加速比，可以看出，其能夠獲得 15%-50% 不等的加速：

六、北大 Centauri

6.1 摘要

對應的論文為：[ASPLOS 24.04] Centauri: Enabling Efficient Scheduling for Communication-Computation Overlap in Large Model Training via Communication Partitioning [5]

高效訓練 LLMs 要求采用混合并行方法，這其中克服通信瓶頸至關(guān)重要，通常通過通信與計算的 Overlap 來實現(xiàn)。然而，現(xiàn)有 Overlap 方法多側(cè)重于細粒度 Kernel 融合或有限的算子（OP）調(diào)度，限制了異構(gòu)訓練環(huán)境下的性能。

本文作者提出 Centauri 框架，其構(gòu)建了一個由三個固有抽象維度組成的切分空間：原語替換、拓撲感知組切分及工作負載切分。這些維度共同構(gòu)成了一個全面的優(yōu)化空間，用于高效 Overlap。為確定通信與計算的高效 Overlap，作者將混合并行訓練中的調(diào)度任務(wù)分解為 OP、Layer 和模型三個層次。通過這些技術(shù)，Centauri 有效 Overlap 了通信時延，提升了硬件利用率。評估結(jié)果表明，在各種并行訓練配置下，Centauri 相較于主流方法可實現(xiàn)高達 1.49x 的加速。

6.2 方法

6.2.1 方案對比

以前的工作在 Overlap 通信和計算時存在一些不足，有些框架專注于優(yōu)化單一并行策略的調(diào)度，未能有效應對混合并行方法中的復雜 Overlap 挑戰(zhàn)。值得注意的是，即便在 Forward 和 Backward 過程中，最優(yōu)的 Overlap 模式也可能存在差異。

• 如下圖Figure 1a所示，其依賴粗粒度方式進行 Graph 級別的 Overlap，可能未能充分利用 GPU 硬件資源。
• 如下圖Figure 1b所示，有些工作依賴復雜的編譯器相關(guān)工作來對集合通信及鄰近計算進行切分，并在算子層面生成融合 Kernel（比如上述的 Microsoft CoCoNet）。然而，細粒度的 Kernel 融合可能忽視了更廣泛的 Graph 級別的調(diào)度計算（上述 Microsoft - CoCoNet 和 Google Intra-layer Overlapping via Kernel）。比如，1b 中的 Matmul B 反而比 1a 中的 Matmul B 慢。
• 如下圖Figure 1c所示，本文方案可以系統(tǒng)且全面地發(fā)掘 Overlap 空間的全部潛力，通過合理切分通信操作，能夠充分擴展通信 Overlap 的優(yōu)化空間。?

6.2.2 方案概覽

如下圖 Figure 3 所示，Centauri 的工作流程包含兩個核心環(huán)節(jié)：通信切分與層次調(diào)度。以 DP 與 FSDP 混合并行訓練為例：

• 通信切分：通過考量三個基本維度，生成潛在切分空間，并為每種集合通信選擇高效策略。
• 層次調(diào)度：在上述全面但較大的切分空間下，優(yōu)化整圖的 Overlap 調(diào)度成為一項復雜的任務(wù)，為了簡化復雜的調(diào)度任務(wù)，作者將復雜的混合并行集合通信分解為三個層次，每個集合通信被分配至特定調(diào)度層級。各層級選取開銷較低的切分與調(diào)度方案，旨在實現(xiàn)整體優(yōu)化 Overlap 方案。?

1. 原語替換：將 AllReduce 拆分為 Reduce-Scatter 和 AllGather。

2. 組切分：Forward 階段中的 AllGather 被切分為節(jié)點間組和節(jié)點內(nèi)組通信。具體來說，集合通信在 Rank Group G 內(nèi)進行，該 Group G 可以進一步細分為若干 Sub Group {???, ???, ???, ...}，以實現(xiàn)更細粒度的通信。在組切分時，應充分考慮網(wǎng)絡(luò)拓撲結(jié)構(gòu)，比如 FSDP 或 DP 等并行方法通常涉及跨節(jié)點的集合通信，導致設(shè)備連接的異質(zhì)性。在帶寬不均衡的子組中，低帶寬鏈路上的瓶頸可能抵消高帶寬鏈路的性能優(yōu)勢，此外，組切分應充分利用本地連接的高帶寬（比如 NVLink+NVSwitch），并盡量限制跨節(jié)點通信量。

3. 任務(wù)切分：以適當粒度切分集合通信與計算任務(wù)。在給定的原始通信與組切分方案下，若通信觸發(fā)的工作負載切分與其依賴的計算鏈之間的 Overlap 不足，則需要進一步切分計算鏈路。比如，F(xiàn)SDP 訓練中，AllGather 可與隨后的 MatMul 及 GeLU 和 Dropout Overlap 執(zhí)行。整個計算鏈路的工作負載切分是各 OP 切分方案的綜合結(jié)果。依賴計算鏈中各 OP 的切分策略會產(chǎn)生多種切分維度的組合，選擇兼容的可行切分組合至關(guān)重要。例如，沿 Batch 維度切分的 MatMul 輸出與隨后沿隱藏層維度切分的逐元素加法 OP 不兼容。

如下圖 Figure 4 展示了通信切分的流程抽象，在混合訓練中的每一個通信都會生成一個樹形結(jié)構(gòu)的切分空間，樹中的每個葉節(jié)點都代表一種可行的切分方案。所選的方案旨在實現(xiàn)最小的調(diào)度成本，所有節(jié)點上的分區(qū)策略構(gòu)成了一個龐大的切分方案森林，適用于混合訓練任務(wù)。

4. OP 級調(diào)度：OP 級別的細粒度調(diào)度旨在有效地 Overlap 每個 Forward Transformer Layer 內(nèi)的通信和計算操作，實現(xiàn)兩個拆分后的集合通信和計算 OP 之間的 Overlap。這個優(yōu)化保證了每個 通信的 Overlap 策略是按順序決定的，從而以貪婪的方式提高整個 Layer 的效率。

對于每個集合通信，基于各種切分模式的不同調(diào)度方案會導致不同的整體性能。

• 過于精細的工作負載切分可能會導致通信和計算幾乎完全 Overlap，但由于多個小型 GPU Kernel 啟動和數(shù)據(jù)移動開銷，它可能會對整體性能產(chǎn)生負面影響，如下圖 Figure 5b 所示。因此，粒度較大的策略更可取。
• 對于組切分，帶寬感知調(diào)度以及節(jié)點間和節(jié)點內(nèi)通信的恰當順序至關(guān)重要，如下圖 Figure 5c 和 5d 的比較所示，適當交錯節(jié)點內(nèi)和節(jié)點間調(diào)度方案，在下圖 Figure  5e 中取得了最大的性能改進。因此，以適當?shù)那蟹至６日_交錯執(zhí)行節(jié)點內(nèi)和節(jié)點間通信至關(guān)重要。?

5. Layer 級調(diào)度：根據(jù) Layer 內(nèi)關(guān)鍵路徑調(diào)整執(zhí)行順序。

與 Forward 階段不同，F(xiàn)orward 階段的高效 Overlap 依賴于切分策略，而 Backward 則具備天然的調(diào)度空間。Backward 包含兩個獨立部分：激活梯度計算與權(quán)重梯度計算。

• 如下圖 Figure 6a 和 6b，傳統(tǒng)方法中，激活計算的輸出作為前一 OP Backward 的輸入，激活計算往往會賦予更高的調(diào)度優(yōu)先級。然而，在混合并行配置中，這兩部分的不同執(zhí)行優(yōu)先級會導致不同的時延。
• 如下圖 Figure 6c，作者區(qū)分了激活梯度計算與權(quán)重梯度計算的兩條關(guān)鍵路徑，在同一個 Layer 內(nèi)，通過不同調(diào)度優(yōu)先級帶來的成本來選擇相應的最優(yōu)策略。（這一部分也可以參考之前我們介紹過的 Zero Bubble）?

6. 模型級調(diào)度：模型級別的 Overlap 旨在隱藏梯度和權(quán)重在 Forward 和 Backward 階段中的通信過程，提升整體訓練效率。

• 在單一數(shù)據(jù)并行（DP）場景中，所有 AllGather 操作與 Forward 階段 Overlap 進行，按塊粒度方式的 ReduceScatter 操作則與 Backward 階段 Overlap。
• 在DP + PP 中，細粒度的流水線調(diào)度策略旨在減少流水線 Bubble，這些 Bubble 影響計算與通信的 Overlap 效果。

通常，Micro-Batch 的模型 Chunk 的計算開銷小于相關(guān)的梯度或權(quán)重通信開銷。啟動多個相同模型 Chunk 的 Micro-Batch 能夠釋放 Overlap 的潛力，但激活內(nèi)存消耗也會隨著同時啟動的 Micro-Batch 數(shù)量增加而增長。內(nèi)存與時間成本是影響 PP 調(diào)度方案設(shè)計的兩大因素。

• 如下圖 Figure 7a 所示，F(xiàn)orward、Backward 和 Weight Update 各階段依次執(zhí)行，每個設(shè)備包含 2 個 PP Stage，每個 Batch 16 個 Micro-Batch，PP 深度為 4。
• 如下圖 Figure 7b 所示，為節(jié)省內(nèi)存消耗，深度優(yōu)先調(diào)度選擇最小數(shù)量的 Micro Batch 同時啟動，數(shù)量等于流水線階段深度，它忽略了為提升 End2End 性能而進行的 Overlap 潛力。
• 如下圖 Figure 7c 所示，廣度優(yōu)先調(diào)度則走向另一個極端，即啟動每個 Batch 中所有大小為 ???? 的 Micro-Batch 以實現(xiàn) Overlap（16 個 Micro-Batch 同時啟動），但伴隨而來的是峰值內(nèi)存消耗的顯著增加。這種權(quán)衡體現(xiàn)在內(nèi)存最小化調(diào)度與 Overlap 最大化調(diào)度之間。
• 如下圖 Figure 7d 所示，AllReduce 拆分為 ReduceScatter 和 AllGather，通過優(yōu)化選擇最優(yōu)策略，同時啟動 8 個 Micro-Batch，內(nèi)存消耗適中，8 倍激活量。?

6.3 結(jié)果

如下圖 Figure 13 和 Figure 14 所示，作者在兩種不同的網(wǎng)絡(luò)環(huán)境中驗證了 Centauri 的可擴展性。

• 集群 A 代表帶寬受限環(huán)境，Centauri 顯著提升了 FSDP/Zero3 配置對應的吞吐量。并且所達到的吞吐量水平與高性能環(huán)境（集群 B）的表現(xiàn)相當，突出了 Centauri 對帶寬變化的不敏感性。
• 盡管在集群 B 中性能提升的潛力有限，但仍能在 256 個 GPU 上將吞吐量提高 5%。
• 在 FSDP + DP 配置中，初始階段，由于 DP 通信開銷的增加，所有 6 種配置的吞吐量均有所下降，然而，Centauri 始終能保持更高的加速比。?

七、字節(jié) Flux

7.1 摘要

對應的論文為：[2406.06858] FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion [6]

對應的開源代碼為：GitHub - bytedance/flux: A fast communication-overlapping library for tensor parallelism on GPUs. [11]

大型模型通常需要分布式訓練和推理，TP 是其中一種常見的技術(shù)，通過在多個設(shè)備間劃分算子或?qū)拥挠嬎?，以克服單個處理器內(nèi)存容量的限制，并/或加速計算以滿足時延要求。然而，TP 引入了額外的通信開銷，可能占據(jù)整體運行時間的重要部分，從而限制了在高速互連設(shè)備（如節(jié)點內(nèi)配備 NVLink 的 GPU）中的可擴展性。

本文作者提出的 Flux 旨在通過依賴計算隱藏 GPU 間的通信延遲。Flux 將通信和計算操作分解為更細粒度的操作，并進一步融合成更大的 Kernel，從而在不損害 Kernel 效率的前提下有效隱藏通信。在融合 Kernel 的情況下，F(xiàn)lux 有望重疊高達 96% 的通信時間。

總體而言，在包含 128 個 GPU（涵蓋不同代際及互連技術(shù)）的集群上，F(xiàn)lux 相較于 Megatron-LM 在訓練中可實現(xiàn) 1.24x 的加速；而在包含 8 個 GPU 的集群上，F(xiàn)lux 在推理的 Prefill 和 Decoding 階段分別比 vLLM 快 1.66x 和1.30x。

7.2 方法

7.2.1 背景

如下圖所示，作者統(tǒng)計了不同訓練任務(wù)、推理任務(wù)在不同 GPU 上的 TP 通信時延，可以看出，在 PCIe 設(shè)備中通信占比很高；而 H800 NVL 相比 A100 NVL 的算力提升更多，通信帶寬提升較少，也就導致通信占比更高。在 PCIe 設(shè)備中 TP 通信占比甚至達到 40%-60%。

7.2.2 ReduceScatter Overlap

在 Flux 中，同樣是將 ReduceScatter 與 GEMM 進行 Overlap 和 Kernel 融合。ReduceScatter 操作可以分解為一個 AlltoAll 操作和一個 local 的 Reduce 操作。這里只有 AlltoAll 需要設(shè)備間通信，因此，將 All2All 融合到 GEMM 的尾部通常足以 Overlap 通信。

該算法要求 GPU 之間支持 P2P 通信，現(xiàn)代 NVIDIA GPU 無論是 NVLink 互聯(lián)還是 PCIe 互聯(lián)，在節(jié)點內(nèi)都已具備此能力，而 NVSHMEM（NVSHMEM | NVIDIA Developer [7]）進一步擴展了 NVIDSIA GPU 在節(jié)點間的 P2P 通信。

如下圖 Algorithm 1 所示為具體的算法：

7.2.3 AllGather Overlap

與 ReduceScatter 不同，AllGather 的實現(xiàn)采用首部融合方式，直接嵌入 GEMM Kernel 中。具體而言，AllGather 的信號檢查功能被融合至 GEMM 內(nèi)核的前序階段。如下圖 Algorithm 2 展示了融合 AllGather 后的 GEMM 偽代碼，用于計算 C = Aagg × B，其中 Aagg 是通過 AllGather 聚合的輸入矩陣 A 的緩沖區(qū)，B 為另一輸入矩陣，C 為輸出矩陣。

在 Kernel 端，GEMM 分塊計算被函數(shù) WaitSignal 阻塞，直至信號值被設(shè)置為真。此處，信號由GetSignal 依據(jù)輸出坐標（m 和 n）以及 TP 中的設(shè)備數(shù)量（NTP）從信號集合（signal_list）中選取。每個通信信號僅在主機端當對應輸入 Tensor 的部分（通信分塊）準備就緒時才被設(shè)置為真，即該部分已在運行融合 Kernel 的設(shè)備上接收完畢后。

如下圖 Algorithm 3 展示了主機端相應的通信過程：主機端（無論是基于 pull 還是 push）執(zhí)行分塊通信操作（DataTransfer），并異步地將相應信號（SetSignal）設(shè)置為真。

• 基于 pull 的方法通過 GetRemotePtr 函數(shù)和 GetLocalPtr 函數(shù)從遠程設(shè)備 pulling 分塊，從分塊 A 矩陣列表（A_list）和聚合矩陣緩沖區(qū)列表（Aagg_list）中選擇正確的指針，然后設(shè)置本地信號。信號由 GetSignalHost 依據(jù)通信分塊索引從信號集合（signal_list）中選取。
• 基于push 的方法則將分塊推送至遠程設(shè)備，隨后設(shè)置遠程信號。
• 需注意的是，在 pull 模式下，signal_list 僅包含本地信號，而在 push 模式下，signal_list 包含遠程設(shè)備的信號。這兩種變體的選擇被視為一個調(diào)優(yōu)參數(shù)。?

值得一提的是，在 AllGather 方法中，作者將通信的等待邏輯融合到 GEMM Kernel 中，而非整個通信操作。因此，AllGather 并不必然依賴 P2P 通信。同時，在 AllGather 中，通信的分塊策略（tilescomm）與 GEMM 計算的分塊策略相互解耦。這一設(shè)計提供了一種靈活的權(quán)衡方式，能夠在不損害 GEMM 效率的前提下，選擇 Overlap 機會與通信效率之間的最佳平衡。

7.2.4 方案對比

如下圖 Figure 5 展示了 ReduceScatter 中 Overlap 之間的主要差異?，F(xiàn)有 Overlap 方案 Tm 理論上可能比原始方法 Tc 執(zhí)行得更快，但通常情況下，Tm 仍慢于原始 GEMM 操作時間 Tg。主要原因在于，將一個 GEMM Kernel 拆分為一系列較小的 GEMM Kernel 會降低 GPU GEMM 的執(zhí)行效率。GEMM 通常需要合理大小的矩陣才能充分利用 GPU 的計算能力。這些具有數(shù)據(jù)依賴性的小型 GEMM 操作序列進一步阻礙了 GEMM Kernel 通過 GPU 多路復用技術(shù)并行運行，因此，Tensor 并行度越高，GPU 上的 GEMM 效率越低。

相比之下，作者提出的技術(shù)不存在上述限制。作者的 Overlap 方案 Tf 能夠在極小開銷下實現(xiàn)與原始 GEMM 操作 Tg 相當?shù)男阅堋Ｆ浼毩６确纸獠呗酝昝榔鹾犀F(xiàn)代 GPU 設(shè)計特性，即通過上下文切換的 Warp 和數(shù)百個在 SM 間并發(fā)活躍的 Warp 來隱藏延遲，如下圖 Figure 5 底部所示。最終，作者的方法在不影響 GEMM 計算效率的前提下，僅在執(zhí)行末尾引入少量通信開銷。

如下圖 Figure 6 展示了 AllGather 中的各種 Overlap 技術(shù)間的關(guān)鍵差異?，F(xiàn)有 Overlap 技術(shù) Tm  雖較原始粗粒度方法 Tc 有所提速，但因 GPU GEMM 效率降低，仍遜于原始 GEMM 操作時間 Tg。而作者的 Overlap 技術(shù) Tf 則能實現(xiàn)與原始 GEMM 操作 Tg 相媲美的性能。

AllGather 中長時延指令源于等待信號，此現(xiàn)象始于每個 Warp 的開端，因 WaitSignal 在起始階段已融合，其時延取決于相應數(shù)據(jù)傳輸?shù)牡竭_時間。對于數(shù)據(jù)已抵達的 Tile，時延近乎為 0；而對于數(shù)據(jù)尚未就緒的 Tile，Warp 間的上下文切換可掩蓋其等待時延。值得一提的是，本地 Tile 的信號始終預設(shè)為真，因此總有部分 Warp 無需等待信號。最終，作者的方法僅在執(zhí)行初期引入少量通信，且未損害 GEMM 計算效率。

7.3 結(jié)果

如下圖 Figure 17 展示了訓練、推理的 Prefill 及 Decoding 階段的性能結(jié)果。Flux 相較于 TransformerEngine：

• 在A100 PCIe上，可實現(xiàn)最高 1.37x 的訓練加速、2.06x 的 Prefill 加速及1.69x 的 Decoding 加速；
• 在A100 NVLink上，則分別可達 1.04x、1.14x 及 2.10x 的加速效果；
• 在H800 NVLink上，訓練、Prefill 及 Decoding 加速分別提升至1.05x、1.18x 及1.76x。

Flux 相較于 Megatron-LM 與 vLLM 基線：

• 在A100 PCIe上可實現(xiàn)1.24x 的訓練加速、1.46x 的 Prefill 加速及1.28x 的 Decoding 加速；
• 在A100 NVLink上，訓練與 Prefill 加速分別達到 1.05x 與 1.45x，Decoding 加速則為 1.30x；
• 在H800 NVLink上，訓練與 Prefill 加速分別提升至 1.10x 與 1.66x，但 Decoding 階段未見顯著加速。?

八、MicroSoft DeepSpeed-Domino

8.1 摘要

對應的論文為：[2409.15241] Domino: Eliminating Communication in LLM Training via Generic Tensor Slicing and Overlapping [8]

對應的開源代碼：DeepSpeedExamples/training/DeepSpeed-Domino/README.md at master [9]

LLM 訓練通常會消耗數(shù)百或數(shù)千個 GPU 來并行化和加速訓練過程，通信開銷也變得更加明顯。為了消除分布式 LLM 訓練中的通信開銷，作者提出了 Domino，提供了一種通用方案來隱藏計算后面的通信。通過將單個 Batch 的訓練數(shù)據(jù)依賴關(guān)系分解為更小的獨立部分，Domino 將這些獨立的部分訓練流水線化，并提供細粒度通信和計算 Overlap 的通用策略。

大量結(jié)果表明，與 Megatron-LM 相比，Domino 在 Nvidia DGX-H100 GPU 上實現(xiàn)了 1.3x 的 LLM 訓練加速。

PS：需要說明的是，本文的 Domino 主要是針對 TP 中的 AllReduce 進行優(yōu)化。此外，本文的方案也有一個比較明顯的約束：要求 Micro-Batch 最小為 4（2 也可以，但是效率很低）

8.2 方法

8.2.1 背景

作者使用 Megatron-LM 框架，在 DGX-H100 集群中，測量了 GPT-3 和 LLaMA-2 系列模型張量并行（TP）訓練中的通信開銷，如下圖 Figure 3 所示，通信時間占到 End2End 訓練時間的 22% 到 47% 不等?？梢钥闯?，即使使用高速的 NVLink + NVSwitch 和 Infiniband 互聯(lián)，通信開銷仍然占據(jù)很大的部分。這主要是對比 V100/A100  GPU，算力的增長更加明顯，通信的開銷就會更加突出。

8.2.2 方案概覽

如下圖 Figure 5 所示，首先是按照輸入數(shù)據(jù)的 Batch 維度進行切分（假設(shè) Tensor 都是 2 維），這樣可以避免按列切分時的通信量激增。由于 Batch 維是完全獨立的，因此不需要在所有 Transformer 層之間進行同步，可以實現(xiàn)層內(nèi)和層間計算和通信 Overlap。

如下圖 Figure 6 所示，同樣可以在 B 的最后一個維度按列切分，也可以實現(xiàn)層內(nèi)的計算和通信的 Overlap，但是在層結(jié)束時需要同步操作，然后才能執(zhí)行下一個注意力或 MLP 計算。

8.2.3 混合切分

對于大型 LLM 的訓練，作者提出了一種混合模型切分方案，同時對輸入 X 和最后的權(quán)重張量 B 進行切分，通過這種切分方式，Domino 能夠?qū)崿F(xiàn)超細粒度的計算與通信 Overlap，并且通信量與原始基線保持一致。

如下圖 Figure 7 中，在 Forward 階段，為了隱藏 SelfAttention 后的 AllReduce 通信，可以首先執(zhí)行 μ-Batch 0 的 SelfAttention，然后異步啟動其 AllReduce 操作（AllReduce(attn 0)），以避免 GPU 在通信過程中的阻塞。隨后立即啟動 μ-Batch 1 的 SelfAttention 計算，其可以與 AllReduce(attn 0) 異步執(zhí)行。而 μ-Batch 1 SelfAttention 后的 AllReduce(attn 1) 可以與隨后的 Dropout，殘差連接及 Noram 操作 Overlap。MLP 中類似。

如下圖 Figure 8 中，在 Backward 階段，首先采用了上述的跨 Micro-Batch 的計算與通信 Overlap 策略。為進一步擴大 Overlap 范圍，作者還采用了同一個 Micro-Batch 內(nèi)的通信與權(quán)重梯度計算的 Overlap。

然而，由于 PyTorch 自動生成了梯度計算圖，精確控制梯度通信以與梯度計算 Overlap 比較有挑戰(zhàn)。為此，作者開發(fā)了一個 no-operation 模塊，在 Forward 階段接收通信句柄，并在 Backward 保留以供使用。其 no-operation 模塊可以與 torch.autograd() 無縫集成。這種方式使得能夠精確控制異步通信的完成時間，而無需復雜的代碼修改。

8.3 結(jié)果

如下圖 Figure 9 所示，提出的 Domino 在不同規(guī)模模型，不同的序列長度下可以有效加快迭代速度：

九、中科大 DHelix

9.1 摘要

對應的論文為：[2411.15871] Hiding Communication Cost in Distributed LLM Training via Micro-batch Co-execution [10]

DHelix，是一種受 DNA 結(jié)構(gòu)啟發(fā)的新型架構(gòu)，可以顯著提升 LLM 訓練的效率。DHelix 的核心是鏈間交錯（Strand Interleaving，SI），它將 GPU 的兩個連續(xù) Micro Batch 流視作兩條鏈。DHelix 將兩條鏈的 Forward 和 Backward 并置，并對 SI 規(guī)劃進行系統(tǒng)優(yōu)化，具體來說，該規(guī)劃基于算子級 Overlap 分析的結(jié)果和基于動態(tài)規(guī)劃的搜索算法實現(xiàn)不同鏈的（Forward 與 Backward）和（Backward 與 Forward）的協(xié)同調(diào)度。同時，DHelix 使兩條鏈能夠共享模型狀態(tài)和激活空間，有效地以不到 3% 的額外內(nèi)存空間容納 2 個 Micro Batch。得益于獨特的模型折疊設(shè)計，DHelix 能夠無縫集成所有形式的現(xiàn)有數(shù)據(jù)/模型并行方案，其中最具挑戰(zhàn)的是 PP（Pipeline Parallelism），該設(shè)計實現(xiàn)了 W 形流水線。

作者在 3 個 GPU 集群（A40、A800 和 H100）上使用 DHelix 評估了常見的 LLaMA 和 GPT 模型以及 Phi MoE 模型。結(jié)果表明，在 64 A40 GPU 和 64 A800 GPU 集群上，DHelix 分別實現(xiàn)了 12-40%（最高 58% MFU）和 2-29%（最高 71% MFU）的提升，顯著優(yōu)于最先進的方案。在 H100 集群上，盡管更快的網(wǎng)絡(luò)削弱了 DHelix 的優(yōu)勢，但依然使得跨節(jié)點的 TP（Tensor Parallelism）更有應用前景。

我們在之前的文章中已經(jīng)詳細介紹過，這里不再贅述，可參考：???DHelix：跨 Micro-Batch 的通信隱藏，SOTA LLM 訓練性能??。

9.2 方法

9.2.1 背景

作者首先分析了 64 卡 A40 GPU 集群上使用 Megatron-LM 框架進行訓練的通信開銷。如下圖 Figure 3 所示，作者展示了多種 Transformer 模型，不同規(guī)模下總訓練時間中計算和 3 種通信操作的分布情況。可以看出，在這個規(guī)模下，通信已經(jīng)占主導地位，尤其是在大模型中。其中主要是模型并行帶來的集合通信開銷，即 TP/SP、CP、EP。另一方面，DP 和 PP 引入的通信則低得多。例如，在 LLaMA-39B 模型中，TP 和 CP 引起的通信占據(jù)執(zhí)行時間的 55%；而 Phi-31B 模型則產(chǎn)生了約 34.3% 的 EP 通信開銷：

9.2.2 方案概覽

DHelix 在算子層面執(zhí)行系統(tǒng)級的交錯處理，以適配兩條并行鏈路，即 α 鏈路 和 β 鏈路，每條鏈路處理一個 Micro Batch，從而最大化 GPU 利用率。作者通過引入時間延遲實現(xiàn) SI 雙鏈路，使得 α 鏈路的 Forward 與 β 鏈路的 Backward 得以協(xié)同調(diào)度，因為它們執(zhí)行過程中呈現(xiàn)出互補的內(nèi)存消耗模式，可以保證總激活內(nèi)存占用量維持在單鏈路的峰值附近。

9.2.3 模型折疊

作者將模型層的原始線性排布在 GPU 間折疊成 U 形布局。如下圖 Figure 7 右側(cè)所示，其包含 32 層，切分在 4 個 GPU 上，每個 GPU 上 8 層。

• 原始線性切分時：L0-7 在 GPU0，L8-15 在 GPU1，L16-23 在 GPU2，L24-31 在 GPU3；
• 本文的 U 形切分為：L0-3 和 L28-31 在 GPU0，L4-7 和 L24-27 在 GPU1，L8-11 和 L20-23 在 GPU2，L12-15 和 L16-19 在 GPU3。

相較于之前的模型復制方案，DHelix 的模型折疊并未改變每個 GPU 上的模型參數(shù)規(guī)模。因此，借助 SI 技術(shù)，同一套模型參數(shù)可以同時執(zhí)行兩條鏈，每個 GPU 上實時處理兩個 Micro Batch，而其消耗的 GPU 內(nèi)存容量幾乎與當前最先進的分布式訓練框架處理單鏈時相當。

9.2.4 調(diào)度優(yōu)化

Dhelix 并不是簡單地釋放兩個 Micro Batch 并讓 GPU 盡力進行協(xié)同調(diào)度，而是精心采用系統(tǒng)化和自適應的方法，在算子級別上有意對齊兩個鏈的執(zhí)行，以盡可能重疊兩個鏈之間的計算和通信操作。如下圖 Figure 9 所示為其整體工作流程：

1. 基于恰當?shù)?DAG 生成所有可能的算子序列。
2. 通過將一堆 Forward 和 Backward 算子劃分為連續(xù) Segment 來生成算子 Segment。
3. 使用動態(tài)規(guī)劃搜索最優(yōu)的 SI 配對方案，通過在執(zhí)行過程中插入 Barrier 來保證兩個鏈的協(xié)同執(zhí)行。?

9.3 結(jié)果

作者在 A800 集群進行了評估，將 LLaMA 模型擴展到 66B，序列長度擴展到 16384 和 32768，相應的 CP 組大小為 2 和 4。

如下圖 Figure 13a 展示了每個 GPU 的訓練吞吐，借助 NVLink，Megatron-LM 獲得了很高的 TFLOPS，在16K 和 32K 序列下分別實現(xiàn) 186 TFLOPS（60% MFU）和 160.9 TFLOPS（52% MFU）。這主要是因為節(jié)點內(nèi) TP 通信成本降低，僅占總訓練時間的 10%；此外，Megatron-LM 能夠部分重疊 CP 相關(guān)通信。相比之下，DHelix 在 Megatron-LM 基礎(chǔ)上仍有 7-24% 的提升，在 CP 為 4 時提升更明顯，這是因為 DHelix 有效隱藏了增加的跨節(jié)點通信，保持了 199.7 TFLOPS（64% MFU）的吞吐量。

十、參考鏈接

1. ???https://arxiv.org/abs/2401.10241???
2. ???https://arxiv.org/abs/2105.05720???
3. ???https://dl.acm.org/doi/pdf/10.1145/3567955.3567959???
4. ???https://arxiv.org/abs/2401.16677???
5. ???https://dl.acm.org/doi/10.1145/3620666.3651379???
6. ???https://arxiv.org/abs/2406.06858???
7. ???https://developer.nvidia.com/nvshmem???
8. ???https://arxiv.org/abs/2409.15241???
9. ???https://github.com/microsoft/DeepSpeedExamples/blob/master/training/DeepSpeed-Domino/README.md???
10. ???https://arxiv.org/abs/2411.15871???
11. ???https://github.com/bytedance/flux???

本文轉(zhuǎn)載自????AI閑談????，作者：AI閑談