文章标题：利用高效并行策略扩展Llama 3训练
作者/机构：Weiwei Chu∗, Xinfeng Xie∗, Jiecao Yu∗, Jie Wang∗, Amar Phanishayee, Chunqiang Tang, Yuchen Hao, Jianyu Huang, Mustafa Ozdal, Jun Wang, Vedanuj Goswami, Naman Goyal, Abhishek Kadian, Andrew Gu, Chris Cai, Feng Tian, Xiaodong Wang, Min Si, Pavan Balaji, Ching-Hsiang Chu, Jongsoo Park (Meta Platforms, Inc.)

A1 主要贡献

本文详细介绍了用于Llama 3预训练的并行技术的设计与实现。为了在数万个GPU上实现高效训练，Llama 3采用了一种四维（4D）并行技术的组合：完全分片数据并行（FSDP）、张量并行（TP）、流水线并行（PP）和上下文并行（CP）。研究的核心问题是在超大规模（hyperscale）下实现高效、灵活且实用的训练。

研究目标是设计并实现一个能够应对Llama 3多阶段、多配置（如变化的批量大小、异构模型架构）训练需求的并行框架，并解决大规模调试中的性能和数值问题。

核心创新点和贡献如下：

  • 优化了适用于16K GPU的4D并行策略：通过优化4D并行，成功将大型模型装入内存，并克服了批量大小的限制。此外，通过将并行策略与模型超参数（如层数和层类型）进行协同设计，平衡了流水线并行中的计算和内存使用。对于在16K GPU上训练的405B Llama 3模型，在8K序列长度下实现了每个GPU 400 TFLOPs的性能，在131K序列长度下实现了每个GPU 380 TFLOPs的性能。
  • 引入了灵活的流水线并行（PP）调度：提出了一种灵活的PP调度方案，支持可变的全局批量大小和异构的层分片。这在Llama 3多模态预训练中得到了验证，实现了均衡的内存使用和高吞吐量。
  • 提出了新颖的上下文并行（CP）解决方案：提出了一种基于all-gather的CP解决方案，该方案促进了模型创新（例如，文档掩码注意力机制），其性能与最先进的基线方法相当，并展现了强大的可扩展性（在四个GPU上相比单个GPU，注意力延迟降低了3.89倍）。
  • 分享了大规模调试的方法论和经验：分享了在规模化场景下调试性能和数值问题的方法与教训。提出的自顶向下跟踪分析方法能够识别跨并行维度的最慢排名（rank）并精确定位根本原因。数值调试方法通过隔离非并行实现来排除软件错误，并识别出需要高精度浮点累积的关键梯度缓冲区。

A3 背景知识

本节概述了Llama 3预训练中使用的并行策略，随后介绍了文本和多模态预训练的各个阶段。

4D并行

4D并行策略概述。鉴于大型语言模型（LLM）规模的不断增长，分布式训练至关重要；因此，我们采用了一种4D并行方法。


图 1: 一个两层的LLM使用4D并行技术被分片到16个GPU上。FSDP和CP对输入数据进行分片，其中FSDP沿批量大小维度分片，CP沿序列维度分片。TP和PP对模型参数进行分片，其中TP在同一层内分片，PP跨层分片。

完全分片数据并行 (FSDP)。传统的数据并行，即分布式数据并行（DDP）【17, PyTorch Distributed: Experiences on Accelerating Data Parallel Training, 2020】，会在工作单元（GPU）之间复制完整的模型权重，并将数据批次分发给它们。这需要在每次训练迭代结束时进行全局梯度同步。Llama 3预训练利用了一种基于Pytorch完全分片数据并行（FSDP）【44, PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel, 2023】的内部实现，该实现通过在工作单元之间分片模型权重、梯度和优化器状态来扩展数据并行。为简化起见，本文后续将DP和FSDP互换使用。我们的FSDP实现支持与DeepSpeed【30, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, 2020】的ZeRO（Zero Redundancy Optimizer）定义一致的三种分片策略（ZeRO-1、ZeRO-2和ZeRO-3），允许对模型参数、梯度和优化器状态进行可选分片。

张量并行 (TP)。TP将Transformer层的线性模块划分到不同的GPU上。我们的实现遵循了Megatron-LM【33, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2020】中引入的方法，沿输入或输出维度分割GEMM（通用矩阵乘法）算子。TP将计算和内存成本分散到各个GPU，但引入了额外的通信开销。序列并行（SP）通常与TP结合使用，以进一步降低激活内存成本【14, Reducing Activation Recomputation in Large Transformer Models, 2022】。SP将序列相关的操作分片到TP ranks上，通常在TP分区的模块周围涉及all-gather和reduce-scatter通信，这以增加通信为代价减少了内存。

流水线并行 (PP)。PP将模型层划分为顺序的阶段，并将这些阶段分布到不同的PP ranks上。然后，跨微批次（micro-batches）的计算以流水线方式进行。图2展示了一个使用交错式1F1B PP调度【33, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2020】的例子。在这种配置中，每个rank管理多个由非连续模型层组成的虚拟阶段（例如，rank 0处理第0层和第3层）。该示例显示了6个微批次在2轮中处理，其中每个虚拟阶段每轮处理 $k = 3$ 个连续的微批次（处理微批次0-2或3-5）。每个PP rank上的激活内存使用量取决于预热（warm-up）微批次的数量。交错式1F1B调度要求总批次大小是PP ranks数量的倍数。


