大模型训练实战
从单GPU到超大规模集群

我们在最多 512 块 GPU 上完成了 4000 多次扩展实验，测量了吞吐量（标记大小）和 GPU 利用率（颜色深浅）。 请注意，在此图中，吞吐量和 GPU 利用率都已按模型大小归一化。
英文原文，译者等信息 英文原文：The Ultra-Scale Playbook: Training LLMs on GPU Clusters
初版翻译：Gemini 2.0 FTE
假装编辑：@Gantrol

数千个 GPU 完美和谐运行。这就是训练当今最强大 AI 模型所需的条件——一种算力的交响，直到不久前，这还只是少数精英研究实验室的专属领域。开源已改变了这一格局，但尚未完全普及。没错，你可以下载最新的 Llama 或 DeepSeek 模型。没错，你可以阅读它们的技术和 实验 报告。但最具挑战性的部分——训练代码，以及协调 GPU 以训练这些庞大系统所需的知识和技术——仍然笼罩在复杂性之中，散落在各种脱节的论文和通常是私有的代码库中。

这本开源书籍旨在改变现状。我们将从基础知识入手，带你了解扩展大型语言模型训练规模的必要知识，从单个 GPU 到数十个、数百个甚至数千个 GPU，并通过实用的代码示例和可复现的基准来阐释理论。

随着用于训练这些模型的集群规模不断扩大，诸如数据并行、张量并行、流水线并行或上下文并行等各种技术，以及 ZeRO 或内核融合等方法应运而生，以确保 GPU 始终保持高利用率。这显著缩短了训练时间，并最大限度地利用了这些昂贵的硬件。更重要的是，随着 AI 训练规模扩展的挑战超越了构建初始模型，团队们发现，在专门数据上微调大型模型通常能产生最佳结果，而这通常也涉及相同的分布式训练技术。本书将逐步介绍所有这些技术——从最简单的到最精细的——同时保持单一的故事主线，以便理解每种方法的由来。

我们假设你对当前的 LLM 架构有一定的基础知识，并且大致熟悉深度学习模型的训练方式，但你可能对分布式训练还比较陌生。如果需要，模型训练的基础知识可以在 DeepLearning.ai 或 PyTorch 教程部分中找到。本书可以看作是三部曲的第二部分，紧随我们关于预训练数据处理的第一篇博客——所谓的“FineWeb 博客文章”。阅读完这两篇博客文章后，你应该掌握构建高性能 LLM 所需的几乎所有核心知识，只差一些关于数据混合和架构选择的最后调味品来完善配方（敬请期待第三部分……）。

本书建立在以下 三个基本原则 之上：

理论和概念的快速入门： 在深入代码和实验之前，我们希望从高层次理解每种方法的工作原理，以及它的优势和局限性。你将了解语言模型的哪些部分会消耗你的内存，以及在训练期间何时发生。你将学习如何通过并行化模型来解决内存约束，以及如何通过扩展 GPU 来提高吞吐量。因此，你将理解以下小部件如何计算 Transformer 模型的内存分解：

内存使用分解

（如果你不清楚这个小部件的作用，别担心。我们就是为此而来。）

虽然这个小部件给出了理论上的分解，但我们也制作了 以下工具，可用于预测训练运行期间的内存使用情况：

Predict Memory Tool

清晰的代码实现： 理论是一回事，但当我们实际实现某些东西时，会发现各种边缘情况和重要细节。因此，我们尽可能链接到实现参考。根据具体情况，我们将使用两个代码参考：

真实的训练效率基准： 最后，*如何* 实际扩展你的 LLM 训练取决于你的基础设施，例如芯片类型、互连方式等，我们无法提供统一的通用方案。但我们将提供一种基准测试多种设置的方法，这正是我们在集群上所做的！我们进行了超过 4100 次分布式实验（包括测试运行，超过 16000 次），最多使用 512 个 GPU，以扫描许多可能的分布式训练布局和模型大小。

正如你所见，我们需要涵盖很多内容。在深入分布式训练的细节之前，让我们先从高层次快速了解本书将涵盖的挑战。

高层概述

本书将涵盖的所有技术都旨在解决以下三个关键挑战中的一个或多个，我们将在本书中不断遇到这些挑战：

  1. 内存使用量：这是一个硬性限制——如果训练步骤无法在内存中容纳，则训练无法继续
  2. 计算效率：我们希望硬件大部分时间都在进行计算，因此我们需要减少花在数据传输或等待其他 GPU 执行工作上的时间。
  3. 通信开销：我们希望最大限度地减少通信开销，因为它会使 GPU 处于空闲状态。为了实现这一点，我们将尝试充分利用节点内（快速）和节点间（较慢）的带宽，并尽可能将通信与计算重叠。

在许多地方，我们将看到我们可以用计算、通信或内存中的一个来换取另一个（例如，重计算或张量并行）。找到正确的平衡是扩展训练的关键。

本书内容广泛，我们制作了一份 速查表，以帮助你浏览本书并掌握要点。在您驾驭这些波涛汹涌的水域时，请将它放在手边！

Cheatsheet

第一步：在单个 GPU 上训练

如果你想在阅读时增加播客的感觉，可以边阅读边听 NotebookLM 主持人讨论本书的第一部分。

在开始扩展到多个 GPU 之前，让我们快速回顾一下模型训练的最基本知识。当模型在单个 GPU 上训练时，训练通常包括三个步骤：

  1. 前向传播：将输入数据传递到模型中，产生输出；
  2. 反向传播：计算梯度；以及
  3. 优化步骤：使用梯度更新模型参数。

它通常看起来像这样：

悬停在网络元素上查看其详细信息

在此图中，顶行和底行的方框可以看作是模型内部的连续层。红色方框是每个层的相关梯度，在反向传播期间计算得出。

批大小 (bs) 是模型训练的重要超参数之一，它同时影响模型的收敛性和吞吐量。

在训练初期，较小的批大小可能有助于快速沿着训练地形移动，达到最佳学习点。然而，在模型训练的后期，较小的批大小会使梯度保持噪声，模型可能无法收敛到最佳的最终性能。另一方面，较大的批大小虽然可以提供非常准确的梯度估计，但往往无法充分利用每个训练样本，导致收敛速度较慢，并可能浪费计算资源。你可以在 OpenAI 关于大批量训练的论文 或 MiniMax-01 技术报告的第 4.2 节中找到对此主题的精彩早期讨论。

批大小也影响在给定文本数据集上进行训练所需的时间：较小的批大小将需要更多的优化器步骤才能在相同数量的样本上进行训练。优化器步骤成本很高（在计算时间方面），因此，与使用较大批大小相比，训练的总时间将会增加。话虽如此，请注意，批大小通常可以在最佳批大小附近进行相当大的调整，而不会对模型的性能产生重大影响，即，最终模型性能对确切批大小值的敏感度通常在最佳批大小附近相当低。

在 LLM 预训练社区中，批大小通常以 tokens 而不是样本数量 (bst = 批大小 Tokens) 报告，这使得训练数值通常独立于训练期间使用的确切输入序列长度。

在最简单的情况下，在单台机器上训练时，bs（以样本为单位）和 bst 可以从模型输入序列长度 (seq) 计算得出，如下所示：

bst=bs *seq

从现在开始，我们将以样本为单位显示批大小的公式，但你始终可以通过将其与序列长度相乘来获得其 token 单位的对应值。

最近 LLM 训练的最佳点通常在每批 4-6 千万个 tokens 的数量级。批大小以及训练语料库多年来一直在稳步增加：Llama 1 使用约 4M tokens 的批大小训练了 1.4 万亿个 tokens，而 DeepSeek 使用约 60M tokens 的批大小训练了 14 万亿个 tokens。

当我们扩展模型训练以适应这些大批量大小时，我们的第一个挑战已经迎面而来：内存不足问题。当我们的 GPU 没有足够的内存来容纳目标批大小的完整批次时，我们应该怎么办？

让我们首先快速了解一下最初导致我们内存不足问题的原因。这将有助于我们对训练模型所需的内存量获得一些有用的直觉。

Transformer 中的内存使用

在训练神经网络模型时，需要在内存中存储几个项目：

📝 注意

你可能会认为对于一个模型，你可以精确地计算出内存需求，但还有一些额外的内存占用者使得精确计算变得困难：

  • CUDA 内核通常需要 1-2 GB 的 GPU 内存，你可以通过运行 import torch; torch.ones((1, 1)).to("cuda")，然后使用 nvidia-smi 检查 GPU 内存来快速验证这一点。
  • 来自缓冲区、中间结果和一些由于碎片而无法使用的内存的剩余内存使用量
我们将忽略最后两个贡献者，因为它们通常是很小的常数因子。

这些项目以张量的形式存储，张量具有不同的 *形状* 和 *精度*。*形状* 由超参数决定，例如批大小、序列长度、模型隐藏层维度、注意力头数、词汇表大小，以及潜在的模型分片（我们稍后将看到）。*精度* 是指 FP32、BF16 或 FP8 等格式，分别需要 4、2 或 1 个字节来存储张量中的每个单个值。我们将在 混合精度训练 部分全面讨论不同的精度及其权衡，现在我们只需记住，这些不同格式的内存需求将有所不同，这将影响我们需要存储的项目的内存使用量。

那么，如何从这些变量中快速确定内存使用量呢？一种简单的方法是通过经验来做到这一点，只需测量它即可。

分析内存使用情况

使用 Pytorch profiler，我们可以了解内存是如何在整个训练过程中分配的。我们可以看到，内存利用率不是静态的，而是在训练过程中以及训练步骤中变化很大：

显然，第一步看起来与随后的步骤非常不同，但让我们首先看看一个步骤的一般结构：首先，当我们进行前向传播时，激活值迅速增加，然后在反向传播期间，梯度累积，并且随着反向传播的进行，用于计算梯度的存储激活值被逐步清除。最后，我们执行优化步骤，在此步骤中，我们需要所有梯度，然后更新优化器状态，然后才能开始下一个前向传播。

为什么第一步看起来不同：激活值迅速增加，然后稳定一段时间。在第一步中，torch 缓存分配器做了很多准备工作，准备内存分配以加速后续步骤，以便它们不需要在之后搜索空闲内存块（参见 Zach 的博客）。在第一步之后，我们还看到了优化器状态的出现，这通常会抵消进一步训练步骤的内存使用量。

现在我们对内存有了初步了解，让我们看看扩展训练规模通常是如何在最大化计算效率的同时，将这些各种项目（激活值、参数、梯度、优化器状态）的内存需求保持在 GPU 的内存约束范围内的。

权重/梯度/优化器状态内存

让我们从列表中的前 3 项开始：模型的权重、梯度和优化器状态。实际上，我们可以很容易地估算出它们所需的内存。

对于一个简单的 Transformer LLM，参数数量由 以下公式 给出：

N = h * v + L * (12 * h^2 + 13 * h) + 2*h

在该等式中，h 是隐藏层维度，v 是词汇表大小，L 是模型中的层数。请注意，查看该等式，我们可以看到，在较大的隐藏层维度下，占主导地位的项将是 h^2 项，因为它是唯一一个随着参数扩展而二次增长的项。

参数和梯度的内存需求只是参数数量乘以每个参数的字节数。在传统的全精度 (FP32) 训练中，参数和梯度都需要 4 个字节，而优化器（如果我们使用 Adam）需要存储动量和方差，这为每个参数增加了另外两个 4 字节。总结如下：

\begin{aligned} & m_{params} = 4 * N \\ & m_{grad} = 4 * N \\ & m_{opt} = (4+4) * N \end{aligned}

现在让我们看看如果我们使用较低的精度，情况会如何变化。出于稳定性的原因（参见下面的 混合精度训练部分），我们通常不使用完全低精度的训练，而是使用更高精度和更低精度的混合，称为“混合精度”。如今，混合精度训练的默认设置通常是对大多数计算使用 BF16——每个参数和梯度需要 2 个字节——以及模型权重和梯度的额外 FP32 副本，因此每个参数总共 12 个字节。除了参数和梯度之外，我们还需要存储优化器状态：对于 Adam 优化器，这需要动量和方差，通常以 FP32 存储以获得数值稳定性，每个使用 4 个字节。

以下是摘要：

\begin{aligned} & m_{params} = 2 * N \\ & m_{grad} = 2 * N \\ & m_{params\_fp32} = 4 * N \\ & m_{opt} = (4+4) * N \end{aligned}

📝 注意

某些库以 fp32 存储梯度，这将需要额外的 m_{params\_fp32} = 4 * N 内存。例如，nanotron 就是这样做的，因为 bf16 对于较小的值是有损的，我们始终优先考虑稳定性。有关更多信息，请参阅 此 DeepSpeed 问题。

📝 注意

参数的 FP32 副本 (m_{params\_fp32}) 有时在文献和代码库中称为“主权重”。

有趣的是，混合精度本身并不能节省总体内存，因为它只是以不同的方式在三个组件之间分配内存，实际上，如果我们在 FP32 中累积梯度，则会比全精度训练增加 4 个字节。它仍然是有利的，因为以半精度计算前向/反向传播允许我们 (1) 使用 GPU 上优化的较低精度运算，这些运算速度更快，并且 (2) 减少前向传播期间的激活内存需求，正如我们在上面和下面的图中看到的那样，这部分内存使用量很大。

让我们了解一下模型通常需要多少内存（全精度和混合精度给出相同的总值）：

模型参数 FP32 或 BF16，不带 FP32 梯度累积 BF16，带 FP32 梯度累积
1B 16 GB 20 GB
7B 112 GB 140 GB
70B 1120 GB 1400 GB
405B 6480 GB 8100 GB

正如我们所见，一旦达到 7B (!)，权重和优化器需求就开始显著增加，并超过了典型 GPU 内存的大小，例如 H100 GPU 的 80GB。

但现在，让我们从仍然适合单个 GPU 的模型开始，看看内存预算的最后一个主要贡献者：激活内存。

激活内存

激活内存比权重、梯度和优化器状态更难计算，部分原因是它取决于模型的输入。如果你不确定为什么我们甚至需要存储激活值以进行反向传播，此参考资料 是一个很好的快速回顾。在仔细检查了反向传播的计算方式后，我们可以估算出混合精度激活值所需的总内存，我们得到以下等式：

m_{act} = L \cdot seq \cdot bs \cdot h \cdot (34 + \frac{5 \cdot n_{heads} \cdot seq}{h})

这里 L 是层数，seq 是序列长度，bs 是样本批大小，h 是模型的隐藏层维度，n_{heads} 是头数。

对于数字的确切推导，你可以参考 NVIDIA 关于重计算的原始论文 ，这基本上需要你对 Transformer 层中每个操作之间的所有中间激活值的大小进行一些核算。

这里一个有趣的观察是，对于给定的模型，内存不是静态的，而是随着序列长度和批大小线性扩展。这意味着激活内存是在我们增加批大小或使用更长序列进行训练时会激增的部分。我们可以使用此等式来查看内存使用量如何随各种序列长度变化，例如对于 Llama 模型 (bs=1)：

这张图讲述了一个引人注目的故事：对于短序列（或类似的小批量大小），激活值几乎可以忽略不计，但从大约 2-4k tokens 开始，它们开始占用大量内存，而参数、梯度和优化器状态使用量（我们稍后将讨论）大致保持独立于序列长度和批大小。

对于大型输入 tokens（也称为大批量大小/序列），激活值成为迄今为止最大的内存负担。

有没有办法驯服这种“激活值爆炸”？好问题，读者！

现在是时候解释我们的第一项技术了——称为 激活重计算——这将帮助我们限制激活内存占用。这是当今大型模型训练工具箱中的一个基本工具。

激活重计算

激活重计算（也称为 *梯度检查点* 或 *重物化*）背后的总体思路是在前向传播期间丢弃一些激活值以节省内存，并在反向传播期间花费一些额外的计算资源来即时重新计算这些激活值。在不进行重计算的情况下，我们存储两个可学习操作（例如，前馈、层归一化等）之间的每个隐藏状态，以便我们可以在反向传播期间使用它们来计算梯度。当我们使用重计算时，我们通常只会在模型架构的几个关键点存储激活值，丢弃其余的激活值，并在反向传播期间从最近保存的激活值即时重新计算它们，基本上再次执行前向传播的子部分，以用计算换取内存。它通常看起来像这样：

悬停在网络元素上查看其详细信息

有几种策略可用于选择要存储的关键激活值：

让我们看看重计算策略在实践中如何显著减少内存占用，以及选择性重计算如何在内存节省和重计算成本之间取得良好的平衡：

这里另一个显而易见的趋势是，对于较小的模型，长序列的激活值扮演着更重要的角色，因此重计算的效果变得更加明显。

📝 注意

当你测量你的训练设置在使用 GPU/TPU/加速器时的效率时，你通常需要考虑重计算来计算总 FLOPS（每秒浮点运算次数），并将其与 GPU/TPU/加速器的理论最大 FLOPS 进行比较。在计算训练步骤的 FLOPS 时考虑重计算会得到一个称为“硬件 FLOPS”的值，这是在加速器上执行的实际运算次数。将此数字除以训练步骤的持续时间和最大加速器 FLOPS，即可得出 硬件 FLOPS 利用率 (HFU)。

然而，最终真正重要的是在给定数据集上训练模型所需的端到端时间。因此，当比较各种 GPU/TPU/加速器时，如果其中一个加速器例如提供足够的内存来跳过重计算，从而执行更少的操作/秒（较低的 HFU），但训练速度更快，则应奖励而不是惩罚它。因此，另一种方法是计算所谓的 模型 FLOPS 利用率 (MFU)，与 HFU 相比，MFU 仅考虑通过模型进行前向 + 反向传播所需的运算，而不将重计算包括在测量的 FLOP 中。因此，此值比训练实现更特定于模型。

如今，大多数训练框架都使用 FlashAttention（我们将在 下文 中介绍），它通过在反向传播中重新计算注意力分数和矩阵而不是存储它们，从而在其优化策略中本地集成了激活重计算。因此，大多数使用 Flash Attention 的人已经在使用选择性重计算。

正如你现在所理解的，激活重计算由于重计算而略微增加了 FLOP 的数量，同时显著减少了内存访问开销。

这种权衡在具有小型高速内存的硬件（如 GPU）上尤其有利，因为访问内存通常比执行计算慢。尽管涉及额外的操作，但总体效果通常是更快的计算，以及更低的内存占用。

现在我们已经了解了重计算，我们可以像在上面的图表中看到的那样驯服激活内存的使用！

然而，激活值仍然线性依赖于批大小，并且上面柱状图中的所有配置文件都使用 bs=1，因此当我们转向更大的批大小时，它可能会再次成为问题。不要灰心，因为我们还有第二个工具——梯度累积——来救援！

梯度累积

梯度累积是一种非常直接的方法，可以避免内存爆炸，它将我们的批次拆分为微批次。我们将依次对每个微批次执行前向和反向传播，计算梯度，并且顾名思义，在执行优化器步骤之前，对所有微批次的梯度求和。在实践中，优化步骤不是在总和上进行的，而是在梯度的平均值上进行的，因此结果与梯度累积步骤的数量无关。

让我们将每次前向传播的批大小称为 微批大小 (mbs)。我们将每次优化器步骤之间的总批大小称为 全局批大小 (gbs)。如果我们对每 8 次前向/反向传播执行一次优化器步骤，则 全局批大小 将是 微批大小 的 8 倍。

我们现在所说的 全局批大小 因此对应于我们迄今为止为简单起见仅称为 批大小 的内容（我们现在使我们的术语更精确以避免歧义）。

通过梯度累积，全局批大小可以简单地计算如下：

bs = gbs = mbs \times grad\_acc

梯度累积使我们能够有效地将批大小增加到无穷大（甚至更大！），同时内存占用保持不变。梯度累积也与激活重计算兼容，以进一步减少内存。

image.png

