一、背景

我們在之前的文章中提到過，在 A100 上進行大規模 LLM 訓練的 MFU（模型浮點運算利用率） 通常可以達到 50%-60%，而在 H100 上往往只有 40%-50%，為什么會存在這樣的現象，能否進一提升對應的性能呢？比如在 H100 中是否可以達到 60% 的 MFU？

今天介紹一篇新的文章，其采用了一種新的雙鏈技術，可以更好實現通信與計算的 Overlap，為實現上述目標提供了更多可能。

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

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

• ???大規模分布式 AI 模型訓練系列——數據并行???
• ???大規模分布式 AI 模型訓練系列——張量并行???
• ???大規模分布式 AI 模型訓練系列——流水線并行???
• ???大規模分布式 AI 模型訓練系列——專家并行???
• ???大規模分布式 AI 模型訓練系列——序列并行???
• ??超長序列 LLM 訓練：DeepSpeed Zero & Ulysses & FPDT??

二、摘要

LLM 的發展促進了大規模分布式訓練的需求，然而，盡管采用高度優化的框架，由于通信量、MFU（通常低于 50%）仍遭受顯著損失。與此同時，作者的全面分析表明，計算密集型和通信密集型算子可以良好 Overlap。

本文中，作者提出了 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）更有應用前景。

PS：針對本文的工作，我們也有一些額外的疑問：

• U 形折疊的另一個影響是開始階段的 Token Embedding 和最后的 LM Head 中的 Embedding 都位于一個 PP Stage 中，進一步增加了負載不均衡的問題；但同時也帶來一個好處，避免了起始 PP Stage 和結束 PP Stage 都需要讀取數據的問題。
• 大部分實驗基于 A40 集群（沒有 NVLink，機間網絡 100Gbps）完成，而大模型訓練通常會使用更高性能的 A100/800 或 H100/H800，往往會有 NVLink + NVSwitch 以及高速 IB 網絡。如果有更多基于這些環境的實驗會更有說服力，比如論文中并沒有具體介紹 A800 最高 71% MFU（312*0.71=221.52 TFLOPS） 的結果是哪個實驗，似乎對應了 Figure 14，此時 Megatron-LM 的 MFU 也非常高。

三、引言

3.1 Intra-batch 調度

計算與通信 Overlap 是模型并行（如 TP 和 PP）中廣泛采用的優化策略，其將計算和通信細化為細粒度任務，實現高效交錯，進而掩蓋通信的成本，實現模型訓練的加速。比如如下這些工作：

• [2406.08756] Optimizing Large Model Training through Overlapped Activation Recomputation [2]
• [2406.06858] FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion [3]
• Centauri: Enabling Efficient Scheduling for Communication-Computation Overlap in Large Model Training via Communication Partitioning [4]
• 字節[2402.15627] MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs [5]
• 微軟 DeepSpeed 團隊[2409.15241] Domino: Eliminating Communication in LLM Training via Generic Tensor Slicing and Overlapping [6]

如下圖 Figure 2 所示，字節 MegaScale 將大型 AllGather 算子及其后續的 GEMM 或 Dgrad（方向數據梯度）分解為細粒度的 Send/Recv 算子和矩陣塊。隨后，單個 Micro Batch 內，MegaScale 在通信與計算之間無數據依賴性，可以重疊這些分塊的通信和計算算子。

類似的，Megatron-LM 和 Ring-Attention 通過用 Send/Recv 替代 CP 中的 ALLGather，實現了對分塊注意力計算的 Overlap。此外，PP 并行優化也致力于重疊流水線通信（如 TP 傳輸）與計算，有效減少 PP Bubble 并提升效率。

然而，上述優化仍受限于單個 Micro Batch，由于數據依賴性，計算和通信無法充分并行。作者也在 Megatron-LM 中實現了 MegaScale 的 Overlap 方法，發現它們僅能重疊集合通信與相鄰的 GEMM 操作，仍有 73.9% 的計算/通信未被 Overlap。如上圖 Figure 2 所示為 MegaScale 調度結果，ALLGather 未被完全 Overlap，主要是因為 GEMM 的執行周期過短，而 Backward 中存在大量計算，難以找到足夠的通信操作與之 Overlap。

3.2 Megatron-LM 通信剖析

作者首先分析了 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 通信開銷：