图 2: 一个6层的LLM被分片到3个PP ranks上，使用1F1B PP调度执行6个微批次。每个rank承载两个虚拟阶段，每个虚拟阶段包含一个模型层。模型层以交错方式分布，因此第0层和第3层在rank 0上，第1层和第4层在rank 1上，依此类推。6个微批次被分为两轮，每个虚拟阶段每轮处理k个连续的微批次，本例中k=3。

上下文并行 (CP)。CP将输入序列数据分片到不同的GPU上。这直接适用于对序列维度不变的模块（例如，前馈网络）。然而，注意力机制需要完整的序列上下文，因此需要通信来重构序列。虽然之前的工作采用环形（ring-style）通信在相邻ranks之间传递键（key）/值（value）张量【21, Ring Attention with Blockwise Transformers for Near-Infinite Context, 2023】，从而重叠通信和计算，但我们提出了一种直接的基于all-gather的CP方法。在我们的方法中，通信延迟（all-gather）是完全暴露的。我们将在第4节详细介绍这种简单但灵活高效的方法。

Llama 3预训练概述

Llama 3预训练流程。Llama 3预训练【10, The Llama 3 Herd of Models, 2024】包括三个主要阶段：短上下文预训练、长上下文文本预训练和多模态预训练。在整个预训练过程中，我们逐步增加了GPU数量、全局批量大小和序列长度。对于多模态预训练，我们通过引入额外的交叉注意力层和可训练的图像编码器来增强文本模型。交叉注意力层将图像编码器和前一个Transformer层的输出作为输入，有效地捕捉图像和文本之间的交互。值得注意的是，在预训练期间，原始的文本模型层（即非交叉注意力层）保持冻结，只有交叉注意力层和图像编码器被训练。

不同阶段的并行策略。为了最大化训练效率，我们对短上下文预训练采用3D并行（FSDP、TP和PP），对长上下文预训练采用4D并行（FSDP、TP、PP和CP）。对于多模态预训练，我们采用混合分片策略，其中图像编码器使用2D并行（FSDP和TP）进行分片，而文本模型则使用3D并行进行分片。

后续内容安排。在接下来的两节中，我们将介绍我们新颖的PP和CP设计，它们实现了高效率和灵活性，以适应多样化的训练工作负载。为清晰起见，我们在表1中总结了本文使用的所有参数及其解释。

表1：本文使用的参数及定义

A2 方法细节

3 流水线并行

本节我们介绍PP的设计及其在多模态训练中的应用。

3.1 设计概述

基于交错式1F1B的优化。我们的PP设计基于交错式1F1B调度【25, Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, 2021】。我们进行了多项优化，并在预训练期间与Llama模型进行了协同设计，以提高效率和灵活性。下面我们详细介绍我们的优化措施。

3.1.1 灵活的PP调度支持任意批量大小

解除批量大小限制。原始的交错式1F1B调度实现将批量大小限制为PP ranks数量的倍数【25, Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, 2021】。在Llama 3训练期间，全局批量大小在多个阶段中会进行调整，这需要一个支持灵活批量大小的PP调度。我们实现了一个灵活的PP调度，取消了对微批次数量的这一限制。

灵活调度的实现与权衡。在1F1B调度中，每个流水线阶段的预热微批次数量为 $(v - 1) \times k + 2 \times (p - rank \times v - 1)$，其中 $v$ 是每个PP rank上的虚拟阶段数， $k$ 是每个阶段连续处理的微批次数， $p$ 是流水线大小，$rank$ 是流水线排名索引。一个PP rank上的总微批次数 $m_p$ 是所有虚拟阶段微批次数的总和，即 $k \times v$，其中 $m$ 是微批次数。原始的1F1B调度要求 $k == m$ 且 $m_p \% p == 0$。需要注意的是，在PP中，存在GPU空闲等待其他PP rank的新微批次或令牌的阶段。我们将这些空闲时间称为PP气泡（bubbles）。PP气泡率定义为PP空闲时间与前向和后向计算时间的比值，计算公式为 $(p - 1) / m_p / v$ 【33, Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2020】。为了最小化PP气泡，我们倾向于选择较小的 $p$、更多的微批次 $m_p$ 和更多的PP虚拟阶段 $v$。

灵活调度的具体行为。我们的灵活PP调度允许 $k$ 设置为1到 $m_p$ 之间的任意数字，而 $m_p$ 可以是任何所需的数量。当 $k$ 超过 $p$ 时，我们在预热阶段为每个虚拟阶段插入 $k - p$ 个额外的微批次。这些额外的微批次有助于重叠点对点通信，如图3所示。然而，这会增加峰值内存使用，因为相比原始的交错式1F1B调度，会有 $(k - p) \times (v - 1)$ 个额外的在途预热微批次。当 $k$ 小于 $p$ 时，调度会退化为“全前向-全后向”（all-forward-all-backward）调度【11, GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, 2019】，即我们在启动后向传递之前执行所有虚拟阶段的前向传递，如图4所示。


图 3: (a) 在1F1B调度中，气泡由暴露的点对点通信（红色箭头）引入。(b) 运行额外的微批次（紫色）有助于减少由暴露的点对点通信引起的气泡。

图 4: 不同PP调度和FSDP ZeRO模式组合下的梯度内存生命周期：(a) 1F1B与ZeRO-1，reduce-scatter仅在最后一个微批次上启动。(b) 全前向全后向（ZeRO-1/2之间行为相同）。(c) 1F1B与ZeRO-2，reduce-scatter在最后一个连续微批次上启动。

3.1.2 通过模型协同设计实现均衡的PP