梯度累积允许我们通过仅计算部分微批次来减少激活内存，激活内存随批大小线性增长。

然而，一个缺点是，梯度累积需要每个优化步骤进行多次连续的前向/反向传播，从而增加计算开销并减慢训练速度。没有免费的午餐！

但是，如果你仔细跟进，你可能已经注意到，每个微批次的前向/反向传播实际上可以并行运行。前向/反向传播是相互独立的，唯一的区别是独立的输入样本。看起来是时候开始将我们的训练扩展到多个 GPU 了！

在此之前，让我们快速了解一下如何使用分布式训练工具箱中最有用的工具之一：分析器 来可视化计算和通信。此工具对于理解和验证 GPU 之间的通信和计算是如何发生的，以及瓶颈在哪里，将非常有用。

分析 GPU 计算和通信

PyTorch 的 分析器 允许我们跟踪和可视化训练期间 CPU 和 GPU 上正在发生的事情。它本地集成在 PyTorch 中。让我们看看如何使用它：

with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, torch.profiler.ProfilerActivity.CUDA, ], schedule=torch.profiler.schedule(wait=1, warmup=1, active=3), on_trace_ready=torch.profiler.tensorboard_trace_handler('./log/profile'), with_stack=True ) as prof: for step in range(steps): train_step() prof.step()

这将生成一个跟踪，我们可以在 TensorBoard 或 Chrome 的跟踪查看器中可视化该跟踪。该跟踪显示：

profile_trace_annotated.png

示例跟踪显示 CPU 线程将内核异步启动到 GPU，其中计算内核和通信在不同的 CUDA 流中并行发生

该跟踪有助于识别瓶颈，例如：

理解这些模式对于优化分布式训练性能至关重要。例如，正如我们稍后将讨论的那样，该跟踪将清楚地显示梯度同步是否与反向计算正确重叠。

现在，让我们找一台更大的工作站 🖥️，配备几个 GPU，并开始研究我们的第一个扩展技术，称为 *数据并行*，正如我们将看到的，它只是梯度累积的并行版本。

数据并行

为了增加阅读的播客感觉，请随意聆听 NotebookLM 主持人讨论本书的以下部分。

​

数据并行 (DP) 背后的思想是在多个 GPU 上复制模型（我们将副本称为“模型实例”），并在每个 GPU 上并行运行不同数据微批次的前向和反向传播，因此得名数据并行。你可能已经在简单的训练示例中见过数据并行，但正如你很快就会看到的，我们将在本节中深入探讨，因此即使你了解一般方法，也请继续关注。

image.png

为每个 GPU 使用不同的微批次意味着每个 GPU 中将有不同的梯度，因此为了使不同 GPU 上的模型实例保持同步，来自模型实例的梯度将使用称为“all-reduce”的操作进行平均，这发生在反向传播期间，在优化器步骤之前。

这涉及我们的第一个“分布式通信”原语：*all-reduce*，它处理 GPU 实例和节点之间的同步和通信。

image.png

一个朴素的 DP 实现只会等待反向传播完成，以便我们拥有所有梯度，然后它会在所有 DP 排名上触发 all-reduce，以同步这些梯度。但是，这种计算步骤后跟通信步骤的顺序是 绝对不可以的！ 因为我们不希望我们的 GPU 在通信发生时保持空闲，就像上图所示。

相反，我们应该尽可能地尝试将通信和计算重叠，以便它们尽可能同时发生。

让我们看看三个优化，它们使我们能够比我们朴素的第一个实现做得更好！

第一个优化： 将梯度同步与反向传播重叠

我们刚刚描述的朴素 DDP 方法的主要缺点是，在反向传播（*计算*）之后，我们必须等待梯度同步（*通信*）才能更新参数。我们能否将此通信与我们的计算重叠？答案是肯定的！

如上图所示，即使在计算出早期层（左侧红色方框）的梯度之前，也可以收集和汇总某层的梯度（红色方框）。例如，一旦最后一层（右侧最后一个方框）的反向传播完成，就可以在早期层的反向计算继续进行（向左移动）时，立即收集和汇总这些梯度。

image.png

这可以通过在 pytorch 中为每个参数附加一个 *all-reduce 钩子函数* 来实现。一旦该参数的梯度准备就绪，就会触发 all-reduce 操作，而其他参数的梯度仍在计算中。这种方法将大部分 all-reduce 操作与梯度计算重叠，从而提高了效率。以下是附加钩子的简单函数：

def register_backward_hook(self, hook): """ 为模型的所有需要梯度的参数注册一个反向钩子。 """ for p in self.module.parameters(): if p.requires_grad is True: p.register_post_accumulate_grad_hook(hook)

重叠计算和通信减少了等待整个模型中梯度同步的时间。梯度同步可以（至少部分地）与反向传播并行发生，从而显著加速数据并行。以下是具有同步重叠的朴素 DP 的完整实现：

👉 Picotron 中具有重叠的朴素 DP 实现（单击展开）

这是我们“*重叠计算和通信*”的第一个示例，我们将在本博客文章中多次讨论它，并且是最大化扩展效率的基本技术。但我们可以进一步提高效率！

第二个优化： 梯度分桶

当 GPU 操作在大型张量上执行而不是在小型张量上运行许多操作时，GPU 操作通常更有效。通信操作也是如此。因此，我们可以有利地将梯度分组到桶中，并为同一桶内的所有梯度启动单个 all-reduce，而不是为每个梯度执行独立的 all-reduce。它通常看起来像这样：

dp_overlap3.svg

可以将其视为在装运前将物品装入箱子。发送几个大箱子比发送许多小箱子更有效。通过为每个桶执行单个 all-reduce 操作，我们可以显著减少通信开销并加速通信操作。

以下是使用分桶的代码实现：

👉 Picotron 中分桶 DP 实现（单击展开）

第三个优化： 与梯度累积的相互作用

最后，正如我们之前看到的，梯度累积的工作原理是在使用 optimizer.step() 更新参数之前执行多个前向和反向传播。当将梯度累积与数据并行结合使用时，我们应该注意何时同步梯度。

在朴素版本中，all-reduce 操作在每次反向传播期间自动触发，这在累积期间是次优的，因为最终步骤之后的单个 reduce 将具有相同的效果，同时减少开销。

在 PyTorch 中，这通常通过在不需要减少的反向传播上添加 model.no_sync() 装饰器来解决，该装饰器禁用梯度同步。

📝 注意

在执行通信操作时，张量必须在内存中是连续的，以避免冗余的内存副本。为了以最佳方式执行此操作，我们通常会预先分配与激活值或模型参数大小相同的连续缓冲区，专门用于通信。虽然这加速了通信，但它也在一定程度上导致了训练期间的峰值内存使用量。

现在让我们看看这对全局批大小意味着什么。

重新审视全局批大小

我们可以使用新添加的数据并行和梯度累积参数来更新我们的批大小等式：

bs = gbs = mbs \times grad\_acc \times dp

这里 grad\_acc 是梯度累积步骤的数量，dp 是用于数据并行的并行实例数。

给定目标全局批大小，我们可以用数据并行进程来换取梯度累积步骤，以加速训练。

在实践中，人们倾向于尽可能地最大化数据并行节点 (DP) 的数量，而不是梯度累积，因为它本质上是并行的，这与梯度累积的顺序性质不同。然后，在数据并行性不足以在 GPU 耗尽之前实现目标全局批大小的情况下，在数据并行性之上添加梯度累积。

能够将训练分布在不同的样本上为我们提供了并行化的第一个维度，从而使这种 1D 并行化（我们将逐步介绍另外 4 个维度）。

我们目前的进展

让我们快速总结一下如何使用最佳数据并行设置的草稿配方来设置我们的第一个 1D 并行训练：

  1. 我们首先应该通过查阅文献或运行实验来测量模型收敛性，从而确定最佳（全局）批大小（以 tokens 为单位）（GBST）。
  2. 然后，我们选择训练的序列长度，再次通过查阅文献或运行实验。通常，对于我们今天的评估来说，2-8k tokens 工作良好（我们不会在此处深入探讨训练配方，但团队通常会在训练结束时增加序列长度，在混合中添加一些更长上下文的数据样本，以达到今天的更长上下文大小）。
  3. 我们现在知道了批大小 (gbs)。我们可以通过增加本地批大小直到内存不足来找到单个 GPU 上的最大本地批大小 (mbs)。
  4. 最后，我们确定目标 DP 的可用 GPU 数量。GBS 与 DP 的比率为我们提供了所需 GBS 的剩余梯度累积步骤数。

如果梯度累积比率低于 1，即我们有太多的 GPU，也就是 GPU 资源丰富 🤑 (!)，我们可以选择不使用我们所有的 GPU，探索更大的全局批大小，或者测试较低的 MBS 是否会加速训练。在后一种情况下，我们将最终优先考虑吞吐量而不是单个 GPU 计算效率，使用比可能更小的 MBS 以加速训练。

现在举一个具体的例子：假设我们要训练一个最新的模型，GBS 为 4M tokens，序列长度为 4k。因此，我们的批大小将为 1024 个样本（我们选择最接近的 2 的幂）。让我们假设我们观察到单个 GPU 只能在内存中容纳 MBS=2，并且我们有 128 个 GPU 可用于训练。这意味着通过 4 个梯度累积步骤，我们将实现每个训练步骤 1024 个样本或 4M tokens 的目标。现在，如果我们突然有 512 个 GPU 可用怎么办？我们可以通过保持 MBS=2 并将梯度累积步骤设置为 1 来实现相同的 GBS，从而实现更快的训练！

📝 注意

请记住，在 512+ GPU 规模下，根据使用的网络，通信操作将开始受到 *环延迟*（信号在环周围传播一次所需的时间）的限制，这意味着我们无法再完全重叠 DP 通信。这将降低我们的计算效率并影响我们的吞吐量。在这种情况下，我们应该开始探索其他维度来进行并行化。

虽然数据并行性很好地将 all-reduce 梯度同步与反向计算重叠以节省时间，但这种优势在大型规模下开始瓦解。为什么？因为随着我们添加越来越多的 GPU（数百个或数千个），协调它们之间的开销会显著增加，并且网络需求变得太大而无法获得收益。因此，随着我们向系统添加更多 GPU，我们的设置将变得越来越低效。

让我们通过一些基准测试在实践中看到这种情况：

我们看到，超过某个限制后，我们的吞吐量开始大幅下降，而每个 GPU 的内存使用量保持不变，并且不受添加更多 DP 排名影响。

数据并行是我们在更多 GPU 上扩展训练的第一个（简单）策略。此技术类似于梯度累积，但并行化了微批次的前向和反向传播，从而提高了吞吐量！

但是，敏锐的读者可能已经注意到，这假设我们可以将至少一个输入样本前向传播 (mbs*=1) 放入我们的 GPU 内存中。情况并非总是如此！正如我们所见，即使激活重计算已激活，更大的模型也无法放入单个 GPU 中：

我们还看到，数据并行性在超过一定程度的扩展后开始出现一些限制性的通信开销。对于这些更大的模型或大批量大小，我们还有其他选择吗？幸运的是，我们有一些解决方案。它们将涉及将一些张量移动到 CPU 或在 GPU 设备之间拆分权重/梯度/优化器状态张量！让我们开始深入研究它们。

有两种主要的分裂方法：并行（张量并行、上下文并行或流水线并行）和共享（DeepSpeed Zero 或 PyTorch FSDP）。这两种方法在某种程度上是正交的，实际上可以结合使用！

共享范式与 DP 密切相关，因此我们将首先通过研究 ZeRO 方法来了解它！

ZeRO (Zero Redundancy Optimizer)

在本节中，我们将介绍 DeepSpeed ZeRO (Zero Redundancy Optimizer)，这是一种旨在减少 LLM 训练中内存冗余的内存优化技术。

虽然数据并行是扩展训练的有效方法，但优化器状态、梯度和参数在每个 DP 排名上的朴素复制引入了显著的内存冗余。ZeRO 通过跨数据并行维度对优化器状态、梯度和参数进行分区来消除内存冗余，同时仍然允许使用完整的参数集进行计算。这有时需要在 DP 排名之间进行更多通信，我们将在接下来看到，这些通信可能或可能无法完全重叠！

此方法分为三个可能的 ZeRO 优化阶段：

你可能会遗漏激活值，因为我们可以在其中分片。由于模型的每个 DP 副本都接收到不同的微批次，因此每个 DP 排名上的激活值也不同，因此它们不会重复，因此无法分片！

让我们仔细看看我们通过每个 ZeRO 阶段的分区可以节省多少内存！

重新审视内存使用量

你可能还记得 我们之前的章节 中，标准训练期间优化器状态、梯度和参数的内存使用量。让我们将模型的参数计数称为 \Psi（之前是 N，但在这里我们使用原始 ZeRO 论文符号）。在 混合精度训练（稍后章节中将有更多详细信息）中使用 Adam 优化器时，我们需要存储的每个项目的内存使用量为：

如果我们不在 fp32 中累积梯度，这将给我们带来 2\Psi + 2\Psi + 12\Psi 的总内存消耗，如果我们累积，则为 2\Psi + 6\Psi + 12\Psi。为了简单起见，我们现在关注不使用 fp32 梯度累积的情况，但你可以只添加额外的字节到受 ZeRO-2 和 3 影响的梯度项。

ZeRO 的思想是在 DP 排名中分片这些对象，每个节点只存储项目的一个切片，这些项目在需要时才会被重建，从而将内存使用量除以数据并行度 N_d：

zero_memory.svg

这里 \Psi 表示参数数量，k 表示优化器状态的内存乘数（正如我们刚刚看到的，对于 Adam 来说 k=12），N_d 表示 DP 度。

让我们通过探索每个 ZeRO 阶段的工作原理来解释此图及其值。我们将从 ZeRO-1 开始。

ZeRO-1：分区优化器状态

在 vanilla DP 中，所有排名在反向传播后收集相同的梯度，并同时执行相同的优化器步骤。这似乎是很多重复的工作。我们能否避免这种情况，同时减少内存使用量？

在 ZeRO-1 中，优化器状态被划分为 N_d 个相等的部分，其中 N_d 是 DP 度。这意味着分布在每个 DP 排名上的每个模型副本仅跟踪 \frac{1}{N_d} 的优化器状态。在优化步骤中，仅更新 \frac{1}{N_d} 的 float32 权重。

然而，在前向传播期间，每个副本都需要所有参数，因此我们需要在优化器步骤后添加额外的 all-gather（我们遇到的第二种类型的集体通信原语！），以便每个模型副本都具有整套更新后的权重。

这解释了我们在上图中看到的 2\Psi + 2\Psi + \frac{k\Psi}{N_d} 的内存公式！以下是单个训练步骤的操作顺序摘要

你可能想知道什么是“reduce-scatter”操作，以及这一切看起来如何，因此让我们尝试使用下图使其更图形化。我们将介绍前向/反向传播周期的所有步骤：

dp_zero1.gif

就实际通信而言，与 vanilla DP 相比，Zero-1 将我们的“all-reduce”梯度通信更改为“reduce-scatter”操作，并在优化器步骤后添加了对所有参数的 all-gather 操作。以下是它的外观：

dp_zero1_overlap.svg

如果你一直关注，你就会记得，从 vanilla DP 中，我们可以将 all-reduce 梯度通信与反向传播计算重叠。在 ZeRO-1 中，我们还可以研究如何有效地重叠新添加的 bf16 参数的 all-gather。对此有两种主要策略：

📝 注意

不幸的是，这些技术不易实现，需要复杂地使用钩子/分桶。在实践中，我们可以只使用 PyTorch 本地 ZeRO-3/FSDP 实现，并将 FSDPUnit 设置为整个模型，稍后将详细介绍。

在 ZeRO-1 中，优化器状态已分区，这意味着每个副本仅更新 \frac{1}{N_d} 的优化器状态。敏锐的读者一定已经注意到，实际上没有必要首先在所有 DP 排名上都拥有所有梯度，因为优化步骤只需要一个子集。认识一下 ZeRO-2！

ZeRO-2：添加 梯度分区

由于在每个副本上，我们只需要拥有与优化器状态分片相对应的梯度分片，因此对梯度进行分片也是有意义的，类似于优化器状态。在反向传播期间，我们不执行对梯度的 all-reduce 操作，而只执行 reduce-scatter 操作！我们只在内存中传播所需的 \frac{1}{N_d} 梯度，从而比 ZeRO-1 节省更多内存。

dp_zero2.gif

现在很容易看出，分片梯度会导致 2\Psi + \frac{2\Psi+k\Psi}{N_d}，并且随着 N_d 的增加，我们可以比基线节省高达 8 倍的内存。在通信方面，与 ZeRO-1 相同的过程适用，唯一的区别是我们即时通信和释放。总的来说，ZeRO-2 因此也等同于 vanilla DP 训练，关于通信。

在通信方面，ZeRO-2 类似于 ZeRO-1，它们都需要 reduce-scatter 用于梯度，all-gather 用于所有参数。

dp_zero2_overlap.svg

现在我们已经分片了梯度，我们是否完成了，或者我们可以继续摆脱困境？嗯，在某种程度上是这样。ZeRO-3 来了！

ZeRO-3：添加 参数分区

对于阶段 3，我们将上述在 DP 副本上分片优化器状态和梯度的方法扩展到分片模型的参数。

📝 注意

此阶段在 PyTorch 本地实现中也称为 FSDP（完全共享数据并行）。在本博客文章中，我们将仅引用 ZeRO-3，但你可以在任何看到它的地方将其视为 FSDP。

那么，如果模型的所有部分都已分布，我们如何在实践中进行前向传播或反向传播？很简单，我们在需要时按需收集它们。在前向传播中，这看起来如下：

dp_zero3_fwd.svg

因此，当我们执行前向传播并按顺序遍历各层时，我们会按需检索必要的参数，并在不再需要它们时立即从内存中刷新它们。反向传播的工作方式相同，只是流程相反，我们生成梯度分片：

dp_zero3_bwd.svg

另一个问题是，我们需要在整个前向和反向步骤中持续执行这些 all-gather，这相当于比 Zero-2 在训练步骤中** 多出 2\cdot \text{num\_layers} -1 次 all-gather，每次都带来较小的 **基本延迟 开销，我们可以在下图中看到：

dp_zero3_overlap.svg

在前向传播期间，我们在需要时对参数执行 all-gather 操作，因此 \Psi 通信税。由于我们在前向传播中需要参数后立即丢弃它们，因此在反向传播期间我们也需要再进行一次 all-gather，从而产生另一个 \Psi 的通信税。最后，我们需要与 ZeRO-2 中相同的 reduce-scatter 用于梯度，这也花费 \Psi 的通信，我们得到的总通信成本为 3\Psi，而 Zero-2 为 2\Psi。

这听起来可能像很多通信开销，但实际上还不错，因为我们可以将下一层参数的通信与当前层的前向传播重叠，这称为 预取。通过预取，我们将在前向传播中“all-gather”*第 n+1 层*的权重，同时为 *第 n 层*执行当前前向传播，类似地，我们将在为 *第 n 层*执行反向传播时“all-gather”*第 n-1 层*的权重。当然，只有在我们不过度扩展 DP 的情况下，这种重叠才成立。（根据经验，DP 不应超过 512）

在内存方面，我们可以看到我们的等式现在达到了它的最终形式 \frac{2\Psi +2\Psi+k\Psi}{N_d}，这意味着如果我们增加 DP 排名，至少对于模型相关参数，我们可以无限期地降低内存使用量。请注意，它对中间激活值没有帮助，为此我们可以使用激活检查点和梯度累积，正如我们在之前的章节中看到的那样。

**让我们总结一下我们迄今为止在 DP 和 ZeRO 中的旅程：我们已经看到，我们可以通过 DP 显著提高训练吞吐量，只需通过添加更多模型副本来扩展训练。通过 ZeRO，我们可以训练即使通常不适合单个 GPU 的模型，方法是在 DP 中分片参数、梯度和优化器状态，同时产生较小的通信成本。 **

