作者/机构: Deepak Narayanan‡★, Mohammad Shoeybi†, Jared Casper†, Patrick LeGresley†, Mostofa Patwary†, Vijay Korthikanti†, Dmitri Vainbrand†, Prethvi Kashinkunti†, Julie Bernauer†, Bryan Catanzaro†, Amar Phanishayee∗, Matei Zaharia‡ (†NVIDIA ‡Stanford University ∗Microsoft Research)

A1 主要贡献

本文旨在解决大规模语言模型训练中的两大挑战：单个GPU内存容量有限无法容纳巨大模型，以及巨大的计算量导致训练时间过长。为应对这些挑战，本文研究了如何高效地组合张量并行、流水线并行和数据并行（PTD-P）以扩展至数千个GPU，同时保持严格的优化器语义。

核心问题与研究目标：
* 核心问题： 单纯使用张量并行或流水线并行在扩展至数千GPU时会遇到瓶颈。张量并行在跨节点通信时效率低下，而流水线并行会因“流水线气泡”（pipeline bubble）导致设备空闲。
* 研究目标： 在给定批量大小（batch size）的情况下，如何组合不同的并行技术（流水线、张量、数据并行），以最大化大型模型的训练吞吐量，并保持严格的优化器语义。

主要创新与贡献：
1. 提出PTD-P并行策略： 本文展示了如何将流水线并行、张量并行和数据并行进行组合，以在数千个GPU上实现高效的大型语言模型训练。该策略利用节点内（intra-server）张量并行、节点间（inter-server）流水线并行，并结合数据并行进行扩展，成功地在3072个A100 GPU上以502 petaFLOP/s的总吞吐量训练了一个万亿参数模型，达到了单个GPU理论峰值性能的52%。
2. 新颖的交错式流水线调度（Interleaved Pipelining Schedule）： 提出了一种新的流水线调度方法，该方法通过将模型的层分成多个块（chunks）并交错分配给流水线中的设备，能够将“流水线气泡”的开销减少一个与块数成正比的因子。实验表明，这种调度策略在内存占用相当的情况下，可将吞吐量提升超过10%。
3. 分布式训练的配置原则： 基于实证和分析研究，本文提出了配置分布式训练的指导原则：
* 张量模型并行在单个多GPU服务器内（利用高速NVLink）非常有效，但对于跨服务器扩展，必须使用流水线模型并行以避免高昂的All-Reduce通信开销。
* 流水线并行的调度方式、微批次大小（microbatch size）等超参数对内存占用、计算效率和流水线气泡大小有显著影响，需要根据具体模型和配置进行优化。
* 大规模训练是通信密集型的，优化的通信策略（如本文提出的Scatter/Gather优化）和高速硬件对于实现良好扩展性至关重要。
4. 性能优势与实践价值：
* 实验证明，PTD-P方法在训练1750亿和5300亿参数模型时，性能比ZeRO-3【索引36，ZeRO: Memory Optimization Towards Training A Trillion Parameter Models，2019，arXiv】高出70%，主要得益于更少的跨节点通信。
* 实现的万亿参数模型训练吞吐量使得端到端训练时间变得现实（估计约3个月），而过去的方法无法训练如此大的模型或需要不切实际的训练时长。
* 本文的软件实现已开源（Megatron-LM），为社区训练大规模NLP模型提供了高效工具。

图1：最先进的自然语言处理（NLP）模型大小随时间变化的趋势。训练这些模型的浮点运算次数正以指数速率增长。

A3 背景知识/关键Observation/设计原则

2.1 数据并行

数据并行 在数据并行【索引25，PyTorch Distributed: Experiences on Accelerating Data Parallel Training，2020，arXiv；索引43，Petuum: A New Platform for Distributed Machine Learning on Big Data，2015，IEEE Transactions on Big Data】中，每个工作节点（worker）都拥有模型的完整副本，输入数据集被分片。工作节点会周期性地聚合它们的梯度，以确保所有节点看到的权重版本保持一致。对于无法在单个工作节点上容纳的大型模型，数据并行可以应用于较小的模型分片上。

2.2 流水线模型并行

流水线并行 在流水线并行中，模型的各层被分片到多个设备上。当应用于具有重复Transformer块的模型时，可以为每个设备分配相等数量的Transformer层。本文不考虑更不对称的模型架构，因为在这种架构中，将层分配到流水线阶段更为困难；我们将这个问题留给相关工作【索引22，Beyond Data and Model Parallelism for Deep Neural Networks，2018，MLSys；索引29，PipeDream: Generalized Pipeline Parallelism for DNN Training，2019，SOSP；索引41，Efficient Algorithms for Device Placement of DNN Graph Operators，2020，NeurIPS】来解决。一个批次（batch）被分割成更小的微批次（microbatches），执行过程在这些微批次上进行流水线化。流水线方案需要确保输入在正向和反向传播过程中看到一致的权重版本，以实现明确的同步权重更新语义。具体来说，简单的流水线化可能导致输入在反向传播时看到正向传播时未见的权重更新。

