一、背景

我們?cè)谥暗奈恼轮刑岬竭^(guò)，在 A100 上進(jìn)行大規(guī)模 LLM 訓(xùn)練的 MFU（模型浮點(diǎn)運(yùn)算利用率） 通?？梢赃_(dá)到 50%-60%，而在 H100 上往往只有 40%-50%，為什么會(huì)存在這樣的現(xiàn)象，能否進(jìn)一提升對(duì)應(yīng)的性能呢？比如在 H100 中是否可以達(dá)到 60% 的 MFU？

今天介紹一篇新的文章，其采用了一種新的雙鏈技術(shù)，可以更好實(shí)現(xiàn)通信與計(jì)算的 Overlap，為實(shí)現(xiàn)上述目標(biāo)提供了更多可能。

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

訓(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??

二、摘要

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

本文中，作者提出了 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)估了常見的 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)用前景。

PS：針對(duì)本文的工作，我們也有一些額外的疑問(wèn)：

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

三、引言

3.1 Intra-batch 調(diào)度

計(jì)算與通信 Overlap 是模型并行（如 TP 和 PP）中廣泛采用的優(yōu)化策略，其將計(jì)算和通信細(xì)化為細(xì)粒度任務(wù)，實(shí)現(xiàn)高效交錯(cuò)，進(jìn)而掩蓋通信的成本，實(shí)現(xiàn)模型訓(xùn)練的加速。比如如下這些工作：

• [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]
• 字節(jié)[2402.15627] MegaScale: Scaling Large Language Model Training to More Than 10,000 GPUs [5]
• 微軟 DeepSpeed 團(tuán)隊(duì)[2409.15241] Domino: Eliminating Communication in LLM Training via Generic Tensor Slicing and Overlapping [6]

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

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

然而，上述優(yōu)化仍受限于單個(gè) Micro Batch，由于數(shù)據(jù)依賴性，計(jì)算和通信無(wú)法充分并行。作者也在 Megatron-LM 中實(shí)現(xiàn)了 MegaScale 的 Overlap 方法，發(fā)現(xiàn)它們僅能重疊集合通信與相鄰的 GEMM 操作，仍有 73.9% 的計(jì)算/通信未被 Overlap。如上圖 Figure 2 所示為 MegaScale 調(diào)度結(jié)果，ALLGather 未被完全 Overlap，主要是因?yàn)?GEMM 的執(zhí)行周期過(guò)短，而 Backward 中存在大量計(jì)算，難以找到足夠的通信操作與之 Overlap。

3.2 Megatron-LM 通信剖析

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

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

如下圖 Table 1 所示，通過(guò)擴(kuò)展 DP 將 GPU 數(shù)量從 8192 擴(kuò)展到 16384 時(shí)，跨節(jié)點(diǎn)通信占總通信時(shí)間的比例從 14.8% 增加到 33%。同時(shí)，如果啟用 CP，由于引發(fā)更多跨節(jié)點(diǎn)的 Send/Recv 操作，跨節(jié)點(diǎn)通信比例進(jìn)一步增加到 49.9%。

3.3 LLM 訓(xùn)練算子 Overlap

目前已經(jīng)有很多研究在推動(dòng)計(jì)算與通信的重疊（Overlap），以掩蓋通信的成本。然而，針對(duì)常見 LLM 訓(xùn)練算子在 Overlap 執(zhí)行時(shí)的性能表現(xiàn)，尚缺乏系統(tǒng)的評(píng)估。本文中，作者進(jìn)行了廣泛的性能剖析，以理解在同一 GPU 上協(xié)同調(diào)度兩個(gè)算子時(shí)的性能影響，并通過(guò)多種 GPU 的完整基準(zhǔn)測(cè)試來(lái)驗(yàn)證。

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

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

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

