万字干货！手把手教你如何训练超大规模集群下的大语言模型

分布式训练的主要难点

简单介绍一下混合并行中经典的三种并行方案。首先是数据并行，简称 DP。正如其名，数据并行是将数据分割到不同的计算设备上，然后由这些设备完成各自的计算任务。第二种是张量并行，简称 TP。张量并行是将模型中某些层的参数分散到不同的设备上，每个设备负责完成部分的计算工作。第三种是流水并行，简称 PP。流水并行是将模型的不同层切分到不同的计算设备上，类似于流水线的工作方式，各个设备协同完成整个模型的计算过程。

大模型训练在超大规模集群下

的挑战与解决方案

随着模型规模和集群规模的扩大，通信在训练过程中的占比越来越大。为了更直观地展示这一现象，我提供了两张时间线图，它们没有应用计算通信重叠技术。第一张图突出显示了在实现 DP 重叠前的数据并行通信状态，第二张图则突出显示了在实现 TP 重叠前的张量并行通信情况。

从图中我们可以看到，在端到端的训练过程中，DP 的通信占比实际上超过了 15%，而 TP 的通信时间占比也超过了 30%。因此，减少通信对训练的影响，对提升训练效率至关重要。

DP Overlap

DP overlap 的方案在理论上看起来非常吸引人，但实际应用中，我们真的能显著提升训练效率吗？在进行 DP overlap 优化时，我们遇到了三个主要问题。首先，是通信和计算资源之间的竞争问题。当通信和计算操作同时进行时，它们会争夺有限的硬件资源，这可能会影响整体的系统性能。其次，在混合并行场景下，DP overlap 还可能带来 PP bubble 的问题。第三，不同通信资源的争抢还可能导致网络拥塞。

DP overlap 引入的 PP bubble 问题。在前面，我们讨论了通信对计算效率的影响。如果我们模仿 ZeRO 的调度策略，由于 overlap 计算的时间会长于 none overlap 计算的时间，这种负载不均衡会导致 PP bubble 的产生。即图中的 Micro batch 2 的前向传播和 Micro batch 1 的反向传播较长的现象，这展示了负载不均的情况。我们提出的解决方案是通信时机的纵向对齐，这样可以极大地缓解 PP bubble 的问题。同时需要强调的是，从计算 overlap 部分移出来的通信都被放在了 PP bubble 上，因此它不会产生任何额外的影响。这种策略有助于平衡负载，减少因通信和计算不匹配而产生的效率损失。

下图展示了我们最终优化后的 timeline。在这个优化版本中，我们实现了 reduce-scatter 与反向传播的 overlap，同时 all-gather 操作与前向传播也实现了 overlap。此外，我们通过分桶通信、网络预测控制、通信 CHANNEL 调优以及通信时机的纵向对齐等方法，大幅优化了 DP 的通信开销。这些优化措施共同作用，提高了整体的训练效率，减少了因通信而产生的延迟和资源浪费。

TP Overlap

然后我们可以将这种策略推广到四个 rank 的场景中。为了简化表述，我们将计算的 stream 都合并到了一起。对于 all-gather overlap GEMM，我们会特别关注第一个 rank。第一步，我们使用自己持有的那部分输入来进行计算，同时将自己持有的内容也发送给其他 rank，并接收其他 rank 中持有的那部分输入。接下来的第二步、第三步、第四步都是按照相同的原理进行。通过这种方式，我们就可以得到一个 all-gather 的 overlap 流程。这样，每个 rank 都在进行本地计算的同时，与其他 rank 进行数据交换，实现了计算与通信的重叠。这种策略可以有效地减少等待时间，提高资源利用率，从而提升整体的并行计算效率。

Reduce scatter 的操作也是类似的。我们可以首先关注 rank 4 在整个计算结果流程中的作用。在第一步中，rank 4 的计算结果被放置在 rank 1 上。rank 1 完成自己的计算后，在第二步中，它会将这个结果发送给 rank 2。rank 2 在接收到来自 rank 1 的结果后，会将其与自己的计算结果进行累加，然后继续进行下一步的计算。接着，在第三步和第四步中，流程与前两步相同。rank 3 和 rank 4 也会按照这个顺序接收之前 rank 传递的结果，并与自己的计算结果进行累加。最终，在流程的最后， rank 4 将拿到汇总后的最终结果。

通过上述步骤，我们得到了一个完整的解决方案，适用于处理通信和计算存在依赖关系时的通信计算重叠问题。