流水线气泡问题 为了精确保留严格的优化器语义，我们引入了周期性的流水线刷新（pipeline flushes），以便在设备之间同步优化器步骤。在每个批次的开始和结束时，设备会处于空闲状态。我们将这种空闲时间称为“流水线气泡”（pipeline bubble），并希望使其尽可能小。像PipeMare、PipeDream和PipeDream-2BW【索引23，Pipelined Backpropagation at Scale: Training Large Models without Batches，2021，MLSys；索引29，PipeDream: Generalized Pipeline Parallelism for DNN Training，2019，SOSP；索引30，Memory-Efficient Pipeline-Parallel DNN Training，2021，ICML；索引45，PipeMare: Asynchronous Pipeline Parallel DNN Training，2021，MLSys】这样的异步和有界陈旧性方法完全取消了刷新，但放宽了权重更新的语义。我们将这类方案的考虑留给未来的工作。

调度方式 跨设备调度正向和反向微批次有多种可能的方式；每种方法在流水线气泡大小、通信量和内存占用之间提供了不同的权衡。本节我们讨论两种这样的方法。

图2：本文工作中用于基于Transformer的模型的张量和流水线模型并行（MP）组合。

2.2.1 默认调度

GPipe调度 GPipe【索引20，GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism，2019，NeurIPS】提出了一种调度方案，即首先执行一个批次中所有微批次的正向传播，然后执行所有微批次的反向传播（如图3所示）。我们可以量化GPipe的流水线气泡大小（$T_{bubble}$）。我们用$m$表示一个批次中的微批次数，用$p$表示流水线阶段数（用于流水线并行的设备数），用$T_{ideal}$表示理想的每次迭代时间（假设完美或理想扩展），用$T_{fwd}$和$T_{bwd}$表示执行单个微批次正向和反向传播的时间。在此调度中，流水线气泡由批次开始时的$p-1$个正向传播和结束时的$p-1$个反向传播组成。因此，在流水线气泡中花费的总时间为$T_{bubble} = (p-1) \cdot (T_{fwd} + T_{bwd})$。该批次的理想处理时间为$T_{ideal} = m \cdot (T_{fwd} + T_{bwd})$。因此，在流水线气泡中花费的理想计算时间的分数为：气泡时间分数（流水线气泡大小）$= \frac{T_{bubble}}{T_{ideal}} = \frac{p-1}{m}$。为了使气泡时间分数较小，我们需要$m \gg p$。然而，对于如此大的$m$，这种方法的内存占用很高，因为它需要在整个训练迭代生命周期内，为所有$m$个微批次保留存储的中间激活值（或在使用激活重计算时，仅保留每个流水线阶段的输入激活值）。

图3：GPipe流水线调度，显示了所有微批次（由数字表示）的正向传播（蓝色）紧接着反向传播（绿色）。灰色区域代表流水线气泡。为简化起见，我们假设反向传播时间是正向传播的两倍。流水线调度的效率不依赖于这个因素。此示例中每个批次包含8个微批次，每个蓝色或绿色框中的数字是对应微批次的唯一标识符（特别是，第一个批次包含微批次1-8，第二个批次包含微批次9-16，依此类推）。优化器在流水线刷新时进行步骤更新和权重参数更新，以确保严格的优化器语义，这导致设备空闲和流水线气泡。

PipeDream-Flush调度 作为替代，我们使用PipeDream-Flush调度【索引30，Memory-Efficient Pipeline-Parallel DNN Training，2021，ICML】。在此调度中，我们首先进入一个预热阶段，其中工作节点执行不同数量的正向传播，如图4（上图）所示。这种调度将正在处理的微批次数（即反向传播尚未完成且需要维持激活值的微批次数）限制在流水线深度之内，而不是批次中的微批次数。预热阶段后，每个工作节点进入稳定状态，执行一次正向传播后跟一次反向传播（简称1F1B）。最后，在一个批次结束时，我们完成所有剩余正在处理的微批次的反向传播。新调度的气泡时间与之前相同，但对于PipeDream-Flush调度，未完成的正向传播数量最多为流水线阶段数$p$。因此，此调度需要存储的激活值最多对应$p$个微批次（而GPipe调度为$m$个微批次）。因此，当$m \gg p$时，PipeDream-Flush比GPipe的内存效率高得多。

2.2.2 带有交错阶段的调度

交错调度设计 为了减小流水线气泡的大小，每个设备可以为多个层的子集（称为模型块，model chunk）执行计算，而不是单个连续的层集合。例如，如果之前每个设备有4个层（即设备1有1-4层，设备2有5-8层，依此类推），我们可以让每个设备为两个模型块（每个块有2层）执行计算，即设备1有1、2、9、10层；设备2有3、4、11、12层；依此类推。通过这种方案，流水线中的每个设备被分配了多个流水线阶段（每个流水线阶段的计算量比之前少）。