解决负载不均衡问题。将模型层均匀地分片到PP ranks上可能导致内存和计算不均衡。这种不均衡是由不同PP ranks上不同数量的预热微批次以及特殊模型结构（如仅位于第一个和最后一个PP ranks上的输入嵌入和输出头）的存在造成的。因此，训练可能会在第一个PP rank上遇到内存不足（OOM）问题，而后续ranks却有大量可用内存；或者由于最后一个PP rank计算负载过重而经历流水线气泡。为了缓解这些问题，我们将PP调度与模型架构进行了协同设计。具体来说，我们从第一个流水线rank减少了一层以降低各PP ranks的峰值内存，并同样从最后一个流水线rank减少了一层以平衡计算工作负载。因此，Llama 3 405B模型被配置为126层（而不是最初的128层）。

3.1.3 PP和FSDP的协同优化

不同ZeRO模式与PP的组合策略。我们研究了将FSDP与ZeRO-1/ZeRO-2和PP相结合。1F1B调度在不同虚拟阶段的执行之间交替进行，需要在同一虚拟阶段的多次执行中累积梯度。FSDP ZeRO-2与PP结合会重新分片梯度以节省内存，但代价是额外的梯度reduce-scatter操作，如图4所示。相比之下，FSDP ZeRO-1与PP结合则在各虚拟阶段间保留未分片的梯度，以增加内存使用为代价来减少通信开销。对于Llama 3训练，当 $k \ge 2 \times p$ 时，我们采用FSDP ZeRO-1与1F1B调度；当 $k < 2 \times p$ 时，我们采用ZeRO-2与“全前向-全后向”调度，以获得更好的性能。我们的实验表明，在较大规模下，FSDP的reduce-scatter可能与其他并行方式导致流量拥塞，从而降低点对点（P2P）性能。

3.2 案例研究：多模态训练

本节我们详细概述Llama 3的多模态训练，重点介绍我们如何调整PP调度以实现高效灵活的训练。

Llama 3多模态模型架构。Llama 3多模态模型结合了预训练的ViT图像编码器【42, Demystifying clip data, 2023】和预训练的Llama 3文本模型。为了整合两者，我们在文本模型的每隔几个原始Transformer层（采用自注意力）之间插入使用交叉注意力的Transformer层（下文称交叉注意力层）。交叉注意力层从图像编码器和自注意力层获取输入。模型架构如图5所示。在预训练期间，自注意力层被冻结，而图像编码器和交叉注意力层则进行训练。更多训练细节请参考Llama 3技术报告【24, Introducing Meta Llama 3: The most capable openly available LLM to date, 2024】。


图 5: Llama 3多模态架构示意图。

多模态训练面临的挑战。由于上述模型架构的差异，我们在扩展多模态预训练时遇到了两个关键挑战：
* 挑战1：图像编码器的分片。除了文本模型，我们还必须考虑图像编码器的分片，它具有独特的计算和内存特性。此外，分片策略必须通用且灵活，以有效扩展到各种图像编码器配置，因为过度拟合特定配置可能导致在其他配置上性能不佳。
* 挑战2：文本模型的工作负载不平衡。自注意力层和交叉注意力层表现出不同的计算和内存特性：（1）文本序列长度（少于200个token）远短于预训练文本模型（8K个token）；（2）大部分模型权重（自注意力层）是冻结的。因此，直接重用文本预训练中每个虚拟PP阶段分配一个Transformer层的PP配置，会导致PP ranks之间严重的工作负载不平衡。

在以下部分，我们将描述如何调整我们的PP设计来应对这两个挑战。

3.2.1 图像编码器的分片

评估三种分片方案。在实现之前，我们评估了三种候选分片方案，如图6所示。
* 方案1：使用PP对整个（图像+文本）模型进行分片。我们将图像编码器放置在第一个PP rank上，并与第一个文本模型虚拟阶段在每个微批次上一起运行。图像编码器的输出与Transformer层的输出一起通过P2P通信传递给所有其他PP ranks。
* 方案2：分离图像和文本模型，仅对文本模型应用PP。我们将图像编码器与文本模型分离，在第一个PP rank上将图像编码器作为文本模型的预处理阶段运行，然后将图像tokens广播到所有流水线阶段并拆分为微批次以输入到使用PP训练的文本模型。文本模型流水线结束后，图像tokens的梯度进行all-reduce，然后我们运行图像编码器的后向传递。
* 方案3：跨PP ranks分片图像模型。我们在所有PP ranks上分片或复制图像编码器。每个图像编码器副本处理一部分输入（$B_{img}/P_p$）。输出在PP阶段间进行all-gather，然后输入到文本模型流水线。


图 6: 使用PP分片编码器的选项。(a) 将编码器放在第一个PP rank上，并在PP ranks之间通信来自图像编码器和Transformer层的输出。(b) 将编码器放在第一个PP rank上，用图像编码器预处理所有输入并广播编码器输出；然后运行文本Transformer模型的训练。(c) 在所有PP ranks上复制/分片编码器，并在PP间分片输入，以便每个编码器副本处理$B_{img}/P_p$部分的输入。在运行Transformer之前对输出进行All-gather。

方案选择与演进。在这三个方案中，方案1需要最少的代码更改，因为我们可以重用文本模型的PP设计，只需为P2P通信打包更多的tokens（包括图像tokens）。然而，这种设计加剧了PP的工作负载平衡问题，因为第一个PP rank被分配了更多的计算，包括图像编码器。由于图像编码器的配置可能在训练期间改变，这种设计不灵活，难以在保持高效率的同时适应这些变化。相比之下，方案2和方案3将图像编码器与文本模型解耦，并在训练期间提供了更大的配置灵活性。在我们的初始实现中，我们因易于实现而采纳了方案2。通过将图像编码器放置在第一个PP rank并减少文本模型Transformer层，我们获得了良好的训练吞吐量。然而，在训练后期，图像分辨率显著增加（从448×448到672×672像素），并且图像编码器中增加了更多的Transformer层。这些因素共同导致图像编码器耗时大大增加（高达图像和文本模型组合训练延迟的33%），从而导致整体训练吞吐量显著下降。为了解决这个问题，我们切换到方案3，在每个PP rank上复制图像编码器，将数据批次拆分为微批次，并让编码器并行计算它们。这一优化将编码器计算比例从33%降低到8%，并恢复了模型变更前的TFLOPs。