• 計(jì)算密集型算子（如 GEMM 和 FA）的重疊效果不佳。
• FA_BWD 算子在與 GEMM 重疊時(shí)尤為不利，這可能是 GEMM 干擾破壞了 FA 在計(jì)算和內(nèi)存 IO 之間的精心編排以實(shí)現(xiàn)的細(xì)粒度交錯(cuò)。
• 兩種通信算子與兩種計(jì)算算子的重疊效果均比較好。
• 無(wú)論是機(jī)內(nèi)通信還是機(jī)間通信，其自身重疊效果均不佳，但由于同時(shí)利用了不同的網(wǎng)絡(luò)資源，它們之間的重疊卻能帶來(lái)收益。

這些結(jié)果啟發(fā)作者構(gòu)建一種跨批次（Batch）執(zhí)行交錯(cuò)，并在算子級(jí)別進(jìn)行系統(tǒng)性的細(xì)粒度折疊優(yōu)化的方案。這使得未來(lái)可以通過(guò)放寬當(dāng)前 TP 約束（當(dāng)前一般僅限單個(gè)節(jié)點(diǎn)內(nèi)的 8 個(gè) GPU）來(lái)實(shí)現(xiàn)模型的擴(kuò)展。

四、DHelix 設(shè)計(jì)與實(shí)現(xiàn)

4.1 鏈交叉概述

DHelix 在算子層面執(zhí)行系統(tǒng)級(jí)的交錯(cuò)處理，以適配兩條并行鏈路，即 α 鏈路 和 β 鏈路，每條鏈路處理一個(gè) Micro Batch，從而最大化 GPU 利用率。

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

引入時(shí)間延遲是實(shí)現(xiàn) SI 雙鏈路的唯一途徑，使得 α 鏈路的 Forward 與 β 鏈路的 Backward 得以協(xié)同調(diào)度，如下圖 Figure 1 所示。這是因?yàn)樗鼈儓?zhí)行過(guò)程中呈現(xiàn)出互補(bǔ)的內(nèi)存消耗模式：Forward 鏈路在逐層推進(jìn)中穩(wěn)定的為激活數(shù)據(jù)分配內(nèi)存，而 Backward 鏈路則以相似速率釋放內(nèi)存（PS：正常來(lái)說(shuō) Backward 的延遲大概是 Forward 的 2 倍，為什么是相似速率？后文中作者給了相關(guān)解釋）。通過(guò)將兩組激活數(shù)據(jù)的“三角形”狀態(tài)緊密結(jié)合，總激活內(nèi)存占用量將維持在單鏈路的峰值附近。

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

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

模型復(fù)制：Wavelet 通過(guò)復(fù)制模型狀態(tài)（參數(shù)、梯度和優(yōu)化器狀態(tài)）來(lái)協(xié)同調(diào)度兩個(gè) Micro Batch 的數(shù)據(jù)訓(xùn)練。值得注意的是，這種方式可能通過(guò)構(gòu)建雙向流水線（如 Figure 6b 所示）來(lái)解決 PP 調(diào)度問(wèn)題，使得 α 鏈路和 β 鏈路 能同向移動(dòng)，比如從 G3 到 G0。然而，在典型的分布式訓(xùn)練中，模型狀態(tài)占據(jù) GPU 內(nèi)存相當(dāng)大的一部分（如上圖 Figure 5 的分析結(jié)果，其超過(guò) 30%）。存儲(chǔ)兩份模型狀態(tài)顯著削弱了 Micro Batch 交錯(cuò)帶來(lái)的收益。

粗粒度交錯(cuò)：此外，Wavelet 協(xié)同調(diào)度 Tick 和 Tock，以執(zhí)行其原始的順序工作流程，未進(jìn)行有意的算子重組以便提供和通信重疊的機(jī)會(huì)。

DHelix 通過(guò)其 SI 機(jī)制消除上述限制，該機(jī)制實(shí)現(xiàn)了在每個(gè) GPU 上對(duì)兩個(gè) Micro Batch 進(jìn)行周期性高效且節(jié)省內(nèi)存的協(xié)同調(diào)度。為便于理解 Dhelix 的工作原理，可以想象二維空間中的雙螺旋 DNA 結(jié)構(gòu)，其兩條鏈耦合成單一的訓(xùn)練計(jì)算流，同時(shí)處理兩個(gè) Micro Batch。

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