交错调度的内存与性能 像之前一样，我们可以使用这种调度的“全正向、全反向”版本，但其内存占用很高（与$m$成正比）。因此，我们开发了一种交错式调度，它改进了之前内存高效的1F1B调度。这个新调度如图4所示，并要求一个批次中的微批次数是流水线并行度（流水线中的设备数）的整数倍。例如，使用4个设备时，一个批次中的微批次数必须是4的倍数。

图4：默认和交错式1F1B流水线调度。上图显示了默认的非交错式1F1B调度。下图显示了交错式1F1B调度，其中每个设备被分配了多个块（本例中为2个）。深色表示第一个块，浅色表示第二个块。流水线气泡的大小更小（在交错式时间线中，流水线刷新发生得更早）。

交错调度的气泡缩减与通信开销 如图4所示，对于相同的批次大小，新调度中的流水线刷新发生得更早。如果每个设备有$v$个阶段（或模型块），那么每个阶段或块的微批次正向和反向时间将变为$T_{fwd}/v$和$T_{bwd}/v$。因此，流水线气泡时间减少到$T_{bubble-int} = (p-1) \cdot \frac{(T_{fwd} + T_{bwd})}{v}$，气泡时间分数则为：气泡时间分数（流水线气泡大小）$= \frac{T_{bubble-int}}{T_{ideal}} = \frac{1}{v} \frac{p-1}{m}$。这意味着新调度将气泡时间减少了$v$倍。然而，这种减小的流水线气泡大小并非没有代价：此调度需要额外的通信。从数量上讲，通信量也增加了$v$倍。在下一节中，我们将讨论如何利用多GPU服务器（例如DGX A100节点）中的8个InfiniBand网卡来减少这种额外通信的影响。

2.3 张量模型并行

张量并行策略 通过张量模型并行，模型的各个层被划分到多个设备上。在本文中，我们采用了Megatron【索引40，Megatron-LM: Training Multi-Billion Parameter Language Models using GPU Model Parallelism，2019，arXiv】中用于Transformer层的特定划分策略，这是语言模型的基础。我们也可以将类似的思想应用于其他类型的模型，如CNN。下面我们简要概述这一策略，如图5所示。一个Transformer层由一个自注意力块和一个两层的多层感知机（MLP）组成。关于Transformer层的更多细节可以在Vaswani等人的论文【索引42，Attention is All You Need，2017，arXiv】中找到。

MLP块的并行化 MLP块包含两个通用矩阵乘法（GEMM）和一个GeLU非线性激活：$Y = \text{GeLU}(XA)B$。我们可以将矩阵$A$按列切分$A = [A_1, A_2]$。这种划分允许GeLU非线性激活独立地应用于每个划分后的GEMM的输出：$Y = \text{GeLU}([XA_1, XA_2])B$。这样做是有利的，因为它避免了同步的需求（如果$A$是按行切分则需要同步，因为GeLU是非线性的）。然后，第二个权重矩阵$B$可以按行切分，以消除GEMM之间的任何通信需求（如图5a所示），如下所示：$Y = \left[ \text{GeLU}(XA_1), \text{GeLU}(XA_2) \right] \begin{bmatrix} B_1 \\ B_2 \end{bmatrix} = \text{GeLU}(XA_1)B_1 + \text{GeLU}(XA_2)B_2$。第二个GEMM的输出在进入dropout层之前，会在所有GPU上进行reduce操作。

自注意力块的并行化 我们利用多头注意力操作中固有的并行性来划分自注意力块（如图5b所示）。键（$K$）、查询（$Q$）和值（$V$）矩阵可以按列并行的方式进行划分。然后，输出的线性层可以直接作用于注意力操作的划分后输出（权重矩阵按行划分）。

通信开销 这种方法将MLP和自注意力块中的GEMM划分到不同的GPU上，同时在正向传播中只需要两次all-reduce操作（$f$算子），在反向传播中也需要两次all-reduce操作（$g$算子）。我们用几行代码实现了$f$和$g$。

图5：使用张量模型并行划分的Transformer模型块（图片来自Megatron【索引40】）。$g$和$f$是共轭的。$f$在正向传播中是恒等算子，在反向传播中是all-reduce；而$g$则相反。

A2 方法细节

3 性能分析与并行配置

本节我们考虑将流水线和张量模型并行与数据并行相结合的性能影响。在给定的GPU预算和批次大小下，可以使用不同程度的PTD-P并行类型来训练模型；每个维度都在内存占用、设备利用率和通信量之间展现出权衡。我们在本节余下部分讨论这些权衡，并在第5.4节展示实证结果。我们在相关处为流水线气泡大小提供了分析模型。我们定性地描述了通信时间的行为，并提出了通信量的成本模型；然而，我们没有为通信时间提供直接的成本模型，因为对于一个分层网络拓扑结构（其中同一服务器上GPU间的互连带宽高于服务器间的互连带宽），建模更加困难。据我们所知，这是第一份分析这些并行维度性能相互作用的工作。