3.2.2 文本模型的工作负载不平衡

交叉注意力与自注意力的差异。交叉注意力层和自注意力层之间有两个关键区别。
* 交叉注意力同时以图像和文本序列为输入，而自注意力只以文本序列为输入。在预训练期间，图像序列远长于文本序列（例如，448×448分辨率为1.2K tokens，672×672分辨率为3K tokens，而文本序列少于200 tokens）。因此，交叉注意力层在前向传递中的计算FLOPs远大于自注意力层，具体取决于批次中图像和文本序列长度的比例。
* 自注意力层在训练中是冻结的。因此，在后向传递期间，自注意力层只计算输入梯度，而交叉注意力层计算权重和输入梯度，这进一步加剧了两层之间的工作负载不平衡。

分片策略选择。在为PP划分文本模型时，我们探索了两种放置自注意力和交叉注意力层的方案：（1）将 $N$ 个自注意力层和一个交叉注意力层包装在一个PP虚拟阶段中，或（2）将 $N$ 个自注意力层或一个交叉注意力层包装在一个虚拟阶段中。方案1在PP ranks之间实现了更均衡的工作负载分布，但导致PP虚拟阶段更少，PP气泡率更大。相比之下，方案2生成更多的PP虚拟阶段，PP气泡率更小。然而，由于上述工作负载差异，实现工作负载平衡具有挑战性。在Llama 3多模态预训练中，我们因其简单性而采用了方案1。我们协同设计了多模态模型，以确定交叉注意力层与自注意力层的最终比例（4:1），这在给定的计算预算下实现了合理的训练吞吐量，并帮助满足了生产训练的截止日期。

4 上下文并行

引入CP的目的。为了给Llama 3启用长上下文训练，我们通过沿序列长度维度分割输入tokens来引入上下文并行（CP）。尽管减小DP大小 $D_p$ 以增加CP大小 $C_p$ 需要增加每个DP组的批次大小 $b$ 以维持相同的全局批次大小 $B_{global}$，但在4D并行中使用PP使得峰值内存使用与 $b$ 无关。因此，当CP沿序列长度分割时，尽管 $b$ 增加，它仍能减少峰值内存使用。关于每种并行方式如何影响内存使用的更多信息，请参考第5节。

设计。在Llama 3中，我们提出并开发了一种基于all-gather的CP注意力机制，以提供一种高效且灵活的解决方案。这是通过在注意力计算之前对键（K）和值（V）张量进行all-gather实现的。尽管现有工作，如RingAttention【21, Ring Attention with Blockwise Transformers for Near-Infinite Context, 2023】，已经提出了重叠token块的点对点通信与计算的解决方案，但我们采用基于all-gather的CP注意力主要有两个原因：


图 7: 不同注意力掩码的CP分片示例。

选择All-gather-based CP的原因。首先，Llama 3模型架构需要灵活性以支持不规则的注意力掩码，而现有工作假设使用完整的因果掩码。具体来说，Llama 3中的注意力掩码强制tokens只关注来自同一文档的其他tokens（即文档掩码），而这个文档边界取决于输入tokens中序列结束ID（eos_ids）的位置。在每个token块上计算掩码容易出错，且不规则的token通信使得充分利用网络带宽变得具有挑战性。其次，与基于环形的方法相比，all-gather方法不会产生显著的性能开销。由于多组注意力（MGA）或分组查询注意力（GQA），KV头的数量小于头的数量，导致K和V张量比注意力中的Q张量小。此外，随着序列长度（表示为 $L_{seq}$）的增加，通信延迟的时间复杂度为 $O(L_{seq})$，而注意力计算的时间复杂度为 $O(L_{seq}^2)$。这使得暴露的all-gather通信延迟在 $L_{seq}$ 增加时占CP注意力总时间的比例越来越小。对于较短的序列长度（例如，$L_{seq} = 8192$ 或 $L_{seq} = 16384$），基于all-gather的注意力仍然可以实现与RingAttention相当的性能，因为RingAttention需要根据softmax的log-sum-exp结果【7, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, 2024】、【8, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022】对部分注意力结果进行缩放和重新缩放合并。关于基于all-gather和基于环形的CP注意力性能比较的更具体结果，见第7.2节。

实现。在我们的CP注意力实现中，我们将输入tokens均匀地分成 $2 \times C_p$ 个块，并分配给每个rank $i$ 来处理第 $i$ 个和第 $(2 \times C_p - i - 1)$ 个token块。这种分片策略确保了CP ranks之间的计算工作负载均衡。例如，图7(a)展示了当 $C_p = 2$ 时将输入tokens分成4个块，其中紫色区域表示使用因果掩码（token $i$ 只关注token $0, 1, ..., i-1$）时token块之间的计算工作负载。图7还显示了Llama 3中文档掩码的一个例子，其中tokens只关注来自同一文档的其他tokens。尽管文档掩码减少了计算工作负载，我们仍然采用对完整因果掩码最优的分片策略，因为一个训练步骤的耗时通常受限于DP组和PP阶段中最慢的rank。在这些情况下，最慢的rank通常处理没有eos_id来分割tokens的完整长序列。