但是，这里有一个限制，DP 仅在模型的层适合单个 GPU 时才有效，ZeRO 只能分区参数、梯度和优化器状态，但不能分区激活内存！我们从 激活内存讨论 中回忆起，这部分内存随序列长度和批大小扩展。自然地，我们可以限制这些，但在实践中，我们不希望受到硬件的限制而只能使用短序列长度进行训练。

为了克服这些问题，是时候探索一个新的、正交的并行轴——张量并行 (TP) 了。与 ZeRO3 依赖于繁重的参数通信不同，TP 建议跨设备分片参数、梯度、优化器状态 *和* 激活值，而无需在 GPU 之间进行任何模型参数通信。

什么？这怎么可能？！让我们一起探索这种看似神奇的方法！🙂

张量并行

为了增加阅读的播客感觉，请随意聆听 NotebookLM 主持人讨论本书的以下部分。

​

因此，我们已经使用 ZeRO 分片了模型的参数、梯度和优化器状态，但是一旦激活内存超过了我们的内存预算，我们就遇到了限制。欢迎使用张量并行 (TP)，这是一种分片权重、梯度和优化器状态以及激活值的方法，并且无需在计算之前收集所有这些值。听起来像是一个梦想！让我们首先看看张量并行如何处理简单的矩阵乘法。

张量并行利用矩阵乘法 A \times B 的数学特性。为了理解它的工作原理，让我们检查使这种并行化成为可能的两个基本方程：

\begin{aligned} &\text{1.} \quad A\cdot B = A \cdot \begin{bmatrix} B_1 & B_2 & \cdots \end{bmatrix} = \begin{bmatrix} AB_1 & AB_2 & \cdots \end{bmatrix} \\ &\text{2.} \quad A\cdot B =\begin{bmatrix} A_1 & A_2 & \cdots \end{bmatrix} \begin{bmatrix} B_1 \\ B_2 \\ \vdots \end{bmatrix} = \sum_{i=1}^n A_i B_i \end{aligned}

这意味着我们可以通过以下方式计算矩阵乘积：1) 单独乘以 B 的每一列，或者 2) 单独乘以每一行并组合结果。在神经网络中，矩阵乘法更常以以下格式表示：X \times W，其中：

在实践中，一个操作的小例子如下所示：

TP diagram

让我们看看如何并行化此操作！在张量并行中，张量将沿特定维度拆分为 N 个分片，并在 N 个 GPU 上分布。矩阵可以按列部分或行部分拆分，从而导致行并行和列并行。我们将在下面看到的一件事是，选择行分片或列分片将需要不同的通信原语。

我们的第一个选择是使用列向分片（也称为 列线性**）：我们将完整的输入矩阵复制到每个工作程序，这需要一个称为 **广播** 的操作，并将权重矩阵拆分为列。然后将输入与部分权重矩阵相乘，最后使用 **all-gather 操作组合结果。

image.png

以下是列向张量并行性的代码实现：

👉 Picotron 中列并行 TP 实现（单击展开）

第二个选项称为行向分片（也称为 行线性**）：正如细心的读者可能猜到的那样，行线性意味着我们将权重矩阵拆分为行块。但是，这也要求我们拆分输入，这需要 **scatter 操作，而不是列线性分片中使用的广播。每个工作程序上的结果已经具有正确的形状，但需要求和才能获得最终结果，因此在这种情况下需要 all-reduce 操作。

我们在这里看到了我们的第四个分布式原语：scatter！

image.png

以下是行向张量并行性的实现：

👉 Picotron 中行并行 TP 实现（单击展开）

现在我们已经掌握了 TP 的基本构建块，让我们看看如何在 Transformer 块内有效地组合它们！

Transformer 块中的张量并行

为了提出要遵循的策略，让我们从一个玩具示例转向一个真正的模型构建块。Transformer 模型由两个主要构建块组成：前馈层 (MLP) 和多头注意力 (MHA)。我们可以对两者都应用张量并行。

前馈部分可以通过“列线性”后跟“行线性”来并行化，这相当于广播以复制输入和前向传播中的 all-reduce。请注意，在实际训练中不需要广播，我们可以确保输入已在 TP 排名之间同步。此设置比从“行线性”开始再到“列线性”更有效，因为我们可以跳过两个拆分操作之间的中间 all-reduce。

image.png

现在我们已经为 Transformer 的前馈部分找到了有效的模式，让我们看一下多头注意力块 (MHA)。

我们通常可以遵循类似的方法，其中 Q、K 和 V 矩阵以列并行方式拆分，输出投影沿行维度拆分。对于多头注意力，列并行方法具有非常自然的解释：每个工作程序计算单个或部分头的注意力。相同的方法也适用于 多查询** (MQA) 或 **分组查询注意力 (GQA)，其中键和值在查询之间共享。

但是值得注意的是，张量并行度不应超过 Q/K/V 头的数量，因为我们需要每个 TP 排名都有完整的头（否则我们无法在每个 GPU 上独立计算注意力，并且我们需要额外的通信操作）。在使用 GQA 的情况下，TP 度实际上应小于 K/V 头的数量。例如，LLaMA-3 8B 有 8 个键/值头，因此张量并行度最好不要超过 8。如果我们为此模型使用 TP=16，我们将需要在每个 GPU 上复制 K/V 头，并确保它们保持同步。

image.png

最后请注意，张量并行仍然不是训练的万能药。我们在模型的计算路径中直接添加了几个分布式通信原语，这些原语很难与计算完全隐藏/重叠（就像我们在 ZeRO 中所做的那样），我们的最终性能将是在计算和内存增益与添加的通信开销之间权衡的结果。让我们举例说明：

Forward pass in Tensor Parallelism

查看张量并行 MLP（同样适用于注意力）中的操作时间线，我们可以更好地理解所涉及的权衡。在每个解码器层的前向传播中，我们使用 AllReduce 操作遇到一个同步点，该操作无法与计算重叠。这种 *公开的通信* 开销是必要的，以便在应用最终 LayerNorm 之前组合跨张量并行排名的部分结果。

张量并行确实有助于减少矩阵乘法的激活内存，因为中间激活值在 GPU 之间分片。但是，我们仍然需要收集 LayerNorm 等操作的完整激活值，这意味着我们没有获得我们可能获得的全部内存优势。此外，TP 引入了显著的通信需求，这些需求在很大程度上取决于网络基础设施。无法完全隐藏这种特定的 AllReduce 在计算背后意味着它直接增加了前向传播的关键路径。

让我们更好地了解权衡，因为我们扩展了 TP 度：

虽然增加 TP 会导致每个 GPU 吞吐量降低（左图），但它能够处理更大的批大小（右图），这说明了分布式训练中计算效率和内存可用性之间的权衡。

在实践中，正如我们在左图中看到的那样，当我们扩展到 8 个 GPU 以上时，张量并行的通信开销变得尤其明显。虽然单个节点内的张量并行可以利用快速 NVLink 互连，但跨节点需要较慢的网络连接。我们观察到从 TP=8 到 TP=16 时吞吐量显著下降，而从 TP=16 到 TP=32 时下降幅度更大。在更高的并行度下，通信开销变得非常高，以至于它迅速占据了计算时间的主导地位。

话虽如此，张量并行通过在 GPU 之间分布模型参数、梯度、优化器状态和激活值（在某种程度上）来提供重要的内存使用量优势。让我们检查一下这种效应对 70B 参数模型的影响：

增加张量并行度会减少每个 GPU 上模型参数、梯度和优化器状态所需的内存，从而使我们能够开始将大型模型放入单个 8 GPU 节点中。

有没有办法从这项技术中获得更多好处？我们已经看到，层归一化和 dropout 仍然需要在每个 GPU 上收集完整的激活值，这在一定程度上抵消了内存节省。我们可以通过找到并行化这些剩余操作的方法来做得更好。

📝 注意

关于张量并行训练中层归一化的一个有趣的注意事项——由于每个 TP 排名在 all-gather 后看到相同的激活值，因此层归一化权重实际上不需要 all-reduce 来在反向传播后同步它们的梯度。它们自然会在各个排名之间保持同步。但是，对于 dropout 操作，我们必须确保在 TP 排名之间同步随机种子以保持确定性行为。

接下来让我们探索张量并行的一个小而自然的扩展，称为 序列并行，它正是这样做的。

序列并行

序列并行 (SP) 涉及沿输入序列维度而不是跨隐藏维度拆分模型的激活值和计算，这些部分不由张量并行 (TP) 处理，例如 Dropout 和 LayerNorm。

📝 注意

术语序列并行有点过载：本节中的序列并行与张量并行紧密耦合，并应用于 dropout 和层归一化操作。但是，当我们移动到更长的序列时，注意力计算将成为瓶颈，这需要诸如环注意力之类的技术，这些技术有时也称为 *序列并行*，但我们将它们称为 *上下文并行* 以区分这两种方法。因此，每次你看到序列并行时，请记住它与张量并行一起使用（与上下文并行相反，上下文并行可以独立使用）。

这是必需的，因为这些操作需要访问完整的隐藏层维度才能正确计算。例如，LayerNorm 需要完整的隐藏层维度来计算均值和方差：

\text{LayerNorm}(x) = \gamma \cdot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta

其中 \mu = \text{mean}(x) 和 \sigma^2 = \text{var}(x) 是跨隐藏层维度 h 计算得出的。

因此，即使这些操作在计算上很便宜，它们仍然需要大量的激活内存，因为它们需要完整的隐藏层维度。SP 允许我们通过沿序列维度拆分来分片此 内存 负担，跨 GPU。

在实践中，我们将从左图转到右图：

in forward: f = no-op ; f* = all-reduce ; g = all-gather ; g* = reduce-scatter in backward: f = all-reduce ; f* = no-op ; g = reduce-scatter ; g* = all-gather SP region needs full hidden_dim

该图显示了我们如何使用不同的集体操作（标记为“f”和“g”）在张量并行和序列并行区域之间转换。关键挑战是如何有效地管理这些转换，同时保持低内存使用量并保持正确性。

在前向传播中：

在反向传播中：

这些操作“f”和“f*”被称为 共轭 对，因为它们相互补充——当一个在前向传播中是 no-op 时，另一个在反向传播中是 all-reduce，反之亦然。

对于序列并行 (SP)，我们使用标记为“g”和“g*”的不同操作。具体来说，我们避免在 SP 区域中使用 all-reduce，因为这将需要收集完整的激活值并增加我们的峰值内存使用量，从而破坏 SP 的目的。

那么这里实际发生了什么？正如一位著名的 LLM 所说，让我们逐步了解一下：

初始 LayerNorm（SP 区域）

  • 输入张量 X1* 和 X2* (b,s/2,h) 进入 LayerNorm，已跨序列维度拆分
  • 每个 GPU 在其序列块上独立计算 LayerNorm，并给出 Y1* 和 Y2*

第一次转换 (SP → TP)

  • “g”操作（all-gather）将 Y1* 和 Y2* 组合回完整序列长度
  • 恢复 Y (b,s,h)，因为列线性需要完整的隐藏层维度 h

第一个线性层（TP 区域）

  • A1 是列线性的，因此它沿隐藏层维度拆分 Y
  • GeLU 独立应用于每个 GPU
  • Z1* 是 (b,s,h/2)

第二个线性层（TP 区域）

  • B1 是行线性的，因此它恢复了隐藏层维度
  • W1 是 (b,s,h)

最终转换 (TP → SP)

  • “g*”操作（reduce-scatter），它为先前的行线性正确性而减少，同时沿序列维度分散
  • W1* 是 (b,s/2,h)
image.png

序列并行的一个关键优势在于，它减少了我们需要存储的最大激活值大小。在仅使用张量并行的情况下，我们必须在各个点存储形状为 (b,s,h) 的激活值。但是，通过序列并行，最大激活值大小减少到 \frac{b \cdot s \cdot h}{tp}，因为我们始终沿序列维度或隐藏层维度拆分。

很难跟踪在 TP 和 TP/SP 中以不同方式分片的所有部分——相信我们，我们也发现很难映射，因此我们制作了这个小表格来总结激活值（又名 hidden_states）形状在正向传播期间如何跨隐藏层维度 h 和序列维度 s 变化：

区域 仅 TP TP 与 SP
进入 TP（列线性） h：分片（weight_out 已分片）
s：完整		 h：分片（weight_out 已分片）
s：all-gather 到完整
TP 区域 h：分片
s：完整		 h：分片
s：完整
退出 TP（行线性） h：完整（weight_out 是完整 + all-reduce 以确保正确性）
s：完整		 h：完整（weight_out 是完整 + reduce-scatter** 以确保正确性）
s：**reduce-scatter 到分片
SP 区域 h：完整
s：完整		 h：完整
s：分片

对于嵌入层：

区域 Vanilla TP TP 与 SP
嵌入层（在词汇表上分片的行线性） h：完整（weight_out 是完整 + all-reduce 以确保正确性）
s：完整		 h：完整（weight_out 是完整 + reduce-scatter** 以确保正确性）
s：**reduce-scatter 到分片

通过使用序列并行，我们可以实现更大的激活内存节省，从而使我们能够比仅使用张量并行时将批大小和序列长度推得更远。让我们看看这对我们之前的 70B 模型示例意味着什么：

正如我们所见，我们再次大大降低了每个 GPU 的最大内存使用量，使我们能够使用 TP/SP=16 来适应 16k tokens 的序列长度，这比 vanilla TP 情况有所改进！（TP=16 仍然有点大，正如我们在上一节中看到的那样，但我们将在下一节中看到如何改进这一点）。

你可能会问自己的一个问题是，使用 TP+SP 是否比 vanilla TP 产生更多的通信？嗯，是的，也不是。在 vanilla TP 的前向传播中，我们每个 Transformer 块有两个 all-reduce，而在 SP 中，我们每个 Transformer 块有两个 all-gather 和两个 reduce-scatter。因此，SP 的通信操作数量是 TP 的两倍。但是，由于 all-reduce 操作可以分解为 all-gather + reduce-scatter（请参阅附录中的 快速关注环 AllReduce 部分），因此它们在通信方面实际上是等效的。反向传播的推理相同，因为我们只使用每个操作的共轭（no-op ↔ allreduce 和 allgather ↔ reducescatter）。

如果你一直密切关注，你就会注意到我们在每个层中讨论 4 个通信操作（注意力为 2 个，MLP 为 2 个）。以下是使用张量 + 序列并行时 MLP 分析的外观：

tp_sp_overlap.svg

与 vanilla TP 类似，TP+SP 也无法轻松地与计算重叠，这使得吞吐量在很大程度上取决于通信带宽。再次强调，与 vanilla TO 类似，TP+SP 通常仅在节点内完成（将 TP 度保持在每个节点 GPU 数量以下，例如 TP≤8）。

我们可以基准测试这种通信开销如何随着我们扩展张量并行而变得越来越成问题。让我们测量吞吐量和内存利用率，因为我们使用 SP 为 3B 模型（4096 seqlen）扩展 TP：

在这里，同样存在计算效率（左图）和内存容量（右图）之间的权衡。虽然更高的并行度通过减少激活内存来支持处理显著更大的批大小，但它们也会降低每个 GPU 的吞吐量，尤其是在超过对应于每个节点 GPU 数量的阈值之上。

让我们总结一下我们的观察结果：

我们已经看到 TP 如何通过沿隐藏层维度拆分注意力和前馈操作来帮助我们在多个 GPU 之间分片激活值，以及 SP 如何通过沿序列维度拆分来自然地补充剩余操作。

📝 注意

由于 SP 区域中的 LayerNorm 对序列的不同部分进行操作，因此它们的梯度在 TP 排名中会有所不同。为了确保权重保持同步，我们需要在反向传播期间 all-reduce 其梯度，类似于 DP 如何确保权重保持同步。然而，这是一种较小的通信开销，因为 LayerNorm 的参数相对较少。

但是，TP 和 SP 有两个限制：1) 如果我们扩展序列长度，激活内存仍然会在 TP 区域中爆炸，并且 2) 如果模型太大而无法使用 TP=8 放入，那么我们将看到由于节点间连接而导致的巨大减速。

我们可以使用上下文并行来解决问题 1)，并使用流水线并行来解决问题 2)。让我们首先看看上下文并行！

上下文并行

通过张量并行和序列并行，我们可以显著降低每个 GPU 的内存需求，因为模型权重和激活值都在 GPU 之间分布。但是，当训练更长序列的模型时（例如，当扩展到每个序列 128k 或更多 tokens 时），我们可能仍然会超出单个节点上的可用内存，因为我们在 TP 区域内仍然需要处理完整的序列长度。

此外，即使我们使用激活值的完全重计算（这会带来约 30% 的沉重计算开销），我们仍然需要在内存中保留一些层边界处的激活值，这些激活值随序列长度线性扩展。让我们看一下上下文并行如何帮助我们：

上下文并行的核心思想是将序列并行方法（又名沿序列长度拆分）应用于我们在其中应用张量并行的模块。因此，我们将沿两个维度拆分这些模块，从而也减少序列长度的影响。在你已经涵盖的所有内容之后，你将发现这种方法非常直观，但是……这里面有一个诀窍，所以请保持清醒！

对于上下文并行；就像序列并行一样，我们将沿序列维度拆分输入，但我们现在将此拆分应用于整个模型，而不是像以前的张量 + 序列并行那样仅应用于模型的序列并行区域。

拆分序列不会影响大多数模块（如 MLP 和 LayerNorm），在这些模块中，每个 token 都是独立处理的。它也不需要像 TP 那样昂贵的通信，因为只有输入被拆分，而不是权重矩阵。与数据并行类似，在计算梯度后，会启动 all-reduce 操作以同步上下文并行组之间的梯度。

但有一个重要的例外，因为我们需要特别注意 注意力块**（哈哈……双关语 :D）。在注意力模块中，每个 token 都需要访问来自 **所有 其他序列 tokens 的键/值对，或者在因果注意力的情况下，至少要注意每个先前的 token。

由于上下文并行跨 GPU 沿序列维度拆分输入，因此注意力模块将需要 GPU 之间的完全通信来交换必要的键/值数据。

如果我们天真地这样做，听起来会非常昂贵。有没有一种方法可以相当高效和快速地做到这一点！谢天谢地，有一种方法：处理键/值对通信的核心技术称为 *环注意力*。

📝 注意

上下文并行与 Flash Attention（稍后会详细介绍）有一些概念上的相似之处——这两种技术都依赖于在线 softmax 计算来减少内存使用量。虽然 Flash Attention 专注于优化单个 GPU 上的注意力计算本身，但上下文并行通过在多个 GPU 之间分布序列来实现内存减少。

发现环注意力

在此注意力机制的实现中，每个 GPU 首先启动异步通信操作，将其键/值对发送到其他 GPU。在等待其他 GPU 数据时，它会计算其内存中已有的数据部分的注意力分数。理想情况下，在计算完成之前，会从另一个 GPU 接收到下一个键/值对，从而允许 GPU 在完成其第一次计算后立即开始下一轮计算。

让我们举例说明。我们将假设我们有 4 个 GPU 和 4 个 tokens 的输入。最初，输入序列沿序列维度均匀拆分，因此每个 GPU 将只有一个 token 及其对应的 Q/K/V 值。假设 Q1、K1 和 V1 表示第一个 token 的查询、键和值，它们位于第一个 GPU 上。注意力计算将需要 4 个时间步骤才能完成。在每个时间步骤中，每个 GPU 执行以下三个连续操作：

  1. 以非阻塞方式将“当前键和值”发送到下一台机器，除了在最后一个时间步骤中，这样我们可以在此步骤完成之前开始后续步骤
  2. 在“当前键和值”上本地计算注意力分数，它通常涉及执行 Softmax(\frac{QK^T}{\sqrt{d}}) * V。
  3. 等待接收来自上一个 GPU 的键和值，然后循环回到步骤 1。其中“当前键和值”现在是从上一个 GPU 接收的键/值。