3.1 符号表示

本节使用的符号
* $(p, t, d)$: 并行化维度。$p$代表流水线模型并行大小，$t$代表张量模型并行大小，$d$代表数据并行大小。
* $N$: GPU数量。我们要求$p \cdot t \cdot d = N$。
* $b$: 全局批次大小（作为输入提供）。
* $\mu$: 微批次大小。
* $m = \frac{1}{p} \cdot \frac{b}{\mu}$: 每个流水线中的批次所含微批次数。

3.2 张量模型并行与流水线模型并行

并行组合对流水线气泡的影响 张量模型并行和流水线模型并行都可以用来将模型的参数划分到多个GPU上。如前所述，使用带有周期性刷新的流水线并行会导致大小为$(p-1)/m$的流水线气泡。假设$d=1$（数据并行大小为1）；因此，$p \cdot t = N$。以$t$表示的流水线气泡大小为：$Bubble\ size = \frac{p-1}{m} = \frac{N/(t \cdot d) - 1}{m}$。因此，当$N$, $b$, 和$\mu$固定时（$m = b/(\mu \cdot p)$也是固定的，此处原文公式有误，应为$m = b/(\mu \cdot p)$，而$p=N/t$，所以$m = (b \cdot t) / (\mu \cdot N)$），随着$t$的增加，$p$减小，流水线气泡也随之减小。

并行组合对通信量的影响 不同GPU之间的通信量也受到$p$和$t$值的影响。流水线模型并行采用成本较低的点对点通信。而张量模型并行则使用all-reduce通信（在正向和反向传播中各有两次all-reduce操作，见2.3节）。对于流水线并行，每个微批次在每对连续设备之间（无论是正向还是反向传播）需要执行的总通信量为$s \cdot h$，其中$s$是序列长度，$h$是隐藏层大小。对于张量模型并行，每个层中总大小为$s \cdot h$的张量需要在$t$个模型副本之间进行两次all-reduce，这导致每个微批次、每个设备、每个层的总通信量为$8sh \frac{t-1}{t}$。每个设备通常有多个层；因此，每个微批次、每个设备的总张量并行通信量为$l_{\text{stage}} \cdot 8sh \frac{t-1}{t}$，其中$l_{\text{stage}}$是流水线阶段中的层数。

配置原则1 因此，我们发现张量模型并行增加了设备间的通信量。所以，当$t$大于单个节点中的GPU数量时，通过较慢的节点间链路执行张量模型并行的开销可能是不切实际的。我们在5.4节的实验中看到了这些结果。要点 #1：当考虑不同形式的模型并行时，张量模型并行通常应在$k$-GPU服务器上使用到$t=k$的程度，然后可以使用流水线模型并行来跨服务器扩展到更大的模型。

3.3 数据并行与模型并行

数据并行与两种模型并行的交互 我们也想考虑数据并行与两种模型并行之间的交互作用。为简化起见，本节我们独立考虑这些交互。

3.3.1 流水线模型并行

流水线与数据并行对气泡的影响 设$t=1$（张量模型并行大小为1）。每个流水线的微批次数为$m = b/(d \cdot \mu) = b'/d$，其中$b' := b/\mu$。在总GPU数量为$N$的情况下，流水线阶段数为$p = N/(t \cdot d) = N/d$。流水线气泡大小为：$Bubble\ size = \frac{p-1}{m} = \frac{N/d - 1}{b'/d} = \frac{N-d}{b'}$。随着$d$变大，$N-d$变小，因此流水线气泡变小。图6显示了在不同$N$, $d$, 和$b'$值下流水线气泡大小的行为。对于所有模型来说，可能无法将$d$一直增加到$N$，因为模型的完整训练内存占用可能大于单个加速器的内存容量。

图6：由于流水线刷新导致的空闲时间比例（流水线气泡大小）与数据并行大小（d）的关系，针对不同的GPU数量（N）和批次大小与微批次大小的比率（b' = b/μ）。

数据并行通信开销 因此，如果数据并行所需的all-reduce通信不会随着$d$的增加而急剧增加，则总吞吐量会增加，这应该是成立的，因为基于环的实现的通信时间与$\frac{d-1}{d} = 1 - \frac{1}{d}$成比例。我们还可以分析增加批次大小$b$的影响。对于给定的并行配置，随着批次大小$b$的增加，$b' = b/\mu$增加，$(N-d)/b'$减少，从而增加吞吐量。数据并行所需的all-reduce通信也变得不那么频繁，进一步增加了吞吐量。

3.3.2 数据并行与张量模型并行

通信模式对比 对于张量模型并行，每个微批次都需要执行all-reduce通信。这在多GPU服务器之间可能代价高昂。另一方面，数据并行每个批次只需要执行一次昂贵的all-reduce通信。此外，对于张量模型并行，每个模型并行等级（rank）在每个模型层中执行一部分计算，因此对于不够大的层，现代GPU可能无法以峰值效率执行这些子矩阵计算。