处理文档掩码。文档掩码对CP注意力提出了挑战，因为文档边界是不规则且依赖于输入的，并且它与tokens的静态CP分片不对齐。在我们基于all-gather的实现中，我们对K和V张量进行all-gather，使得对每个Q token的计算是独立的。例如，图7(c)中的块1只需要块0和块1的K和V张量来计算输出结果。然而，块1中的一些tokens仍然需要关注来自块0的tokens，因为它们属于跨越CP分片边界的同一文档。例如，考虑图7中16个tokens和文档长度为[3, 3, 8, 2]的例子，块1中的前两个tokens需要关注来自同一文档的所有三个tokens。为了解决这个挑战，我们用前导零填充Q序列长度，用于在此token块之外的注意力计算。然而，我们保留K和V的序列长度信息，因为我们在all-gather后有完整的K和V张量。这种方法使我们能够实现一个支持文档掩码的高效且准确的CP注意力机制。

集成。在将CP集成到端到端训练系统时，我们需要考虑其对几个组件的影响：
* 数据并行（DP）组：尽管CP沿序列长度分割输入tokens，但由同一CP组处理的tokens仍然共享同一组模型参数。因此，在通信模型参数时（例如，在参数的all-gather操作和参数梯度的reduce-scatter期间），CP可以被视为DP的扩展。
* CP ranks：同一CP组内的ranks选择它们的本地tokens，并使用整个序列计算它们自己的注意力掩码。正如我们的CP注意力实现中详细说明的，每个CP rank需要一个完整的序列来准确计算KV序列长度并调整填充的Q序列长度。本地token选择遵循我们的分片方法，其中rank $i$ 获取第 $i$ 个和第 $(2 \times C_p - i - 1)$ 个token块。此外，位置编码应被适当地选择【41, Effective Long-Context Scaling of Foundation Models, 2023】。
* 数据加载器（Dataloaders）：数据加载器继续向原始的DP组提供不同批次的输入训练数据。序列长度的分割对分词器（tokenizer）是不可见的，因为每个CP rank需要完整的序列信息来准确计算注意力掩码。

5 Llama 3的4D并行

并行配置概览。我们在表2中提供了大规模训练Llama 3的并行配置概览。在16K个GPU上训练Llama 3模型是一项能力计算挑战，其目标是最大化效率以减少总训练时间。在此规模下，每个训练步骤16M tokens的固定批次大小限制了跨训练样本可实现的并行度。本节我们详细阐述最终确定并行配置（即确定每个并行维度的大小）的推理过程，以在全局数据批次大小为16M tokens的情况下充分利用16K个GPU训练Llama 3模型。

表2：在16K GPU上使用16M tokens全局批次大小预训练405B Llama 3模型的各并行维度大小（更多细节见Llama 3技术报告[10]的表4）。

5.1 并行维度的大小

并行维度确定的步骤。我们将多个并行维度大小的确定分解为几个关键步骤：首先，我们考虑有限的全局批次大小，确定最小的TP大小。其次，我们探讨在Llama 3中使用3D并行而非2D并行的理由。最后，我们解释在长上下文训练中引入CP的优势，以及我们4D并行的配置。本节使用的符号定义见表1。

TP大小。对于16M的全局token预算和8K的序列长度 $L_{seq}$，全局批次大小 $B_{global}$ 为2048。每个GPU的批次大小 $b = B_{global} / D_p$，其中 $D_p$ 是数据并行组的数量，计算为 $D_p = N_{gpus} / T_p = N_{gpus} / T_p / P_p / C_p$。使用2D并行，$b$ 变为 $B_{global} / D_p = 2K / (16K / T_p / 1 / 1) = T_p / 8$。为确保 $b \ge 1$，我们需要 $T_p \ge 8$。对于3D并行，$b$ 为 $T_p \times P_p / 8$。为了实现最小气泡的高效PP，我们倾向于 $b \ge P_p$，因此 $T_p \ge 8$。因此，在16K GPU上每步训练16M tokens时，无论2D还是3D都需要 $T_p \ge 8$。在我们的训练集群设置中，每个节点有8个GPU。设置 $T_p \le 8$ 将TP限制在节点内通信（即NVLink），其带宽远高于节点间通信。综上，考虑到批次大小和分层网络带宽的限制，$T_p=8$ 是最优的TP大小。

2D还是3D并行。为了将模型装入内存，我们考虑2D并行（FSDP ZeRO-3 + TP）或3D并行（FSDP ZeRO-1/2 + TP + PP）。对于2D和3D，当 $T_p=8$ 时，效率主要取决于FSDP和PP的通信开销。对于 $b=1$ 的2D并行，计算延迟不足以隐藏FSDP通信，而3D并行具有更廉价且更稳定的P2P通信；因此，我们选择了3D并行。例如，每个模型参数在FSDP ZeRO-3中贡献2字节的通信数据（假设为BF16数据类型），并在前向传递中为每个token贡献2个计算FLOPs。当使用8K tokens（$b=1$ 且 $L_{seq}=8192$）进行训练时，计算与通信的算术强度为 $(2 \times 8K) / 2$ FLOPs/通信字节，这远低于硬件的峰值计算FLOPs与网络带宽的比率，即Nvidia H100 GPU【26, NVIDIA Hopper Architecture In-Depth, 2022】上BF16的989 TFLOPs与50 GB/s RoCE网络带宽【24, Introducing Meta Llama 3: The most capable openly available LLM to date, 2024】的比率 ($989T / 50G = 19.78K$)。