當(dāng)深入觀察 α 鏈路與 β 鏈路的耦合時(shí)，兩條鏈路上下相對(duì)移動(dòng)，交替進(jìn)行 Forward 和 Backward。在每次傳播中，鏈會(huì)重新排列其算子以找到重疊良好的兼容算子（例如，計(jì)算和通信的重疊），創(chuàng)建一個(gè)這些耦合點(diǎn)為錨點(diǎn)的調(diào)度，類似 DNA 鏈中的堿基對(duì)。DHelix 利用離線分析結(jié)果測(cè)量算子間重疊兼容性，并采用動(dòng)態(tài)規(guī)劃搜索耦合鏈的高效協(xié)同調(diào)度。

因此，DHelix 的 SI 設(shè)計(jì)通過(guò)使訓(xùn)練路徑能夠同時(shí)容納兩個(gè)相鄰 Micro Batch，有效隱藏了 LLM 訓(xùn)練關(guān)鍵路徑中的通信開銷，顯著提升整體性能。同時(shí)，SI 在現(xiàn)有并行級(jí)別之下運(yùn)行，可以無(wú)縫集成于 TP、SP、CP 和 EP。

4.2 模型折疊

這里，作者具體介紹了其模型折疊（Folding）技術(shù)。這一關(guān)鍵的 DHelix 技術(shù)使得 PP 得以實(shí)現(xiàn)，具體而言，作者將模型層的原始線性排布在 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(xiàn)orward 和 Backward 在兩個(gè) GPU 上的流動(dòng)方向完全一致。如下圖 Figure 7 左側(cè)所示，熟悉的 1F1B 的 V 形結(jié)構(gòu)變成了 W 形，其中 Forward 和 Backward 各形成一個(gè)相同的 V，以此可以實(shí)現(xiàn)  α 鏈路與 β 鏈路在跨 GPU 時(shí)的完美重疊，以及 Forward 和 Backward 在各 GPU 上的協(xié)同調(diào)度。

相較于之前的模型復(fù)制方案（如 Figure 6 ），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)。

如上 Figure 7 簡(jiǎn)化了雙鏈交錯(cuò)傳播調(diào)度的時(shí)間表示，藍(lán)色和綠色塊具有相同的時(shí)間寬度。眾所周知，Backward 更慢，大約是 Forward 2 倍，因此綠色塊的寬度幾乎是藍(lán)色的 2 倍。當(dāng) SI 在同一 GPU 上協(xié)同調(diào)度 α 鏈路與 β 鏈路時(shí)，α 鏈路的 Backward 內(nèi)存釋放是否能及時(shí)滿足 β 鏈路的 Forward 內(nèi)存需求，從而引發(fā)兩條鏈路之間的空間依賴問(wèn)題。實(shí)際上，作者發(fā)現(xiàn)這并不會(huì)引發(fā)問(wèn)題，至多需要額外數(shù)百 MB 的閑置內(nèi)存，因?yàn)樽髡呤窃趯蛹?jí)別而非整個(gè)模型層面重疊 Forward 和 Backward。

為了提供更全面的展示，如下圖 Figure 8 所示，作者展示了包含 12 個(gè) Micro Batch 的完整 W 形流水線調(diào)度。它與經(jīng)典的 1F1B 流水線調(diào)度一樣，經(jīng)歷了預(yù)熱（Warmup），穩(wěn)定（Steady）和冷卻（Cooldown） 3 個(gè)階段。

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

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

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

4.3 算子配對(duì)下的鏈耦合

現(xiàn)在，可以進(jìn)一步聚焦到 Figure 7 中兩個(gè)鏈（α 鏈路與 β 鏈路）耦合的 Forward + Backward 過(guò)程。盡管 U 形模型折疊機(jī)制使得這兩個(gè)鏈的協(xié)同執(zhí)行成為可能，但接下來(lái)進(jìn)行的算子級(jí)重疊則為 DHelix 帶來(lái)了顯著的性能提升。