这是 TP overlap 的整体解决方案，对于计算通信没有依赖的情况，这里是指 column-wise linear 的反向传播。由于这部分操作没有数据依赖关系，我们采用了 bulk overlap 技术。对于其余的通信和计算，因为它们之间存在依赖关系，我们采用了 split pipeline overlap 的方法。

下图展示了实现 TP overlap 后的 timeline，我们可以看到 TP 的通信和计算重叠在了一起。同时，我们进行了两项优化措施：第一项是使用了 peer-to-peer memory copy，以此来减轻通信对 SM 的消耗。第二项优化是将计算分散到不同的 stream 上，这样计算也可以实现部分的重叠。

超长文本场景解决方案

针对 TP 作为通信换显存的两大弊端——在 h 维度上切分导致的不可扩展性以及方案本身的通信量大，我们希望找到一种在 s 维度上可以切分并且通信量相比 TP 小一些的方案。为此，我们实现了上下文并行（context parallel，简称 CP）。

在 CP 场景下，整个模型的 activation 从始至终都在 s 维度上保持着切分状态。之前无法解决的问题，通过 CP=4 就可以解决。我们可以计算这个方案的通信开销，CP 引入的通信开销仅有 KV 前向时的 all-gather 和反向时的 all-gather 以及 reduce-scatter。同时，我们改变了 QKV 的计算顺序，使得 K 的通信可以与 V 的计算重叠，V 的计算可以与 Q 的计算重叠。因此，我们可以得出下述两个结论。

CP 的通信量与 KV 的 activation 大小成正比。在混合并行场景下，利用了 TP 可以减少 activation 大小的特点，使得 CP 的通信量相比于直接扩大 TP 可以减少 TP 倍。
由于 CP 的通信可以与计算进行重叠，因此进一步减少了对训练的影响。同时，由于 CP 的切分维度在 s 上，理论上如果有足够的机器，CP 可以解决任意大小的上下文窗口问题。

CP 与其他技术结合时，会带来一些额外的好处和挑战。首先是计算负载均衡问题，这个问题的背景是大语言模型采用了 Decoder Only 架构，并且在 attention 中使用了 causal mask，这导致 CP 会引入计算负载不均的问题。从下面的左图中可以看到，rank 0 的计算负载明显低于 rank 1。

为了解决这个问题，我们采用了类似高斯求和的方法，让每个设备负责一大一小两个 attention 的计算，以此来缓解负载不均的问题。由于同一个设备上的这两个 attention 计算之间不存在依赖关系，为了进一步提升硬件利用率，我们仿照 TP overlap，使用了不同的 CUDA stream 来 launch 两个 kernel。借助 CUDA 的 runtime 调度，我们实现了更高效的并行计算。

结合 CP 还有一些额外的好处。GQA（Grouped Query Attention）是在长上下文场景下几乎必选的技术。与原来的 Multihead attention 相比，GQA 将多个 query 作为一个 group，每个 group 对应一个 K 和 V。可以发现，GQA 可以极大地减少 KV activation 的大小。正如之前提到的，CP 的通信量与 KV 的 activation 大小成正比。因此在 GQA 的场景下，我们可以进一步减少 CP 的通信量，这是结合使用 CP 和 GQA 技术的一个显著优势。

下面是关于计算换显存的方案，其中 recomputing 是一个非常经典的技术。首先，让我们对 recomputing 做一个介绍。下图展示了一个正常的模型训练过程中的数据流。由于反向传播计算对前向传播计算结果存在数据依赖，因此在前向计算完成后，计算结果并不会立即释放，而是要等到反向计算完成后才释放。

右侧的图展示了使用 recomputing 方案的情况。可以看到，在 0 到 3 层的中间结果被释放了，只有 recomputing block 的输入，也就是 layer 0 的 input 被保存了下来。在反向传播过程中，我们会使用保存下来的 input 重新计算 0 到 3 层的前向传播结果，然后再进行反向计算，从而达到节省显存的目的。这个方案从理论上看起来非常理想，但在实际应用中也会遇到一些问题。

首先，主流的框架都采用了 full computing，这导致每次反向计算都会执行一次完整的 forward pass，引入了大量的无效计算。在大模型时代，这种情况是不可接受的。其次，目前的开源框架 Megatron-LM 对 attention 部分实现了 selective recomputing。然而，在 flash attention 时代，这个方案的效率已经不如以前了。