3D并行配置。在配置3D并行时，我们选择 $P_p=16$ 以将模型装入内存。我们没有考虑将FSDP ZeRO-3与PP结合，原因有三：首先，模型并行维度（$P_p=16$ 和 $T_p=8$）足够大，可以容纳405B模型。其次，FSDP ZeRO-3在每个PP阶段的前向和后向都有额外的通信开销。第三，FSDP通信与PP重叠时会产生性能干扰。特别是，PP的跨主机P2P通信由于PP和FSDP通信之间对网络硬件带宽的资源争用而变慢。

用于长上下文的4D并行。对于Llama 3预训练的长上下文阶段，随着序列长度增加到128K，全局批次大小从2K减少到128，而每个全局批次的tokens数量保持不变。如果不改变并行配置，这意味着 $b$ 下降到1，这将由于PP中难以忍受的气泡而完全破坏性能。减少DP大小以换取更大的TP或PP来增大 $b$ 也是不可行的。因为以相同速率增加PP并不能解决流水线气泡问题，而将TP增加到8以上会在关键路径上引入昂贵的跨主机TP通信。鉴于此，CP通过在每个训练样本内分片序列维度，成为完美的解决方案。

4D并行配置。当我们将CP引入现有并行时，我们首先需要考虑用CP替换哪个并行维度。当全局token预算保持16M但序列长度增加到131K时，$B_{global}$ 变为128。在 $b \ge 1$ 的批次大小约束下，我们不能用CP替换TP或PP；因此只能用CP替换DP。如前所述，对于 $T_p=8$ 和 $P_p=16$，为了PP效率，强烈推荐 $B_{global} / (N_{gpus} / T_p / P_p / C_p) \ge P_p$，这意味着 $C_p \ge 16$。我们使用 $C_p=16$ 来最小化CP通信开销。总的来说，CP使我们能够轻松扩展到长上下文训练阶段，同时保持相同的 $T_p$ 和 $P_p$ 配置，并通过分片序列长度维度额外减少了激活内存的使用。

5.2 并行维度的顺序

并行维度排序原则。为了最大化我们训练集群网络带宽的效率，我们仔细地对并行级别进行排序。我们的训练集群具有分层网络拓扑，从作为最内层同一主机内GPU的高带宽NVLink到作为外层的较低带宽跨节点网络。作为一个指导原则，我们将通信需求更高的并行维度（即更高的通信数据量、更高的通信频率，和/或更难隐藏的通信延迟）放在并行的内层：

TP通信。TP在每个线性模块上涉及激活或梯度张量的all-gather和reduce-scatter。这些通信完全暴露在关键路径上，并在每个Transformer层中发生四次，两次用于注意力模块，两次用于前馈网络（FFN）模块。

CP通信。CP在内部注意力模块上涉及KV张量的all-gather或KV张量梯度的reduce-scatter。通信延迟是完全暴露的，并且在每个Transformer层中发生一次。尽管CP的通信数据量与PP相似，但它涉及 $C_p$ 个ranks的集体通信，而PP涉及两个ranks之间的P2P。因此，由于CP ranks之间的同步，CP通信延迟更长。

PP通信。PP在每个虚拟流水线阶段进行通信。由于解耦的异步P2P发送和接收，两个P2P ranks之间没有同步。当 $k=p$ 时，所有P2P通信都是完全暴露的。尽管如此，在分析CP通信时，由于前述原因，我们优先考虑CP而非PP。

DP通信。使用ZeRO-1/ZeRO-2的DP每个训练步骤只通信一次，即all-gather模型参数和reduce-scatter梯度。尽管其通信量与PP相当，但我们可以潜在地用前向/后向计算隐藏其通信延迟（即将all-gather参数与模型前向重叠，reduce-scatter梯度与模型后向重叠）。因此，我们将DP作为4D并行中最外层的级别。

最终排序。综上所述，考虑到每个并行维度的通信，我们得到的并行顺序从最内层到最外层是 [TP, CP, PP, DP]。

6 大规模并行调试

在利用高效的多维并行技术扩展Llama 3训练时，调试性能、内存使用和数值问题成为一项至关重要的任务。我们分享我们的调试过程和从中吸取的经验教训，以促进未来的研究和开发。

6.1 性能：识别慢速Rank

识别慢速Rank的挑战。在多维并行中，大规模调试性能问题可能具有挑战性，特别是当试图确定问题的根本原因时。各种并行之间的交互使追踪一个缓慢的通信集合操作到其根本原因的过程变得复杂。例如，图8展示了一个8个GPU和（$C_p = 2$, $T_p = 4$）的配置，以及一个TP组内4个GPU的TP通信集合操作的堆叠性能轨迹。轨迹显示Rank 2是该组中最慢的rank，因为它的通信集合操作最短，表明其他ranks正在等待Rank 2加入。然而，尚不清楚Rank 2是否是整个系统的瓶颈，因为它的缓慢可能是由其CP通信集合操作引起的，而其在CP组中的对等rank（Rank 6）可能是实际的瓶颈。


图 8: 在进程组中识别慢速ranks。

自顶向下的分析方法。为了应对这一挑战，我们的性能轨迹分析采用自顶向下的方法，从最外层的并行级别开始。如第5.2节所述，我们的并行顺序从内到外是 [TP, CP, PP, DP]。我们首先分析DP组以识别最慢的一个，然后迭代地对PP、CP和TP组重复此过程，以缩小慢速ranks的范围。一旦识别出慢速rank，我们就可以检查CPU、GPU计算和GPU通信的详细分析轨迹，以调查根本原因，无论是软件还是硬件问题（例如，有故障的GPU）。这种方法类似于分布式系统中的故障定位【39, A survey on software fault localization, 2016】、【43, The inflection point hypothesis: a principled debugging approach for locating the root cause of a failure, 2019】，其中有问题的宿主不一定是第一个崩溃并报告错误的主机。一个用于分析性能轨迹并识别最慢rank根本原因的自动化工具，将是Llama训练系统中性能调试的宝贵资产。