通常的做法是將大通信量操作保留在機器內部（比如 TP 和 EP 經常會安排在機內，以便使用高速的 NVLink + NVSwitch）。然而，隨著模型規模變大，所需要的 GPU 數量增加，跨節點通信的比例必然增加；與此同時，機器內部 GPU 間的帶寬也在不斷提高，機內通信速度加快，這兩者的結合可能削弱層內通信的主導地位。為了理解這一點，作者估算了在 10000+ GPU 集群上進行 LLM 訓練時機內通信和跨節點通信的分布情況。具體而言，作者遵循 LLaMA 3.1 405B 技術報告（The Llama 3 Herd of Models | Research - AI at Meta [7]）中的分布式訓練配置和模型結構，以及 8*H100 NVSwitch 全互聯，外加 8*400 Gbps NIC 的硬件配置。

如下圖 Table 1 所示，通過擴展 DP 將 GPU 數量從 8192 擴展到 16384 時，跨節點通信占總通信時間的比例從 14.8% 增加到 33%。同時，如果啟用 CP，由于引發更多跨節點的 Send/Recv 操作，跨節點通信比例進一步增加到 49.9%。

3.3 LLM 訓練算子 Overlap

目前已經有很多研究在推動計算與通信的重疊（Overlap），以掩蓋通信的成本。然而，針對常見 LLM 訓練算子在 Overlap 執行時的性能表現，尚缺乏系統的評估。本文中，作者進行了廣泛的性能剖析，以理解在同一 GPU 上協同調度兩個算子時的性能影響，并通過多種 GPU 的完整基準測試來驗證。

作者采用重疊效果因子（Overlap Effectiveness Factor，OEF）來衡量各種算子的重疊行為，該因子的定義為：

其中 Ti 表示一個單一算子 opi 的順序執行時間，而 Pi,j 是 opi 和 opj 的重疊執行時間。換言之，OEF 衡量了重疊執行能夠較短算子執行時間的程度。

如下圖 Figure 4 所示，為 4 個代表性算子的結果，包括兩個計算密集型算子（GEMM：矩陣乘和 FA_BWD：FlashAttention  的反向操作）以及兩個通信通信密集型算子（C1：節點內 AllGather 和 C2：節點間 All-to-All）。其中分別對應 3 種 GPU（A40：GPU 間通過 PCIe 連接，A800 和 H100 兩者均通過 NVLink + NVSwitch 連接），主要觀察結果如下：

• 計算密集型算子（如 GEMM 和 FA）的重疊效果不佳。
• FA_BWD 算子在與 GEMM 重疊時尤為不利，這可能是 GEMM 干擾破壞了 FA 在計算和內存 IO 之間的精心編排以實現的細粒度交錯。
• 兩種通信算子與兩種計算算子的重疊效果均比較好。
• 無論是機內通信還是機間通信，其自身重疊效果均不佳，但由于同時利用了不同的網絡資源，它們之間的重疊卻能帶來收益。

這些結果啟發作者構建一種跨批次（Batch）執行交錯，并在算子級別進行系統性的細粒度折疊優化的方案。這使得未來可以通過放寬當前 TP 約束（當前一般僅限單個節點內的 8 個 GPU）來實現模型的擴展。

四、DHelix 設計與實現

4.1 鏈交叉概述

DHelix 在算子層面執行系統級的交錯處理，以適配兩條并行鏈路，即 α 鏈路 和 β 鏈路，每條鏈路處理一個 Micro Batch，從而最大化 GPU 利用率。

值得注意的是，當前常見的做法是處理單一鏈路，已通過盡可能增大 Micro Batch 來實現最大訓練吞吐，從而最大化利用 GPU 內存。如下圖 Figure 5 展示 展示了針對若干 Micro Batch 的樣本訓練時的內存使用情況。在給定模型規模和并行訓練配置（如 LLaMA 25B，DP=8，TP=8）的情況下，用戶通常會增大 Micro Batch（比如加 1），直到系統內存即將耗盡時為止。