這種提升源自于兩個(gè)鏈的協(xié)同執(zhí)行，盡可能的利用硬件并行性。如前所示，當(dāng)前 LLM 訓(xùn)練流程中的各個(gè)算子，無(wú)論是計(jì)算還是通信，都已得到深度優(yōu)化，使得這兩種操作能夠重疊。直接結(jié)果是，由于數(shù)據(jù)依賴性迫使其順序調(diào)度，當(dāng)前的算子幾乎沒(méi)有機(jī)會(huì)在每個(gè)鏈中進(jìn)一步折疊。

然而，由于 α 鏈路與 β 鏈路之間沒(méi)有數(shù)據(jù)依賴性，SI 也就為 α 鏈路的通信操作與 β 鏈路的計(jì)算操作，以及反之的計(jì)算和通信操作之間的交叉重疊開辟了新的可能性。除了交錯(cuò)兩個(gè)鏈的 Forward 和 Backward 所帶來(lái)的內(nèi)存收益外，還存在兩個(gè)額外的性能優(yōu)勢(shì)：

• Forward 和 Backward 具有不同的算子布局，為計(jì)算算子和通信算子的系統(tǒng)調(diào)度留出了更多空間（相較于 Forward 與 Forward，以及 Backward 與 Backward 的交錯(cuò)）。
• 在某些情況下，Backward 往往是計(jì)算密集型的，這提供了更多機(jī)會(huì)來(lái)隱藏 Forward 中較高比例的通信操作。

Dhelix 并不是簡(jiǎn)單地釋放兩個(gè) Micro Batch 并讓 GPU 盡力進(jìn)行協(xié)同調(diào)度（如 Wavelet 那樣），而是精心采用系統(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í)行。?

算子基礎(chǔ)特性：再次借用 DNA 結(jié)構(gòu)類比，其中算子“配對(duì)”與互補(bǔ)鏈上的相應(yīng)伙伴，形成“結(jié)合（Bond）”，在此情景下，通過(guò)利用異構(gòu) GPU 資源實(shí)現(xiàn)協(xié)同執(zhí)行，從而顯著提升性能。與 DNA 鏈中 4 種核苷酸基不同，從概念上可將 Transformer 算子歸為兩種類型：計(jì)算和通信，如下圖 Table 2 所示。

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

算子排序預(yù)切分：Dhelix 無(wú)法控制單個(gè)算子的 CUDA 調(diào)度。然而，它可以通過(guò)在 α 鏈路與 β 鏈路的執(zhí)行過(guò)程中策略性地設(shè)置 Barrier 來(lái)強(qiáng)制它們同步。這樣，通信算子可以與另一個(gè)鏈中的對(duì)等算子協(xié)同執(zhí)行，以期最大化被計(jì)算隱藏的機(jī)會(huì)。換言之，單鏈的線性算子執(zhí)行順序被劃分為如干個(gè) Segment，這些 Segment 通過(guò) DHelix 注入的 Barrier 在鏈間配對(duì)。DHelix 通過(guò)構(gòu)建 Forward 和 Backward 中單層計(jì)算的 DAG 開始鏈配對(duì)過(guò)程，如啟用 TP/CP，則一個(gè) Transformer 層由 14 或 18 個(gè)算子組成。如上圖 Figure 9 所示，基于這些 DAG，DHelix 首先通過(guò)枚舉其拓?fù)漤樞蛏?Forward/Backward 算子序列。

自然地，下一步是將每對(duì)候選序列劃分為算子 Segment，并在兩序列間協(xié)同調(diào)度這些 Segment。給定一堆候選 Forward/Backward 序列，有效地放置跨鏈 Barrier 即生成一個(gè)候選配對(duì)規(guī)劃。

這兩類 DAG 在當(dāng)前的 Transformer 工作流程中較為簡(jiǎn)單，包含少數(shù)可在結(jié)果序列中移動(dòng)的算子。這些算子包含反向權(quán)重梯度計(jì)算（wgrad，執(zhí)行 GEMM 操作）以及在 MoE 中的路由計(jì)算。監(jiān)管如此，候選序列的數(shù)量及其切分也導(dǎo)致搜索空間過(guò)于龐大，難以通過(guò)人工方式進(jìn)行搜索。幸運(yùn)的是，尋找最優(yōu)配對(duì)方案的問(wèn)題可以形式化為一個(gè)動(dòng)態(tài)規(guī)劃問(wèn)題。