我们执行这 4 个步骤四次以完成注意力计算。

以下动画显示了具有 4 个 GPU 的整个过程：

ring-attention.gif

你可能在此动画中清楚地看到了为什么作者选择将此方法称为环注意力。

但是，这里有一个大问题，即环注意力的朴素实现导致 GPU 之间存在一些严重的失衡，这来自因果注意力矩阵的形状。让我们通过考虑带有因果注意力掩码的注意力分数矩阵来查看 SoftMax 计算：

cp_attnmask.svg

SoftMax 是按行计算的，这意味着只要 GPU 收到一行的所有 tokens，就可以计算它。我们看到 GPU1 可以立即计算它，因为它从 tokens 1-4 开始，而 GPU1 实际上不需要从任何其他 GPU 接收任何信息。但是，GPU2 将需要等待第二轮才能也接收 1-4，从而拥有 tokens 1-8 的所有值。此外，GPU1 似乎执行的工作比所有其他 GPU 少得多。

让我们看看我们是否可以更好地平衡我们的计算：

Zig-Zag 环注意力 – 一种均衡计算的实现

我们需要一种更好的方法来分布输入序列。这可以通过将 tokens 不是纯粹顺序地分配给 GPU，而是通过稍微混合排序来实现，这样我们在每个 GPU 上都有早期和后期 tokens 的良好组合。这种方法称为 Zig-Zag 注意力，在这种新排列中，注意力掩码将显示均匀的计算分布，但如果你计算彩色方块的数量，你将看到计算现在在所有 GPU 之间是平衡的。

cp_zigzagmask.svg

同时，我们还将看到，为了完成所有行，每个 GPU 将需要来自所有其他 GPU 的信息。

我们有两种通用方法来重叠计算和通信，要么执行通用 all-gather，同时在每个 GPU 上重新组合所有 KV（以 Zero-3 类型的方式），要么根据需要从每个 GPU 逐个收集到每个 GPU：

cp_overlap_allgather.svg

cp_overlap_all2all.svg

这两种实现之间的主要区别在于它们的通信模式和内存使用量：

1. AllGather 实现：

2. All-to-All（环）实现：

All-to-All 方法通常提供更好的内存效率，但代价是通信模式稍微复杂一些，而 AllGather 方法更简单，但在注意力计算期间需要更多临时内存。

我们现在已经看到了如何使用 TP 在一个节点上拆分模型以驯服大型模型，以及如何使用 CP 来驯服长序列的激活值爆炸。好的，让我们继续。

但是，我们仍然知道 TP 在跨节点时无法很好地扩展，那么如果模型权重不容易放入 1 个节点怎么办？这就是另一个并行维度，我们的第四个维度，称为 流水线并行，来救援！

流水线并行

为了增加阅读的播客感觉，请随意聆听 NotebookLM 主持人讨论本书的以下部分。

​

在 张量并行 部分中，我们看到尝试将张量并行扩展到超过每个单节点 GPU 数量（通常为 4 个或 8 个）的限制时，会遇到带宽较低的网络，称为“节点间连接”，这会严重损害我们的性能。我们可以在 all-reduce 操作中清楚地看到这一点，当我们在集群中的多个节点（每个节点有 8 个 GPU）上对其进行基准测试时：

跨不同节点计数的节点间通信带宽测量值，显示 AllReduce、AllGather 和 ReduceScatter 操作的中位数（线）和第 5-95 百分位数范围（阴影区域）。

序列并行和上下文并行可以帮助处理长序列，但如果序列长度不是我们内存问题的根本原因，而是模型本身的大小，则帮助不大。对于大型模型 (70B+)，仅权重的大小就已经突破了单个节点上 4-8 个 GPU 的限制。我们可以通过调用第四个（也是最后一个）并行维度：“流水线并行”来解决此问题。

流水线并行是一种简单但功能强大的技术——我们将模型的层拆分到多个 GPU 上！例如，如果我们有 8 个 GPU，我们可以将第 1-4 层放在 GPU 1 上，将第 5-8 层放在 GPU 2 上，依此类推。这样，每个 GPU 只需要存储和处理模型层的一部分，从而显著降低了每个 GPU 的内存需求。让我们看看流水线并行在 8B 模型的内存使用量方面的实际效果：

查看上图，我们注意到一些有趣的事情：虽然模型参数在 GPU 之间得到了很好的拆分，但每个 GPU 上的激活内存保持不变！这是因为每个 GPU 仍然需要处理完整批次的数据，只是使用不同的层。来自一个 GPU 层的激活值将发送到下一个 GPU 以继续前向传播。

这引入了一种新的通信模式：与我们在数据并行中使用 ZeRO-3 通信参数不同，我们现在正在“流水线”中在 GPU 之间顺序传递激活张量。虽然概念上很简单，但有效地实现此技术非常棘手。让我们直接深入细节！

在各个节点上拆分层——全前向，全反向

因此，假设我们只是将层分散在多个设备上，例如，第一个 GPU 将采用前几层，第二个 GPU 将采用模型的第二部分，依此类推。我们模型的前向传播现在只是简单地沿着模型顺序传递数据批次，从而连续使用每个计算设备。

我们有一个直接的第一个优势：所需的互连带宽保持相当低，因为我们仅在模型深度的少量位置发送中等大小的激活值。与例如张量并行中的通信相比，这可能会产生巨大的差异，张量并行中的通信在每一层中发生多次。

但是，也许你开始感觉到即将到来的麻烦：“顺序地”** 和 **“连续地” ？！在我们讨论了计算和通信重叠之后，这在并行计算的世界中听起来效率不高。

确实如此，读者！流水线并行中的主要挑战将是如何有效地规避 PP 的顺序性质，以使我们的 GPU 始终处于忙碌状态，并避免一个 GPU 计算而其他 GPU 等待的情况。以下是当通过模型执行朴素而简单的前向和反向传播时，我们的 GPU 利用率的外观（这里的数字表示模型层）：

image.png

模型具有 16 层并分布在 4 个 GPU 上的流水线并行的示例。数字对应于层 ID。

剩余的空闲时间以灰色表示，通常称为“气泡”，在花费这么多时间优化吞吐量之后，看到这种情况可能会让你心碎。

我们可以通过查看我们因气泡而损失的时间来量化流水线设置的效率。假设 t_f 和 t_b 分别是前向传播和反向传播的时间，如针对一个微批次和一个流水线阶段测量的那样（一个简单的假设通常是 t_b \approx 2 \times t_f，你可以在上图中看到）。如果我们能够完美地并行化，理想的总时间将是 t_{id}=t_f + t_b。但是，我们可以在图上计算出，由于流水线气泡，存在 t_{pb}=(p-1)*(t_f+t_b) 的额外时间（其中 p 是流水线并行的程度，即上图中的 GPU 数量），即每个 GPU 在其他 GPU 计算时等待的时间。

我们可以计算额外气泡时间与理想时间的比率：

r_{bubble} = \frac{(p-1)*(t_f+t_b)}{t_f+t_b} = p-1

随着我们添加更多阶段，气泡时间因此增加，利用率下降。正如我们所见，在朴素的实现中，气泡可能非常大！

值得庆幸的是，已经设计了各种流水线并行方案来 减少气泡的大小。

让我们从我们的工具箱中取出第一个工具，并考虑将我们的批次拆分为可以并行或几乎并行处理的较小位大小部分，就像我们之前在数据并行中所做的那样。现在，当第二个 GPU 忙于处理微批次 1 时，第一个 GPU 已经可以开始处理微批次 2。以下是使用 8 个微批次的计划：

pp_afab2.svg

上述计划称为 全前向全反向 (AFAB) 计划，因为我们首先执行所有前向传播，然后才执行所有反向传播。优点是前向和反向步骤仍然通常是顺序的，因此我们保留了模型训练代码的一般组织。这使得此 PP 实现成为最容易实现的实现之一。

你可以在 picotron 中找到 AFAB 流水线的完整实现：

👉 Picotron 中的 AFAB PP 实现（单击展开）

让我们估计一下此示例中的气泡。与我们的第一个示例的区别在于，处理 m 微批次的理想时间现在是 t_{id} = m*(t_f+t_b)：

r_{bubble} = \frac{(p-1)*(t_f+t_b)}{m*(t_f+t_b)} = \frac{p-1}{m}

正如我们所见，我们可以通过添加更多微批次来对抗流水线阶段的一些低效率，从而将气泡的大小减少 m 倍。

然而，与气泡一样烦人的是，存储所有激活值所需的内存存储。我们需要将所有激活值保存在内存中，直到我们达到反向阶段，这会导致 PP 的这些实现中的内存快速爆炸。我们能否做得更好并避免这种内存爆炸？

由于内存爆炸是由我们为反向传播存储的激活值触发的，因此让我们尝试看看是否可以在我们仍在执行其他前向传播部分计算的同时开始执行反向传播。这将使我们能够在尽可能早的时候删除反向传播所需的一些激活值。

单次前向单次反向和 Llama 3.1 方案

此计划称为 单次前向单次反向 (1F1B)，因为中间/稳定状态涉及交替执行一次前向传播和一次反向传播。总的想法是尽快开始执行反向传播。该计划如下所示：

image.png

如果你仔细计算，你会发现气泡仍然具有相同的大小，因此我们的训练效率并没有显著提高。但是，我们只需要存储 p 微批次的激活值（其中 p 是流水线并行的程度），而不是 m（其中 m 是微批次的数量），这可以减少我们在 AFAB 计划中遇到的激活内存爆炸。因此，我们可以添加更多微批次，从而实际减少气泡。

此设置的主要复杂性（在上图中可见）是前向传播和反向传播不再是完全顺序的，而是在设备之间并行执行和交错执行。这意味着我们将不得不独立于每个设备上的前向传播到反向传播的转换进行计划，而不是像往常一样在简单而常见的中央训练循环中进行。

这是实现流水线并行通常需要对训练代码以及建模代码进行相当广泛修改的原因之一。

你也可以在 picotron 中找到 1F1B 的完整实现：

👉 Picotron 中的 1F1B PP 实现（单击展开）

让我们看一下 1F1B 流水线并行计划在实践中如何随我们在集群上进行的一些基准测试进行扩展：

Throughput scaling of Pipeline Parallelism with varying microbatch sizes

在左侧，当微批次数量等于或小于 PP 度减一时 (m = p - 1)，我们看到流水线气泡的危害有多大——性能很低，甚至随着我们扩展 PP 而下降。右图显示，使用比 PP 度多得多的微批次 (m = 32 \gg p - 1) 有助于提高低 PP 度性能，同时在非常大的 PP 度下仍然受到限制。在实践中，不可能任意增加微批次的数量来保持 m \gg p - 1 的比率，因为我们最终会受到目标全局批大小的约束。由于最大可能的微批次数量，当我们添加更多 PP 度时，我们将最终不得不根据 r_{bubble} = \frac{p - 1}{m} 增加气泡大小。

有趣的是，在少量微批次下，当从一个节点 (p = 8) 扩展到两个节点 (p = 16) 时，性能仅下降 14%——比张量并行性好得多，张量并行性通常在类似的跨节点场景中看到约 43% 的性能下降。当遇到较低带宽的节点间网络时，这种类型的行为使流水线并行对于跨多个节点的分布式训练特别有吸引力。

虽然 1F1B 显著减少了我们的激活内存占用，但我们在此最后一张图中看到，流水线气泡仍然是主要的效率瓶颈。由于气泡大小仍然与流水线阶段的数量成正比，我们让宝贵的 GPU 计算处于空闲状态。我们能否设计一个更智能的计划来最大限度地减少这种浪费的计算时间？

交错阶段

1F1B 计划让我们改进了内存使用量，但对空闲气泡的大小没有太大改进。我们还能以某种方式进一步推动这一界限吗？

事实证明，如果我们愿意引入一些额外的通信操作，这是可能的。是时候讨论 交错阶段 了。

到目前为止，我们已经沿着模型深度维度朴素地切分了我们的模型，例如，将第 1-4 层托管在第一个 GPU 上，将第 5-8 层托管在第二个 GPU 上。但是，我们还可以通过其他方式来考虑切分我们的层，例如，让奇数层 1、3、5、7 在第一个 GPU 上，偶数层 2、4、6、8 在第二个 GPU 上。

这通常可以看作是一种“循环流水线”，其中微批次将在 GPU 之间循环移动，因为它会通过模型进行前向传播。让我们以图形方式看一下它是如何工作的：

pp_1f1b_interleaved.svg

模型层分布在 4 个 GPU 上的交错流水线并行的示例。数字仍然对应于微批次 ID，但为了清楚起见，我们以不同颜色对模型的第一个和最后一个层进行了着色，以说明层如何在 GPU 之间分布。

因此，我们看到发生了额外的通信，因为模型多次通过每个 GPU，以进行之前仅需一次传递的相同计算。但是，每个前向和反向传播都除以一个因子 v，其中 v 是每个 GPU 的阶段或模型块的数量，因为我们能够更好地交错前向和反向传播。

\begin{aligned} &t_{pb} = \frac{(p-1)*(t_f+t_b)}{v} \\ &r_{bubble} = \frac{1}{v}\frac{(p-1)*(t_f+t_b)}{m*(t_f+t_b)} = \frac{p-1}{v*m} \end{aligned}

因此，我们现在可以通过添加微批次和交错阶段来减少气泡，但请注意，定量地，通信量也增加了 v，因此这是一种权衡。在下图中，你可以看到 PP 设置的几种配置，其中 p=8，其中 m=1, v=1 的特殊情况对应于朴素流水线并行，配置为 v=1 的配置是 AFAB 或 1F1B 设置，而 v \neq 1 是交错配置。

计划也变得更加复杂，因为我们必须决定在给定的 GPU 上以及在给定的时刻，我们是优先处理通过后续层的早期微批次——这意味着我们尽快关闭前向和反向循环（所谓的“深度优先”，即优先考虑尽快将批次从模型中取出）——还是我们优先处理首先通过早期层的后期微批次（所谓的“广度优先”，即优先考虑尽可能多地填充流水线）。在优秀的“广度优先流水线”论文中详细解释了这种选择。

你现在拥有所有元素来理解 Llama 3.1 中的流水线并行方法，该方法使用单次前向单次反向设置，其中包含交错阶段和在深度优先和广度优先之间可调的优先级设置。

pp_llama3.1_schedule.png

但是，我们尚未达到可能的流水线计划的尽头，最近提出了一些将气泡 减少到几乎为零 的方法！这些技术例如在 DeepSeek V3/R1 实现 中使用。激发了你的好奇心？在我们离开流水线并行的世界之前，让我们最后快速浏览一下这些神奇的计划！

零气泡和 DualPipe

最近提出了更复杂的方法来减少气泡，这些方法接近于达到“零气泡”状态。这里的秘密是以更细粒度的级别拆分所涉及的操作，以便以最有效的方式交错它们。例如，DeepSeek V3/R1 中的流水线实现方法（称为 DualPipe）接近于达到零气泡状态。

让我们简要地了解一下这是如何工作的，方法是总结 ZeroBubble 工作，这是 DualPipe 的先驱。ZeroBubble 的基本观察是，通过矩阵乘法的反向传播实际上涉及两个分离的操作：输入的反向操作 (B) 和权重的反向操作 (W)：

虽然 B 的输出（输入的反向传播）对于执行较低层的反向传播是必要的，但权重的反向传播 W 对于其余的反向传播不是必需的，通常只需要在优化器步骤之前执行。我们可以在下图中看到这一点：

image.png

这意味着 W 可以灵活地安排在同一阶段的相应 B 之后的任何位置。这允许战略性地放置 W 以填充流水线气泡。右上角的 ZB-H2 计划是（理论上的）零气泡计划的示例，该计划利用了这种细粒度的分解。

image.png

在顶部（来自 ZeroBubble 论文的图 2）：经典的 1F1B 计划，交错前向和反向传播，但保持粗粒度的反向传播。在底部两张图（来自 ZeroBubble 论文的图 3）中，ZeroBubble 计划的两种变体，将反向操作拆分为“B”和“W”更细粒度的操作。最后一个计划（所谓的“ZB-H2”）是（理论上的）零气泡计划的示例，该计划利用了这种细粒度的分解。

DeepSeek 的 DualPipe 在其 V3 技术报告 中引入了这种分解方法的扩展，以适应从 PP 维度的两端传播的额外两个流的情况，这些流被交错以进一步最小化 GPU 中的空闲时间。此计划显示在以下计划图中，并且比以前的计划更复杂：

image.png

一般来说，完全优化此类复杂计划涉及仔细测量各种细粒度操作的持续时间，并求解 ILP 以最小化最终气泡时间。例如，有关执行此类计划的启发式方法和算法的讨论，请参阅 ZeroBubble 论文。因此，ZeroBubble 和 DualPipe 计划过于复杂，我们无法在此处提供代码片段，但你应该开始对所涉及的概念有一个大致的了解。

这结束了我们对流水线计划和气泡世界的游览。我们希望你喜欢这次导览！

现在是时候转向我们将详细介绍的最后一种并行方法了，我们可以使用该方法有效地训练大型模型：专家并行。

专家并行

这是我们要讨论的最后一种并行方法。在处理它之前，如果你不了解混合专家模型，请随意阅读我们在不久前发布的 这篇较短的博客文章，它应该可以帮助你更好地理解混合专家模型 (MoE) 架构。

混合专家模型最近通过 GPT-4、Mixtral 或最近的 DeepSeek-V3/R1 等模型获得了新的关注和可见性。基本思想是，对于每一层，我们都可以有几个并行模块，并将 tokens 路由到其中一个或另一个以进行不同的处理，而不是每个层只有一个前馈模块。

ep_moe.png

来自 Switch Transformers 论文 的 MoE 层插图

MoE 层的设计实际上可以轻松地实现跨专家维度进行并行化，我们将其称为 专家并行 (EP)。由于前馈层是完全独立的，我们可以简单地将每个专家的前馈层放在不同的工作程序上。与 TP 相比，它要轻量得多，因为我们不需要拆分矩阵乘法，我们只需要将 token 的隐藏状态路由到正确的专家。

在实践中，EP 通常会与其他形式的并行性结合使用——例如数据并行性。这是因为 EP 仅影响 MoE 层，而不分片输入 tokens（与沿序列长度维度分片 tokens 的上下文并行不同）。这意味着，如果我们仅使用 EP，我们的 GPU 将为所有非 MoE 块执行冗余计算。通过将 EP 与 DP 结合使用，我们可以有效地在 GPU 之间分片专家和输入批次，正如我们在简化的图表中看到的那样：

ep_schema.png

来源：混合专家调查

但我们不要操之过急——我们的 以下部分 将专门讨论不同并行策略之间的所有交互，因此，如果你还不理解最后一张图，请不要担心。

在实践中，有一些技巧可以使 EP 有效地工作，并且它们与模型设计紧密相关。例如，DeepSeek-V3 在路由器中强制执行约束，确保每个 token 最多发送到 M 个节点（在他们的情况下为 4 个），以将 tokens 保留在单个节点上并减少通信开销。虽然专家并行已经存在一段时间了，但随着 MoE 架构越来越受欢迎，它现在才获得新的关注。

我们计划很快在 picotron/nanotron 中添加一个更完整的 EP 示例，敬请关注更多信息！

5D 并行概览

恭喜读者，你现在已经了解了可用于扩展模型训练的 5 种并行策略：

  1. 数据并行 (DP) – 沿批次维度
  2. 张量并行 (TP) - 沿隐藏层维度
  3. 序列并行和上下文并行 (SP/CP) - 沿序列维度
  4. 流水线并行 (PP) - 沿模型层
  5. 专家并行 (EP) - 沿模型专家