引入時間延遲是實現 SI 雙鏈路的唯一途徑，使得 α 鏈路的 Forward 與 β 鏈路的 Backward 得以協同調度，如下圖 Figure 1 所示。這是因為它們執行過程中呈現出互補的內存消耗模式：Forward 鏈路在逐層推進中穩定的為激活數據分配內存，而 Backward 鏈路則以相似速率釋放內存（PS：正常來說 Backward 的延遲大概是 Forward 的 2 倍，為什么是相似速率？后文中作者給了相關解釋）。通過將兩組激活數據的“三角形”狀態緊密結合，總激活內存占用量將維持在單鏈路的峰值附近。

盡管 Wavelet 方法（Wavelet: Efficient DNN Training with Tick-Tock Scheduling [8]）已經探討過這一理念，其提出的 Tick-Tock 調度策略在數據并行中考慮了兩鏈路內存使用高峰和低谷的重疊。然而，在應用于 LLM 訓練環境中該方法仍有不足：

與流水線并行（PP）不兼容：PP 的計算松耦合且通信量低，成為擴展 LLM 訓練的關鍵機制（PS：Transformer 各層參數量一致，也更適合按層切分）。然而，在 PP 模式下，相鄰 Micro Batch 的 Forward 和 Backward（Wavelet 中的 Tick 和 Tock）呈反向移動。如下圖 Figure 6a 所示，α 鏈路的 Forward（藍色）從 GPU G0 開始， 而 β 鏈路的 Backward（綠色）從 GPU3 開始，兩條鏈在傳遞過程中僅交叉一次，使得 GPU 上協同執行的機會甚微。

模型復制：Wavelet 通過復制模型狀態（參數、梯度和優化器狀態）來協同調度兩個 Micro Batch 的數據訓練。值得注意的是，這種方式可能通過構建雙向流水線（如 Figure 6b 所示）來解決 PP 調度問題，使得 α 鏈路和 β 鏈路 能同向移動，比如從 G3 到 G0。然而，在典型的分布式訓練中，模型狀態占據 GPU 內存相當大的一部分（如上圖 Figure 5 的分析結果，其超過 30%）。存儲兩份模型狀態顯著削弱了 Micro Batch 交錯帶來的收益。

粗粒度交錯：此外，Wavelet 協同調度 Tick 和 Tock，以執行其原始的順序工作流程，未進行有意的算子重組以便提供和通信重疊的機會。

DHelix 通過其 SI 機制消除上述限制，該機制實現了在每個 GPU 上對兩個 Micro Batch 進行周期性高效且節省內存的協同調度。為便于理解 Dhelix 的工作原理，可以想象二維空間中的雙螺旋 DNA 結構，其兩條鏈耦合成單一的訓練計算流，同時處理兩個 Micro Batch。

耦合的 “DNA 鏈”作為一個整體，被折疊成 U 形，模型層在參與 PP 的 GPU 之間來回翻轉。通過這種排布，解決了上述 PP 不兼容和模型復制的問題。后文會詳細介紹，α 鏈路的 Forward 與 β 鏈路的 Backward在其生命周期內始終朝著同一方向移動，使得它們在 PP 方案中從一個設備到下一個設備內完美協同執行，同時，U 形折疊使得每個 GPU 上的模型層能夠被 α 鏈路與 β 鏈路共享，以單一的模型參數副本同時服務于 Forward 和 Backward。

當深入觀察 α 鏈路與 β 鏈路的耦合時，兩條鏈路上下相對移動，交替進行 Forward 和 Backward。在每次傳播中，鏈會重新排列其算子以找到重疊良好的兼容算子（例如，計算和通信的重疊），創建一個這些耦合點為錨點的調度，類似 DNA 鏈中的堿基對。DHelix 利用離線分析結果測量算子間重疊兼容性，并采用動態規劃搜索耦合鏈的高效協同調度。

因此，DHelix 的 SI 設計通過使訓練路徑能夠同時容納兩個相鄰 Micro Batch，有效隱藏了 LLM 訓練關鍵路徑中的通信開銷，顯著提升整體性能。同時，SI 在現有并行級別之下運行，可以無縫集成于 TP、SP、CP 和 EP。

4.2 模型折疊

這里，作者具體介紹了其模型折疊（Folding）技術。這一關鍵的 DHelix 技術使得 PP 得以實現，具體而言，作者將模型層的原始線性排布在 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。