動(dòng)態(tài)規(guī)劃配對(duì)法：給定一對(duì) Forward 和 Backward 候選算子序列，Sf 和 Sb，各自被切分為 Nf 和 Nb 個(gè) Segment。Dhelix 尋求產(chǎn)生最短執(zhí)行時(shí)間跨度 Topt(Nf,Nb) 的最優(yōu)算子配對(duì)方案。

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

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

4.4 討論

兼容性：從之前提供的系統(tǒng)設(shè)計(jì)描述中可以看出，DHelix 的 SI 方案具有很高的通用性，不依賴于實(shí)際的分布式訓(xùn)練框架及底層硬件。它可以通過(guò)重復(fù)執(zhí)行離線分析過(guò)程，以適配新的訓(xùn)練框架或硬件平臺(tái)。如下圖 Figure 10 所示，展示了同一訓(xùn)練工作流在兩個(gè)不同平臺(tái)（A40 40GB 集群）和 （A800 80GB 集群）上的搜索結(jié)果，顯示出了兩種截然不同的 SI 規(guī)劃。需要說(shuō)明的是，當(dāng)分布式訓(xùn)練參數(shù)（如 TP、PP、SP 等）改變時(shí)，即使在同一硬件上，配對(duì)調(diào)度也會(huì)有所不同，因?yàn)橛?jì)算和通信算子的權(quán)重會(huì)發(fā)生變化。

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

4.5 實(shí)現(xiàn)細(xì)節(jié)

作者基于 NVIDIA 的 Megatron-LM 實(shí)現(xiàn)了 DHelix，涉及 5000 行 Python 代碼。Megatron-LM 原有的 Transformer 實(shí)現(xiàn)包含 3 個(gè)模塊：預(yù)處理、Transformer 和后處理。在 DHelix 中，作者將 Transformer Block 替換為本文中的 SI-enabled Transformer Block，以便與 Megatron-LM 的預(yù)處理和后處理模塊結(jié)合，實(shí)現(xiàn)雙鏈執(zhí)行。所有 DHelix 組件均基于 torch.nn.Module 構(gòu)建，使用戶能夠利用標(biāo)準(zhǔn) PyTorch API 創(chuàng)建自定義 Transformer 模型，并享受 SI 帶來(lái)的優(yōu)勢(shì)，且無(wú)需對(duì)用戶級(jí)代碼進(jìn)行任何修改。

在運(yùn)行時(shí)，DHelix 根據(jù)生成的配對(duì)計(jì)劃依次處理算子對(duì)。通過(guò)在不同的 Segment 之間使用 torch.cuda.synchronize（也就是 barrier），確保這種順序執(zhí)行語(yǔ)義。在單個(gè) Segment 的執(zhí)行過(guò)程中，算子進(jìn)一步分為 3 類：計(jì)算、節(jié)點(diǎn)內(nèi)通信和節(jié)點(diǎn)間通信。作者啟動(dòng)了 3 個(gè) CUDA Stream，并將每類算子分配到專屬的 Stream 進(jìn)行處理。此外，PyTorch 允許不同的 CUDA Stream 獨(dú)立分配和釋放內(nèi)存，這可能會(huì)導(dǎo)致內(nèi)存碎片或內(nèi)存不足錯(cuò)誤。為了解決此問(wèn)題，DHelix 僅使用默認(rèn) CUDA Stream 進(jìn)行所有內(nèi)存分配和釋放操作。

此外，作者也遵循現(xiàn)有實(shí)踐，通過(guò)調(diào)整 NCCL 環(huán)境變量 NCCL_NTHREADS 和 NCCL_MAX_NCHANNELS 來(lái)緩解計(jì)算與通信內(nèi)核之間的競(jìng)爭(zhēng)（Liger: Interleaving Intra- and Inter-Operator Parallelism for Distributed Large Model Inference [9]）。

五、實(shí)驗(yàn)&評(píng)估

5.1 實(shí)驗(yàn)設(shè)置