以及可以与数据并行结合使用的用于减少内存的 3 种 ZeRO 策略：

  1. ZeRO-1 – 在 DP 副本之间分片优化器状态
  2. ZeRO-2 – 在 DP 副本之间分片优化器状态和梯度
  3. ZeRO-3 – 在 DP 副本之间分片优化器状态、梯度和参数

在这个阶段，你可能很好奇的一个方面是，所有这些并行和 ZeRO 策略如何相互比较以及如何相互作用。换句话说，我们应该使用哪些策略并有效地组合在一起，以及我们应该将哪些策略分开？

让我们看一下相似之处和相互作用。我们将首先并排比较流水线并行和 ZeRO-3，因为它们有一些非常相似之处，但也存在重要差异。

流水线并行与 ZeRO-3 - PP 和 ZeRO-3 都是在多个 GPU 上分区模型权重并在模型深度轴上执行通信/计算的方法（例如，在 ZeRO-3 中，我们在计算时预取下一层）。这意味着在这两种情况下，完整的层操作都在每个设备上计算，这与 TP 或 EP 不同，在 TP 或 EP 中，计算是在子层单元上执行的。

但是，PP 和 ZeRO-3 方法之间存在一些主要差异：

ZeRO-3 流水线并行
每个计算单元存储 仅一层的一小部分 一层完整层
通信用于传输 权重 激活值
编排 模型无关 模型无关
实施挑战 难以处理模型分区和通信 难以处理高效的 PP 计划
扩展考虑因素 首选大 mbs 和 seq\_len 以隐藏通信 首选大 \text{grad\_acc} 以隐藏气泡

正如你所见，ZeRO-3 和 PP 解决了相同的挑战，但涉及不同的方法，两者之间的选择将取决于你是否决定将通信重点放在权重还是激活值上。虽然它们可以结合使用，但在实践中并不常见，因为这样做需要显著增加全局批大小以分摊通信成本，从而在全局批大小、模型大小、网络带宽和训练效率之间创建权衡。如果你决定将它们结合使用，则应将 ZeRO-3 配置为在 PP 微批次序列期间将权重保留在内存中，以尽可能减少不必要的通信开销。

另一方面，ZeRO-1 和 ZeRO-2 专注于优化器状态和梯度，可以轻松地与流水线并行结合使用，并且是对流水线并行的补充。将它们结合使用不会引起任何特殊的新挑战。例如，DeepSeek-v3 的训练使用了 PP 与 ZeRO-1 的组合（原文如此）。

张量并行（与序列并行一起使用）自然是互补的，可以与流水线并行和 ZeRO-3 结合使用，因为它依赖于矩阵乘法的分配律，这允许权重和激活值被分片和独立计算，然后再组合。

TP & SP diagram
​

我们不希望仅将 TP 用于并行化的主要原因是，在实践中，TP 有我们在之前的章节中讨论过的两个限制：首先，由于其通信操作是计算关键路径的一部分，因此很难扩展到一定程度以上，此时通信开销开始占主导地位。其次，与 ZeRO 和 PP 不同，ZeRO 和 PP 与模型无关，TP 需要仔细处理激活分片——有时沿隐藏层维度（在 TP 区域中），有时沿序列维度（在 SP 区域中）——这使得正确实现起来更加繁琐，并且需要模型特定的知识来确保在整个过程中的正确分片模式。

因此，在组合并行策略时，TP 通常会保留用于高速节点内通信，而 ZeRO-3 或 PP 可用于跨越较低速度节点间通信的并行组，因为它们的通信模式需要更少的带宽（对于 PP）或可以更容易地与计算重叠（对于 ZeRO-3）。在组合这些技术时，主要考虑因素是以高效的方式将 GPU 组织成每个并行维度的组，以最大化吞吐量并最小化通信开销，同时注意 TP 的扩展限制。例如，为 TP 通信的 GPU 组应保留在节点内。

上下文并行** 和 **专家并行 也有助于我们分片激活值，并且可以被视为对 TP 的补充。第一个处理长序列，而第二个支持分布式混合专家训练，它们可以组合在一起而不会出现任何特殊问题。

上下文并行 (CP) 专门针对使用非常长序列进行训练的挑战，方法是跨 GPU 沿序列维度分片激活值。虽然大多数操作（如 MLP 和 LayerNorm）可以独立处理这些分片的序列，但注意力层需要通信，因为每个 token 都需要访问来自完整序列的键/值。正如我们在 CP 部分 中看到的那样，这可以通过环注意力模式有效地处理，该模式将计算和通信重叠。当扩展到极端序列长度（128k+ tokens）时，CP 尤其有价值，即使使用完全激活重计算，注意力的内存需求在单个 GPU 上也是令人望而却步的。

CP diagram

专家并行 (EP) 专门针对训练混合专家 (MoE) 模型的挑战，方法是在 GPU 之间分片专门的“专家”，并在计算期间动态地将 tokens 路由到相关专家。EP 中的关键通信操作是将 tokens 路由到其指定专家并将结果收集回来的 `all-to-all` 操作。虽然此操作引入了一些通信开销，但它显著提高了模型容量的扩展能力，因为在推理（和训练）期间，每个 token 仅由总参数中一小部分参数处理。在分布式训练/推理方面，当模型扩展到大量专家时，跨 GPU 分区专家变得相关。

EP diagram

📝 注意

EP 和 DP 在输入处理方面的这种相似性是为什么某些实现将专家并行视为数据并行的一个子组，主要区别在于 EP 使用专门的专家路由，而不是让所有 GPU 通过相同的模型副本处理输入。

范围和重点 让我们快速总结一下这些不同并行策略对模型的哪些子部分影响最大：

张量 + 序列并行 上下文并行 专家并行
沿隐藏层/序列维度分片权重和激活值 沿序列维度分片激活值 分片专门的专家权重和激活值
矩阵乘法运算（列/行线性）的通信 注意力键/值的通信 token 路由到专家的通信
需要模型特定的实现 模型无关，注意力除外 模型无关，MoE 层除外
首选高带宽节点内通信 首选长序列长度 需要 MoE

总结所有内容—— 现在如何在一个图中收集和组合我们所看到的所有技术。是的，我们准备迎接挑战！

在这个摘要图中，你将找到一个 Transformer 层的激活值和模块的图示——以其 MoE 变体为例。我们还说明了各种并行方向以及我们在之前所有章节中讨论过的通信操作。

image.png

我们还可以并排表示所有这些策略的 完整概述，以了解内存节省情况。我们将绘制它们，其中包含不同的序列长度以及选择性（顶部）和完全（底部）重计算，以便你可以看到它们如何与激活值一起发挥作用：

5Dparallelism_8Bmemoryusage.svg

让我们在本节的结尾处以高层次的视角来看待所有这些技术、它们的主要基本思想和主要瓶颈：

方法 内存节省专门应用于 并行/分片维度 缺点
DP 激活值（减少本地批大小） 批次 受最大批大小限制
PP 模型参数 模型层 空闲气泡和复杂的计划
TP/SP 模型参数和激活值 隐藏层维度/序列长度 需要高带宽通信
CP 激活值 序列长度 在注意力模块中增加通信开销
EP 专家参数 专家维度 需要 MoE 层，增加路由通信开销
ZeRO-1 优化器状态 在 DP 副本之间分片 参数通信开销
ZeRO-2 优化器状态和梯度 在 DP 副本之间分片 参数通信开销
ZeRO-3 优化器状态、梯度和模型参数 在 DP 副本之间分片 参数通信开销

显然，这些技术都不是神奇扩展的万能药，我们通常不得不以某种方式将它们组合起来。我们能否实际提出一些规则，帮助我们找到一个好的起点来选择和组合它们？这将是我们下一节的主题。

寻找最佳训练配置

我们现在已经介绍了所有用于分布式训练和训练更大模型的并行技术，以及如何以及为何可以将它们组合在一起。仍然存在一个普遍的问题：最终我们应该选择哪些技术，以及如何决定特定的组合？

我们在上一节中稍微谈到了这一点，但现在让我们详细地逐步了解一个可能的决策过程，请记住，你始终必须运行一些实验，才能为你的计算集群找到最终的最佳设置，因为它的各种物理属性、网络带宽、每个节点的 GPU、每个 GPU 的内存等。

步骤 1：使训练步骤适应内存

首先，我们需要弄清楚如何使完整的模型实例适应我们的 GPU。通常有两种情况。

GPU 资源丰富的情况 🤑 - 当你有大量 GPU 可用时：

特殊注意事项：

GPU 资源匮乏的情况 😭 - 当你的 GPU 资源可能不足时：

现在我们有了第一个模型实例训练，我们需要确保我们有正确的批大小。

步骤 2：实现目标全局批大小

根据步骤 1 在微批大小和 DP 方面留给我们的结果，我们当前的批大小可能太小或太大。现在是时候达到我们的目标批大小了。

要增加我们当前的全局批大小：

要减小我们当前的全局批大小：

好的，现在我们的模型在模型大小和批大小方面以我们想要的一般配置运行，但是我们是否以最快的方式训练它？现在让我们开始尽可能地优化吞吐量。

步骤 3：优化训练吞吐量

因此，我们希望确保训练尽可能快地运行，以便我们所有宝贵的 GPU 始终得到充分利用。只要内存和通信不是瓶颈，我们就可以尝试以下操作：

基准测试数千种配置

现在我们已经介绍了逐步操作，让我们在现实生活中实现此搜索过程。

你将在 nanotron 仓库中找到几个脚本，你可以使用这些脚本来运行我们上面讨论的所有实验，并能够在现实生活中基准测试你自己的模型和集群。

实际上，我们自己在 数千个分布式配置 上运行了基准测试，涵盖了我们上面讨论的每个模型大小，以及我们可以在其中尝试的大量集群配置（即 8xH100 的 1-64 个节点），以便生成我们迄今为止在本书中介绍的结果。

现在让我们退后一步，收集和分析我们所有基准测试的结果，看看除了理论之外，我们是否可以在真实世界的数据中发现各种配置如何相互竞争。

以下所有基准测试均以 4096 的序列长度和 1M tokens 的全局批大小进行。我们收集了每个模型和集群大小的所有顶级配置，并将它们绘制在以下热图中：

image.png

热图可视化显示了不同模型大小和计算节点计数（我们每个节点有 8 个 GPU）的最佳训练配置。对于每种组合，配置详细信息包括数据并行 (DP)、张量并行 (TP)、流水线并行 (PP)、梯度累积步骤 (GAS)、微批大小 (MBS) 和 ZeRO 优化阶段。颜色强度表示模型 FLOPS 利用率 (MFU)，颜色越亮表示效率越高。

从这种高层次的可视化中，我们可以得出几个重要的见解：

首先，随着我们增加节点数量（更高的并行性），我们观察到效率下降。对于较小的模型，这种效果尤其明显，这些模型的计算与模型大小之比较低。虽然我们通常可以通过增加批大小来补偿小模型大小，但我们受到 1M 的全局批大小限制的约束。

其次，较大的模型提出了不同的挑战。随着模型大小的增加，内存需求也大幅增长。这在节点较少的情况下创建了两种情况：模型根本不适合，或者勉强适合，但由于在 GPU 内存限制附近运行而效率低下（例如，在 4 个节点上训练 80B 参数模型）。

最后，我们的基准测试表明，性能在很大程度上取决于实现质量。当我们首次实施这两种并行策略时，张量并行 (TP) 的性能优于流水线并行 (PP)。在优化了我们的 PP 代码后，它成为更快的选择。现在，当我们改进 TP 实现中的通信重叠时，我们期望它重新获得性能领先地位。

基准测试中学到的经验

本书的目标不仅是讨论理论和实现，还要提供实际数据点。因此，计划很简单：让我们为每个模型和多个集群大小（即 8xH100 的 1-64 个节点）运行每种可能的分布式配置。即使在排除不可能的配置后，我们仍然需要运行数千个实验。

从理论上讲，这听起来很容易：我们可以轻松地在我们的集群上启动大型作业数组。然而，当我们启动第一批实验时，麻烦开始了：

在有限的时间内运行所有实验需要额外的工程设计，我们最终花费了大量时间在以下方面：

这些挑战值得单独讲述，但它们教会了我们关于分布式训练基础设施复杂性的宝贵经验。理论上看起来很简单的事情，在实践中通常需要仔细关注许多可移动的部件。

在实践中重现理论结果具有挑战性，尤其是在生产训练代码的可用性有限的情况下。通过像 nanotron 和 picotron 这样的开源项目，我们希望能够帮助使分布式训练技术更易于访问，并协作开发简单高效的代码库，以帮助研究人员和从业人员最大限度地利用其硬件资源。

这结束了我们对 5D 并行分布方法的深入探讨。

退后一步，我们到目前为止的讨论经常依赖于一个关键假设——计算和通信可以在 GPU 上有效地重叠，而不会对计算吞吐量产生任何影响。现实情况更为细致。当使用常见的通信原语（如 NCCL 发送/接收）时，我们面临着计算和通信资源之间隐藏的争用，因为通信内核通常会使用与计算相同的 GPU 流式多处理器 (SM)，从而导致当通信与计算重叠时吞吐量降低。为了真正优化我们的分布式训练，我们需要更深入地了解 GPU 架构本身。

是时候关灯并激活 CUDA 模式了！

深入 GPU – 融合、线程、混合

为了增加阅读的播客感觉，请随意聆听 NotebookLM 主持人讨论本书的以下部分。

到目前为止，我们的讨论一直集中在模型操作的高层组织上。我们已经在各种加速器上移动了计算，同时考虑了一般的内存约束和计算单元的高层调度。

但这忽略了我们可以在更低的层面上进行的优化，方法是仔细了解我们的模型操作是如何在每个 GPU 上调度和执行的。

本节将深入探讨 GPU 架构的更多细节，尤其是 NVIDIA 的 GPU 架构，但一般思想（通常情况下）可以重用于类似的加速器单元。

我们将简要解释 GPU 的组织方式，然后再介绍 Flash-Attention 革命、如何在 GPU 上高效调度工作负载，最后解释如何在 GPU 上高效使用各种精度。

GPU 入门

通常，GPU 具有非常分层的组织结构。在本入门中，我们将讨论保持在概念层面，这对于我们演示的其余部分是必要的。

在计算方面，GPU 由称为 流式多处理器 (SM) 的计算单元阵列组成。每个 SM 包含并控制一组流式处理器，也称为内核。例如，Nvidia H100 GPU 具有 132 个 SM，每个 SM 有 128 个内核，总共有 16,896 个内核（有关详细信息，请参阅 张量内核文档），每个内核都能够同时处理多个线程。

image.png

来源：https://blog.codingconfessions.com/p/gpu-computing

内存方面也是高度分层的，具有多层缓存和内存：寄存器** 是最小的单元，在执行期间对线程是私有的，**共享内存** 和 **L1 缓存** 在单个 SM 上运行的线程之间共享，更高级的是所有 SM 共享的 **L2 缓存**，最后是 **全局内存，它是 GPU 上最大的内存（对于 H100，广告宣传为 80 GB），但也是访问和查询速度最慢的内存。

image.png

来源：https://www.youtube.com/watch?v=ZQKMZIP3Fzg

GPU 的目标是在 GPU 内核上并行运行尽可能多的工作负载，方法是利用计算/内存的分层组织结构。

在 GPU 内核上运行的一段代码称为 内核**。它可以用高级语言（例如 **CUDA** 或 **Triton）编写，然后编译为并行线程执行 PTX，即 NVIDIA GPU 使用的低级汇编。

要运行内核，你还需要一个特定的代码部分，称为 主机代码**，它在 **CPU/主机 上执行，并将负责准备数据分配和加载数据和代码。

// 主机代码 void vecAdd(float* h_A, float *h_B, float *h_c, int n) { // 在设备内存中分配向量 int size = n * sizeof(float); float *d_A, *d_B, *d_C; cudaMalloc(&d_A, size); cudaMalloc(&d_B, size); cudaMalloc(&d_C, size); // 将向量从主机内存复制到设备内存 cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice); cudaMemcpy(d_B, h_B, size, cudaMemcpyHostToDevice); // 调用内核 int threadsPerBlock = 256; int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock; VecAdd<<>>(d_A, d_B, d_C, N); // 将结果从设备内存复制到主机内存 // h_C 包含主机内存中的结果 cudaMemcpy(h_C, d_C, size, cudaMemcpyDeviceToHost); // 释放设备内存 cudaFree(d_A); cudaFree(d_B); cudaFree(d_C); }

用于添加两个向量的 CUDA 内核的主机代码。改编自 https://docs.nvidia.com/cuda/cuda-c-programming-guide/ 和 https://blog.codingconfessions.com/p/gpu-computing

// 设备代码 __global__ void VecAdd(float* A, float* B, float* C, int N) { int i = blockDim.x * blockIdx.x + threadIdx.x; if (i < N) C[i] = A[i] + B[i]; }

包含向量加法内核定义的设备代码，改编自 https://docs.nvidia.com/cuda/cuda-c-programming-guide/ 和 https://blog.codingconfessions.com/p/gpu-computing

内核通常按以下方式调度：

从这些细节中要记住的主要内容是，存在各种大小调整和分配约束（各种内存的大小、并发块的数量以及 warp 中线程的数量），需要考虑这些约束才能以最有效的方式使用 GPU 架构。

大多数时候，你不需要深入到这种精确程度，你可以幸运地重用社区其他成员准备的内核和代码。但在任何情况下，我们都想让你入门了解如何开始使用内核！

如何使用内核提高性能？

如果你希望添加缺少优化内核的新操作或加速现有的 PyTorch 函数，从头开始编写内核似乎是最直接的途径。但是，从头开始创建高性能 CUDA 内核需要丰富的经验和陡峭的学习曲线。通常，更好的入门方法是利用 torch.compile，它通过捕获你的操作并在 triton 中生成低级、高性能内核来动态优化 PyTorch 代码。

假设你想为一个名为指数线性单元的激活函数编写内核：

\text{ELU}(x) = \begin{cases} e^x - 1 & \text{if } x < 0 \\ x & \text{if } x \geq 0 \end{cases}

你可以从一个简单的 pytorch 实现开始，然后在顶部添加 @torch.compile 装饰器：

@torch.compile def elu(x, alpha=1.0): return torch.where(x < 0, alpha * (torch.exp(x) - 1), x)

编译版本和非编译版本之间的区别非常显著，特别是考虑到我们只添加了一个装饰器。下图说明了这种显著的差异（N 是列数）：

image.png

但是，如果这种性能提升不足，你可以考虑实现 Triton 内核。作为起点，你可以查看由 @torch.compile 生成的 triton 内核。为此，你只需将环境变量 TORCH_LOGS 设置为 "output_code"：

export TORCH_LOGS="output_code"

一旦你运行带有 @torch.compile 装饰器的 Python 脚本，它将生成并输出相应的 Triton 内核，在本例中为：

@triton.jit def triton_(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr): xnumel = 100000000 xoffset = tl.program_id(0) * XBLOCK xindex = xoffset + tl.arange(0, XBLOCK)[:] xmask = xindex < xnumel x0 = xindex tmp0 = tl.load(in_ptr0 + (x0), xmask) tmp1 = 0.0 tmp2 = tmp0 < tmp1 tmp3 = tl_math.exp(tmp0) tmp4 = 1.0 tmp5 = tmp3 - tmp4 tmp6 = tl.where(tmp2, tmp5, tmp0) tl.store(out_ptr0 + (x0), tmp6, xmask)

为了提高可读性，我们可以修改变量名称、添加注释并进行细微调整（或让 LLM 为我们做），如下所示：