這樣操作之后，Forward 和 Backward 在兩個 GPU 上的流動方向完全一致。如下圖 Figure 7 左側所示，熟悉的 1F1B 的 V 形結構變成了 W 形，其中 Forward 和 Backward 各形成一個相同的 V，以此可以實現  α 鏈路與 β 鏈路在跨 GPU 時的完美重疊，以及 Forward 和 Backward 在各 GPU 上的協同調度。

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

如上 Figure 7 簡化了雙鏈交錯傳播調度的時間表示，藍色和綠色塊具有相同的時間寬度。眾所周知，Backward 更慢，大約是 Forward 2 倍，因此綠色塊的寬度幾乎是藍色的 2 倍。當 SI 在同一 GPU 上協同調度 α 鏈路與 β 鏈路時，α 鏈路的 Backward 內存釋放是否能及時滿足 β 鏈路的 Forward 內存需求，從而引發兩條鏈路之間的空間依賴問題。實際上，作者發現這并不會引發問題，至多需要額外數百 MB 的閑置內存，因為作者是在層級別而非整個模型層面重疊 Forward 和 Backward。

為了提供更全面的展示，如下圖 Figure 8 所示，作者展示了包含 12 個 Micro Batch 的完整 W 形流水線調度。它與經典的 1F1B 流水線調度一樣，經歷了預熱（Warmup），穩定（Steady）和冷卻（Cooldown） 3 個階段。

• 在預熱階段，DHelix 將α 鏈路中 α1-α4（藍色）填充到流水線中。一旦它們完成 Forward，DHelix 便注入β 鏈路中 β5-β8，通過 SI 塊（頂部藍色和底部綠色）開始預填充流水行，此時α 鏈路中的 Forward 和β 鏈路中 Backward 融合在一起。
• 當發出 SI 塊 [β5α1] 時，DHelix 進入穩定階段。在此階段，DHelix 啟動α 鏈路中 α9-α12，且 SI 塊占滿所有 GPU。
• 最后，在冷卻階段，DHelix 耗盡用來創建 SI 塊的 Forward 塊，并隨著流水線情況回收 Backward 塊（綠色）。?

無論是原始的 1F1B 還是 DHelix 的 W 形調度，在穩定階段都能使 GPU 完全占用。此外，在預熱和冷卻階段，DHelix 僅執行原始的單鏈操作（無 SI），因此，DHelix 并未改變整體調度中 Bubble 率，仍然為 p/(m-1)。DHelix 的性能優勢在于：首先，能夠啟用 SI 的流水線執行（不同于 Wavelet）；其次，通過重疊兩條鏈的算子時間來縮短整體執行時間。換言之，每個 SI 塊 [βiαj] 完成所需的時間小于兩個算子獨立執行的時間。

此外，如上圖所示，W 形流水線對每個 Micro Batch 需要進行兩次上行和下行，從而使得 Send/Recv 通信量比原始的 1F1B 調度增加一倍。然而，在常見的分布式 LLM 訓練中，這種 PP Send/Recv 的通信量占比極小，對性能影響微乎其微。

4.3 算子配對下的鏈耦合

現在，可以進一步聚焦到 Figure 7 中兩個鏈（α 鏈路與 β 鏈路）耦合的 Forward + Backward 過程。盡管 U 形模型折疊機制使得這兩個鏈的協同執行成為可能，但接下來進行的算子級重疊則為 DHelix 帶來了顯著的性能提升。

這種提升源自于兩個鏈的協同執行，盡可能的利用硬件并行性。如前所示，當前 LLM 訓練流程中的各個算子，無論是計算還是通信，都已得到深度優化，使得這兩種操作能夠重疊。直接結果是，由于數據依賴性迫使其順序調度，當前的算子幾乎沒有機會在每個鏈中進一步折疊。

然而，由于 α 鏈路與 β 鏈路之間沒有數據依賴性，SI 也就為 α 鏈路的通信操作與 β 鏈路的計算操作，以及反之的計算和通信操作之間的交叉重疊開辟了新的可能性。除了交錯兩個鏈的 Forward 和 Backward 所帶來的內存收益外，還存在兩個額外的性能優勢：

• Forward 和 Backward 具有不同的算子布局，為計算算子和通信算子的系統調度留出了更多空間（相較于 Forward 與 Forward，以及 Backward 與 Backward 的交錯）。
• 在某些情況下，Backward 往往是計算密集型的，這提供了更多機會來隱藏 Forward 中較高比例的通信操作。

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

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

