DHelix：跨 Micro-Batch 的通信隐藏，SOTA LLM 训练性能

一、背景

我们在之前的文章中提到过，在 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闲谈 ​​​​