一、背景

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

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

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

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

訓練相關可以參考我們之前的文章：

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

二、引言

2.1 AllReduce

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

對于常見的基于 Ring 的 AllReduce 實現中，通常可以把 AllReduce 操作看成為一個 ReduceScatter 和一個 AllGather 操作，如下圖所示：

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

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

2.2 LLM Tensor Parallelism AllReduce

當前 LLM 推理通常會采用 Tensor Parallelism（TP）模型并行，以便在多個 GPU 上實現較大 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 的優化，可以等效為一個 All2All 操作實現數據的重排，然后在 Local 進行 Reduce 操作。此過程只有一個 All2All 的整體通信操作，雖然實際上與 Ring 實現的方式的通信量和計算量沒有變化，但可以避免 K-1 個 Ring Step 的同步，進而可以有效降低時延。

2.4 ZeroBubble

 [2401.10241] Zero Bubble Pipeline Parallelism [1]  中作者提出 Zero Bubble，核心思路是將 Backward 分為兩個部分，一部分計算輸入的梯度，另一部分計算參數的梯度，如下圖 Figure 1 所示。這里計算輸入的梯度有明確的依賴關系，也是鏈式法則不斷傳遞的基礎；而計算權重的梯度卻沒有明確的依賴，甚至可以滯后很多。此外，三個紅色部分計算量相當，這也就是為什么常見的流水線并行（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 包含：

• 一種領域特定語言（DSL），用于以計算和通信操作的形式表示分布式機器學習程序。
• 一組保持語義的轉換規則，以優化程序。
• 一個編譯器，可以生成聯合優化通信和計算的 GPU Kernel。

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

3.2 方法

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

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

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

具體的幾個轉換過程如下：

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

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

3. 算子融合：

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

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

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

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

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

3.3 結果

如下圖 Figure 1 所示，MatMul 與 AllReduce 的細粒度 Overlap 執行可掩蓋 80% 的 MatMul 執行時間，并帶來 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]

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

該技術將識別出的原始集合通信與依賴的計算操作分解為一系列更細粒度的操作，通過創造更多 Overlap 機會并并行執行新生成的細粒度通信與計算操作，可以有效隱藏數據傳輸時延，實現更優的系統利用率。

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

4.2 方法

4.2.1 背景

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

4.2.2 方案概覽

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

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

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

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

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

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

4.3 結果

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

五、AMD - T3

5.1 摘要

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

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

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

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

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

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

5.2 方法

5.2.1 背景

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

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

5.2.2 對比

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

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

5.2.3 方案概覽

現代 GPU 首先執行生產者 GEMM 并將結果存儲于本地內存中。隨后，啟動集合通信操作。相比之下，T3 系統在 GEMM 生成數據的同時立即啟動集合通信操作，以實現細粒度的 Overlap 執行。如下圖 Figure 1 所示：

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

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

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

5.2.4 硬件設計

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

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

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

5.3 結果

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

如下圖 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 要求采用混合并行方法，這其中克服通信瓶頸至關重要，通常通過通信與計算的 Overlap 來實現。然而，現有 Overlap 方法多側重于細粒度 Kernel 融合或有限的算子（OP）調度，限制了異構訓練環境下的性能。

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

6.2 方法

6.2.1 方案對比

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

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

6.2.2 方案概覽

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

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

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

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

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

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

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

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

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

5. Layer 級調度：根據 Layer 內關鍵路徑調整執行順序。

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

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

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

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

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

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

6.3 結果

如下圖 Figure 13 和 Figure 14 所示，作者在兩種不同的網絡環境中驗證了 Centauri 的可擴展性。

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

七、字節 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 是其中一種常見的技術，通過在多個設備間劃分算子或層的計算，以克服單個處理器內存容量的限制，并/或加速計算以滿足時延要求。然而，TP 引入了額外的通信開銷，可能占據整體運行時間的重要部分，從而限制了在高速互連設備（如節點內配備 NVLink 的 GPU）中的可擴展性。

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

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

7.2 方法

7.2.1 背景

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

7.2.2 ReduceScatter Overlap

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

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

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

7.2.3 AllGather Overlap

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

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

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

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

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

7.2.4 方案對比

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

相比之下，作者提出的技術不存在上述限制。作者的 Overlap 方案 Tf 能夠在極小開銷下實現與原始 GEMM 操作 Tg 相當的性能。其細粒度分解策略完美契合現代 GPU 設計特性，即通過上下文切換的 Warp 和數百個在 SM 間并發活躍的 Warp 來隱藏延遲，如下圖 Figure 5 底部所示。最終，作者的方法在不影響 GEMM 計算效率的前提下，僅在執行末尾引入少量通信開銷。

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

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

7.3 結果

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

• 在A100 PCIe上，可實現最高 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上可實現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 訓練通常會消耗數百或數千個 GPU 來并行化和加速訓練過程，通信開銷也變得更加明顯。為了消除分布式 LLM 訓練中的通信開銷，作者提出了 Domino，提供了一種通用方案來隱藏計算后面的通信。通過將單個 Batch 的訓練數據依賴關系分解為更小的獨立部分，Domino 將這些獨立的部分訓練流水線化，并提供細粒度通信和計算 Overlap 的通用策略。

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

PS：需要說明的是，本文的 Domino 主要是針對 TP 中的 AllReduce 進行優化。此外，本文的方案也有一個比較明顯的約束：要求 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 互聯，通信開銷仍然占據很大的部分。這主要是對比 V100/A100  GPU，算力的增長更加明顯，通信的開銷就會更加突出。

8.2.2 方案概覽

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

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

8.2.3 混合切分

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

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

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

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

8.3 結果

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

九、中科大 DHelix

9.1 摘要

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

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

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

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

9.2 方法

9.2.1 背景

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

9.2.2 方案概覽

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

9.2.3 模型折疊

作者將模型層的原始線性排布在 GPU 間折疊成 U 形布局。如下圖 Figure 7 右側所示，其包含 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 上的模型參數規模。因此，借助 SI 技術，同一套模型參數可以同時執行兩條鏈，每個 GPU 上實時處理兩個 Micro Batch，而其消耗的 GPU 內存容量幾乎與當前最先進的分布式訓練框架處理單鏈時相當。

9.2.4 調度優化

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

1. 基于恰當的 DAG 生成所有可能的算子序列。
2. 通過將一堆 Forward 和 Backward 算子劃分為連續 Segment 來生成算子 Segment。
3. 使用動態規劃搜索最優的 SI 配對方案，通過在執行過程中插入 Barrier 來保證兩個鏈的協同執行。?

9.3 結果

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

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

本文轉載自????AI閑談????，作者：AI閑談