算子基礎特性：再次借用 DNA 結構類比，其中算子“配對”與互補鏈上的相應伙伴，形成“結合（Bond）”，在此情景下，通過利用異構 GPU 資源實現協同執行，從而顯著提升性能。與 DNA 鏈中 4 種核苷酸基不同，從概念上可將 Transformer 算子歸為兩種類型：計算和通信，如下圖 Table 2 所示。

從概念上講，“堿基配對”僅發生在可獲利的堿基類型之間：comp-comm 或 comm-comm。在 DHelix 設計中，為了考慮算子之間復雜的交錯形式，作者對 Table 2 中列出的算子進行詳細的成對測試。這種離線的預分析必須在新的硬件配置上重新進行，或者訓練流程被修改（如算子更新）時重新執行，耗時 10-30 分鐘。

算子排序預切分：Dhelix 無法控制單個算子的 CUDA 調度。然而，它可以通過在 α 鏈路與 β 鏈路的執行過程中策略性地設置 Barrier 來強制它們同步。這樣，通信算子可以與另一個鏈中的對等算子協同執行，以期最大化被計算隱藏的機會。換言之，單鏈的線性算子執行順序被劃分為如干個 Segment，這些 Segment 通過 DHelix 注入的 Barrier 在鏈間配對。DHelix 通過構建 Forward 和 Backward 中單層計算的 DAG 開始鏈配對過程，如啟用 TP/CP，則一個 Transformer 層由 14 或 18 個算子組成。如上圖 Figure 9 所示，基于這些 DAG，DHelix 首先通過枚舉其拓撲順序生成 Forward/Backward 算子序列。

自然地，下一步是將每對候選序列劃分為算子 Segment，并在兩序列間協同調度這些 Segment。給定一堆候選 Forward/Backward 序列，有效地放置跨鏈 Barrier 即生成一個候選配對規劃。

這兩類 DAG 在當前的 Transformer 工作流程中較為簡單，包含少數可在結果序列中移動的算子。這些算子包含反向權重梯度計算（wgrad，執行 GEMM 操作）以及在 MoE 中的路由計算。監管如此，候選序列的數量及其切分也導致搜索空間過于龐大，難以通過人工方式進行搜索。幸運的是，尋找最優配對方案的問題可以形式化為一個動態規劃問題。

動態規劃配對法：給定一對 Forward 和 Backward 候選算子序列，Sf 和 Sb，各自被切分為 Nf 和 Nb 個 Segment。Dhelix 尋求產生最短執行時間跨度 Topt(Nf,Nb) 的最優算子配對方案。

在最優解中，一對前綴序列的配對子解包含的來自 Sf 和 Sb 的 i(0<=i<Nf) 和 j(0<=j<Nb) 算子 Segment，也必須是最優的。

這里，P(i, j) 表示從 Forward 序列的第 i 個算子 Segment與 Backward 序列的第 j 個算子 Segment 重疊執行時的時間重疊量，該值由離線分析結果計算得出。P(i, ?) 則給出了第 i 個算子 Segment 單獨執行的時間，因其被安排獨立執行。

4.4 討論

兼容性：從之前提供的系統設計描述中可以看出，DHelix 的 SI 方案具有很高的通用性，不依賴于實際的分布式訓練框架及底層硬件。它可以通過重復執行離線分析過程，以適配新的訓練框架或硬件平臺。如下圖 Figure 10 所示，展示了同一訓練工作流在兩個不同平臺（A40 40GB 集群）和 （A800 80GB 集群）上的搜索結果，顯示出了兩種截然不同的 SI 規劃。需要說明的是，當分布式訓練參數（如 TP、PP、SP 等）改變時，即使在同一硬件上，配對調度也會有所不同，因為計算和通信算子的權重會發生變化。

可擴展性：最后，作者強調，除了算子特征化部分（包括離線分析），SI 對模型訓練用戶保持透明。若某一框架設置了多維并行配置，相同的配置可以在不改變參數或引入新的 SI 特定參數的情況下啟用 SI。因此，這使得 SI 能夠在對影響極小的情況下優化整體訓練吞吐量。本文測試受硬件資源限制，在擴展到更大規模時可以將其視為最小單元，通過在 PP、SP/CP 和 EP 維度上進一步擴展。