经过观察，我发现某些 kernel，例如 GEMM，其反向计算实际上并不依赖于前向传播的输出结果。例如，对于公式 𝑌=𝑋𝑊，𝑑𝑋 和 𝑑𝑊 的计算并不依赖于前向传播的结果 𝑌。如果我们将这类算子作为 recomputing block 中的最后一个算子，就无需对它们进行重计算。

大家可以看下图右侧。假设层 3 是一个 GEMM 操作，那么 layer 3 的反向计算只依赖于层 3 的输入，而不是层 3 的输出。这样，在重计算时，我们可以省去 layer 3 的前向计算。我将这种重算策略称为 GEMM last recomputing。

我们将 GEMM last recomputing 策略实施到大语言模型的训练中，发现只需要对计算量较小的算子进行重算。相比于没有采用 recomputing 的方案，我们的策略在增加了不到 1.5% 的计算量的情况下，减少了 40% 的显存开销。这是一个在保持计算效率的同时显著减少显存需求的有效方法。

接下来是内存换显存的方案。我最初产生这个想法的原因是，在训练过程中，显存资源已经非常紧张，然而内存资源在训练状态下却几乎处于闲置状态，这为我们提供了一定的操作空间。其次，随着硬件的升级，PCIe 已经升级到第五代，每张卡分配到的 x16 带宽达到了 64GB/s。同时，由于 H2D（Host to Device）和 D2H（Device to Host）是 memory copy 操作，它们对计算的影响几乎可以忽略不计。在混合并行场景下，每次前向计算产生的 activation 并不会立即被使用，而是至少要间隔一个完整的虚拟 pipeline stage 计算，因此混合并行架构也为我们提供了足够的时间窗口。

我们的解决方案是，将每个虚拟 pipeline stage 前一个 micro batch 的 activation H2D 和 D2H 的通信操作与下一个 micro batch 的计算进行 overlap，这样可以极大减少 offload 对关键路径上计算的影响。通过这个 offload 方案，我们能够在几乎不影响计算性能的情况下实现内存换显存的效果，上下文窗口大小提升了 2.5 倍。

接下来展示的是这个解决方案的整体成果。我们在 H800 集群上进行的测试显示，在吞吐量上，与现有的最先进开源方案相比，我们在** 任意上下文窗口下都能实现超过 30% 的性能提升。** 能达到这样的性能提升主要归功于两点原因：

第一，我们采用了通信代价更小的 CP 来替代 TP，从而降低了为解决显存问题而引入的通信开销；
第二，我们采用了 GEMM last recomputing 和 pipeline aware offloading 这两种更具成本效益的显存问题解决方案，减少了以通信换取显存的需求，进一步实现了训练吞吐量的提高。

在支持的序列长度上限方面，首先，我们通过内存换显存、通信换显存、计算换显存的方法，大幅提升了单个设备支持的上下文窗口。同时，由于该方案还具有极强的可扩展性，因此在设备资源充足的情况下，我们可以支持无限大的上下文窗口。

接下来是 cost model（成本模型）的介绍。在进行大模型训练时，参数调整是一个非常痛苦的过程，因为模型有大量的参数，并且这些参数之间相互影响，比如 TP、CP、DP 的大小，以及 offload 的比例，还有网络设置中的 CTS。如果对所有参数都进行实际运行测试，将会消耗大量的计算资源。然而，如果不进行实际运行，仅仅通过比例和一些基于 FLOPs 理论算力的简单折算来预测，会导致预测极其不准确。因此，这样的成本模型是不可行的。

为了解决这个问题，我们对 TP、CP、PP 等一系列可能影响性能的因素进行了细致的建模。我们将所需信息分为与模型相关的信息，比如不同组合下单层前向和反向传播的时间；以及与集群相关的信息，比如跨机器的集群通信带宽或者 H2D 的带宽等。整体的测量耗时可以在一个小时内完成，并且这些信息可以多次复用。

在 175b 的案例中，我们建模的预测值和实测值之间的误差控制在 2% 以内。在实际使用过程中，我们的成本模型的误差与实测值的对比也不超过 5%，其中大部分误差来源于网络的不稳定性。下图右边展示了我们的成本模型给出的参数配置表。通常情况下，搜索完成后，我们可以根据 MFU 的前 5 名进行实际测试，最终得到我们的训练配置。这种方法大大提高了参数调整的效率和准确性。