震撼！HuggingFace 超大规模训练手册来袭，解锁 GPU 集群训 LLM 密码

• 作者：Nouamane Tazi, Ferdinand Mom, Haojun Zhao, Phuc Nguyen, Mohamed Mekkouri, Leandro Werra, Thomas Wolf
• 发布时间：2023年02月19日
• 阅读时长：2 - 4 天

“

HuggingFace: 我们在多达 512 个 GPU 上运行了 4000 多次缩放实验，并测量了吞吐量（tokens的大小）和 GPU 利用率（tokens的颜色）。请注意，在此可视化中，两者都按模型大小进行标准化

数千个 GPU 完美和谐地协同工作，这就是训练当今最强大的人工智能模型所需要的——一场计算能力的交响乐，直到最近，这还只是精英研究实验室的专属领域。开源已经改变了这一局面，但尚未完全改变。

是的，你可以下载最新的 Llama 或 DeepSeek 模型，阅读它们的技术和实验报告。但最具挑战性的部分——训练代码、协调 GPU 训练这些大规模系统所需的知识和技术——仍然笼罩在复杂性之中，分散在一系列不相关的论文和通常私有的代码库中。

这本开源书籍旨在改变这一现状。从基础开始，我们将带你了解将大型语言模型的训练从一个 GPU 扩展到数十、数百甚至数千个 GPU 所需的知识，用实际代码示例和可复现的基准测试来说明理论。

随着用于训练这些模型的集群规模不断扩大，诸如数据并行、张量并行、流水线并行或上下文并行以及 ZeRO 或内核融合等各种技术被发明出来，以确保 GPU 始终得到高效利用。这大大缩短了训练时间，并充分利用了这种昂贵的硬件。更重要的是，随着扩展人工智能训练的挑战不仅仅局限于构建初始模型，还包括对大规模模型进行微调，这些技术通常也适用于微调结果，通常涉及相同的分布式训练技术。在本书中，我们将逐步介绍所有这些技术——从最简单到最精细的技术——同时保持一个单一的故事线，以了解每种方法的来源。

我们假设你对当前的大语言模型架构有一些简单的基础知识，并且大致熟悉深度学习模型的训练，但你可能对分布式训练比较陌生。如果需要，模型训练的基础知识可以在 DeepLearning.ai 课程或 PyTorch 教程部分找到。本书可视为我们关于预处理数据的第一篇博客（即所谓的“FineWeb 博客文章”）的三部曲的第二部分。阅读了这两篇博客文章后，你应该几乎掌握了全面了解当今高性能大语言模型构建所需的所有核心知识，只缺少一些关于数据混合和架构选择的最后细节来完成这个配方（请期待第三部分……）。

本书基于以下三个一般基础：

• 快速介绍理论和概念：在深入研究代码和实验之前，我们希望从高层次上了解每种方法的工作原理及其优缺点。你将了解语言模型的哪些部分会占用内存以及在训练过程中何时发生。你将学习如何通过并行化模型来解决内存限制问题，并通过扩展 GPU 来提高吞吐量。因此，你将理解以下用于计算 transformer模型内存分解的小部件是如何工作的：

  

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

    

• 清晰的代码实现：理论是一回事，但在实现过程中我们会发现各种边界情况和重要细节。这就是为什么我们尽可能提供实现参考。根据情况，我们将使用两种代码参考：

• • 如果您想了解分布式训练的概念，picotron 存储库是为教育目的而构建的，因此它通常在单个、自包含的短文件中实现概念。
  • 另一方面，要查看生产就绪代码，我们将参考 nanotron 的实现，这是 Hugging Face 使用的生产训练代码库。

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

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

1. 高层次概述

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

• 内存使用：这是一个硬限制——如果一个训练步骤无法放入内存，训练就无法进行。
• 计算效率：我们希望硬件将大部分时间用于计算，因此需要减少在数据传输或等待其他 GPU 执行工作上花费的时间。
• 通信开销：我们希望尽量减少通信开销，因为它会使 GPU 闲置。为了实现这一点，我们将尝试充分利用节点内（快速）和节点间（较慢）的带宽，并尽可能使通信与计算重叠。

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

由于本书内容广泛，我们制作了一个速查表来帮助你浏览本书并获得主要要点。在你探索这些复杂内容时，请务必牢记！

2. First Steps: 在单个 GPU 上训练

如果你想在阅读过程中增加播客的感觉，可以随时收听 NotebookLM 主持人讨论本书的前几节内容。

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

• 前向传播：将输入通过模型以生成输出。
• 反向传播：计算梯度。
• 优化步骤：使用梯度更新参数。

通常看起来像这样：

将鼠标悬停在网络元素上可以查看其详细信息。

在这个图中，顶行的框可以看作是模型内部的连续层（底行也是如此）。红色框是在反向传播过程中计算的每个层的相关梯度。

批次大小（batch size,bs）是模型训练的重要超参数之一，它会影响模型的收敛和吞吐量。

例如，在 DeepSeek-V3/R1 训练中，“在训练的前 469B 个tokens中，批次大小从 3072 个输入序列逐渐增加到 15360，然后在剩余的训练中保持在 15360 个输入样本”。

在训练初期，小批次大小有助于快速在训练空间中移动，以达到最佳学习点。然而，在模型训练的后期，小批次大小会使梯度嘈杂，模型可能无法收敛到最优的最终性能。另一方面，大批次大小虽然能给出非常准确的梯度估计，但会减少对每个训练tokens的利用，导致收敛变慢，并可能浪费计算资源。你可以在 OpenAI 关于大批量训练的论文或 MiniMax-01 技术报告的 4.2 节中找到关于这个主题的早期讨论。

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

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

最近大语言模型训练的一个合适的批次大小通常在每批次 400 万到 6000 万tokens的数量级。

多年来，批次大小和训练语料库一直在稳步增加：Llama 1 使用约 400 万tokens的批次大小训练了 1.4 万亿tokens，而 DeepSeek 使用约 6000 万tokens的批次大小训练了 14 万亿tokens。

当我们将模型训练扩展到这些大批次大小时，

• 第一个挑战已经出现：内存不足问题。当我们的 GPU 没有足够的内存来容纳目标批次大小的完整批次时，我们该怎么办？

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

2.1. Transformer 中的内存使用

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

• 模型权重
• 模型梯度
• 优化器状态
• 计算梯度所需的激活值

“