4.5 實現細節

作者基于 NVIDIA 的 Megatron-LM 實現了 DHelix，涉及 5000 行 Python 代碼。Megatron-LM 原有的 Transformer 實現包含 3 個模塊：預處理、Transformer 和后處理。在 DHelix 中，作者將 Transformer Block 替換為本文中的 SI-enabled Transformer Block，以便與 Megatron-LM 的預處理和后處理模塊結合，實現雙鏈執行。所有 DHelix 組件均基于 torch.nn.Module 構建，使用戶能夠利用標準 PyTorch API 創建自定義 Transformer 模型，并享受 SI 帶來的優勢，且無需對用戶級代碼進行任何修改。

在運行時，DHelix 根據生成的配對計劃依次處理算子對。通過在不同的 Segment 之間使用 torch.cuda.synchronize（也就是 barrier），確保這種順序執行語義。在單個 Segment 的執行過程中，算子進一步分為 3 類：計算、節點內通信和節點間通信。作者啟動了 3 個 CUDA Stream，并將每類算子分配到專屬的 Stream 進行處理。此外，PyTorch 允許不同的 CUDA Stream 獨立分配和釋放內存，這可能會導致內存碎片或內存不足錯誤。為了解決此問題，DHelix 僅使用默認 CUDA Stream 進行所有內存分配和釋放操作。

此外，作者也遵循現有實踐，通過調整 NCCL 環境變量 NCCL_NTHREADS 和 NCCL_MAX_NCHANNELS 來緩解計算與通信內核之間的競爭（Liger: Interleaving Intra- and Inter-Operator Parallelism for Distributed Large Model Inference [9]）。

五、實驗&評估

5.1 實驗設置

實驗平臺：包括 3 中 NVIDIA GPU 集群：

• 8 臺 A40（48GB） 機器的集群，機器間通過 100 Gbps 的 IB 網絡連接，每臺機器 8 張 A40 PCIe GPU，PCIe 4 帶寬為 32 GB/s。
• 8 臺 A800（80GB） 機器的集群，機器間通過 4x200 Gbps IB 網絡連接，機器內通過 NVLink（400 GB/s） + NVSwitch 互聯。
• 4 臺 H100（80GB） 機器的集群，機間通過 8x400 Gbps IB 網絡連接，每臺集群 8 張 H100 GPU，機器內通過 NVLink（900 GB/s） + NVSwitch 互聯。

CUDA 版本為 12.2，由于高端 GPU 資源申請難度大，A800/H100 集群僅在極短時間內可供使用。因此在所有平臺重復進行了整體性能測試，其余則在 A40 集群完成。

LLM/MoE 模型。如下圖 Table 3,4,5 所示，作者測試了 3 種主流開源 LLM，涵蓋 9 種不同規模。其中 Phi 為 MoE 模型，同樣 MoE 中的 top-k 參數為 2，也就是每個 Token 選擇兩個專家。

在具體實驗中，作者根據原始模型結構調整模型層數，以適應集群總的 GPU 內存容量。考慮到 PP 維度擴展具有相對較小的 PP 通信量，性能結果預計與更改模型規模時非常接近。

基線系統：作者將 DHelix 與多個基線系統進行對比，其主要是廣泛采用的 Megatron-LM。該框架支持所有形式的數據/模型并行，包括 DP、TP/SP、CP、PP 和 EP。

• Megatron-LM 基線已包含了對 CP 的通信 Overlap 優化，即注意力計算中重疊 Send/Recv 操作。
• Intra-batch 基線是指遵循 MegaScale 的設計理念，引入 Batch 內的通信重疊方案。主要是對 GEMM 算子進行切分，并在單個 Batch 內將這些算子與通信算子交錯執行，適用于 TP 和 SP。
• Wavelet+ 基線是對 Wavelet 的擴展，據作者了解，這是唯一探索 Micro Batch 之間交叉的方案。其僅應用于 DP，并且采用模型復制策略。為公平起見，作者在 DHelix 中實現了 Wavelet 的簡單輪詢調度，以交錯 DHelix 中的算子，形成流水線化的 Wavelet+。

