一、背景

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

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

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

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

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

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

二、引言

2.1 AllReduce

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

對(duì)于常見(jiàn)的基于 Ring 的 AllReduce 實(shí)現(xiàn)中，通常可以把 AllReduce 操作看成為一個(gè) ReduceScatter 和一個(gè) AllGather 操作，如下圖所示：

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

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

2.2 LLM Tensor Parallelism AllReduce

當(dāng)前 LLM 推理通常會(huì)采用 Tensor Parallelism（TP）模型并行，以便在多個(gè) GPU 上實(shí)現(xiàn)較大 LLM 的推理。對(duì)于標(biāo)準(zhǔn)的 Transformer Decoder Only 模型，通常會(huì)在每個(gè) Transformer Block 中采用如下的 TP 切分方式：

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

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

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

2.3 ReduceScatter = All2All + Local Reduce

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

2.4 ZeroBubble

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

三、Microsoft - CoCoNet

3.1 摘要

對(duì)應(yīng)的 Paper 為：[ASPLOS 22] [2105.05720] Breaking the Computation and Communication Abstraction Barrier in Distributed Machine Learning Workloads [2]

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

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

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

3.2 方法

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

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

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

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

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

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

3. 算子融合：

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

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

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

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

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

3.3 結(jié)果

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

四、Google - Intra-layer Overlapping via Kernel Fusion

4.1 摘要

對(duì)應(yīng)的論文為：[ASPLOS 23] Overlap Communication with Dependent Computation via Decomposition in Large Deep Learning Models [3]

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

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

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

4.2 方法

4.2.1 背景

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

4.2.2 方案概覽

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

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

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

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

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

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

4.3 結(jié)果

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

五、AMD - T3

5.1 摘要

對(duì)應(yīng)的論文為：[2401.16677] T3: Transparent Tracking & Triggering for Fine-grained Overlap of Compute & Collectives [4]

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

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

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

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

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

5.2 方法

5.2.1 背景

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

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

5.2.2 對(duì)比

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

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

5.2.3 方案概覽

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

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

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

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

5.2.4 硬件設(shè)計(jì)

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

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

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

5.3 結(jié)果

由于需要新的硬件支持，因此作者只是使用 Accel-Sim 來(lái)模擬評(píng)估 T3 的性能，當(dāng)然，作者也對(duì) Accel-Sim 進(jìn)行了擴(kuò)展，以支持多 GPU 系統(tǒng)。

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

六、北大 Centauri

6.1 摘要

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

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

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

6.2 方法

6.2.1 方案對(duì)比

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

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

6.2.2 方案概覽

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

6.3 結(jié)果

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

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

七、字節(jié) Flux

7.1 摘要

對(duì)應(yīng)的論文為：[2406.06858] FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion [6]

對(duì)應(yīng)的開(kāi)源代碼為：GitHub - bytedance/flux: A fast communication-overlapping library for tensor parallelism on GPUs. [11]

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

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

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

7.2 方法

7.2.1 背景

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

7.2.2 ReduceScatter Overlap

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

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

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

7.2.3 AllGather Overlap

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

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

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

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

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

7.2.4 方案對(duì)比

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

相比之下，作者提出的技術(shù)不存在上述限制。作者的 Overlap 方案 Tf 能夠在極小開(kāi)銷下實(shí)現(xiàn)與原始 GEMM 操作 Tg 相當(dāng)?shù)男阅?。其?xì)粒度分解策略完美契合現(xiàn)代 GPU 設(shè)計(jì)特性，即通過(guò)上下文切換的 Warp 和數(shù)百個(gè)在 SM 間并發(fā)活躍的 Warp 來(lái)隱藏延遲，如下圖 Figure 5 底部所示。最終，作者的方法在不影響 GEMM 計(jì)算效率的前提下，僅在執(zhí)行末尾引入少量通信開(kāi)銷。

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

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

7.3 結(jié)果

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

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

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

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

八、MicroSoft DeepSpeed-Domino

8.1 摘要

對(duì)應(yīng)的論文為：[2409.15241] Domino: Eliminating Communication in LLM Training via Generic Tensor Slicing and Overlapping [8]

對(duì)應(yīng)的開(kāi)源代碼：DeepSpeedExamples/training/DeepSpeed-Domino/README.md at master [9]

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

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

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

8.2 方法

8.2.1 背景

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

8.2.2 方案概覽

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

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

8.2.3 混合切分

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

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

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

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

8.3 結(jié)果

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

九、中科大 DHelix

9.1 摘要

對(duì)應(yīng)的論文為：[2411.15871] Hiding Communication Cost in Distributed LLM Training via Micro-batch Co-execution [10]

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

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

我們?cè)谥暗奈恼轮幸呀?jīng)詳細(xì)介紹過(guò)，這里不再贅述，可參考：???DHelix：跨 Micro-Batch 的通信隱藏，SOTA LLM 訓(xùn)練性能??。

9.2 方法

9.2.1 背景

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

9.2.2 方案概覽

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

9.2.3 模型折疊

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

• 原始線性切分時(shí)：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。

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

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

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

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

9.3 結(jié)果

作者在 A800 集群進(jìn)行了評(píng)估，將 LLaMA 模型擴(kuò)展到 66B，序列長(zhǎng)度擴(kuò)展到 16384 和 32768，相應(yīng)的 CP 組大小為 2 和 4。

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