@triton.jit def elu_kernel(input_ptr, output_ptr, num_elements, BLOCK_SIZE: tl.constexpr): # 计算此块的起始索引 block_start = tl.program_id(0) * BLOCK_SIZE # 为此块创建索引数组 block_indices = block_start + tl.arange(0, BLOCK_SIZE)[:] # 创建一个掩码，以确保仅处理有效索引 valid_mask = block_indices < num_elements # 根据有效索引从输入指针加载输入值 input_values = tl.load(input_ptr + block_indices, valid_mask) # 定义 ELU 参数 zero_value = 0.0 # ELU 激活的阈值 negative_mask = input_values < zero_value exp_values = tl.math.exp(input_values) # 定义 ELU 输出偏移 one_value = 1.0 shifted_exp_values = exp_values - one_value output_values = tl.where(negative_mask, shifted_exp_values, input_values) # 将计算出的输出值存储回输出指针 tl.store(output_ptr + block_indices, output_values, valid_mask)

在此，tl.program_id(0) 提供了一个唯一的块 ID，我们使用它来确定该块将处理的数据部分。使用此块 ID，block_start 计算每个块部分的起始索引，而 block_indices 指定该部分内的索引范围。valid_mask 确保仅处理 num_elements 内的索引，从而使用 tl.load 安全地加载数据。然后应用 ELU 函数，根据值是否为负数来修改值，并将结果写回内存，使用 tl.store。

当我们使用 triton.testing.Benchmark 对生成的内核进行基准测试时，我们具有以下性能：

image.png

与 @torch.compile 相比，此独立内核甚至在较小尺寸下也表现出卓越的性能，但这可能只是 torch.compile 编译时间的伪影。在任何情况下，与其从头开始，请记住你可以从这种生成的内核开始，并将注意力集中在优化其性能上，从而为你节省大量时间。

即使在 Triton 中，有时由于语言限制而无法处理共享内存和流式多处理器 (SM) 内调度等底层细节，我们也无法完全实现设备的峰值性能。Triton 功能仅限于块以及跨 SM 的块调度。为了获得更深层次的控制，你需要在 CUDA 中直接实现内核，你将在其中访问所有底层细节。

转移到 CUDA 后，可以使用各种技术来提高内核的效率。我们在这里仅介绍几个：优化内存访问模式以减少延迟，使用共享内存来存储频繁访问的数据，以及管理线程工作负载以最大限度地减少空闲时间。

在我们深入研究 CUDA 示例之前，让我们总结一下我们已经看到的工具，这些工具使我们能够编写内核代码以在 GPU 上执行指令：

  1. Pytorch：简单但速度慢
  2. torch.compile：简单、快速，但不灵活
  3. triton：更难、更快、更灵活
  4. CUDA：最难、最快、最灵活（如果你做得对）

让我们讨论一下我们可以在 CUDA 中使用的一种最常见的技术：优化内存访问。GPU 中的全局内存（我们上面图表中最大的内存）与缓存相比具有较长的延迟和较低的带宽，这通常会为大多数应用程序创建主要瓶颈。有效地从全局内存访问数据可以大大提高性能。

内存合并

为了有效地利用全局内存的带宽，必须了解其架构。在 CUDA 设备中，全局内存是使用 DRAM 实现的。

内存合并利用了 DRAM 在每次访问内存地址时以突发或连续内存位置范围提供数据的方式。每次访问 DRAM 位置时，都会通过 DRAM 芯片中的多个传感器并行读取一系列连续位置（包括请求的位置）。一旦读取，此数据就可以快速地作为突发传输到处理器。在 CUDA 中，合并使用此突发行为来最大限度地提高内存访问效率，方法是确保 warp 中的线程（32 个线程，以锁步方式执行相同的指令（SIMD））访问连续的内存位置。例如，如果线程 0 访问位置 M，线程 1 访问 M + 1，线程 2 访问 M + 2，依此类推，则 GPU 硬件将这些请求合并或组合为一个大型高效的 DRAM 突发访问请求，而不是单独处理每个访问。

让我们以矩阵乘法为例。一个简单、直接的实现将让每个线程计算输出矩阵的单个元素，如下所示：

__global__ void matmul_naive(int M, int N, int K, const float *A, const float *B, float *C) { const uint x = blockIdx.x * blockDim.x + threadIdx.x; const uint y = blockIdx.y * blockDim.y + threadIdx.y; if (x < M && y < N) { float tmp = 0.0; for (int i = 0; i < K; ++i) { tmp += A[x * K + i] * B[i * N + y]; } C[x * N + y] = tmp; } }

以下是来自 精彩博客文章 的内核的绝佳可视化：

image.png

但是，当使用像 ncu 这样的工具分析此内核时，我们可以看到问题，包括内存吞吐量低和未合并的内存访问。

image.png image.png

原因是，在此内核中，线程 ID 为 (0, 0) 和 (1, 0) 的同一块中的两个线程（最终将进入同一 warp）都将从矩阵 B 的同一列加载，但从矩阵 A 的不同行加载。由于矩阵元素以行优先顺序存储（意味着行元素位于连续的内存地址中，如下图所示），线程 (0, 0) 将加载 A_{0,0}，线程 (1, 0) 将在第一次迭代 i = 0 时加载 A_{1,0}。这些元素没有紧密地存储在内存中，并且这种未对齐将在每次迭代中都存在，从而阻止内存访问被合并。

image.png

​

为了提高内核的性能，我们可以更改坐标 x 和 y 的计算方式，如下所示：

const int x = blockIdx.x * BLOCKSIZE + (threadIdx.x / BLOCKSIZE); const int y = blockIdx.y * BLOCKSIZE + (threadIdx.x % BLOCKSIZE); if (x < M && y < N) { float tmp = 0.0; for (int i = 0; i < K; ++i) { tmp += A[x * K + i] * B[i * N + y]; } C[x * N + y] = tmp; }

我们没有使用 2D 块，而是切换到 1D 块，并重新定义我们如何确定 x 和 y 的值。在这种新方法中，同一 warp 中的线程（具有接近 threadIdx.x 值）将共享相同的 x 值，但具有不同的 y 值。这意味着它们将加载矩阵 A 的同一行，但加载矩阵 B 的不同列。因此，对于行优先矩阵，内存访问可以合并。

当我们分析我们的新内核时，我们注意到关于未合并内存访问的警告已消失，并且 GPU 的内存吞吐量提高了大约 10 倍。

image.png

我们还注意到内核的执行时间 减少了 10 倍！太棒了。

现在让我们介绍一下你在文献中经常看到的另一种技术：平铺。

平铺

​

平铺是一种利用 *共享内存* 来优化内存访问模式的技术。正如我们上面提到的，共享内存是一种小型、快速的内存，块内的所有线程都可以访问它。它允许数据被多个线程重用，从而减少了重复从较慢的全局内存加载数据的需要。

例如，在矩阵乘法中，块中的每个线程可能需要来自两个矩阵（例如 A 和 B）的元素。如果每个线程独立地从全局内存加载它需要的行和列，我们最终会得到许多冗余加载，因为块中的多个线程将访问重叠的数据。相反，我们可以使用平铺将 A 和 B 的块（或平铺）加载到共享内存中一次，从而允许该块中的所有线程重用相同的共享数据。

在平铺方法中，每次迭代都涉及块中的所有线程协同加载两个平铺——一个来自矩阵 A，另一个来自矩阵 B——到共享内存中。具体而言，线程加载矩阵 A 的平铺（大小为 BLOCK_SIZE_M 乘以 BLOCK_SIZE_K）和矩阵 B 的平铺（大小为 BLOCK_SIZE_K 乘以 BLOCK_SIZE_N）。一旦平铺位于共享内存中，线程就会对这些平铺执行矩阵乘法，从而实现高效计算，因为所有必要的数据都可以快速访问。平铺乘法的结果存储在一个累积矩阵中，该矩阵保存中间结果。在每次迭代之后，来自当前平铺乘法的结果都会添加到此累积矩阵中，并继续进行，直到矩阵 A 和矩阵 B 中的所有平铺都已处理完毕。

image.png

来自 https://cnugteren.github.io/tutorial/pages/page4.html

让我们看一下你需要从实现中理解的重要部分：

// 设置指向起始元素的指针 A += blockRow * TILE_SIZE * K; // 从行 = blockRow，列 = 0 开始 B += blockCol * TILE_SIZE; // 从行 = 0，列 = blockCol 开始 C += blockRow * TILE_SIZE * N + blockCol * TILE_SIZE; // 从行 = blockRow，列 = blockCol 开始 float sum = 0.0; // 外循环遍历 A 的平铺（跨列）和 B 的平铺（向下行） for (int tileIdx = 0; tileIdx < K; tileIdx += TILE_SIZE) { sharedA[localRow * TILE_SIZE + localCol] = A[localRow * K + localCol]; sharedB[localRow * TILE_SIZE + localCol] = B[localRow * N + localCol]; // 确保块中的所有线程都已完成数据加载 __syncthreads(); // 将指针移动到下一个平铺 A += TILE_SIZE; B += TILE_SIZE * N; // 计算此平铺的部分点积 for (int i = 0; i < TILE_SIZE; ++i) { sum += sharedA[localRow * TILE_SIZE + i] * sharedB[i * TILE_SIZE + localCol]; } // 再次同步以防止任何线程在其他线程完成计算之前将新数据加载到共享内存中 // 之前 __syncthreads(); } C[localRow * N + localCol] = sum;

每个线程首先从 矩阵 A** 和 **矩阵 B** 加载一个元素到共享内存中。在这种情况下，实现合并的内存访问非常简单，通过分配 threadIdx.x 作为 **本地列索引 (localCol)**，同一 warp 中的线程将访问两个矩阵的相邻元素。在块中的每个线程完成将其元素加载到共享内存中后（通过调用 __syncthreads() 确保），它们将继续计算两个平铺的点积。一旦线程遍历了所有平铺（水平遍历 **矩阵 A**，垂直遍历 **矩阵 B**），结果总和将存储在 **矩阵 C 的相应位置中。

当使用 ncu 对此内核进行基准测试时，我们注意到内存吞吐量增加到 410 Gb/s，内核执行时间减少了约 43%，实现了约 6.6 TFLOP 的性能

线程粗化

平铺技术显著提高了我们内核的性能。但是，当分析 warp 状态（量化每个状态花费的周期数）时，我们观察到以下情况：

image.png

这些神秘状态名称的含义可以在 NVidia 的分析指南 的 Warp Stall Reasons 部分中找到。在那里我们可以读到：

"smsp__pcsamp_warps_issue_stalled_mio_throttle：Warp 正在等待 MIO（内存输入/输出）指令队列未满而停顿。在 MIO 管道的极端利用率情况下，此停顿原因是高的，其中包括特殊数学指令、动态分支以及共享内存指令。当由共享内存访问引起时，尝试使用更少但更宽的加载可以减轻管道压力。"

因此，warp 似乎正在停顿，等待共享内存访问返回！为了解决这个问题，我们可以应用一种称为 线程粗化 的技术，该技术涉及将多个线程合并为一个粗化的线程。这将显著减少共享内存访问，因为每个粗化的线程可以处理多个输出元素。

让我们简要提及最后一个在编写或改进自定义内核时需要考虑的重要事项：最大限度地减少控制发散。

最大限度地减少控制发散

流式多处理器 (SM) 的构建目的是使用单指令多数据 (SIMD) 模型执行 warp 中的所有线程。这意味着在任何给定时刻，都会提取一条指令并为 warp 内的所有线程同时执行。当执行 warp 时，其中的线程对不同的数据段进行操作，但遵循相同的指令，因此得名单指令多数据。SIMD 的主要优点是其效率；负责指令提取和分派的控制硬件在多个执行单元之间共享。这种设计最大限度地减少了与控制功能相关的硬件开销，从而使硬件的更大部分能够专注于提高算术吞吐量。

当同一 warp 中的线程采用不同的执行路径时，就会发生控制发散。例如，如果条件语句（如 if 语句）导致某些线程执行一个代码块，而其他线程执行不同的代码块，则 warp 必须串行化这些执行，从而导致空闲线程等待其他线程完成。 为了最大限度地减少控制发散，我们需要设计内核以确保同一 warp 中的线程遵循相同的执行路径。这可以通过重组代码以减少分支、使用确保所有线程遵循相似执行路径的数据结构或采用诸如谓词之类的技术来实现。

我们已经介绍了编写自定义内核以及提高 GPU 操作的性能和内存占用的主要注意事项。但在我们转向一个真实的示例之前，还有一个重要的概念是“融合内核”。

融合内核

在多个地方，我们都提到 GPU 和 CPU 操作如何异步。特别是，CPU 上的主机代码可以以非阻塞方式在 GPU 上调度工作负载。

非阻塞对于重叠通信和计算非常有用——正如我们在整个旅程中多次看到的那样——但可以扩展到更一般的想法，即尝试不惜一切代价避免在主机和 GPU 内核命令之间来回切换。

这个想法由 Horace He 在这些图中精美地说明：

image.png

需要全局内存和计算单元之间来回切换的一系列内核

image.png

与其将我们的三角形发送回全局内存只是为了再次读取它，不如我们一次性完成所有操作。

我们如何避免这种来回切换？最好的方法是使我们的 GPU 尽可能自主。这是通过将尽可能多的连续计算操作打包在一个内核中以供 GPU 运行来实现的，称为“融合内核”。

对于在一系列点状操作上执行的独立于每个输入 tokens 的操作，融合内核尤其高效且易于编写。在这种情况下，没有必要在将计算值移动到 SM 内存并启动新内核之前将其带回全局内存。最好将所有值本地保留，直到执行完一系列计算。

在 Transformer 模型中有许多地方可以应用这种“融合”方法：每次我们有一系列点状操作时，例如在层归一化中涉及的计算中。

我们现在已经掌握了所有必要的理解，可以惊叹于内核工程的真正杰作：Flash Attention

Flash Attention 1-3

Flash Attention 由 Tri Dao 引入，旨在通过编写自定义 CUDA 内核来优化注意力计算，从而使其更快 *且* 内存效率更高。Flash Attention 背后的思想是有效利用 GPU 的各种内存，以避免过多地依赖最慢的内存：GPU 的全局内存。

注意机制的基本实现涉及内存和工作程序之间的大量传输。它需要物化 HBM 中的 S 和 P 矩阵，这意味着结果需要发送到 HBM，然后再返回 SRAM 以进行下一次计算：

image.png

由于 HBM 中的带宽要低得多，因此这会在注意力计算中引入严重的瓶颈。我们能做得更好吗？Tri Dao 说是的！

关键要素是以小块计算 S 矩阵，这些小块可以放入 SM 的较小共享内存中。但我们可以做得更好，避免将非常大的 S 矩阵完全物化，而是只保留计算 softmax 归一化因子所需的统计信息。因此，我们可以直接在 SRAM 中的一次计算中计算部分 O，而不是来回移动中间结果。在这种情况下，我们不仅利用了共享内存，还释放了由模型中最大的激活矩阵之一（在长上下文长度下）物化（注意力矩阵）引起的内存瓶颈。

image.png

来源：FlashAttention 论文

Flash Attention 的思想解决了模型训练中的许多瓶颈，因此它已迅速成为所有 Transformer 中执行注意力的默认方式：

因此，线性能量和二次线性方法的所有变体都用于近似注意力——在 Transformer 架构发明后不久开发——主要被搁置，转而支持这种精确且快速的 Flash Attention 实现和机制。

在 Flash-attention 1 之后，同一实验室发布了两个连续的改进版本：Flash-attention 2 和 3。与 Flash-attention 1 相比，Flash-attention 2 和 3 中的改进更多地是关于根据 GPU 定制其低级实现，方法是 (1) 尽可能减少非 matmul 操作的数量，(2) 仔细划分 warp 和线程块之间的工作负载（对于 Flash Attention 2），以及仔细优化 FP8 和最新 Hopper (H100) 架构上的张量内核支持（对于 Flash Attention 3）。

Flash-Attention 是一个大师级的演示，展示了当你考虑到当前 GPU 加速器的内部内存/计算设计时，可以取得突破性的改进。

到目前为止，本操作融合部分中描述的技术要求我们实现建模代码更改，并为某些操作编写自定义内核，以加速训练。

在我们对计算操作本身进行低级别深入研究的最后一节中，我们将了解一系列与建模代码无关且可用于任何模型的方法，并且这些方法被广泛使用，以至于已成为行业标准：混合精度训练！

混合精度训练

在本书的多个部分中，我们讨论了较低精度格式及其对存储激活值、参数和优化器状态的内存需求的影响。现在是时候深入了解这些格式的细节，并更好地了解它们的权衡、优势和局限性了。

混合精度训练，顾名思义，涉及在训练时混合不同的精度。PyTorch 张量的默认数值精度是单精度浮点格式，也称为 FP32 或 float32，这意味着存储的每个数字占用 32 位或 4 个字节。用于表示数字的可用位分为 3 个部分：

sign-mantissa-exponent.svg

浮点数的原理可以通过回忆数字的科学计数法轻松说明，例如 - 5.734 \times 10^{7}，其中我们首先有符号，然后是尾数和指数。因此，我们可以使用自适应精度表示各种量级的数字。虽然 float32 是默认值，但 PyTorch 中有一系列可用的浮点格式：

格式 总位数 符号 指数 尾数
float32 32 1 8 23
float16 16 1 5 10
bfloat16 16 1 8 7
float8 (e4m3) 8 1 4 3
float8 (e5m2) 8 1 5 2

减少总位数是有代价的（这里也没有免费的午餐），但我们可以控制如何支付。我们可以牺牲尾数或指数上更多的位。因此，也存在两种 float8 格式，根据指数和尾数命名，以灵活地选择最合适的格式。我们可以查看每种格式的可能数字范围：

image.png

我们可以看到，float32 跨越 80 个数量级，float16 牺牲了很多范围，而 bfloat16 保持了完整范围。两种 float8 格式进一步缩小了范围，其中 e5e2 可以保持 float16 范围，而 e4m3 的范围更小。

为什么有些格式能够保持范围，而其他格式不能？让我们通过绘制 1 和 2 之间的 10,000 个点来研究分辨率。每个点将四舍五入到每种格式中最接近的可表示数字：

image.png

我们可以在这里看到，bfloat16 比 float16 保持了 float32 的范围，但这以牺牲更多精度为代价。在 float8 的情况下，情况更加糟糕，因为 e4m3 可以表示 7 个，而 e5m2 只能表示 1-2 区间上的 3 个数字。

衡量格式分辨率的一个常见指标是 epsilon：1.00 之后的第一个可表示数字。我们可以看到，对于 float32 格式，10^{-4} 是上限（实际上是 1.19^{-7}）。对于 float16，它是 \tilde 10^{-3}，对于 bfloat，它仍然高出 10 倍。

混合精度训练的思想是使用其中一些较低精度的格式，同时保持全精度训练的性能。

事实证明，我们 不能** 完全放弃 float32，通常需要将某些部分保持在全精度下。这就是为什么较低精度训练通常被称为 **混合精度 训练的原因。

现在让我们看一下使用 16 位训练模型，然后看看我们是否可以将训练进一步推进到 8 位。

​

FP16 和 BF16 训练

遗憾的是，朴素地将所有张量和运算切换到 float16 无法正常工作，结果通常是损失发散。但是，最初的混合精度训练论文提出了三个技巧来匹配 float32 训练：

  1. FP32 权重副本：float16 权重可能存在两个问题。在训练期间，某些权重可能会变得非常小，并将四舍五入为 0。但是，即使权重本身不接近于零，如果更新非常小，则幅度差异可能会导致权重在加法期间下溢。一旦权重为零，它们将在训练的剩余时间内保持为 0，因为不再有梯度信号通过。
  2. 损失缩放：我们在梯度方面也存在类似的问题，因为梯度往往远小于 1，因此有下溢的风险。一个简单但有效的方法是在反向传播之前缩放损失，并在反向传播之后取消缩放梯度。这确保了反向传播期间不会发生下溢，并且缩放不会影响训练，因为我们在进一步处理梯度（例如，裁剪）和优化步骤之前取消缩放。
  3. 累积：最后，当以 16 位精度执行某些算术运算（例如平均值或总和）时，我们也可能面临下溢或溢出。解决方案是在运算期间以 float32 累积中间结果，并且仅将最终结果转换回 16 位精度。