指標：依照之前的慣例，對于 LLaMA 和 GPT 系列模型，主要通過單個 GPU 上作業的總浮點數除以 End2End 時間來代表每個 GPU 的訓練性能（TFLOPS/GPU）；對于 Phi MoE 模型，主要使用 Tokens/s 來衡量，以指示生成速度。最后，因為 DHelix 并未改變分布式訓練的語義，因此不影響模型的收斂性和準確性。

5.2 整體性能：A40 集群

作者通過 3 項任務評估 DHelix 的性能：Dense 模型訓練、長序列及上下文并行（CP）的 Dense 模型訓練、MoE 模型訓練。對于 LLaMA/GPT 模型，Global Batch 對應 800 萬 Token，也就是 Sequence Length 為 8K 時 Global Batch Size 為 1K。而 Micro Batch Size 則會相應調整以充分利用 GPU 內存。同樣，Phi 模型的 Global Batch Size 為 1600。

5.2.1 正常序列長度的 Dense 模型

如下圖 Figure 11 所示，在 LLaMA/GPT 模型，8192 序列長度，保持 TP 和 PP 相同配置下，DHelix 相比 Megatron-LM，Intra-Batch、Wavelet+ 分別提升 27-40%，15-28% 以及 21-25%。

與基準 Megatron-LM 相比，Intra-batch 的性提升顯著較低，只有 5%-15%。分析發現，由于每個 Micro Batch 內的順序依賴性，Intra-batch 通過與切分的 GEMM 算子重疊，僅隱藏了 26.1% 的 TP 通信開銷。相比之下，DHelix 幾乎隱藏了 83% 的通信開銷。

5.2.2 長序列的 Dense 模型

考慮到長序列的需求越來越多，作者進一步評估了長序列訓練性能，將序列長度從 8192 翻倍到 16384，并加入序列并行（CP），需要說明的是，加入 CP 會導致 DP 維度的降低。

如下圖 Figure 12 所示，DHelix 再次一致性的優于 Megatron-LM，Intra-batch 和 Wavelet+，提升幅度分別為22%-39%、13%-30%和11%-30%。

5.2.3 MoE 模型

在 Phi-31B MoE 模型上，使用專家并行（EP），且 DP 組的大小需要能被 EP 組大小整除。這里作者將 EP 大小設置為 8，DP 與 PP 大小分別為 16 和 4。此外，每個專家的容量因子設置為 6，全局并行組大小為 16x4，因為 EP 組為 DP 組的子集。如下圖 Table 6 所示，DHelix 可以有效實現 EP 中 All-to-All 通信的 Overlap，實現 27% 的性能提升。

5.3 整體性能：高性能集群

為了評估 DHelix 性能的普適性，作者進一步在 A800 集群進行評估。需要說明的是，作者這里忽略了 Intra-batch 實驗，因為 A800 機內有高速 NVLink 互聯，Intra-batch 帶來的提升幾乎可以忽略。

5.3.1 長序列 Dense 模型

進一步將 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）的吞吐量。

5.3.2 更大的 MoE 模型

得益于內存容量變大，A800 集群能夠容納完整的 Phi-42B 模型，并且可以將每個專家的容量因子設置為 8。如上圖 Figure 13b 所示，相比 Megatron-LM，DHelix 同樣實現了 15% 的提升（PS：這里如果也列下 MFU 就好了）。

5.3.3 增加 TP 組規模

大規模訓練中 TP 通常限制在 8 以內（單節點內）。然而隨著模型規模擴大和網絡帶寬提升，跨節點 TP 逐漸受到關注。為此，作者在 4 節點 H100 集群進行了短暫測試，以評估 DHelix 在解鎖跨節點 TP 方面的潛力。

如下圖 Figure 14 列出了 TP 規模從 8 增加到 32 的結果。值得注意的是，更高的 TP 規模支持更大的模型（PS：增加 PP 也可以，這樣相比 PP 的明顯優勢是什么？），但會犧牲每個 GPU 的吞吐量。在 H100 上，DHelix 相比 Megatron-LM 的收益較低（最高 17%），而在 A800 上則高達 29%，這主要是 H100 的互聯通信帶寬更高。不過，DHelix 在兩種平臺展現了相同的優勢，即 TP 越大，DHelix 相比 Megatron-LM TFLOPS 損失的越小。