配置原则2 要点 #2：当同时使用数据和模型并行时，应使用一个总模型并行大小$p \cdot t$，使得模型的参数和中间元数据能够装入GPU内存；数据并行可用于将训练扩展到更多的GPU。

3.4 微批次大小

微批次大小对吞吐量的影响 微批次大小$\mu$的选择也影响模型训练的吞吐量。例如，我们在图7中看到，在单个GPU上，随着微批次大小的增加，每个GPU的吞吐量最多增加1.3倍。我们现在想确定在给定并行配置$(p, t, d)$和批次大小$b$的情况下，最优的微批次大小$\mu$。无论微批次大小如何，数据并行的通信量将保持不变。给定将微批次大小映射到单个微批次正向和反向计算时间的函数$T_{fwd}(\mu)$和$T_{bwd}(\mu)$，计算一个批次的总时间（忽略通信成本）为（与之前一样，定义$b'$为$b/\mu$）：
$T_{comp} = (p-1) \cdot (T_{fwd}(\mu) + T_{bwd}(\mu)) + m \cdot (T_{fwd}(\mu) + T_{bwd}(\mu)) = (\frac{b'}{d} + p - 1) \cdot (T_{fwd}(\mu) + T_{bwd}(\mu))$

图7：对于一个十亿参数的GPT模型（128个注意力头，隐藏层大小4096，4个Transformer层），每个GPU的吞吐量与微批次大小的关系。

微批次大小的权衡 因此，微批次大小既影响操作的算术强度，也影响流水线气泡的大小（通过影响$m$）。图8显示了对于一个十亿参数的GPT模型和$(p, t) = (8, 8)$配置下的估计吞吐量（使用公式(1)估计处理时间）。对于两种批次大小，最优的$\mu$都是4。

图8：对于与图7相同的GPT模型，归一化估计吞吐量（时间计算为 T = (b'/d + p - 1) · (T_fwd(μ) + T_bwd(μ))）相对于微批次大小μ的行为。

配置原则3 要点 #3：最优微批次大小$\mu$取决于模型的吞吐量和内存占用特性，以及流水线深度$p$、数据并行大小$d$和批次大小$b$。

3.5 激活重计算

激活重计算的原理与作用 激活重计算【索引12，Training Deep Nets with Sublinear Memory Cost，2016，arXiv；索引18，Revolve: An Implementation of Checkpointing for the Reverse or Adjoint Mode of Computational Differentiation，2000，ACM Transactions on Mathematical Software；索引20，GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism，2019，NeurIPS；索引21，Breaking the Memory Wall with Optimal Tensor Rematerialization，2020，MLSys】是一种可选技术，它通过在反向传播之前再次运行正向传播来换取额外的内存占用，从而增加了执行的计算操作数量（只存储给定流水线阶段的输入激活值，而不是整个中间激活值集合，后者要大得多）。为了以可接受的低内存占用训练相当大的模型，使用流水线并行时需要激活重计算。之前的工作如PipeDream-2BW【索引30，Memory-Efficient Pipeline-Parallel DNN Training，2021，ICML】已经研究了激活重计算的性能影响。

最优检查点数量 激活检查点的数量不影响吞吐量，但影响内存占用。设$A_{input}$为一层输入激活值的大小，$A_{intermediate}$为每层中间激活值的大小。如果一个模型阶段有$l$层，并且$c$是检查点的数量，那么总内存占用将是$c \cdot A_{input} + \frac{l}{c} \cdot A_{intermediate}$。该函数的最小值在$c = \sqrt{l \cdot \frac{A_{intermediate}}{A_{input}}}$时取得。在实践中，我们凭经验测量$A_{intermediate}$。在大多数情况下，每1或2个Transformer层设置一个检查点是最佳的。

其他内存优化技术 其他技术如激活分区【索引36，ZeRO: Memory Optimization Towards Training A Trillion Parameter Models，2019，arXiv】也可以与张量模型并行结合使用，以进一步减少因激活值产生的内存占用。

4 实现

实现细节 我们将PTD-P作为Megatron-LM代码库的扩展来实现。我们的实现基于PyTorch【索引32，PyTorch: An Imperative Style, High-Performance Deep Learning Library，2019，NeurIPS】构建。我们使用NCCL【索引7，NVIDIA Collective Communication Library (NCCL)】在设备之间进行通信。为了获得良好性能，我们实现了针对通信和计算的优化，具体如下。

4.1 通信优化

流水线通信并行化 在使用流水线并行时，我们希望在正向和反向传播中并行地发送和接收张量。每个DGX A100都配备了8个InfiniBand（IB）网卡。不幸的是，发送和接收是点对点的，并且只发生在两个服务器上的一对GPU之间，这使得在流水线内的单个通信调用中很难利用所有8个网卡。

Scatter/Gather优化 然而，我们可以利用我们同时使用张量模型并行和流水线模型并行这一事实来减少跨节点通信的开销。特别地，我们注意到每个Transformer层的输出在张量并行等级（ranks）之间是复制的（在MLP块中的$f$之后，见图5a）。因此，在两个连续的流水线阶段中执行张量模型并行的等级会发送和接收完全相同的张量集合（图9a）。对于足够大的模型，我们使用大小为8的张量模型并行。这意味着我们在相邻多GPU服务器上的相应GPU之间发送相同的张量集合8次。为了减少这种冗余，我们可以在发送端将张量分割成大小相等的块，然后只使用该等级自己的InfiniBand卡将一个块发送到下一个节点上相应的等级（例如，图9中等级1发送到等级3，等级2发送到等级4）。使用8个张量模型并行等级，每个块会小八分之一。然后，在接收端，我们可以通过NVLink执行all-gather操作，这比InfiniBand互连快得多，以重新物化完整的张量。这在图9b中显示。我们称之为scatter/gather通信优化。这项优化有助于更好地利用DGX A100服务器上的多个IB卡，并使得像交错式调度这样通信密集型的调度变得可行。

（a）无scatter/gather优化。（b）有scatter/gather优化。图9：Scatter/gather通信优化。浅蓝色块是第一个流水线阶段的层，深蓝色块是第二个流水线阶段的层。没有scatter/gather优化时，相同的张量会通过节点间的InfiniBand链路冗余地发送。相反，在发送端，我们可以将张量分散成更小的块，减少通过Infini-Band链路发送的张量大小。然后，最终的张量可以在接收端使用gather操作重新物化。

通信量化分析 从数量上看，通过scatter-gather通信优化，每对连续阶段之间需要执行的总通信量减少到$\frac{sh}{t}$，其中$t$是张量模型并行的大小，$s$是序列长度，$h$是隐藏层大小（在我们的实验中$t=8$）。

4.2 计算优化

模型特定优化 我们对计算图实施了三项模型特定的优化以获得高性能。首先，我们更改了Transformer层中的数据布局以避免内存密集的转置操作，并启用跨步批处理GEMM（strided batched GEMM）核函数。具体来说，我们将数据布局从$[b, s, a, h]$更改为$[s, b, a, h]$，其中$b, s, a, h$分别是批次、序列、注意力头和隐藏层大小的维度。其次，我们使用PyTorch JIT【索引10，PyTorch JIT】为一系列逐元素操作（bias + GeLU 和 bias + dropout + add）生成了融合核函数。第三，我们创建了两个自定义核函数来融合缩放、掩码和softmax（归约）操作：一个支持通用掩码（用于BERT等模型），另一个支持隐式因果掩码（用于GPT等自回归模型）。我们将在下一节量化这些优化的效果。

A4 实验环境

模型架构： 实验使用不同规模的GPT模型，参数量从10亿到1万亿不等。模型采用标准架构，如GPT-3【索引11，Language Models are Few-Shot Learners，2020，arXiv】。关键参数包括词汇表大小（51,200）、序列长度（2048），并根据模型规模调整隐藏层大小、注意力头数和层数。
数据集： 论文未明确指定数据集名称，但根据模型类型（语言模型），使用的是大规模文本语料库。训练量以token计，例如GPT-3（1750亿参数）模型在3000亿token上进行训练，万亿参数模型估计需要4500亿token。
硬件配置：
- 平台： Selene超级计算机【索引8，NVIDIA Selene Supercomputer】。
- GPU： 每个节点包含8个NVIDIA 80GB A100 GPU【索引6，NVIDIA A100 Tensor Core GPU】，节点内GPU通过NVLink和NVSwitch【索引9，NVLink and NVSwitch】互联。
- 网络： 每个节点配备8个NVIDIA Mellanox 200Gbps HDR InfiniBand HCA用于应用通信，采用三级胖树拓扑结构。
- 存储： 全NVME共享并行文件系统。
软件配置：
- 代码实现： 基于PyTorch【索引32，PyTorch: An Imperative Style, High-Performance Deep Learning Library，2019，NeurIPS】和Megatron-LM【索引40，Megatron-LM: Training Multi-Billion Parameter Language Models using GPU Model Parallelism，2019，arXiv】代码库进行扩展。
- 依赖库： 使用NVIDIA NCCL【索引7，NVIDIA Collective Communication Library (NCCL)】进行通信，使用PyTorch JIT【索引10，PyTorch JIT】进行算子融合。
- 训练精度： 所有实验均采用混合精度训练。

A4 实验结果

5.1 端到端性能

实验内容： 在10亿到1万亿参数的GPT模型上进行弱扩展（weak-scaling）测试，即模型规模与GPU数量同步增加。
实验结果：
- 在3072个A100 GPU上训练万亿参数模型时，实现了502 PetaFLOP/s的总吞吐量，单GPU吞吐量为163 TeraFLOP/s，达到A100理论峰值性能（312 TFLOP/s）的52%。
- 随着模型增大，GPU利用率提高，展现出超线性扩展趋势。
分析结论：
- 基于此吞吐量，估计训练1750亿参数的GPT-3模型需34天，训练万亿参数模型需84天，证明了该方法在实际应用中的可行性。
- 详细数据见表1。

表1：GPT模型从10亿到1万亿参数的弱扩展吞吐量。

5.2 与ZeRO-3的比较

实验内容： 在1750亿和5300亿参数的GPT模型上，比较PTD-P与ZeRO-3【索引36,37，ZeRO: Memory Optimization Towards Training A Trillion Parameter Models, 2019; ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning, 2021】的性能。
实验结果：
- 在少量GPU上，PTD-P的吞吐量比ZeRO-3高6%到24%。
- 随着GPU数量增加，PTD-P的扩展性更好，在GPU数量翻倍时，性能比ZeRO-3高出70%。
分析结论： PTD-P由于跨节点通信量更少，因此扩展性优于单独使用ZeRO-3（未结合模型并行）。结果见表2和图10。

表2：PTD并行与ZeRO-3（无模型并行）的比较。

图10：PTD-P和ZeRO-3在两种不同GPT模型上的每GPU吞吐量（175B GPT-3模型用虚线表示，530B模型用实线表示）。全局批次大小固定，ZeRO-3不使用任何模型并行。

5.3 流水线并行性能

实验内容： 1) 评估默认流水线并行的弱扩展性能。2) 比较交错式（interleaved）和非交错式（non-interleaved）流水线调度的性能。
实验结果：
- 弱扩展（图11）： 流水线并行表现出良好的扩展性，尤其是在大批次下，因为更多的微批次可以摊销流水线气泡的开销。
- 调度比较（图12）： 在1750亿参数模型上，交错式调度结合scatter/gather优化，吞吐量比默认调度高出10%。随着批次增大，两者差距缩小。
分析结论： 交错式调度能有效减少流水线气泡，提升性能，尤其适合小批次场景。

图11：在弱扩展实验设置中（模型大小随流水线并行规模增加），使用两种不同批次大小的流水线并行的每GPU吞吐量。

图12：在96个GPU上，针对一个GPT模型（1750亿参数）的交错式和非交错式调度的每GPU吞吐量。

5.4 并行配置比较

实验内容： 在固定GPU数量下，比较不同并行维度组合（张量vs流水线，流水线vs数据，张量vs数据）的性能。
实验结果：
- 张量 vs 流水线（图13）： 最佳性能出现在张量并行规模等于单节点GPU数（t=8）时，此时跨节点通信由开销较低的流水线并行承担。这证实了应在节点内使用张量并行，节点间使用流水线并行的原则。
- 流水线 vs 数据（图14）： 在模型能装入内存的前提下，增加数据并行度通常比增加流水线并行度更有效，因为前者可以减小流水线气泡。
- 张量 vs 数据（图15）： 数据并行比跨节点张量并行扩展性更好，因其通信更少（每批次一次all-reduce vs 每微批次多次all-reduce）且计算效率更高（更大的矩阵乘法）。
分析结论： 各种并行策略的组合存在复杂的性能权衡。最佳策略是：首先使用足够的模型并行（张量+流水线）将模型装入GPU内存，然后使用数据并行进行扩展。

图13：使用一个1622亿参数的GPT模型和64个A100 GPU，结合流水线和张量模型并行的各种并行配置的每GPU吞吐量。

图14：使用一个59亿参数的GPT模型、三种不同批次大小、微批次大小为1和64个A100 GPU，结合数据和流水线模型并行的各种并行配置的每GPU吞吐量。

图15：使用一个59亿参数的GPT模型、三种不同批次大小、微批次大小为1和64个A100 GPU，结合数据和张量模型并行的各种并行配置的每GPU吞吐量。

5.5 微批次大小

实验内容： 评估不同微批次大小对模型训练吞吐量的影响。
实验结果（图16）： 对于一个910亿参数的模型，最佳微批次大小为2。增加微批次大小会提高GPU计算效率，但同时会减少微批次总数，从而增大流水线气泡。
分析结论： 最佳微批次大小是计算效率和流水线气泡之间权衡的结果，需要根据模型和配置进行选择。

图16：在两种不同批次大小下，使用64个A100 GPU，针对一个910亿参数的GPT模型，(p, t) = (8, 8)并行配置下不同微批次大小的每GPU吞吐量。

5.6 激活重计算

实验内容： 比较使用和不使用激活重计算时的吞吐量。
实验结果（图17）： 在小批次下，激活重计算因额外的正向传播开销导致吞吐量降低最多33%。但它能支持更大的批次，从而通过减小流水线气泡，使最大吞吐量比不使用时高出2倍。
分析结论： 激活重计算是训练大模型时实现高吞吐量的关键技术，因为它能以计算换内存，从而支持更大的、更高效的批次大小。

图17：使用128个A100 GPU（(p, t) = (8, 16)）的1450亿参数GPT模型，在使用和不使用激活重计算时的吞吐量（以每秒序列数计）。

5.7 Scatter-Gather优化

实验内容： 比较使用和不使用Scatter/Gather通信优化时的吞吐量。
实验结果（图18）： 对于通信密集的交错式调度，该优化通过减少跨节点链路上的通信量，将吞吐量提升了高达11%。
分析结论： Scatter/Gather优化能有效利用多网卡硬件，对提升跨节点流水线并行的性能至关重要。

图18：使用96个A100 GPU和交错式调度，针对一个1750亿参数的GPT模型，在使用和不使用scatter/gather优化时的每GPU吞吐量。

5.8 算子融合

实验内容： 评估第4.2节中描述的算子融合（fused operators）的性能影响。
实验结果： 对于1750亿参数模型，吞吐量提升了19%（从113 TFLOP/s提升到135 TFLOP/s）；对于5300亿参数模型，提升了11%（从133 TFLOP/s提升到148 TFLOP/s）。
分析结论： 算子融合通过减少内存访问和kernel启动开销，显著提升了计算效率。

5.9 节点间通信带宽

实验内容： 测量万亿参数模型在3072个GPU上训练时的实际通信带宽。
实验结果： 观察到流水线阶段间的点对点通信有效对剖带宽为892 GB/s，数据并行副本间的all-reduce操作有效对剖带宽为12.9 TB/s。
分析结论： 高性能来自于软硬件栈的协同优化。若算子划分不当导致更多跨节点通信，扩展性能将受损。

5.10 检查点加载与保存

实验内容： 评估万亿参数模型检查点（13.8TB）的I/O性能。
实验结果： 初始加载时，384个节点（3072个GPU）的峰值读取带宽达到1TB/s（文件系统读取上限）。保存检查点时，写入带宽达到峰值的40%（273 GB/s）。
分析结论： 高效的并行文件系统对于处理大规模模型的检查点至关重要，是实际训练中的一个重要考量。

A5 结论

本文展示了如何通过组合使用流水线并行（节点间）、张量并行（节点内）和数据并行（PTD-P），在训练万亿参数级别的大型模型时实现高聚合吞吐量（502 petaFLOP/s）。这一性能使得在合理的时间内（例如，万亿参数模型约3个月）完成端到端训练成为可能。

论文深入探讨了各种并行策略之间的权衡与相互作用，并强调了在组合使用时需要仔细考虑这些因素。尽管本文的实现和评估是以GPU为中心，但其核心思想具有普适性，同样适用于其他类型的加速器。这些通用思想包括：
1. 智能地划分模型训练图以最小化通信量，同时保持设备的高利用率。
2. 通过算子融合和优化的数据布局来最小化内存密集型核函数的数量。
3. 应用领域特定的优化（如本文的scatter-gather优化）来充分利用硬件特性。

A6 附录

附录：浮点运算

FLOPs计算方法 本节描述了我们如何计算模型中的浮点运算（FLOPs）次数。我们考虑一个语言模型，它有$l$个Transformer层，隐藏层大小为$h$，序列长度为$s$，词汇表大小为$v$，训练批次大小为$b$。一个$m \times k \times k \times n$的矩阵乘法需要$2m \times k \times n$个FLOPs（因子2用于计算乘法和加法）。

Transformer层FLOPs 一个Transformer层由一个注意力块和一个2层前馈网络组成。对于注意力块，主要的FLOPs贡献者是键、查询、值的变换（$6bsh^2$次操作），注意力矩阵计算（$2bs^2h$次操作），值上的注意力（$2bs^2h$次操作），以及注意力后的线性投影（$2bsh^2$次操作）。前馈网络将隐藏层大小增加到$4h$然后再减少回$h$；这需要$16bsh^2$个FLOPs。将这些加在一起，每个Transformer层在正向传播中产生$24bsh^2 + 4bs^2h$个FLOPs。反向传播需要的FLOPs是正向传播的两倍，因为我们需要计算关于输入和权重张量的梯度。此外，我们使用了激活重计算，这需要在反向传播之前进行一次额外的正向传播。因此，每个Transformer层的总FLOPs数为$4 \times (24bsh^2 + 4bs^2h) = 96bsh^2(1 + \frac{s}{6h})$。

Logit层FLOPs FLOPs计数的另一个主要贡献者是语言模型头部的logit层，它将维度为$h$的特征转换为词汇表维度$v$。此操作在正向传播中需要$2bshv$个FLOPs，在反向传播中需要$4bshv$个FLOPs，总共为$6bshv$个FLOPs。

总FLOPs 因此，对于一个有$l$个Transformer层的Transformer模型，总的浮点运算次数为：
$C \approx 6bsv + l \left( 96bsh^2 \left( 1 + \frac{s}{6h} \right) \right)$

Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM