超大规模集群下大语言模型训练实操指南

超长文本场景解决方案

针对 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 的计算重叠。因此，我们可以得出下述两个结论：

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

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 做一个介绍。下图展示了一个正常的模型训练过程中的数据流。由于反向传播计算对前向传播计算结果存在数据依赖，因此在前向计算完成后，计算结果并不会立即释放，而是要等到反向计算完成后才释放。

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

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

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

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

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

Pipeline Aware Offloading

接下来是内存换显存的方案。我最初产生这个想法的原因是，在训练过程中，显存资源已经非常紧张，然而内存资源在训练状态下却几乎处于闲置状态，这为我们提供了一定的操作空间。其次，随着硬件的升级，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% 的性能提升。能达到这样的性能提升主要归功于两点原因：

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

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

Cost Model 成本模型

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

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

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

未来展望

未来在训练引擎方面我们会专注于五个主要方向：

1. 万亿参数规模的 MoE 模型：我们期望能够训练具有万亿参数的 MoE 模型，这将推动模型容量和性能的显著提升。
2. 继续扩大序列长度：我们希望能够支持达到百万级别的序列长度，这将极大地扩展模型处理长文本数据的能力。
3. RLHF 框架：目前还没有看到非常高效的 RLHF 框架实现，这将是未来研究的一个重要领域。
4. 低精度训练：随着 Hopper 系列架构的推广以及 FP8、FP6 等多精度配置训练，我们将需要关注低精度训练技术的发展。
5. 异构算力的引入：我们需要考虑引入异构算力来增强训练引擎的灵活性和健壮性。