通过这些技术，我们可以获得稳定的训练，同时受益于更高的吞吐量，这是因为更快、更低精度的算术运算。自然而然地，作为一个好奇的读者——并且现在有点沉迷于最大化吞吐量——你可能会问一个问题：我们能否比 16 位精度走得更远更快？

也许可以！

FP8 预训练

即使我们完美地将通信与计算重叠，我们总是最终会遇到硬件本身的低级别理论 FLOPS 限制，即硬件上每个单独操作的效率。这就是数值精度变得至关重要的地方。例如，在 NVIDIA 的 H100 GPU 上，FP8 矩阵乘法（GEMM 操作）实现了 bfloat16 两倍的理论 FLOPS，这使得低精度训练成为进一步优化的有吸引力的途径。

最近的研究——包括 FP8-LM、torchao 和 DeepSeek-V3——已经证明了 FP8 训练在大型模型中的潜力。尽管如此，FP8 预训练还是带来了巨大的挑战：稳定性。在较低精度下，数值不稳定性通常会导致损失发散，从而难以匹配更高精度训练的准确性。

我们知道，对于固定模型大小，不稳定性会随着学习率的升高而增加，这使得 FP8 预训练尤其棘手。

以下是 FP8 训练的典型发散损失曲线示例：

首次公开发布的成功的超大规模 FP8 混合精度训练是在 DeepSeek-V3 上进行的。作者仔细分析了前向传播 (Fprop) 的每个操作以及激活值 (Dgrad) 和权重 (Wgrad) 反向传播。与 BF16 混合精度训练类似，一些聚合和主权重保持在更高的精度，而操作本身则以 FP8 执行。

image.png

为了从高精度（例如 FP32 或 BF16）切换到范围较小的低精度（例如 FP16 或 FP8），我们需要标准化激活值的范围，例如通过计算其绝对最大值。DeepSeek-V3 进一步引入了一种特定的量化方案，其中范围按平铺进行标准化：输入/激活值为 1x128，权重和缩放元素为 128x128。这使得标准化受激活值中异常值的影响较小。他们提出了许多额外的技巧，以进一步减少内存和通信占用，你可以在 DeepSeek-V3 技术报告 的 3.3 节中找到这些技巧的后续内容。

以下是一些已知的 FP8 训练方法摘要：

GEMM 的精度 主模型权重 累积梯度 模型权重 梯度 优化器状态 总内存
带 fp32 混合精度基线的 bfloat16 bf16 fp32 fp32 bf16 bf16 fp32 + fp32 4 + 4 + 2 + 2 + 4 + 4 = 20 字节
以上不带 FP32 梯度累积 bf16 fp32 不适用 bf16 bf16 fp32 + fp32 4 + 2 + 2 + 4 + 4 = 16 字节
Transformer 引擎 fp8 不适用 不适用 fp32 fp32 fp32 + fp32 4 + 4 + 4 + 4 = 16 字节（减少 20%）
FP8-LM 的 O3 级别 fp8 fp16 fp16 fp8 fp8 fp8 + fp16 2 + 2 + 1 + 1 + 1 + 2 = 9 字节（55%）
DeepSeek-V3 fp8 fp32 fp32 fp8 bf16 bf16 + bf16 4+4+1+2+2+2 = 15 (25%)
nanotron 的 FP8 fp8 bf16 fp32 fp8 fp8 fp8 + fp8 2 + 4 + 1 + 1 + 1 + 1 = 10 字节（50%）

总的来说，FP8 仍然是——在 2025 年初——一种实验性技术，方法仍在不断发展。鉴于其显而易见的好处，它很可能成为标准，并很快取代 bf16 混合精度。要关注 FP8 训练技术的开源实现，请访问 此 PR 中的 nanotron 实现。

展望未来，下一代 NVIDIA 芯片 Blackwell 已宣布 支持 FP4 训练，进一步加快了训练速度，但毫无疑问也引入了新的训练稳定性挑战。

最后一部分结束了我们对 GPU 集群上快速和大型模型训练领域的漫长旅程。是时候慢慢让我们的 GPU 集群休息一下，并退后一步，总结一下我们一路学到的所有知识。

结论

​

亲爱的读者，恭喜你坚持到了最后！我们完成了一段相当长的旅程：我们从了解如何在单个 GPU 上训练一个简单的模型开始，一直到掌握用于在数千个 GPU 上高效训练 Llama-405B 和 DeepSeek-V3 等大型语言模型的所有复杂技术。到现在为止，你可以（相对）轻松地阅读像 Llama-3 的 4D 并行设置这样的图表：

image.png

协调大型 GPU 集群以高效训练 LLM 并非易事。我们学习了如何优化 GPU 之间的计算和通信，以便它们始终以最大利用率运行。这包括为给定的模型和集群大小选择正确的并行化策略、尽可能重叠通信和计算，以及编写考虑到硬件布局的自定义内核，以便在 GPU 上尽可能快地执行操作。

你可能仍然认为，这种知识有点小众，只涉及少数预训练 LLM 的人。从历史上看，情况可能确实如此，但随着 AI 构建者社区 和模型大小都在快速增长，使用分布式技术进行推理、微调和训练的人员社区也在呈指数级增长，这使得分布式训练设置越来越普遍。因此，深入研究所有分布式事物可能会被证明是非常及时的。

这是一段漫长的学习之旅，但不仅仅是对你而言！在 GPU 集群上运行数千个基准测试比我们预期的更具挑战性，我们想分享一些我们自己的学习体验的亮点。

那么，接下来是什么？

你现在对主要的分布式训练概念有了很好的概述，但与此同时，我们只是触及了其中几种工具和技术的表面。有很多方法可以深入研究一个主题，但以下是我们推荐的一些步骤：

我们希望本书能帮助你入门分布式训练，并且你将训练下一代出色的模型，伴随着你的 GPU 集群的嗡嗡声！

对第一批读者说最后一句话。我们对这篇写作作品感到非常满意，因此我们决定分发数量有限的实体印刷版，作为送给我们第一批读者的礼物。

如果你是前 50 名填写以下电子邮件地址的人之一，我们将在今年晚些时候与你联系，以便在我们将博客文章格式化为印刷副本后向你发送真正的实体版。

我们预计本书约为 100-150 页，内容与博客文章相同，但我们也可能决定根据印刷品的意义来缩短或延长它。

要获得你的实体副本，请在以下 google 表单 中填写你的电子邮件地址。

无论你是我们的第一批读者还是稍后才看到这篇博客文章，我们都很高兴看到你喜欢这种知识分享。愿开源和开放科学的力量永远与你同在。

致谢

我们感谢 Elie 进行深入的审查并使用 NotebookLM 创建音频组件。特别感谢 Hynek 优化前端性能。我们还要感谢 Simon 解决了一些 Hub 上的问题。

讨论页面

如果你想讨论这篇博客文章的内容、提出问题、提出更改或只是打个招呼，请在 讨论页面 上打开一个主题。

参考文献

具有里程碑意义的 LLM 扩展论文

Megatron-LM

介绍了张量并行和高效的模型并行技术，用于训练大型语言模型。

Megatron-Turing NLG 530B

描述了使用 DeepSpeed 和 Megatron-LM 框架的组合训练 530B 参数模型。

PaLM

介绍了 Google 的 Pathways 语言模型，展示了在数百个语言任务和推理能力方面的强大性能。

Gemini

介绍了 Google 的多模态模型架构，该架构能够处理文本、图像、音频和视频输入。

Llama 3

Llama 3 模型家族

DeepSeek-V3

DeepSeek 关于 DeepSeek-V3 模型的架构和训练的报告。

训练框架

Nanotron

我们用于训练大型语言模型的框架，具有各种并行策略

Megatron-LM

NVIDIA 用于训练大型语言模型的框架，具有各种并行策略。

DeepSpeed

Microsoft 的深度学习优化库，具有 ZeRO 优化阶段和各种并行策略。

FairScale

PyTorch 扩展库，用于大规模训练，提供各种并行和优化技术。

ColossalAI

集成的超大规模模型训练系统，具有各种优化技术。

torchtitan

用于大型模型训练的 PyTorch 本地库。

GPT-NeoX

EleutherAI 用于训练大型语言模型的框架，用于训练 GPT-NeoX-20B。

LitGPT

Lightning AI 的最新开源 LLM 实现，重点是可重复性。

DiLoco

使用 DiLoco 跨计算集群训练语言模型。

torchgpipe

PyTorch 中的 GPipe 实现。

OSLO

OSLO：用于大规模优化的开源。

调试

速度分析

关于使用分析器分析模型性能和瓶颈的官方 PyTorch 教程。

内存分析

全面指南，用于理解和优化 PyTorch 中的 GPU 内存使用情况。

简单示例中的内存分析演练

可视化和理解 PyTorch 中的 GPU 内存。

TensorBoard 分析器教程

使用 TensorBoard 的分析工具分析 PyTorch 模型的指南。

分布技术

数据并行

深度学习中数据并行训练的全面解释。

ZeRO

介绍了零冗余优化器，用于使用内存优化训练大型模型。

FSDP

PyTorch 中完全分片数据并行训练的实现。

张量和序列并行 + 选择性重计算

用于高效大规模模型训练的先进技术，结合了不同的并行策略。

流水线并行

NVIDIA 关于实现大型模型训练的流水线并行的指南。

广度优先流水线并行

包括围绕 PP 计划的广泛讨论。

All-reduce

环形 all-reduce 算法的详细解释，该算法用于分布式训练。

环形闪存注意力

环形注意力机制与闪存注意力的结合实现，用于高效训练。

环注意力教程

教程，解释了环注意力的概念和实现。

ZeRO 和 3D

DeepSpeed 指南，用于理解 ZeRO 和 3D 并行策略之间的权衡。

混合精度训练

介绍了深度学习模型的混合精度训练技术。

可视化 6D 网格并行

解释了 6D 并行网格中涉及的集体通信。

硬件

Fire-Flyer - 一个 10,000 个 PCI 芯片的集群

DeepSeek 关于设计具有 1 万个 PCI GPU 的集群的报告。

Meta 的 24k H100 Pod

Meta 对其使用 NVIDIA H100 GPU 构建的大规模 AI 基础设施的详细概述。

Semianalysis - 10 万个 H100 集群

对超大规模 H100 GPU 集群及其对 AI 基础设施的影响的分析。

Modal GPU 词汇表

人类可读的 CUDA 文档

其他

Stas Bekman 的手册

涵盖 LLM 训练各个方面的综合手册。

Bloom 训练编年史

BLOOM 模型训练过程和挑战的详细文档。

OPT 日志

Meta 记录 OPT-175B 模型训练过程的详细日志。

Harm 关于更长时间训练小型模型的定律

调查模型大小与训练开销之间的关系。

Harm 关于长上下文的博客

从数据和训练成本方面调查长上下文训练。

GPU 模式

GPU 阅读小组和社区。

EleutherAI Youtube 频道

ML 可扩展性和性能阅读小组

Google Jax 扩展书

如何扩展你的模型

@fvsmassa & @TimDarcet FSDP

独立的约 500 行 LoC FSDP 实现

thonking.ai

Horace He 的一些博客文章 - 让 GPU 飞速运行……

Aleksa 的 ELI5 Flash Attention

Flash Attention 的简单解释

TunibAI 的 3D 并行教程

使用 PyTorch 进行大规模语言建模教程。

附录

A0：并行编程速成课程

在整个博客文章中，我们将 LLM 训练从一个 GPU 扩展到数百个 GPU。这将需要在所有机器之间通信和同步权重、梯度和数据。有一组分布式模式可以完全实现这一点，称为 集体操作。在本节中，我们将简要介绍所有操作，如 *广播、AllReduce、Scatter* 等。让我们深入了解一下！

一般设置是我们有许多独立的节点，这些节点可以是 CPU 内核、GPU 或计算节点。每个节点执行一些计算，然后我们希望将结果或部分结果通信给其他节点，以进行下一步计算（t+1）。

image.png

也许我们需要将结果从一个节点发送到所有其他节点，或者我们需要对每个节点的所有中间结果求和，以报告总体结果。通常，有一个具有提升状态的节点，它在中心角色中发挥作用，这里用 root 表示，它是某些操作的目标或来源。让我们从最简单的原语之一开始：广播操作。

广播

一个非常常见的模式是，你在节点 1 上有一些数据，并且你想与其他所有节点共享它，以便它们可以使用该数据进行一些计算。广播操作正是如此：

image.png

集体操作由 PyTorch 本地提供，因此我们可以轻松编写一个小例子来演示广播的工作原理。我们首先需要使用 dist.initi_process_group 初始化进程组，该进程组设置通信后端（我们稍后将讨论 NCCL），它确定存在多少工作程序（又名节点），并为每个工作程序分配一个排名（我们可以使用 dist.get_rank 获取排名）。最后，它在工作程序之间建立连接。

为了展示 dist.broadcast 操作，让我们在 rank=0 上创建一个具有非零值的张量，并在其他工作程序上创建充满零的张量。然后，我们使用 dist.broadcast(tensor, src=0) 将 rank=0 张量分发给所有其他排名：

import torch import torch.distributed as dist def init_process(): dist.init_process_group(backend='nccl') torch.cuda.set_device(dist.get_rank()) def example_broadcast(): if dist.get_rank() == 0: tensor = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32).cuda() else: tensor = torch.zeros(5, dtype=torch.float32).cuda() print(f"Before broadcast on rank {dist.get_rank()}: {tensor}") dist.broadcast(tensor, src=0) print(f"After broadcast on rank {dist.get_rank()}: {tensor}") init_process() example_broadcast()

你可以使用 torchrun --nproc_per_node=3 dist_op.py 运行上面的脚本（你需要 3 个 GPU，否则请相应地更改 nproc_per_node），你应该会看到以下输出：

Before broadcast on rank 0: tensor([1., 2., 3., 4., 5.], device='cuda:0') Before broadcast on rank 1: tensor([0., 0., 0., 0., 0.], device='cuda:1') Before broadcast on rank 2: tensor([0., 0., 0., 0., 0.], device='cuda:2') After broadcast on rank 0: tensor([1., 2., 3., 4., 5.], device='cuda:0') After broadcast on rank 1: tensor([1., 2., 3., 4., 5.], device='cuda:1') After broadcast on rank 2: tensor([1., 2., 3., 4., 5.], device='cuda:2')

太棒了，似乎可以按预期工作。请注意，排名消息可能会无序打印，因为我们无法控制哪个打印语句首先执行（我们在此处对它们进行了排序以提高可读性）。现在让我们继续介绍 Reduce 和 AllReduce 模式！

Reduce 和 AllReduce

Reduce 模式是分布式数据处理中最基本的模式之一。其思想是，你想通过一个函数 f() 组合每个节点上存在的数据，该函数可以是求和或求平均值。在 Reduce 范例中，结果仅发送到根节点，而在 AllReduce 情况下，结果会广播到所有节点：

image.png

当然，没有神奇的“自由飞行”节点可以执行此操作，通常每个节点都在节点环或树结构中执行部分计算。这是一个简单的例子：假设我们需要计算每个节点上数字的总和，并且我们的节点以环状模式连接。第一个节点将其数字发送给邻居，邻居将其数字添加到接收到的数字中，然后再转发给下一个邻居。在一轮沿着节点环之后，第一个节点将收到总和。

以下是运行简单 Reduce 操作（对张量求和）的代码，我们使用 op=dist.ReduceOp.SUM 指定要使用的操作（你可以在 Pytorch 文档 中找到有关支持操作的更多信息）：

def example_reduce(): tensor = torch.tensor([dist.get_rank() + 1] * 5, dtype=torch.float32).cuda() print(f"Before reduce on rank {dist.get_rank()}: {tensor}") dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM) print(f"After reduce on rank {rank}: {tensor}") init_process() example_reduce()

请注意，在 Reduce 操作中，只有 dst 节点上的张量会更新：

Before reduce on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0') Before reduce on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1') Before reduce on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2') After reduce on rank 0: tensor([6., 6., 6., 6., 6.], device='cuda:0') After reduce on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1') After reduce on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')

类似地，我们可以执行 AllReduce（在这种情况下我们不需要指定目标）：

def example_all_reduce(): tensor = torch.tensor([dist.get_rank() + 1] * 5, dtype=torch.float32).cuda() print(f"Before all_reduce on rank {dist.get_rank()}: {tensor}") dist.all_reduce(tensor, op=dist.ReduceOp.SUM) print(f"After all_reduce on rank {dist.get_rank()}: {tensor}") init_process() example_all_reduce()

在这种情况下，结果在所有节点上都可用：

Before all_reduce on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0') Before all_reduce on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1') Before all_reduce on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2') After all_reduce on rank 0: tensor([6., 6., 6., 6., 6.], device='cuda:0') After all_reduce on rank 1: tensor([6., 6., 6., 6., 6.], device='cuda:1') After all_reduce on rank 2: tensor([6., 6., 6., 6., 6.], device='cuda:2')

现在让我们转向我们的下一个分布式通信操作。在许多实际情况下，每个节点单独执行许多复杂的计算，我们需要在节点之间共享最终结果。在这种情况下，Gather 和 AllGather 是我们要使用的操作。让我们看一下！

Gather 和 AllGather

Gather 和 AllGather 与 Broadcast 非常相似，因为它们允许在节点之间分发数据而不进行修改。与 Broadcast 的主要区别在于，没有一个值我们需要从一个节点共享到所有其他节点，而是每个节点都有一个单独的数据块，我们希望将所有数据收集到一个节点（在 Gather 的情况下）或将所有数据收集到所有节点（在 AllGather 的情况下）。一图胜千言，让我们看一下：

image.png

请注意，虚线表示某些数据实际上根本没有移动（因为它已经存在于节点上）。

在 gather 操作的情况下，我们需要准备一个容器对象，用于存储收集的张量，在本例中为 gather_list：

def example_gather(): tensor = torch.tensor([dist.get_rank() + 1] * 5, dtype=torch.float32).cuda() if dist.get_rank() == 0: gather_list = [ torch.zeros(5, dtype=torch.float32).cuda() for _ in range(dist.get_world_size()) ] else: gather_list = None print(f"Before gather on rank {dist.get_rank()}: {tensor}") dist.gather(tensor, gather_list, dst=0) if dist.get_rank() == 0: print(f"After gather on rank 0: {gather_list}") init_process() example_gather()

我们看到 gather_list 确实包含所有排名的张量：

Before gather on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0') Before gather on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1') Before gather on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2') After gather on rank 0: [tensor([1., 1., 1., 1., 1.], device='cuda:0'), tensor([2., 2., 2., 2., 2.], device='cuda:0'), tensor([3., 3., 3., 3., 3.], device='cuda:0')]

对于 AllGather 示例，我们需要更改的唯一内容是，每个节点都需要结果的占位符：

def example_all_gather(): tensor = torch.tensor([dist.get_rank() + 1] * 5, dtype=torch.float32).cuda() gather_list = [ torch.zeros(5, dtype=torch.float32).cuda() for _ in range(dist.get_world_size()) ] print(f"Before all_gather on rank {dist.get_rank()}: {tensor}") dist.all_gather(gather_list, tensor) print(f"After all_gather on rank {dist.get_rank()}: {gather_list}") init_process() example_all_gather()

实际上，我们可以看到，现在每个节点都具有所有数据：

Before all_gather on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0') Before all_gather on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1') Before all_gather on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2') After all_gather on rank 0: [tensor([1., 1., 1., 1., 1.], device='cuda:0'), tensor([2., 2., 2., 2., 2.], device='cuda:0'), tensor([3., 3., 3., 3., 3.], device='cuda:0')] After all_gather on rank 1: [tensor([1., 1., 1., 1., 1.], device='cuda:1'), tensor([2., 2., 2., 2., 2.], device='cuda:2'), tensor([3., 3., 3., 3., 3.], device='cuda:2')] After all_gather on rank 2: [tensor([1., 1., 1., 1., 1.], device='cuda:2'), tensor([2., 2., 2., 2., 2.], device='cuda:2'), tensor([3., 3., 3., 3., 3.], device='cuda:2')]