實(shí)驗(yàn)平臺(tái)：包括 3 中 NVIDIA GPU 集群：

• 8 臺(tái) A40（48GB） 機(jī)器的集群，機(jī)器間通過(guò) 100 Gbps 的 IB 網(wǎng)絡(luò)連接，每臺(tái)機(jī)器 8 張 A40 PCIe GPU，PCIe 4 帶寬為 32 GB/s。
• 8 臺(tái) A800（80GB） 機(jī)器的集群，機(jī)器間通過(guò) 4x200 Gbps IB 網(wǎng)絡(luò)連接，機(jī)器內(nèi)通過(guò) NVLink（400 GB/s） + NVSwitch 互聯(lián)。
• 4 臺(tái) H100（80GB） 機(jī)器的集群，機(jī)間通過(guò) 8x400 Gbps IB 網(wǎng)絡(luò)連接，每臺(tái)集群 8 張 H100 GPU，機(jī)器內(nèi)通過(guò) NVLink（900 GB/s） + NVSwitch 互聯(lián)。

CUDA 版本為 12.2，由于高端 GPU 資源申請(qǐng)難度大，A800/H100 集群僅在極短時(shí)間內(nèi)可供使用。因此在所有平臺(tái)重復(fù)進(jìn)行了整體性能測(cè)試，其余則在 A40 集群完成。

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

在具體實(shí)驗(yàn)中，作者根據(jù)原始模型結(jié)構(gòu)調(diào)整模型層數(shù)，以適應(yīng)集群總的 GPU 內(nèi)存容量?？紤]到 PP 維度擴(kuò)展具有相對(duì)較小的 PP 通信量，性能結(jié)果預(yù)計(jì)與更改模型規(guī)模時(shí)非常接近。

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

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

指標(biāo)：依照之前的慣例，對(duì)于 LLaMA 和 GPT 系列模型，主要通過(guò)單個(gè) GPU 上作業(yè)的總浮點(diǎn)數(shù)除以 End2End 時(shí)間來(lái)代表每個(gè) GPU 的訓(xùn)練性能（TFLOPS/GPU）；對(duì)于 Phi MoE 模型，主要使用 Tokens/s 來(lái)衡量，以指示生成速度。最后，因?yàn)?DHelix 并未改變分布式訓(xùn)練的語(yǔ)義，因此不影響模型的收斂性和準(zhǔn)確性。

5.2 整體性能：A40 集群

作者通過(guò) 3 項(xiàng)任務(wù)評(píng)估 DHelix 的性能：Dense 模型訓(xùn)練、長(zhǎng)序列及上下文并行（CP）的 Dense 模型訓(xùn)練、MoE 模型訓(xùn)練。對(duì)于 LLaMA/GPT 模型，Global Batch 對(duì)應(yīng) 800 萬(wàn) Token，也就是 Sequence Length 為 8K 時(shí) Global Batch Size 為 1K。而 Micro Batch Size 則會(huì)相應(yīng)調(diào)整以充分利用 GPU 內(nèi)存。同樣，Phi 模型的 Global Batch Size 為 1600。

5.2.1 正常序列長(zhǎng)度的 Dense 模型

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

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

5.2.2 長(zhǎng)序列的 Dense 模型

考慮到長(zhǎng)序列的需求越來(lái)越多，作者進(jìn)一步評(píng)估了長(zhǎng)序列訓(xùn)練性能，將序列長(zhǎng)度從 8192 翻倍到 16384，并加入序列并行（CP），需要說(shuō)明的是，加入 CP 會(huì)導(dǎo)致 DP 維度的降低。

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

5.2.3 MoE 模型

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

5.3 整體性能：高性能集群

為了評(píng)估 DHelix 性能的普適性，作者進(jìn)一步在 A800 集群進(jìn)行評(píng)估。需要說(shuō)明的是，作者這里忽略了 Intra-batch 實(shí)驗(yàn)，因?yàn)?A800 機(jī)內(nèi)有高速 NVLink 互聯(lián)，Intra-batch 帶來(lái)的提升幾乎可以忽略。

5.3.1 長(zhǎng)序列 Dense 模型