需要注意的是，对于一个模型，你可能认为可以准确计算内存需求，但实际上有一些额外的内存占用因素使得这很难做到：

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

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

那么如何根据这些变量快速确定内存使用情况呢？一种简单的方法是通过实验测量它。

2.1.1. 分析内存使用情况

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

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

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

现在我们已经对内存有了初步了解，让我们看看如何在将训练扩展到多个 GPU 时，在保持这些各种项目（激活值、参数、梯度、优化器状态）的内存需求在 GPU 内存限制内的同时，最大限度地提高计算效率。

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

(WEIGHTS/GRADS/OPTIMIZER STATES MEMORY) 让我们从列表中的前三个项目开始：模型的权重、梯度和优化器状态。我们实际上可以很容易地估计它们所需的内存。

对于一个简单的 Transformer 大语言模型，参数数量由以下公式给出：

在传统的全精度（FP32）训练中，参数和梯度的内存需求是参数数量乘以每个参数的字节数，全精度（FP32） 训练参数和梯度两者都需要 4 字节，而对于 Adam 优化器，需要存储动量和方差，这又为每个参数增加了 4 字节。总之：

• m_{optimizer} = (4 + 4) * N

现在让我们看看如果使用较低精度会发生什么。出于稳定性原因（在下面的混合精度训练部分会详细介绍），我们通常不使用全低精度训练，而是使用一种称为“混合精度”的高低精度混合方法。如今混合精度训练的默认做法通常是在大多数计算中使用 BF16（每个参数和梯度需要 2 字节），以及额外的 FP32 模型权重和梯度副本，因此每个参数总共需要 12 字节。除了参数和梯度，我们还需要存储优化器状态：对于 Adam 优化器，这需要存储通常为 FP32 的动量和方差以保证数值稳定性，每个使用 4 字节。

总结如下：

有趣的是，混合精度本身并没有总体上节省内存，因为它只是在三个组件之间不同地分配内存，实际上如果我们在 FP32 中累积梯度，比全精度训练还多了 4 字节。它仍然有优势，因为在前向/反向传播中使用半精度计算允许我们:

• （1）在 GPU 上使用优化的低精度操作，这些操作更快，
• （2）减少前向传播期间的激活内存需求，正如我们在上面和下面的图表中看到的，这是内存使用的很大一部分。

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

• 使用 FP8 训练代替 BF16 将进一步减少内存使用，但它不太稳定，并且是一个非常活跃的研究课题（见此推文），我们将在后面更详细地介绍。

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

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

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

2.1.3. 激活内存(ACTIVATIONS MEMORY)

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

在这个公式中，是层数，是序列长度，是样本批次大小，是模型的隐藏维度，是前馈维度。

关于这些数字的精确推导，你可以参考 NVIDIA 关于重计算的原始论文，它本质上要求你对transformer layer中每个操作之间的所有中间激活的大小进行核算。

这里一个有趣的观察结果是，对于给定的模型，内存不是静态的，它与序列长度和批次大小呈线性比例关系。这意味着当我们增加批次大小或使用更长的序列进行训练时，激活内存将会大幅增加。我们可以使用这个公式来查看不同序列长度下内存使用的变化情况，例如对于 Llama 模型（）：

这个图表揭示了一个惊人的事实：对于短序列（或小批次大小情况类似），激活内存几乎可以忽略不计，但从大约 2 - 4k 个标记开始，它们会占用大量内存，而参数、梯度和优化器状态的使用（我们稍后会讨论）大致与序列长度和批次大小无关。

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

那么，有没有办法控制这种“激活爆炸”呢？读者朋友们，这是个好问题！接下来是时候介绍我们的第一种技术——激活重计算了，它将帮助我们控制激活内存占用。这是当今大型模型训练工具箱中的一个重要工具。

2.2. Activation recomputation

激活重计算（也称为梯度检查点或重新计算）的基本思想是在正向传播过程中丢弃一些激活值以节省内存，并在反向传播过程中花费一些额外的计算来即时重新计算这些值。在不进行重计算的情况下，我们会存储两个可学习操作（例如前馈、层归一化等）之间的每个隐藏状态，以便在反向传播时用于计算梯度。

当我们使用重计算时，通常只会在模型架构的几个关键点存储激活值，丢弃其余的激活值，并在反向传播时从最近保存的激活值即时重新计算它们，基本上是在正向传播的一个子部分上进行计算，以用计算换取内存。通常情况如下所示：

将鼠标悬停在网络元素上可以查看其详细信息。

有几种选择关键激活值进行存储的策略：

• 完全策略：我们在 Transformer 模型的每一层之间的过渡点对激活值进行检查点操作。这通常被称为完全策略，因为它需要在反向传播时对每一层进行一次正向传播，本质上是在反向传播期间添加了一个完整的正向传播。这种策略节省的内存最多，但在计算方面也是最昂贵的。它通常会使激活计算的计算成本和时间增加多达 ，这是非常显著的。
• 选择性策略：通常我们可以做得比完全策略更好。重计算论文的作者进行了详细的分析，研究了哪些激活值增长最大且重新计算成本（以 FLOPs 为单位）最低。结果发现注意力计算属于这一类，因此我们通常可以丢弃它们，而专注于对昂贵的前馈计算进行检查点操作。对于一个 GPT - 3（175B）模型，这意味着在 2.7%的计算成本下减少 70%的激活内存。

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

这里另一个明显的趋势是，对于较小的模型，长序列的激活值起着更大的作用，因此重计算的效果更加显著。

“

注意，当你测量训练设置在使用 GPU/TPU 加速器时的效率时，通常需要考虑重计算来计算总 FLOPS（每秒浮点运算次数），并将其与 GPU/TPU 加速器的理论最大 FLOPS 进行比较。在计算训练步骤的 FLOPS 时考虑重计算会得到一个称为“硬件 FLOPS”的值，它是在加速器上实际执行的操作数量。将这个数字除以训练步骤的持续时间和最大加速器 FLOPS，得到硬件 FLOPS 利用率（HFU）。 然而，归根结底，真正重要的是在给定数据集上训练模型所需的端到端时间。因此，当比较不同的 GPU/TPU 加速器时，如果其中一个加速器提供了足够的内存来跳过重计算，从而每秒执行较少的操作（较低的 HFU）但训练速度更快，那么它应该受到奖励而不是惩罚。因此，另一种方法是计算所谓的模型 FLOPS 利用率（MFU），与 HFU 不同，它只考虑模型正向/反向传播所需的操作，不包括重计算的实现。

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