现在，反过来考虑一下 gather？在这种情况下，我们将所有数据都放在一个节点上，并希望在节点之间分发/切片它，可能在中间进行一些处理？我们可以使用 Scatter，或者在操作之间使用 Reduce Scatter 模式：

Scatter 和 ReduceScatter

顾名思义，Scatter 操作的目标是获取一个节点上的数据，并将其切片分发给所有其他节点。因此，它与 Broadcast 操作不同，Broadcast 操作复制数据而不进行切片，并且它是 Gather 操作的逻辑逆运算。

ReduceScatter 模式稍微复杂一些：想象一下，你应用一个操作，就像在 Reduce 情况中一样，但我们不是将结果仅移动到一个节点，而是均匀地将其分发给所有节点：

image.png

Scatter 操作在代码中的编写方式与 Gather 相反：我们不是准备张量列表作为目标，而是准备源数据作为我们想要分发的张量列表。我们还需要指定 src：

def example_scatter(): if dist.get_rank() == 0: scatter_list = [ torch.tensor([i + 1] * 5, dtype=torch.float32).cuda() for i in range(dist.get_world_size()) ] print(f"Rank 0: Tensor to scatter: {scatter_list}") else: scatter_list = None tensor = torch.zeros(5, dtype=torch.float32).cuda() print(f"Before scatter on rank {dist.get_rank()}: {tensor}") dist.scatter(tensor, scatter_list, src=0) print(f"After scatter on rank {dist.get_rank()}: {tensor}") init_process() example_scatter()

因此，我们可以看到空张量如何填充了 scatter_list 的内容

Rank 0: Tensor to scatter: [tensor([1., 1., 1., 1., 1.], device='cuda:0'), tensor([2., 2., 2., 2., 2.], device='cuda:0'), tensor([3., 3., 3., 3., 3.], device='cuda:0')] Before scatter on rank 0: tensor([0., 0., 0., 0., 0.], device='cuda:0') Before scatter on rank 1: tensor([0., 0., 0., 0., 0.], device='cuda:1') Before scatter on rank 2: tensor([0., 0., 0., 0., 0.], device='cuda:2') After scatter on rank 0: tensor([1., 1., 1., 1., 1.], device='cuda:0') After scatter on rank 1: tensor([2., 2., 2., 2., 2.], device='cuda:1') After scatter on rank 2: tensor([3., 3., 3., 3., 3.], device='cuda:2')

让我们创建更多有趣的数据来演示 ReduceScatter 逻辑：在每个节点上，我们创建了一个包含 2 元素向量的列表，每个向量都有一个幂指数和一个节点排名偏移函数（这有点难以想象，所以只需查看下面的示例）：

def example_reduce_scatter(): rank = dist.get_rank() world_size = dist.get_world_size() input_tensor = [ torch.tensor([(rank + 1) * i for i in range(1, 3)], dtype=torch.float32).cuda()**(j+1) for j in range(world_size) ] output_tensor = torch.zeros(2, dtype=torch.float32).cuda() print(f"Before ReduceScatter on rank {rank}: {input_tensor}") dist.reduce_scatter(output_tensor, input_tensor, op=dist.ReduceOp.SUM) print(f"After ReduceScatter on rank {rank}: {output_tensor}") init_process() example_reduce_scatter()

让我们打印出我们创建的数据模式。我们还可以立即看到 ReduceScatter 模式：第一个排名接收到每个节点的第一个张量的总和，第二个排名包含每个节点的第二个张量的总和，依此类推：

Before ReduceScatter on rank 0: [tensor([1., 2.], device='cuda:0'), tensor([1., 4.], device='cuda:0'), tensor([1., 8.], device='cuda:0')] Before ReduceScatter on rank 1: [tensor([2., 4.], device='cuda:1'), tensor([ 4., 16.], device='cuda:1'), tensor([ 8., 64.], device='cuda:1')] Before ReduceScatter on rank 2: [tensor([3., 6.], device='cuda:2'), tensor([ 9., 36.], device='cuda:2'), tensor([ 27., 216.], device='cuda:2')] After ReduceScatter on rank 0: tensor([ 6., 12.], device='cuda:0') After ReduceScatter on rank 1: tensor([14., 56.], device='cuda:1') After ReduceScatter on rank 2: tensor([ 36., 288.], device='cuda:2')

让我们快速了解一下 AllReduce 的常见实现，它使用 ReduceScatter 和 AllGather：环 AllReduce。

快速关注环 AllReduce

环 AllReduce 是 AllReduce 的一种特定实现，针对可扩展性进行了优化。与其让所有设备直接相互通信（这可能会造成通信瓶颈），不如将环 All-Reduce 分解为两个关键步骤：ReduceScatter 和 AllGather。以下是它的工作方式：

  1. ReduceScatter
    • 每个设备将其数据（例如，梯度）拆分为块，并将一个块发送给其邻居。同时，每个设备都会从其另一个邻居接收一个块。
    • 当每个设备接收到一个块时，它会将其对应的块添加到接收到的块中（减少）。
    • 此过程在环周围继续进行，直到每个设备都拥有一个部分减少的块，该块表示所有设备上该块的梯度总和。
  2. AllGather
    • 现在，每个设备都需要从其他设备收集完全减少的块。
    • 设备开始将其减少的块发送给邻居。
    • 每个设备转发它接收到的块，直到每个设备都拥有所有完全减少的块，从而使每个设备都获得完整的、汇总的梯度。

让我们用以下 gif 来演示这一点，其中我们有 5 个 GPU，每个 GPU 都有一个长度为 5 的张量。第一个动画显示了 ReduceScatter 步骤，在步骤结束时，每个 GPU 接收到特定数据块（橙色矩形）的减少结果。

image.png

下一个动画显示了 AllGather 步骤，在步骤结束时，每个 GPU 都获得了 AllReduce 操作的完整结果：

image.png

你可能已经注意到，每个 N GPU 在 reduce-scatter 和 all-gather 步骤中都发送和接收值 N-1 次。每个 GPU 每次传输发送 \frac{K}{N} 个值，其中 K 是在 GPU 之间求和的数组中值的总数。因此，传输到每个 GPU 和从每个 GPU 传输的数据总量约为 2 \times (N-1) \times \frac{K}{N}。当 N（GPU 的数量）很大时，传输到每个 GPU 和从每个 GPU 传输的数据总量约为 2 \times K，其中 K 是参数的总数。

关于 AllReduce，需要记住两件关键的事情：

  1. 当 N（GPU 的数量）很大时，AllReduce 的通信成本约为 2xK。
  2. AllReduce 操作可以分解为 reduce-scatter，后跟 all-gather。这两个操作的通信成本约为 AllReduce 的一半，约为 K。

正如我们所见，此实现可以有效利用即使是有限的节点间带宽。

我们现在已经看到了分布式操作的主要构建块，但在我们看到它们在实际应用之前，让我们看一下用于同步的特殊操作：Barrier。

Barrier

Barrier 是一个简单的操作，用于同步所有节点。只有当所有节点都到达 Barrier 时，Barrier 才会解除。然后，它们才被允许继续进行进一步的计算：

image.png

我们可以通过在每个节点上设置不同的睡眠时间来模拟延迟节点，并查看所有节点通过 Barrier 需要多长时间：

def example_barrier(): rank = dist.get_rank() t_start = time.time() print(f"Rank {rank} sleeps {rank} seconds.") time.sleep(rank) # 模拟不同的处理时间 dist.barrier() print(f"Rank {rank} after barrier time delta: {time.time()-t_start:.4f}") init_process() example_barrier()

我们可以看到，虽然第一个排名根本没有睡眠，但它也花费了 2 秒钟才通过 Barrier：

Rank 0 sleeps 0 seconds. Rank 1 sleeps 1 seconds. Rank 2 sleeps 2 seconds. Rank 0 after barrier time delta: 2.0025 Rank 1 after barrier time delta: 2.0025 Rank 2 after barrier time delta: 2.0024

我们需要小心地同步所有节点，因为这会破坏并行独立操作的目的，因此可能会减慢整个处理速度。在许多情况下，如果快速节点已经开始处理下一个作业，这可能就足够了，因为快速节点在下一次迭代中可能会变慢，因此可以平衡整个过程中的延迟。

在转向实际的分布式训练实现之前，让我们首先解决一个谜团：NCCL 到底是什么？

NCCL：NVIDIA 集体通信库

当在许多 GPU 上训练大型模型时，我们有时可能会走运，但我们总是会遇到镍（或 NCCL 🥁）！那是什么？

有几个库实现了集体通信，并且 PyTorch 支持：有经典的 MPI**（消息传递接口）、Meta 的 **Gloo，最后是 `NCCL`（NVIDIA 集体通信库）。它们都提供类似的功能（在集体通信模式方面），但针对不同的硬件设置进行了优化；NCCL 旨在有效地服务于 GPU-GPU 通信，而 MPI 和 Gloo 设置用于 CPU-CPU 或 CPU-GPU 通信。PyTorch 提供了一个 很好的指南 来决定使用哪个：

在决策树中还有一些更精细的点，我们留给读者在上面引用的 PyTorch 指南中探索。

现在我们已经介绍了分布式操作的基本操作，你现在应该准备好轻松地阅读博客文章了。

A1：分布式训练分析

内核

让我们首先假设内核已集成到 PyTorch 中。作为一个简单的示例，我们可以查看在 PyTorch 中实现的层归一化函数，即 torch.nn.functional.layer_norm。有几种方法可以分析作为此函数基础的内核。最直接的方法可能是使用 Python time 模块。但是，由于 CUDA 操作是异步的，因此使用此方法测量时间只会捕获与在 Python 中启动内核相关的开销，而不是内核本身的实际执行时间。

为了解决这个问题，我们可以使用 torch.cuda.Event 进行精确计时，并使用 torch.cuda.synchronize() 指令来确保我们等待内核执行完成。以下代码段演示了此方法：

def profile_pytorch(func, input): # 创建 CUDA 事件以跟踪时间。CUDA 操作是异步的， start = torch.cuda.Event(enable_timing=True) # 标记开始时间的事件 end = torch.cuda.Event(enable_timing=True) # 标记结束时间的事件 # 预热以消除首次运行的任何开销，首次运行可能无法反映 # 实际性能。 for _ in range(10): func(input) # 在执行要分析的函数之前记录开始时间 start.record() func(input) # 调用我们要分析的函数 # 在函数完成后记录结束时间 end.record() # 同步 CUDA 操作以确保所有操作都已完成 # 在测量经过的时间之前。 torch.cuda.synchronize() # 计算并返回以毫秒为单位的经过时间。 return start.elapsed_time(end)

一种更有效的分析方法是使用 PyTorch 分析器，如前所述。例如，考虑以下代码：

import torch import torch.nn.functional as F def pytorch_layer_norm(input): return F.layer_norm(input, input.size()[1:]) a = torch.randn(10000, 10000).cuda() with torch.profiler.profile( activities=[ torch.profiler.ProfilerActivity.CPU, # 分析 CPU 活动 torch.profiler.ProfilerActivity.CUDA, # 分析 CUDA 活动 ], # 定义分析器的计划 schedule=torch.profiler.schedule( wait=1, # 在开始分析之前等待 1 次迭代 warmup=3, # 预热 3 次迭代以稳定性能 active=2, # 分析 2 次活动迭代 repeat=1, # 重复分析计划一次 ), on_trace_ready=torch.profiler.tensorboard_trace_handler('.'), ) as p: for iter in range(10): pytorch_layer_norm(a) p.step() # 打印分析结果表，按 CUDA 总时间排序，仅限前 10 个条目 print(p.key_averages().table(sort_by="cuda_time_total", row_limit=8))

这将打印按 CUDA 总时间排序的聚合分析结果，输出将是：

image.png

你还可以尝试检查跟踪，正如我们之前在 chrome://tracing/ 上提到的那样

💡 提示

如果你不熟悉此工具，可以使用向右和向左箭头键导航跟踪。此外，你可以按住 Alt 键，同时使用鼠标左右滚动来放大和缩小。

放大后，你可以观察在调用 layer_norm 时操作的流程，在此跟踪中：

image.png

该序列从 CPU（上部）中的 aten::layer_norm 开始，进行到 aten::native_layer_norm，然后转换为 cudaLaunchKernel。从那里，我们继续前进到 GPU，在 GPU 中调用 vectorized_layer_norm_kernel 内核。

📝 注意

你可以通过在分析器中将 profile_memory 设置为 True 来启用内存分析。但是，这可能会导致更复杂的跟踪。

虽然 PyTorch 分析器提供了快速的性能概述，但 NVIDIA Nsight Compute (ncu) 提供了对 GPU 性能的更深入了解，包括每个内核的详细执行时间和内存使用情况。要运行分析器，这非常简单：

ncu --set full python layer_norm.py

其中 layer_norm.py 是一个直接的文件，用于执行层归一化函数。此命令将生成日志输出，但可视化结果的更有效方法是设置输出标志：

ncu --set full -o output python layer_norm.py

并使用 Nsight Compute 打开文件 output.ncu-rep，你将获得如下所示的视图：

image.png

其中包含关于计算和内存利用率的明确警告，以及如何使内核更好地平衡计算和内存并实现最大占用率。

CPP 扩展

如果 你要分析的内核尚未集成到 PyTorch 中，你可以使用 PyTorch 的 cpp_extension 模块轻松编译和运行自定义 CUDA 代码。该过程非常简单——只需在 .cu 文件中创建你的 CUDA 内核，并使用来自 cpp_extension 模块的 load 函数在 Python 中加载它。

对于一个简单的 add 内核，.cu 文件将如下所示：

#include #include #include __global__ void add_kernel(float* x, float* y, float* output, int size) { int index = blockIdx.x * blockDim.x + threadIdx.x; if (index < size) { output[index] = x[index] + y[index]; } } void add_cuda(torch::Tensor x, torch::Tensor y, torch::Tensor output) { int threads = 1024; int blocks = (x.size(0) + threads - 1) / threads; add_kernel<<>>(x.data_ptr(), y.data_ptr(), output.data_ptr(), x.size(0)); } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("add_cuda", &add_cuda, "Vector addition (CUDA)"); }

以及用于加载内核的 python 文件：

import torch from torch.utils.cpp_extension import load # 加载和编译 CUDA 扩展 vector_add = load( name="vector_add", sources=["add_kernel.cu"], verbose=True ) # 定义输入张量 size = 10000 x = torch.randn(size, device='cuda') y = torch.randn(size, device='cuda') output = torch.empty(size, device='cuda') # 运行 CUDA 内核 vector_add.add_cuda(x, y, output)

使用此方法，你可以像我们之前使用 PyTorch 的分析器或 NVIDIA 工具演示的那样，分析自定义 CUDA 内核。

A2：LLM 训练中的典型规模

让我们感受一下 LLM 训练中事物的典型大小。当我们谈论内存或计算时，我们通常会计算“元素”——将它们视为张量中的数字。要获得以字节为单位的实际内存，你需要乘以每个数字的大小（例如，bf16 为 2 个字节，fp32 为 4 个字节）。

以下是一些快速的粗略数字：

A3：计算/通信重叠的数学原理

使用上一节中的公式，我们可以估算计算和通信何时可以在分布式训练中有效重叠。让我们以数据并行（Zero-0）为例。

数据并行通信分析

需要通信的梯度总大小为：

在反向传播期间，这些梯度在桶中通信（默认 25MB）。all-reduce 每个桶的通信时间为：

t_{comm} = t_{comm\_bucket} = \frac{bucket\_size \cdot 2(DP-1)}{DP \cdot peak\_bw}

📝 注意

对于带宽计算，我们使用 NCCL 文档 中的总线带宽公式。这些公式在计算 GPU 之间的有效带宽时考虑了特定的通信模式。

反向传播的计算时间为：

t_{compute} = \frac{4 \cdot num\_tokens \cdot num\_params}{peak\_flops}

为了有效重叠，我们需要：

\frac{t_{comm}}{t_{compute}} = \frac{num\_params}{2 \cdot num\_tokens} \cdot \frac{DP-1}{DP} \cdot \frac{peak\_flops}{peak\_bw} \leq 1

此比率有助于确定通信是否会成为训练中的瓶颈。当比率小于 1 时，通信可以与计算完全重叠。

ZeRO-3 (FSDP) 通信分析

对于 ZeRO-3，参数和梯度在 GPU 之间分片。让我们分析一个模型（Transformer 块大小为 16h^2 参数）的通信模式：

allgather 操作的通信时间为：

t_{comm} = 16h^2 \cdot \frac{DP-1}{DP \cdot peak\_bw}

一个解码器层的前向传播计算时间为：

t_{compute} = \frac{32 \cdot seq\_len \cdot mbs \cdot h^2}{peak\_flops}

为了使计算和通信有效重叠，我们需要：

\frac{t_{comm}}{t_{compute}} = \frac{1}{2 \cdot seq\_len \cdot mbs} \cdot \frac{DP-1}{DP} \cdot \frac{peak\_flops}{peak\_bw} \leq 1

当此比率小于 1 时，下一层的参数通信可以隐藏在当前层的计算背后。

`

TP 通信分析

对于张量并行 (TP)，激活值在线性层期间在 GPU 之间分片。让我们分析一下通信模式：

让我们分析一下我们是否可以将一层的 allgather 通信与下一个线性层的计算重叠。allgather 操作的通信时间为：

t_{comm} = \frac{seq \cdot mbs \cdot h \cdot (TP-1)}{TP \cdot peak\_bw}

而下一个线性层（参数为 h^2）的计算时间为：

t_{compute} = \frac{2 \cdot seq \cdot mbs \cdot h^2}{TP \cdot peak\_flops}

为了有效重叠，我们希望通信时间小于计算时间：

\frac{t_{comm}}{t_{compute}} = \frac{TP-1}{2 \cdot h} \cdot \frac{peak\_flops}{peak\_bw} \leq 1

此比率告诉我们，我们是否可以成功地将 allgather 通信隐藏在下一个线性层的计算背后。有趣的是，该比率仅取决于隐藏层大小 h 和张量并行度 TP，而不取决于序列长度或批大小。

PP 通信分析

对于流水线并行 (PP)，激活值和梯度在流水线阶段之间通信。让我们分析一下通信模式：

让我们分析一下我们是否可以将激活值/梯度的通信与下一个 Transformer 块的计算重叠。下一个流水线阶段中 Transformer 块的计算时间为：

t_{compute} = \frac{32 \cdot seq \cdot mbs \cdot h^2 \cdot num\_layers\_in\_next\_pp}{peak\_flops}

而 P2P 传输的通信时间为：

t_{comm} = \frac{seq \cdot mbs \cdot h}{peak\_bw}

为了有效重叠，我们希望：

\frac{t_{comm}}{t_{compute}} = \frac{peak\_flops}{32 \cdot h \cdot num\_layers\_in\_next\_pp \cdot peak\_bw} \leq 1

与 TP 类似，此比率独立于序列长度和批大小。它取决于隐藏层大小 h、下一个流水线阶段中的层数以及硬件的计算与 P2P 带宽能力之比。

引用

为了在学术环境中署名，请将此作品引用为

Tazi et al., "The Ultra-Scale Playbook: Training LLMs on GPU Clusters", 2025.

BibTeX 引用

@misc{ultrascale_playbook,
      title={The Ultra-Scale Playbook: Training LLMs on GPU Clusters},
      author={Nouamane Tazi, Ferdinand Mom, Haojun Zhao, Phuc Nguyen, Mohamed Mekkouri, Leandro Werra, Thomas Wolf},
      year={2025},
}