進(jìn)一步將 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）的吞吐量。

5.3.2 更大的 MoE 模型

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

5.3.3 增加 TP 組規(guī)模

大規(guī)模訓(xùn)練中 TP 通常限制在 8 以內(nèi)（單節(jié)點(diǎn)內(nèi)）。然而隨著模型規(guī)模擴(kuò)大和網(wǎng)絡(luò)帶寬提升，跨節(jié)點(diǎn) TP 逐漸受到關(guān)注。為此，作者在 4 節(jié)點(diǎn) H100 集群進(jìn)行了短暫測(cè)試，以評(píng)估 DHelix 在解鎖跨節(jié)點(diǎn) TP 方面的潛力。

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

為了評(píng)估 DHelix 在 H100 上的算子重疊效果，作者深入分析了使用跨節(jié)點(diǎn) TP 進(jìn)行訓(xùn)練時(shí)的通信開銷。盡管已經(jīng)隱藏了 50%-70% 的可見通信成本，但仍有進(jìn)一步優(yōu)化的空間，如下圖 Table 7 展示了限制這種重疊的 2 個(gè)重要因素：

• 內(nèi)核降速：在多個(gè) CUDA Stream 上同時(shí)運(yùn)行計(jì)算和通信 Kernel 會(huì)導(dǎo)致性能下降 20%-30%，限制了潛在的性能提升。
• 啟動(dòng)間隔：每個(gè)配對(duì) Segment 開始時(shí)計(jì)算和通信 Kernel 啟動(dòng)之間的小間隔會(huì)產(chǎn)生相當(dāng)于通信成本 10%-20% 的開銷。

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

5.4 性能提升統(tǒng)計(jì)分析

這個(gè)部分，作者分析了 DHelix 相對(duì)于 Megatron-LM 基線在性能提升方面的各個(gè)來(lái)源，這些提升主要來(lái)源于 DHelix 引入的各種通信 Overlap 策略。作者在 A40 集群上使用 LLaMA-39B 模型進(jìn)行了兩組統(tǒng)計(jì)分析實(shí)驗(yàn)：

• a：CP 并行度為 2，序列長(zhǎng)度為 16384。
• b：CP 并行度為 4，序列長(zhǎng)度為 32768。

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

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

CP Overlap：進(jìn)一步啟用 CP Overlap，進(jìn)一步隱藏了由 CP 引起的 Send/Recv 開銷，分別使 TP Overlap 后的性能再提升 22.5% 和 14.4%。這一步 DHelix 帶來(lái)了最顯著的性能提升。因?yàn)樵?a 和 b 中，CP 引起的層內(nèi)通信分別占到總訓(xùn)練時(shí)間的 28% 和 53%。同樣，過(guò)多的通信也降低了 b 中 CP Overlap 的收益

通信與通信 Overlap：又格外新增了節(jié)點(diǎn)內(nèi)與節(jié)點(diǎn)間通信的重疊，在 b 中帶來(lái)了 7% 的提升，而在 a 中，大部分通信已經(jīng)被重疊，因此沒(méi)什么收益。

權(quán)重梯度 Overlap：最后，進(jìn)一步啟用權(quán)重梯度（Wgrad）算子在 DAG 排序允許的范圍內(nèi)移動(dòng)，展示了算子重排的影響。在 a 上進(jìn)一步提升 5%，有助于減少數(shù)據(jù)依賴性導(dǎo)致的 “Overlap Bubble”。

5.5 SI 內(nèi)存空間利用率

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

通過(guò)逐步增加模型層數(shù)，以確定系統(tǒng)支持的最大模型參數(shù)規(guī)模。實(shí)驗(yàn)結(jié)果表明，DHelix 支持的最大模型尺寸為 39B，非常接近 Megatron-LM 的 40B，僅減少了 2.5% 的模型尺寸，卻換來(lái)了 40% 的訓(xùn)練吞吐提升（PS：這里應(yīng)該是 A40 GPU）。相比執(zhí)行，Wavelet 最大只支持 32B 參數(shù)的模型。

六、參考鏈接

• ??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???

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