如你现在所理解的，激活重计算由于重新计算会略微增加 FLOPs 的数量，但它显著减少了内存访问开销。这种权衡在具有小高速内存的硬件（如 GPU）上特别有利，因为访问内存通常比执行计算慢。尽管涉及额外的操作，但总体效果通常是更快的计算，以及更低的内存占用。

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

然而，激活内存仍然与批次大小呈线性依赖关系，并且我们在上面的柱状图中的所有分析都是使用 进行的，所以当我们转向更大的批次大小时，它可能会再次成为一个问题。别担心，我们还有第二个工具——梯度累积来解决这个问题！

2.3. Gradient accumulation

梯度累积是一种避免内存爆炸的非常直接的方法，它包括将我们的批次分割为微批次(micro-batches)。我们将在每个微批次上依次执行前向和反向传播，计算梯度，并且，正如其名称所示，在执行优化器步骤之前对所有微批次的梯度求和。在实践中，优化器步骤不是对总和进行操作，而是对梯度的平均值进行操作，这样结果就与梯度累积的步数无关。

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

“

使用梯度累积意味着我们需要保留缓冲区来累积在整个训练步骤中持续存在的梯度。 在不进行梯度累积的情况下，在计算激活内存时，前向传播是在没有累积梯度的情况下计算的，这意味着内存峰值较低。

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

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

但是如果你仔细阅读了前面的内容，你可能已经注意到每个微批次的前向/反向传播实际上可以并行运行。前向/反向传播彼此独立，唯一的区别是独立的输入样本。似乎是时候开始将我们的训练扩展到多个 GPU 了！

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

2.3.1. 分析 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 的跟踪查看器中可视化它。跟踪显示：

• CPU 线程异步向 GPU 启动内核
• 多个 CUDA 流并行处理计算和通信
• 内核执行时间和内存分配

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

跟踪有助于识别以下瓶颈：

• 可以重叠的顺序计算和通信
• GPU 等待数据传输的空闲时间
• CPU 和 GPU 之间的内存移动
• CPU 的内核启动开销

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

现在让我们使用几个 GPU 来搭建一个更大的工作站，并开始研究我们的第一种扩展技术——数据并行，正如我们将看到的，它只是梯度累积的并行版本。

3. Data Parallelism

为了给你的阅读体验增添播客的感觉，你可以在阅读过程中收听 NotebookLM 主持人讨论本书的以下章节。

数据并行（DP）的理念是在多个 GPU 上复制模型（我们将这些副本称为“模型实例”），并针对每个 GPU 并行地在不同微批次数据上运行前向和反向传播，这就是它被称为数据并行的原因。你可能已经在简单的训练示例中见过数据并行，但正如你很快就会看到的，在本节中我们将进行更深入的探讨，所以即使你知道一般的方法，也请继续关注。

“