為了評估 DHelix 在 H100 上的算子重疊效果，作者深入分析了使用跨節點 TP 進行訓練時的通信開銷。盡管已經隱藏了 50%-70% 的可見通信成本，但仍有進一步優化的空間，如下圖 Table 7 展示了限制這種重疊的 2 個重要因素：

• 內核降速：在多個 CUDA Stream 上同時運行計算和通信 Kernel 會導致性能下降 20%-30%，限制了潛在的性能提升。
• 啟動間隔：每個配對 Segment 開始時計算和通信 Kernel 啟動之間的小間隔會產生相當于通信成本 10%-20% 的開銷。

一個有前景的方法是將這些 Kernel 融合成一個單一的、整體的 Kernel，以更精確地管理 GPU 資源，如 LIBRA: Contention-Aware GPU Thread Allocation for Data Parallel Training in High Speed Networks [10] 所示。

5.4 性能提升統計分析

這個部分，作者分析了 DHelix 相對于 Megatron-LM 基線在性能提升方面的各個來源，這些提升主要來源于 DHelix 引入的各種通信 Overlap 策略。作者在 A40 集群上使用 LLaMA-39B 模型進行了兩組統計分析實驗：

• a：CP 并行度為 2，序列長度為 16384。
• b：CP 并行度為 4，序列長度為 32768。

結果如下圖 Figure 15a 和 15b 所示，測試 15a 采用了與 Figure 12 中 LLaMA-39B 相同配置，而測試 15b 則評估了 DHelix 在更長序列下的延遲情況：

TP Overlap：首先加入 TP Overlap，使得 15a 和 15b 的性能提升為 11.18% 和 5.37%。通常來說，更高的通信成本更利于 DHelix。然而，在 15b 中，32K 序列下的通信開銷過高，無法完全 Overlap，這反而成了不利因素。

CP Overlap：進一步啟用 CP Overlap，進一步隱藏了由 CP 引起的 Send/Recv 開銷，分別使 TP Overlap 后的性能再提升 22.5% 和 14.4%。這一步 DHelix 帶來了最顯著的性能提升。因為在 a 和 b 中，CP 引起的層內通信分別占到總訓練時間的 28% 和 53%。同樣，過多的通信也降低了 b 中 CP Overlap 的收益

通信與通信 Overlap：又格外新增了節點內與節點間通信的重疊，在 b 中帶來了 7% 的提升，而在 a 中，大部分通信已經被重疊，因此沒什么收益。

權重梯度 Overlap：最后，進一步啟用權重梯度（Wgrad）算子在 DAG 排序允許的范圍內移動，展示了算子重排的影響。在 a 上進一步提升 5%，有助于減少數據依賴性導致的 “Overlap Bubble”。

5.5 SI 內存空間利用率

最后，作者也進一步驗證了 DHelix 的內存優化能力，這對于支持盡可能大的模型訓練任務至關重要。作者將 DHelix 的 SI 方案與 Megatron-LM（單鏈）以及另外一種直接但內存不友好的雙鏈實現（Wavelet）進行了比較。實驗中采用 LLaMA-25B 模型，在 DP=4，TP=8 和 PP=2 下，Micro Batch 設置為 1。

通過逐步增加模型層數，以確定系統支持的最大模型參數規模。實驗結果表明，DHelix 支持的最大模型尺寸為 39B，非常接近 Megatron-LM 的 40B，僅減少了 2.5% 的模型尺寸，卻換來了 40% 的訓練吞吐提升（PS：這里應該是 A40 GPU）。相比執行，Wavelet 最大只支持 32B 參數的模型。

六、參考鏈接

• ??https://arxiv.org/abs/2411.15871??
• ??https://arxiv.org/abs/2406.08756??
• ??https://arxiv.org/abs/2406.06858??
• ??https://dl.acm.org/doi/10.1145/3620666.3651379??
• ??https://arxiv.org/abs/2402.15627??
• ??https://arxiv.org/abs/2409.15241??
• ??https://ai.meta.com/research/publications/the-llama-3-herd-of-models/????????https://proceedings.mlsys.org/paper_files/paper/2021/hash/099268c3121d49937a67a052c51f865d-Abstract.html??
• ??https://dl.acm.org/doi/10.1145/3627535.3638466??
• ???https://jhc.sjtu.edu.cn/~shizhenzhao/infocom2023.pdf???

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