6.2 4D并行中的数值问题

数值问题的来源。在4D并行中，训练数据和模型参数的划分不可避免地会改变数值行为，这是由于浮点加法的非交换和非结合性。使用低精度数据类型，如BFloat16（BF16）【13, A Study of BFLOAT16 for Deep Learning Training, 2019】，进一步加剧了这个问题。因此，识别和缓解数值差距对于确保训练稳定性至关重要，并且是训练系统的一个重要设计目标。

区分数值问题与实现错误。在开发4D并行时，区分训练损失行为的差异是由于数值问题（例如，不同的累积顺序）还是实现错误至关重要。前者可以通过在某些并行实现中使用更高精度的累积顺序来缓解，而后者需要进一步调查以修复根本原因。由于并行将计算分成块并归约部分结果，它无法实现与顺序版本逐位相匹配的结果。为了区分这两种不同的原因，我们采用一种方法，将顺序版本分割成与并行版本相同的累积顺序，并检查是否逐位完全匹配。例如，我们维护一个带有微批处理的2D并行（DP和TP）设计，以模拟PP微批处理的累积顺序，作为确认数值差距是由于PP实现错误还是累积顺序差异的参考基线。

使用FP32累积梯度。通过上述调试过程识别数值问题后，我们使用FP32累积梯度和优化器状态来弥合数值差距，同时为模型计算和通信保持BF16格式。具体来说，我们对DP组的梯度reduce-scatter和在PP后向传递中累积微批次的梯度使用FP32。这种累积精度与硬件单元相一致，其中GEMM核在计算两个BF16输入矩阵时以FP32累积部分结果。在后向计算中，累积发生在批量大小维度上，DP将其拆分为小批量并进行梯度reduce-scatter，而PP进一步将小批量拆分为微批次并在PP后向传递期间累积梯度。在多模态训练中，我们进一步将所有交叉注意力层分片的图像tokens转换为FP32，以便在后向传递期间，梯度在所有交叉注意力层之间以FP32精度进行归约。

6.3 内存优化

利用4D并行进行内存优化。在大型系统中运行4D并行时，我们发现4D并行本身除了分割模型参数和输入训练数据外，还为提高内存效率带来了独特的机会。例如，PP阶段只需要前向输出张量的元数据（即张量形状）来启动后向传递，但传统的PyTorch自动求导引擎在通过引用计数释放内存方面是保守的。为了优化Llama 3训练系统中的内存使用，我们首先使用内存快照工具【1, Understanding GPU Memory 1: Visualizing All Allocations over Time, 2023】分析内存成本，以获得详细的内存分配轨迹。然后，我们或者开发一个定制的自动求导算子在前向传递期间保存张量检查点，或者利用现有的自动求导引擎但通过手动调整张量存储的大小来释放底层张量数据。我们注意到，这些优化在我们的并行配置中有助于消除激活重计算，并避免为了将训练放入内存而增加PP或TP，从而有助于提高训练效率。

A4 实验环境

  • 模型架构：实验主要基于Llama 3 405B模型，其词汇量为128K，包含126个Transformer层。在一些小规模实验中，使用了层数较少的缩减版模型。评估的序列长度包括8192（8K）和131072（131K）。
  • 硬件配置：
    • GPU：训练在多达16,000个NVIDIA H100 GPU上进行，每个GPU的TDP为700W，配备80GB HBM3内存。部分上下文并行实验在配备HBM2e内存的H100上进行，以评估在较低内存带宽下的可扩展性。
    • 平台：使用了Meta的Grand Teton AI服务器平台。
    • 互联：节点内GPU通过高带宽的NVLink连接，节点间通过RoCE网络连接。
  • 软件配置：
    • 框架：使用了基于PyTorch的内部训练框架。
    • 并行库：并行实现基于PyTorch FSDP，并进行了大量内部优化和定制。
    • 基线：在上下文并行评估中，与Flash-Attention V2和NVIDIA TransformerEngine（TE）进行了性能比较。
  • 并行配置：
    • 3D并行（8K序列长度）: 在16K GPU上，配置为 $T_p=8$, $P_p=16$, $C_p=1$, $D_p=128$。
    • 4D并行（131K序列长度）: 在16K GPU上，配置为 $T_p=8$, $P_p=16$, $C_p=16$, $D_p=8$。

A4 实验结果

流水线并行

  • 不同PP调度的性能与内存对比：使用一个26层的缩减模型进行测试，比较了“全前向-全后向”、1F1B和灵活PP三种调度。结果（图9）显示：

    • 1F1B调度内存占用最低，但由于暴露了P2P通信，TFLOPs最低。
    • “全前向-全后向”调度通过处理更多微批次隐藏了P2P通信，TFLOPs最高，但内存占用也最高。
    • 灵活PP调度在内存使用和训练吞吐量之间取得了良好的平衡。

      图 9: (a) 全前向-全后向、1F1B和灵活PP的训练TFLOPs。(b) 最大已分配内存使用量。

  • 均衡PP的有效性：由于Llama 3的词汇量（128K）导致嵌入和输出模块很大，在第一个和最后一个PP rank上造成了计算和内存不平衡。通过从首尾PP阶段各移除一层来实现负载均衡，实验结果（图10）表明：

    • 该优化将最大内存分配减少了5GB。
    • 训练TFLOPs提高了6.5%。
    • 更重要的是，该优化使得可以关闭激活重计算，从而在缩减模型上额外带来了17.5%的TFLOPs提升。

      图 10: (a) 有无工作负载均衡时，PP ranks上的最大已分配内存使用量。(b) 使用均衡PP的训练吞吐量。