如果你不熟悉分布式通信模式，如广播、收集或归约，[我们准备了一个简短的分布式编程速成课程(https://nanotron-ultrascale-playbook.static.hf.space/dist/index.html#a0%3A_parallel_programming_crash_course)]。

其思路是为每个 GPU 使用不同的微批次数据，这意味着每个 GPU 中的梯度会有所不同。因此，为了使不同 GPU 上的模型实例保持同步，在反向传播期间，在优化器步骤之前，会使用一个名为“归约（all-reduce）”的操作来平均模型实例的梯度。

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

一个简单的实现方式可能是在反向传播完成后等待所有梯度计算完成，然后在所有 DP 等级上触发归约操作来同步这些梯度。但这种先计算后通信的顺序步骤是绝对不可取的！因为我们不希望在通信发生时 GPU 处于空闲状态，就像上图所示那样。

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

让我们看看三种优化方法，它们能让我们比最初的简单实现做得更好！

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

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

如上图所示，在一层的梯度（红色框）中，即使前面层（左边的红色框）的梯度尚未计算完成，也可以收集和求和。例如，一旦最后一层的反向传播完成（最右边的框），就可以收集和求和这些梯度，同时前面层的反向计算继续向左进行。在 PyTorch 中，可以通过为每个参数附加一个归约钩子函数来实现这一点。一旦某个参数的梯度准备好，就会触发归约操作，而其他参数的梯度仍在计算中。这种方法将大部分归约操作与梯度计算重叠，从而提高了效率。以下是一个附加钩子的简单函数：

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 实现（带有重叠）：

class DataParallelNaive(nn.Module):
    """
    Naive Data Parallelism. Not used in practice. But it is a good starting point to understand how data parallelism works.
    It implements a simple all-reduce operation to synchronize gradients across multiple processes.
    And `no_sync` context manager to disable gradient synchronization.
    """
    def __init__(self, module):
        """
        Initializes the DataParallel wrapper for a given module.

        Args:
            module (nn.Module): The model to be wrapped for data parallelism.
            process_group (torch.distributed.ProcessGroup): The process group used for gradient synchronization. 
                                                            It could be a data parallel or context parallel group.
        """
        super().__init__()
        self.module = module
        self.require_backward_grad_sync = True# whether to synchronize gradients during backward pass. Set to False when using gradient accumulation
        self.register_backward_hook(self._allreduce_grads)
    
    def forward(self, *inputs, **kwargs):
        return self.module(*inputs, **kwargs)
    
    def register_backward_hook(self, hook):
        """
        Registers a backward hook for all parameters of the model that require gradients.    
        """
        for p in self.module.parameters():
            if p.requires_grad isTrue:
                p.register_hook(hook)
                
    def _allreduce_grads(self, grad):
        """
        Performs an all-reduce operation to synchronize gradients across multiple processes.    
        """
        # No synchronization needed during gradient accumulation, except at the final accumulation step.
        if self.require_backward_grad_sync:
            dist.all_reduce(grad, op=dist.ReduceOp.SUM, group=pgm.process_group_manager.cp_dp_group)
            grad /= pgm.process_group_manager.cp_dp_world_size
        return grad 
    
    @contextlib.contextmanager
    def no_sync(self):
        """
        A context manager to temporarily disable gradient synchronization. 
        This is useful for performing multiple backward passes during gradient accumulation without synchronizing 
        gradients in between.
        """
        self.require_backward_grad_sync = False
        yield
        self.require_backward_grad_sync = True

这是我们在本文中多次讨论的“重叠计算和通信”的第一个示例，它是实现最大缩放效率的关键技术。但我们可以进一步提高效率！

优化二：梯度分桶

GPU 操作在处理大型张量时通常比处理多个较小张量更高效。通信操作也是如此。

因此，我们可以有利地将梯度分组到桶中，并为同一桶内的所有梯度启动单个归约操作，而不是为每个梯度执行独立的归约。它通常看起来像这样：

可以把它想象成在发货前将物品打包成箱子。发送几个大箱子比发送许多小箱子更高效。通过为每个桶执行单个归约操作，我们可以显著减少通信开销并加快通信操作。

以下是带有分桶的代码实现：

class DataParallelBucket(nn.Module):
    """
    Data Parallelism with gradient grouped into buckets to reduce the communication overhead.
    """
    def __init__(self, module, bucket_cap_mb=25, grad_type = torch.float32):
        """
        Initialize the DataParallelBucket module.
        
        Args:
            module (nn.Module): The model to be parallelized.
            process_group: The process group for gradient synchronization, which can be either 
                           a data parallel group or a context parallel group.
            bucket_cap_mb (int, optional): The maximum size of each gradient synchronization bucket in megabytes. 
                                           Defaults to 25 MB.
            grad_type (torch.dtype, optional): The data type of gradients, defaulting to float32.
        """
        super().__init__()
        self.module = module
        self.require_backward_grad_sync = True# whether to synchronize gradients during backward pass. Set to False when using gradient accumulation
        grad_size = 2if grad_type == torch.bfloat16 else4# float32 gradient: 4 bytes
        bucket_size = bucket_cap_mb * 1024 * 1024 // grad_size # number of gradients in one bucket
        self.bucket_manager = BucketManager(module.parameters(), pgm.process_group_manager.cp_dp_group, bucket_size, grad_type)
        self.register_backward_hook()
        self._post_backward_callback_set = False# whether the callback for wait gradient synchronization is set
        
    def forward(self, *inputs, **kwargs):
        return self.module(*inputs, **kwargs)

    def backward(self, input_tensor, output_tensor, output_tensor_grad):
        return self.module.backward(input_tensor, output_tensor, output_tensor_grad)
    
    def register_backward_hook(self):
        """
        Registers a backward hook to manually accumulate and synchronize gradients.
        
        This hook serves two main purposes:
        1. PyTorch does not natively support gradient accumulation with mixed precision.
        2. After gradient accumulation, it flags parameters as ready for synchronization.
        
        The gradient accumulation functions are stored to prevent them from going out of scope.
        
        References:
        - https://github.com/NVIDIA/Megatron-LM/issues/690
        - https://pytorch.org/docs/stable/generated/torch.autograd.graph.Node.register_hook.html
        - https://arxiv.org/abs/2006.15704 (page 5)
        """
        self.grad_accs = []
        for param in self.module.parameters():
            if param.requires_grad:
                # Expand so we get access to grad_fn.
                param_tmp = param.expand_as(param)
                # Get the gradient accumulator function.
                grad_acc_fn = param_tmp.grad_fn.next_functions[0][0]
                grad_acc_fn.register_hook(self._make_param_hook(param, self.bucket_manager))
                self.grad_accs.append(grad_acc_fn)
                
    def _make_param_hook(self, param: torch.nn.Parameter,bucket_manager: BucketManager):
        """
        Creates the a hook for each parameter to handle gradient accumulation and synchronization.
        """
        def param_hook(*unused):
            """
            The hook called after the gradient is ready. It performs the following:
            1. Accumulates the gradient into the main gradient.
            2. Adds a post-backward callback to wait for gradient synchronization completion.
            3. Marks the parameter as ready for synchronization.
            """
            if param.requires_grad:
                assert param.grad isnotNone
                param.main_grad.add_(param.grad.data) # accumulate the gradients
                param.grad = None
                
                # skip the gradient synchronization (gradient accumulation/PP micro batches)
                if self.require_backward_grad_sync:
                    # Add a callback to wait for gradient synchronization. Ensures the callback is added only once.
                    # Callback is executed after the backward pass. It should be added per backward pass.
                    ifnot self._post_backward_callback_set:
                        Variable._execution_engine.queue_callback(self._post_backward)
                        self._post_backward_callback_set = True
                        
                    # mark the parameter as ready for gradient synchronization. 
                    bucket_manager.mark_param_as_ready(param) 
        return param_hook
    
    @contextlib.contextmanager
    def no_sync(self):
        """A context manager to disable gradient synchronization."""
        self.require_backward_grad_sync = False
        yield
        self.require_backward_grad_sync = True
        
    def _post_backward(self):
        """
        A post-backward callback that waits for gradient synchronization to finish, then copies 
        the synchronized gradients back to the parameters' grad attribute.
        
        This method is called after the backward pass and before the optimizer step.
        """
        self.bucket_manager.wait()
        self._post_backward_callback_set = False
        # copy to params.grad so we can use the optimizer to update the parameters
        for p in self.module.parameters():
            if p.requires_grad:
                p.grad = p.main_grad.to(p.dtype) # In PyTorch, you cannot assign a gradient with one data type to a tensor of another data type.

    def reset(self):
        """
        Reset the bucket manager and zero out gradients in the model
        """
        self.bucket_manager.reset() 

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

最后，如我们之前所见，梯度累积通过在执行 optimizer.step() 之前执行多次前向和反向传播来工作。当将梯度累积与数据并行结合时，我们在同步梯度时需要小心。

在一个朴素的版本中，在累积过程中的每次反向传播后都会自动触发归约操作，这不是最优的，因为在最后一步进行单个归约会产生相同的效果，同时减少开销。

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

“

需要注意的是，在执行通信操作时，张量必须在内存中连续，以避免冗余的内存复制。为了实现这一点，我们通常会预先分配与激活值或模型参数大小相同的连续缓冲区，专门用于通信。虽然这加快了通信速度，但也在一定程度上增加了训练期间的峰值内存使用。

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

重新审视全局批次大小(global batch size)

其中 是梯度累积的步数， 是用于数据并行的并行实例数量。

给定目标全局批次大小，我们可以通过调整梯度累积步数和数据并行进程来加快训练速度。

“

进一步阅读数据并行的一个好资源是 https://siboehm.com/articles/22/data-parallel-training。

在实践中，人们倾向于尽可能最大化数据并行节点（DP）的数量，而不是梯度累积，因为数据并行本身是并行的，不像梯度累积具有顺序性。当仅通过扩展数据并行无法达到目标全局批次大小时，再添加梯度累积来实现目标。

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

我们目前的进展

让我们快速总结如何设置我们的第一个 1D 并行训练，并给出一个最佳数据并行设置的初步方案：

“

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

虽然数据并行通过将归约梯度同步与反向计算重叠节省了时间，但在大规模情况下，这种优势开始减弱。为什么呢？因为随着我们添加越来越多的 GPU（成百上千个），协调它们的开销显著增加，并且网络需求变得过大，以至于超过了其带来的好处。结果，随着我们向系统中添加更多的 GPU，我们的设置将变得越来越低效。

让我们通过一些基准测试来看看实际情况：

从图表中我们可以看到，在超过一定限制后，我们的吞吐量开始显著下降，而每个 GPU 的内存使用保持不变，并且不受添加更多 DP 等级的影响。

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

“

提示：你可以通过将模型参数所需的最小内存乘以 2 来快速估算你的模型所需的内存。例如，对于 70B 模型，70B + 70B = 140GB（约为 133GB）。

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

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

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

共享范式与 DP 密切相关，所以我们先来研究 Zero 方法！

ZeRO (Zero Redundancy Optimizer)

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

虽然数据并行是扩展训练的一种有效方式，但在每个数据并行等级上天真地复制优化器状态、梯度和参数会引入大量内存冗余。ZeRO 通过在数据并行维度上对优化器状态、梯度和参数进行分区来消除内存冗余，同时仍允许使用完整的参数集进行计算。这有时需要在数据并行等级之间进行更多的通信，这些通信可能会或可能不会完全重叠，我们接下来会看到！

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

• ZeRO-1：优化器状态分区
• ZeRO-2：优化器状态 + 梯度分区
• ZeRO-3（也称为 FSDP，即“完全分片数据并行”）：优化器状态 + 梯度 + 参数分区

当我们说分区时，是指沿着数据并行维度进行划分。我们将主要关注 ZeRO-1 到 ZeRO-3，因为这应该能让我们对其如何减少内存以及其中的权衡有一个广泛的了解。你可以在 DeepSpeed 文档中找到更多 ZeRO 的变体。

你可能会注意到在可分片的内容中缺少激活值。由于模型的每个数据并行副本接收不同的微批次，每个数据并行等级上的激活值也不同，所以它们不会被复制，因此也不能被分片！

让我们更仔细地看看在 ZeRO 的每个阶段进行分区可以节省多少内存！

重新审视内存使用

ZeRO-1：优化器状态分区

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

这是 ZeRO 中的一个新操作，在普通的数据并行中不使用。

你可能会想知道这个“reduce-scatter”操作是什么以及整个过程是如何工作的，让我们通过下面的图表使其更直观。我们将详细介绍前向/反向传播周期的所有步骤：在实际通信方面，与普通的数据并行相比，ZeRO-1 将我们的“归约”梯度通信更改为“reduce-scatter”操作，并在优化器步骤之后添加了对所有参数的all-gather操作。如下所示：

如果你一直在跟进，你会记得在普通的数据并行中，我们可以将归约梯度通信与反向传播计算重叠。在 ZeRO-1 中，我们也可以研究如何有效地重叠新添加的 bf16 参数的all-gather操作。有两种主要策略：

• 在优化器步骤期间：我们可以在优化器更新部分参数后立即启动all-gather操作。这允许通信有可能与其他参数更新重叠。
• 在前向传播期间：我们可以将每个层的参数的all-gather操作与前向传播重叠。

“

注意，不幸的是，这些技术并不容易实现，需要巧妙地使用钩子和分桶。在实践中，我们可以直接使用 PyTorch 原生的 ZeRO-3/FSDP 实现，并将 FSDPUnit 设置为整个模型，稍后会详细介绍。

在 ZeRO-1 中，优化器状态已分区，这意味着每个副本仅更新优化器状态。敏锐的读者一定已经注意到，一开始就没有必要在所有 DP 等级上拥有所有梯度，因为优化步骤只需要一个子集。认识 ZeRO-2！

ZeRO-2：添加梯度分区

在通信方面，ZeRO-2 与 ZeRO-1 类似，它们都需要对梯度进行reduce-scatter操作，并对所有参数进行all-gather操作。

“

注意，你可能会发现使用 ZeRO-2 相对于 ZeRO-1 并没有真正的通信开销增加，实际上 ZeRO-2 通常是最佳选择。

现在我们已经对梯度进行了分片，我们是否已经完成了所有工作，还是可以继续改进呢？嗯，差不多了。接下来是 ZeRO-3！

ZeRO-3：添加参数分区

在第三阶段，我们将上述在数据并行副本上对优化器状态和梯度进行分片的方法扩展到对模型的参数进行分片。

“

注意，这个阶段在 PyTorch 原生实现中也称为 FSDP（完全共享数据并行）。在本文中我们将只提及 ZeRO-3，但你可以在看到 FSDP 的地方将其视为相同的概念。

那么在实践中，如果模型的所有部分都分布在不同的 GPU 上，我们如何进行前向或反向传播呢？很简单，我们在需要时按需收集它们。在前向传播中，这看起来如下所示：因此，当我们执行前向传播并依次遍历层时，我们按需检索必要的参数，并在不再需要它们时立即从内存中刷新它们。反向传播的工作方式相同，只是流程相反，并且我们生成梯度分片：

另一个问题是，我们需要在整个前向和反向步骤中持续进行这些all-gather操作，与 ZeRO-2 相比，在一个训练步骤中总共会增加 2 * num_layers - 1 个额外的all-gather操作，每个操作都会带来一个小的基本延迟开销，如下所示：

让我们总结一下到目前为止我们在数据并行和 ZeRO 方面的探索：我们已经看到，通过数据并行，我们可以通过添加更多的模型副本显著提高训练吞吐量。通过 ZeRO，我们可以训练那些通常无法放入单个 GPU 的模型，通过在数据并行等级上对参数、梯度和优化器状态进行分片，同时只产生少量的通信成本。

然而，这里存在一个限制，DP仅在模型的一层能够放入单个 GPU 时才有效，而 ZeRO 只能对参数、梯度和优化器状态进行分区，无法处理激活内存！

我们从激活内存的讨论中记得，这部分内存与序列长度和批次大小成比例增长。当然，我们可以限制这些因素，但在实践中，我们不想因为硬件限制而只能使用短序列长度进行训练。

为了克服这个问题，是时候探索一种新的正交平行轴——张量了并行性（TP）。与依赖重参数通信的ZeRO3不同，TP建议分片参数、渐变、优化器状态和跨设备激活，无需GPU之间的任何模型参数通信

让我们看看如何对这个操作进行并行化！在张量并行中，张量将沿着特定维度被分割成 个分片并分布在 个 GPU 上。矩阵可以按列或行进行分割，从而导致列并行和行并行。我们将在下面看到，选择行或列分片将需要不同的通信原语。

4. Tensor Parallelism

因此，我们使用 ZeRO 对模型的参数、梯度和优化器状态进行了分片，但一旦激活内存超过我们的内存预算，我们就会达到极限。欢迎使用张量并行 （TP），这是一种对权重、梯度和优化器状态以及激活进行分片的方法，无需在计算之前收集它们。好像一场梦！首先，我们来看看 Tensor Parallel 如何处理简单的矩阵乘法。

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

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

我们的第一个选择是使用列向分片（也称为列线性）：我们将完整的输入矩阵复制到每个工作节点，这需要一个名为广播(broadcast)的操作，然后将权重矩阵按列分割。输入与部分权重矩阵相乘，结果最终使用all-gather操作进行组合。

这里是列向分片的代码实现，Picotron 中的列并行 TP 实现：

class ColumnParallelLinear(torch.nn.Module):
    """Column Parallel Linear layer
    Y = XW + b, where weight matrix W is parallelized along its second dimension. W = [W_1, ..., W_p]
    This module returns the results of Y_i = XW_i + b_i in the forward method, Y_i is parallelized in the second dimension.
    Arguments:
        in_features: first dimension of weight matrix W.
        out_features: second dimension of weight matrix W.
        bias: If true, add bias
        init_method: method to initialize weights
        gather_output: If true, gather the output from all the partitions. This is used for the last linear layer
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool = False,
        gather_output: bool = False,
        async_all_reduce: bool = False,
    ) -> None:
        super(ColumnParallelLinear, self).__init__()

        self.tp_world_size = pgm.process_group_manager.tp_world_size
        self.tp_rank = pgm.process_group_manager.tp_rank 

        self.in_features = in_features
        self.out_features = out_features
        assert out_features % self.tp_world_size == 0, "Hidden dimension must be divisible by the tensor parallel world size"
        self.output_size_per_partition = out_features // self.tp_world_size
        self.gather_output = gather_output
        self.async_all_reduce = async_all_reduce
        # Allocate space for the weight and bias
        # Note: torch.nn.functional.linear performs XW^T + b so we exchange the order of dimensions
        self.weight = nn.Parameter(torch.Tensor(self.output_size_per_partition, self.in_features)) # W_i
        if bias:
            self.bias = nn.Parameter(torch.Tensor(self.output_size_per_partition))
            with torch.no_grad():
                self.bias.zero_()
        else:
            self.register_parameter("bias", None)

        self.reset_parameters()

    def reset_parameters(self):
        # Initialize weight tensor with the default initialization method used for nn.Linear in PyTorch
        master_weight = torch.empty(
            self.out_features, 
            self.in_features, 
            dtype=self.weight.dtype,
            device=self.weight.device,
            requires_grad=False
        )
        
        # Calculate bound based on master weight's input dimension
        k = 1 / master_weight.size(1)
        bound = math.sqrt(k)
        torch.nn.init.uniform_(master_weight, -bound, bound)
        
        # Split the model into size of self.output_size_per_partition
        weight_list = torch.split(master_weight, self.output_size_per_partition, dim=0)
        self.weight.data = weight_list[self.tp_rank].contiguous()
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.async_all_reduce:
            output = linear_with_async_all_reduce(x, self.weight, self.bias) 
        else:
            output = linear_with_all_reduce(x, self.weight, self.bias) 
        if self.gather_output:
            output = GatherFromModelParallelRegion.apply(output)
        return output

第二种选择称为行向分片（也称为行线性）：正如细心的读者可能猜到的，行线性意味着我们将权重矩阵分割成行块。然而，这也需要我们分割输入，这需要一个scatter操作，而不是列线性分片中使用的广播操作。每个工作节点上的结果已经是正确的形状，但需要求和以得到最终结果，因此在这种情况下需要一个归约操作。这里我们看到了第四个分布式原语：散射！

这是行并行张量并行的实现，Picotron 中的行并行 TP 实现：

class RowParallelLinear(nn.Module):
    """Linear layer with row parallelism.
    Y = XW + b. W is parallelized along its first dimension and X along its second dimension as:
               -   -
              | W_1 |
              | .   |
          W = | .   |        X = [X_1, ..., X_p]
              | .   |
              | W_p |
               -   -
    We assume that X is already parallelized. This is the case after ColumnParallelLinear.
    This module returns the results of Y = sum(X_i * W_i + b_i) in the forward method.
    Arguments:
        in_features: first dimension of matrix W.
        out_features: second dimension of matrix W.
        bias: If true, add bias
        init_method: method to initialize weights.
    """
    def __init__(self, in_features: int, out_features: int, bias: bool):
        super(RowParallelLinear, self).__init__()

        self.tp_world_size = pgm.process_group_manager.tp_world_size
        self.tp_rank = pgm.process_group_manager.tp_rank 

        self.in_features = in_features
        self.out_features = out_features
        assert in_features % self.tp_world_size == 0, "Hidden dimension must be divisible by the tensor parallel world size"
        self.input_size_per_partition = in_features // self.tp_world_size

        self.weight = nn.Parameter(torch.Tensor(self.out_features, self.input_size_per_partition))
        if bias:
            self.bias = nn.Parameter(torch.Tensor(self.out_features))
            # Always initialize bias to zero.
            with torch.no_grad():
                self.bias.zero_()
        else:
            self.register_parameter("bias", None)

        self.reset_parameters()

    def reset_parameters(self):
        # Initialize weight tensor with same dtype and device as self.weight
        master_weight = torch.empty(
            self.out_features, 
            self.in_features, 
            dtype=self.weight.dtype,
            device=self.weight.device,
            requires_grad=False
        )
        
        # Calculate bound based on master weight's input dimension
        k = 1 / master_weight.size(1)
        bound = math.sqrt(k)    
        torch.nn.init.uniform_(master_weight, -bound, bound)
        
        # Split the model into size of self.input_size_per_partition
        weight_list = torch.split(master_weight, self.input_size_per_partition, dim=1)
        self.weight.data = weight_list[self.tp_rank].contiguous()

    def forward(self, x):
        # X_i * W_i^T + b
        output_parallel = F.linear(x, self.weight)
        # All-reduce across all the partitions.
        output = ReduceFromModelParallelRegion.apply(output_parallel)
        return output if self.bias isNoneelse output + self.bias

现在我们已经有了张量并行的基本构建块，让我们看看如何在一个变压器层内有效地组合它们！

Transformer 块中的张量并行

为了制定一个策略，让我们从一个简单的示例转移到一个真实的模型构建块。一个变压器模型由两个主要构建块组成：前馈层（MLP）和多头注意力（MHA）。我们可以对两者应用张量并行。

前馈部分可以通过先进行“Column linear”然后进行“Row Linear”来并行化，这相当于在正向传播中进行广播以复制输入并进行归约。需要注意的是，在实际训练中，广播操作可能不是必需的，因为我们可以确保输入在张量并行等级之间已经同步。这种设置比先进行“Row Linear”然后进行“Column linear”更有效，因为我们可以跳过两个分割操作之间的中间归约。

现在我们已经为变压器的前馈部分找到了一个有效的方案，让我们看看多头注意力块（MHA）。

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

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

最后需要注意的是，张量并行并不是训练的万能解决方案。我们在模型的计算路径中直接添加了几个分布式通信原语，因此很难完全隐藏（像在 ZeRO 中那样与计算重叠），我们的最终性能将是计算和内存收益与增加的通信开销之间的权衡结果。让我们通过缩放张量并行度来说明这一点：

(通过执行块矩阵乘法和异步通信/计算，可以部分隐藏这种通信。)

查看张量并行 MLP 中的作时间表（同样适用于 Attention ），我们可以更好地理解所涉及的权衡。

在每个解码器层的前面，我们用 AllReduce作打到一个不能与计算重叠的同步点。在应用最终的 LayerNorm 之前，这种公开的通信开销对于跨张量并行排名组合部分结果是必要的。

“

例如， Megatron-LM/Nanotron 实现了全聚集与 FC1 计算的部分重叠，其中矩阵乘法结果的一部分将开始发送到另一个 GPU，而另一部分仍在计算中。这个研究领域仍然是一个活跃的研究领域，最近的工作如 Domino 探索了新技术以最大限度地利用这种重叠。

张量并行确实有助于减少矩阵乘法的激活内存，因为中间激活是在 GPU 之间分片的。但是，我们仍然需要为 LayerNorm 等作收集完全激活，这意味着我们无法获得完全内存优势。此外，TP 还引入了大量依赖于网络基础设施的通信要求。无法将这个特定的 AllReduce 完全隐藏在计算后面意味着它直接增加了前向传播的关键路径。

让我们更好地看一下在缩放 TP 度数时的权衡：虽然增加 TP 会导致每个 GPU 的吞吐量降低（左），但它可以处理更大的批量（右），说明了分布式训练中计算效率和内存可用性之间的权衡。

在实践中，正如我们在左图中看到的，当我们扩展到超过 8 个 GPU 时，张量并行的通信开销变得特别明显。虽然在单个节点内的张量并行可以利用快速的 NVLink 互连，但跨节点的通信需要更慢的网络连接。我们观察到从 TP = 8 到 TP = 16 有显著的性能下降，从 TP = 16 到 TP = 32 下降更为陡峭。在更高的并行度下，通信开销变得如此之高，以至于它很快主导了计算时间。

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

增加张量并行度减少了每个 GPU 上模型参数、梯度和优化器状态所需的内存，以至于我们可以开始在一个 8 个 GPU 的单个节点上拟合一个大型模型。

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

“

注意，关于张量并行训练中的层归一化有一个有趣的点——由于每个张量并行等级在all-gather后看到相同的激活值，层归一化权重在反向传播后实际上不需要归约来同步它们的梯度。它们自然地在等级之间保持同步。然而，对于随机失活操作，我们必须确保在张量并行等级之间同步随机种子以保持确定性行为。

接下来让我们探索张量并行的一个小而自然的扩展，称为Sequence Parallelism，它正是为此而设计的。

5. Sequence Parallelism

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

“

注意 术语 Sequence Parallelism 有点超载： 本节中的 Sequence Parallelism 与 Tensor Parallelism 紧密耦合，适用于 dropout 和 layer norm作。然而，当我们转向更长的序列时，注意力计算将成为一个瓶颈，这需要诸如 Ring-Attention 之类的技术，有时也称为序列并行，但我们将它们称为上下文并行来区分这两种方法。因此，每次您看到 sequence parallelism 时，请记住它与 tensor parallelism 一起使用（与上下文并行相反，后者可以独立使用）。

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

所以尽管这些操作在计算上很简单，但它们仍然需要大量的激活内存，因为它们需要完整的隐藏维度。序列并行允许我们通过沿着序列维度分割来分担GPU 之间对此内存负担。

在实践中，我们将从左侧的图示转换到右侧的图示：

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

在正向传播中：

• “f”是一个无操作（no operation, no - op），因为激活值已经在等级之间复制。
• “f*”是一个全收集(all-reduce)操作，用于同步激活值并确保正确性。

在反向传播中：

• “f*”是一个无操作，因为梯度已经在等级之间复制。
• “f”是一个(all-reduce)操作，用于同步梯度。

这些操作“f”和“f*”被称为共轭对，因为它们在正向传播中一个是无操作时，在反向传播中另一个是all-gather操作，反之亦然。

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

那么实际情况是怎样的呢？正如一个著名的大语言模型可能会说的，让我们一步一步来看：

在 TP 和 TP/SP 中，不同部分的激活值（又名隐藏状态）在隐藏维度和序列维度上的形状变化有点难以跟踪——相信我们，我们也觉得很难映射，所以我们制作了这个小表格来总结在正向传播期间的变化：

区域	仅 TP	TP 与 SP
进入 TP（列线性）	：分片（权重输出分片） ：完整	：分片（权重输出分片） ：全收集到完整
TP 区域	：分片 ：完整	：分片 ：完整
退出 TP（行线性）	：完整（权重输出完整 + 归约以保证正确性） ：完整	：完整（权重输出完整 + reduce-scatter以保证正确性） ：reduce-scatter到分片
SP 区域	：完整 ：完整	：完整 ：分片

对于embedding层：

区域	普通 TP	TP 与 SP
嵌入层（行线性，在词汇表上分片）	：完整（权重输出完整 + 归约以保证正确性） ：完整	：完整（权重输出完整 + reduce-scatter以保证正确性） ：reduce-scatter到分片

通过使用序列并行，我们可以实现更大的激活内存节省，从而允许我们进一步增加批次大小和序列长度，这比仅使用张量并行是可能的。让我们看看这对我们之前的 70B 模型示例意味着什么：

如我们所见，我们再次显著降低了每个 GPU 的最大内存使用量，允许我们在 TP/SP = 16 时适应 16k 个标记的序列长度，这比普通的 TP 情况有所改进！（TP = 16 仍然有点大，正如我们在上一节中看到的，但我们将在下一节中看到如何改进这一点）。

你可能会问自己，使用 TP + SP 是否比普通的 TP 会产生更多的通信？嗯，是也不是。在普通的 TP 正向传播中，每个变压器块有两个归约操作，而在 SP 中，每个变压器块有两个全收集和两个reduce-scatter操作。所以从操作数量上看，SP 的通信操作数量是 TP 的两倍。但由于归约操作可以分解为全收集 + reduce-scatter（见附录中的“快速聚焦环归约”部分），所以它们在通信量上实际上是等效的。在反向传播中也是同样的道理，因为我们只是使用每个操作的共轭（无操作与归约，全收集与reduce-scatter）。

如果你一直在密切关注，你会注意到我们在每一层中谈论 4 个通信操作（2 个用于注意力，2 个用于 MLP）。这是使用张量 + 序列并行时 MLP 的分析情况：

就像普通的 TP 一样，TP + SP 不容易与计算重叠，这使得吞吐量在很大程度上依赖于通信带宽。同样，像普通的 TP 一样，TP + SP 通常只在节点内完成（保持 TP 度低于每个节点的 GPU 数量，例如 TP ≤ 8）。

我们可以通过测量在 3B 模型上随着 TP 与 SP 一起缩放时的吞吐量和内存使用情况来基准测试这种通信开销如何变得越来越成问题，序列长度为 4096：

在这里，我们再次看到计算效率（左）和内存容量（右）之间的权衡。虽然更高的并行度通过减少激活内存能够处理明显更大的批次大小，但它们也会降低每个 GPU 的吞吐量，特别是在超过每个节点的 GPU 数量对应的阈值时。

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

• 对于这两种方法，我们注意到从 TP = 8 到 TP = 16 时性能下降最大，因为这是我们从仅在单个节点（NVLink）内通信转换到跨节点（EFA）通信的时候。
• 使用 TP 与 SP 时激活值的内存节省帮助我们比单独使用 TP 适应更大的批次。
• 使用 TP 与 SP 时激活值的内存节省帮助我们比单独使用 TP 适应更大的批次。

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

“

注意，由于 SP 区域中的 LayerNorms 在序列的不同部分上操作，它们的梯度在 TP 等级之间会有所不同。为了确保权重保持同步，我们需要在反向传播期间对它们的梯度进行归约，类似于数据并行（DP）确保权重同步的方式。不过，这是一个较小的通信开销，因为 LayerNorm 具有相对较少的参数。

然而，TP 和 SP 有两个限制：1）如果我们扩展序列长度，TP 区域中的激活内存仍然会爆炸；2）如果模型太大以至于 TP = 8 无法适应，由于节点间的连接性，我们将看到严重的减速。

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

如何学习AI大模型？

大模型时代，火爆出圈的LLM大模型让程序员们开始重新评估自己的本领。 “AI会取代那些行业？”“谁的饭碗又将不保了？”等问题热议不断。

不如成为「掌握AI工具的技术人」，毕竟AI时代，谁先尝试，谁就能占得先机！

想正式转到一些新兴的 AI 行业，不仅需要系统的学习AI大模型。同时也要跟已有的技能结合，辅助编程提效，或上手实操应用，增加自己的职场竞争力。

但是LLM相关的内容很多，现在网上的老课程老教材关于LLM又太少。所以现在小白入门就只能靠自学，学习成本和门槛很高

那么我作为一名热心肠的互联网老兵，我意识到有很多经验和知识值得分享给大家，希望可以帮助到更多学习大模型的人！至于能学习到多少就看你的学习毅力和能力了 。我已将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

这份完整版的大模型 AI 学习资料已经上传优快云，朋友们如果需要可以微信扫描下方优快云官方认证二维码免费领取【保证100%免费】

👉 福利来袭优快云大礼包：《2025最全AI大模型学习资源包》免费分享，安全可点 👈

全套AGI大模型学习大纲+路线

AI大模型时代的学习之旅：从基础到前沿，掌握人工智能的核心技能！

640套AI大模型报告合集

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。

👉学会后的收获：👈
• 基于大模型全栈工程实现（前端、后端、产品经理、设计、数据分析等），通过这门课可获得不同能力；

• 能够利用大模型解决相关实际项目需求： 大数据时代，越来越多的企业和机构需要处理海量数据，利用大模型技术可以更好地处理这些数据，提高数据分析和决策的准确性。因此，掌握大模型应用开发技能，可以让程序员更好地应对实际项目需求；

• 基于大模型和企业数据AI应用开发，实现大模型理论、掌握GPU算力、硬件、LangChain开发框架和项目实战技能， 学会Fine-tuning垂直训练大模型（数据准备、数据蒸馏、大模型部署）一站式掌握；

• 能够完成时下热门大模型垂直领域模型训练能力，提高程序员的编码能力： 大模型应用开发需要掌握机器学习算法、深度学习框架等技术，这些技术的掌握可以提高程序员的编码能力和分析能力，让程序员更加熟练地编写高质量的代码。

👉 福利来袭优快云大礼包：《2025最全AI大模型学习资源包》免费分享，安全可点 👈

这份完整版的大模型 AI 学习资料已经上传优快云，朋友们如果需要可以微信扫描下方优快云官方认证二维码免费领取【保证100%免费】

作为普通人，入局大模型时代需要持续学习和实践，不断提高自己的技能和认知水平，同时也需要有责任感和伦理意识，为人工智能的健康发展贡献力量。