上下文并行

  • CP注意力的效率与可扩展性：将本文提出的all-gather CP与单GPU上的Flash-Attention V2进行比较，评估其硬件FLOPs利用率（HFU）。

    • 结果（图11）显示，随着序列长度增加，相对HFU也随之提高，在128K序列长度时可达95%。这验证了计算开销（$O(L_{seq}^2)$）的增长速度快于通信开销（$O(L_{seq})$）的理论分析。
    • 使用块状因果掩码（文档掩码）时的相对HFU较低。这是由于静态的token分片与动态的文档边界不匹配，导致了工作负载不平衡，而非网络带宽瓶颈（图12）。

      图 11: 相对于单GPU上注意力的硬件FLOPs利用率（HFU）。

      图 12: 上下文并行all-gather实现的GPU间通信带宽。

  • 与RingAttention（TE实现）的比较：在生产硬件（H100 with HBM3）上，将本文的CP方案与TransformerEngine中类似RingAttention的方案进行比较。

    • 结果（图13）显示，当 $C_p=2$ 时，TE的HFU略高。但当序列长度超过64K时，两者性能都超过95%。
    • 当 $C_p=4$ 时，本文的CP方案显著优于TE，尤其是在4K和8K序列长度下，相对HFU高出13.53%。这是因为Ring-style方法在序列短、并行度高时，会因计算分片过细和合并部分结果的开销而降低效率。

      图 13: 上下文并行注意力（CP Attn）与TransformerEngine注意力（TE Attn）的相对硬件FLOPs利用率（HFU）比较。

端到端性能

  • 最终性能：在16K H100 GPU上训练405B Llama 3模型，取得了优异的性能：
    • 8K序列长度（3D并行）：400 TFLOPs/GPU。
    • 131K序列长度（4D并行）：380 TFLOPs/GPU。
  • 3D并行性能分析：FSDP的通信大部分被计算和其它通信所掩盖。PP的气泡率很低（当每个DP组的微批次数量是PP大小的两倍时为5%，相等时为12%），实现了高效的流水线执行。
  • 4D并行性能分析：在长上下文训练中，CP通信（all-gather和reduce-scatter）的暴露延迟占总时间的7.64%。然而，深入分析发现，其中65.75%的延迟是由于等待CP组中最慢的rank造成的。根本原因是Llama 3训练中使用的文档掩码导致了跨GPU的工作负载不平衡（图14），最慢rank的计算时间是最快rank的1.44倍。这表明，即使采用通信计算重叠的CP算法，性能提升的上限也仅为2.62%，证明了all-gather方案的有效性。

    图 14: (a) 所有GPU上总计算时间的分布。(b) 所有GPU上注意力核总时间的分布。

A5 结论与未来硬件建议

结论

本文详细介绍了Llama 3文本和多模态预训练的训练系统。通过采用4D并行技术，该系统成功扩展至多达16K个GPU，并通过众多优化在批量大小受限的情况下实现了高效率。该系统设计高度灵活，能够支持不同训练阶段的动态工作负载和异构模型架构。此外，本文还提供了一套大规模调试性能和数值问题的方法论。基于Llama 3的训练经验，作者为未来的训练节点和集群设计提供了建议。尽管高效的Llama训练需要模型架构、学习算法和训练基础设施（如容错）的整体协同设计，但本文分享的细节和见解有望为未来模型开发和软硬件协同设计指明方向。

对未来硬件的建议

节点级建议
* 优化各种形状的计算效率：硬件加速器不仅要在大矩阵尺寸上提供高计算吞吐量，也要在并行化导致GEMM维度减小的情况下保持高效率，这意味着需要提供与计算吞吐量相匹配的充足内存带宽。
* 更高的HBM容量可以提升性能：更大的HBM容量可以扩展多维并行的超参数空间，从而获得更高的整体性能。例如，减少张量维度的分片会增加内存使用，但通过更好地分摊通信开销来减少TP通信，从而提升性能。
* 确保充足的CPU性能：随着加速器性能的增长远超CPU，未来大规模LLM训练可能会受限于CPU。这源于更小的GEMM尺寸和模型中复杂的轻量级算子，这些都会增加CPU启动内核的开销。
* 最小化性能差异并使DVFS确定化：并行计算中的同步使得整个集群的性能取决于最慢的加速器。动态电压频率缩放（DVFS）策略应在所有加速器上保持确定性，以避免因瞬时降速累积导致的整体性能下降。

训练集群级建议
* 优化网络层次结构：将LLM训练扩展到10万或更多加速器需要高效的多级交换机网络。设计分层网络时，上层交换机可以采用较低或超额认购的带宽，以实现成本和功耗效率。网络参数应与预期的工作负载需求（包括模型超参数和所有并行维度）协同设计。
* 确保稳健的网络性能：由于大规模并行中的细粒度同步，任意两个rank之间的网络降速（如拥塞或丢包）都会影响整个集群的性能。必须确保整个网络以一致的性能运行，没有瞬时降速。
* 优先考虑能效：未来的LLM训练集群规模巨大，其瓶颈将是数据中心的总功率而非加速器数量。因此，加速器的每瓦性能（Perf/Watt）与绝对性能同等重要，甚至更重要。