终极训练指南：
在大规模 GPU 集群上训练大语言模型

我们在多达 512 个 GPU 上进行了超过 4000 次扩展实验，并测量了吞吐量（标记大小）和 GPU 利用率（标记颜色）。需要注意的是，在本可视化中，两者都根据模型大小进行归一化。

 原文：The Ultra-Scale Playbook: Training LLMs on GPU Clusters
翻译工具：ChatGPT
校对者：@Ki_Seki_here

数千块 GPU 协同工作，完美协调——这正是当今最强大 AI 模型所需的训练规模。直到最近，这种规模还只属于顶尖研究实验室的专有领域。开源让这块领域发生了翻天覆地的变化，但并没有彻底改变现状。是的，你可以下载最新的 Llama 或 DeepSeek 模型。是的，你能阅读它们的 技术和 实验报告。但最具挑战性的部分——用于协调 GPU 训练这些超大模型的训练代码、必要的知识和技术——仍然极其复杂，零散地分布于多篇论文或私有代码库中。

本开源书将尝试改变这一点。我们将从基础开始，循序渐进地带你了解如何将大型语言模型的训练从单卡 GPU，扩展到数十、数百甚至数千个 GPU，并用实践代码示例和可复现的基准来演示理论。

随着用于训练这些模型的集群规模不断增大，出现了如数据并行、张量并行、流水线并行、上下文并行，以及 ZeRO 或核融合等技术，以确保 GPU 始终保持高利用率，从而显著缩短训练时间并最大化昂贵硬件的使用效率。更进一步地，随着对大型模型的微调在真实应用中越来越受重视，人们发现微调在特定领域数据上能带来更好的效果，这通常也需要相同的分布式训练技术。在本书中，我们将从最简单的方法到最精巧的方法，依次介绍所有这些技术，同时围绕单一的故事线说明每种方法的来龙去脉。

本书假设你对当前 LLM 架构有一些基本了解，并大致熟悉深度学习模型的训练原理，但在分布式训练方面可以是新手。如果有需要，可在 DeepLearning.ai 或 PyTorch 教程等地方学习模型训练的基础知识。可以将本书视为我们第一篇关于预训练数据处理博客——即“FineWeb 博客”——的后续篇章。读完这两篇博客后，你几乎掌握了当今高性能 LLM 构建所需的核心知识，只差一些关于数据混合和模型架构选择等“最后的调味料”就可以完整配方了（敬请期待第三部分……）。

本书基于以下 三个核心支柱：

简明的原理与概念简介：在开始编码和实验前，我们希望先高层次地理解每种方法如何工作，及其优劣势何在。你会了解到语言模型在训练时哪些部分会占用最多显存，以及何时占用显存。你会了解到如何通过对模型进行并行拆分来缓解显存的限制，并如何通过扩展 GPU 集群来提高吞吐量。于是你就能理解下面这个小工具是如何计算 Transformer 模型显存占用的：

显存使用分析

（如果你对这个小工具里发生的事情还摸不着头脑，不用担心——这正是我们要讲的内容。）

虽然上面的小工具给出的是理论分析，但我们还做了一个 在线工具，可用来预测训练过程中的实际显存用量：

预测显存占用工具截图

清晰的代码实现：理论是一方面，但在真正实现过程中会遇到各种边界情况以及重要的细节。因此我们会在可能的地方关联到实现参考。我们会根据情况引用两份代码：

真实的训练效率基准：最后，如何在实际环境中真正扩展 LLM 训练还要取决于你的硬件环境，比如 GPU/TPU 芯片种类、网络带宽等。因此我们不可能给出一个适用于所有硬件的统一方案，但我们会介绍评测方法，而我们也已经在自己的集群上进行了此类评测！我们在最多 512 块 GPU 上运行了超过 4100 个分布式实验（加上测试一共做了 16000 多次），以探索众多可能的分布式训练布局和模型规模。

如你所见，还有很多内容需要讨论。在深入了解分布式训练的细节之前，让我们先从更高层面回顾一下在本书中会涉及到的挑战。

高层概览

本书中所涉及到的所有技术都是为了应对如下三个关键挑战，而这三个挑战会在全书中反复出现：

  1. 显存使用：这是训练的硬性限制——如果某个训练步骤超过了显存容量，训练便无法进行。
  2. 计算效率：我们希望硬件大部分时间都在计算，而不是花在数据传输或等待其他 GPU 的工作上。
  3. 通信开销：我们希望最小化通信开销，因为通信会让 GPU 处于空闲。为此，需要尽量利用节点内（更快）和节点间（更慢）的带宽，并将通信尽量与计算重叠。

在很多地方，我们会看到可以在（计算、通信、显存）三者之间做权衡（例如重计算或者 Tensor 并行）。找到合适的平衡是成功扩展训练的关键。

因为本书内容非常详实，我们做了一个 快速参考表，以便在阅读本书时辅助理解并帮你提炼要点。可以在阅读过程中随时翻看它！

快速参考表

单卡训练

如果你想添加“播客”式的阅读体验，可以在阅读本节时收听 NotebookLM 主播关于本书第一部分的讨论。

在扩展到多 GPU 之前，让我们先快速回顾在单卡 GPU 上训练模型的基本流程。单卡训练通常包括三步：

  1. 前向传播（forward pass）：输入数据通过模型，得到输出
  2. 反向传播（backward pass）：计算梯度
  3. 优化步骤（optimizer step）：利用梯度更新参数

通常可以表示为下图：

鼠标悬停于网络元素可查看详情

在此图中，顶部每个彩色框代表模型的某一层（最后一行同理）；红色框是对应层在反向传播时的梯度。

批大小（bs）是模型训练中非常重要的超参数，会同时影响模型收敛和吞吐量。

小批量（batch size 小）有利于训练初期快速更新权重，且梯度带有随机性。但在训练后期，小 batch size 可能使梯度噪音过大，模型无法收敛到最优的点。而相反，如果 batch size 太大，虽然可以让梯度估计更精确，但会使每次更新投入的计算代价增大，从而降低了整体效率。关于这个话题可参见 OpenAI 的大 batch 训练论文，或 MiniMax-01 技术报告 第 4.2 节中的讨论。

batch size 同时也影响训练相同数据集所需的训练时间：较小的 batch size 会导致需要更多次的优化步骤（optimizer steps），而优化步骤通常是最耗算力的。因此小 batch size 往往会拉长整体训练时间。话虽如此，在实际训练中，batch size 的最终性能对其附近的一些取值并没有那么敏感，通常存在一个相对平缓的区间。

在 LLM 预训练社区中，batch size 通常会用 token 而不是样本量来表示（bst = Batch Size Tokens），这样就不必关心具体输入序列长度，用 tokens 来衡量可以更直观地比较训练代价。

在最简单的单机训练场景下，bs（以样本计）和 bst（以 token 计）可以通过序列长度（seq）相互转化：

bst=bs *seq

近来对于 LLM 预训练而言，较常见的 token batch 范围在几百万到几千万之间不等。随着预训练规模的不断增长：Llama 1 约在 4M tokens batch，对 1.4T token 进行训练；DeepSeek 则用到了 60M tokens batch，训练了 14T token。

而当我们将模型训练扩展到如此大的批次时，第一个挑战便出现了：内存不足问题。当我们的GPU没有足够的内存来容纳目标批次大小的完整数据时，我们该怎么办？

让我们先快速了解一下最初导致内存不足问题的原因。这将帮助我们对训练模型所需的内存有一些有用的直观认识。

Transformer 中的显存使用

在训练神经网络模型时，一般需要存储：

📝 注意

你可能认为对于一个模型，我们可以精确地计算出需要多少显存，但实际上显存里还会有以下内容：

  • CUDA Kernel 通常会占用 1-2 GB 的 GPU 显存。可以用 import torch; torch.ones((1, 1)).to("cuda") 并配合 nvidia-smi 观测 GPU 显存来验证。
  • 一些缓存、临时 buffer，以及由于碎片化而无法充分使用的显存空间
在分析时，我们往往忽略最后两点，因为它们一般是较小且固定的开销。

这些需要存储的项目以张量形式存在，不同张量有不同的 形状（shape） 和 精度（precision）。形状由例如 batch size、序列长度、模型隐层维度、注意力头数、词表大小，以及是否进行模型切分等超参数决定；精度则对应 FP32、BF16 或 FP8 等格式，会影响每个元素所占的字节数（4、2 或 1 字节）。我们稍后会在 混合精度训练部分更详细讨论精度，这里只需知道不同精度会影响存储需求即可。

那如何快速确定模型的实际显存占用？一个简单的方法是直接在真实环境中测量。

分析显存使用

借助 PyTorch 的 profiler 工具，我们可以查看训练过程中不同阶段的显存分配。可以发现显存使用并非静态，而是在训练过程（尤其是单个 step 内）不断变化：

显然，第一步看起来和后续步骤有些不同；但先让我们看看一次完整训练 step 的内存使用模式：前向传播时，会随着激活值的产生，显存占用快速上涨；随后在反向传播时，梯度逐渐累加，而用于计算梯度的激活值也在此过程中逐步释放。最后，我们进行优化步骤，此时需要所有梯度，然后更新优化器状态，接着进入下一次前向传播。

为什么第一步看上去与众不同：激活值先快速增加，然后会有一段时间保持在高位。这是因为第一步里 Torch 的缓存分配器（allocator）要初始化分配显存块，以便随后的训练步骤中不必再频繁搜索可用显存（见 Zach 的博客）。在第一步结束后，我们还要为优化器状态分配显存，这通常会让后续步骤的起始显存占用有所提高。

现在我们已经对显存变化有了初步认识，下面我们就来看看在扩展训练规模的同时，如何在保证计算效率的前提下，让模型的各种存储需求（激活值、参数、梯度、优化器状态）都在显存范围内。

权重、梯度和优化器状态的显存

让我们先看一下前面提到的前三项：模型权重、梯度和优化器状态。我们可以相对容易地对它们的显存需求做出估算。

对一个简化的 Transformer LLM，其参数量可按 如下公式计算：

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

其中 h 是隐层维度，v 是词表大小，L 是层数。显然，当我们增大 h 时，h^2 这一项会主导参数量。

对于这些参数以及对应的梯度，我们在显存中需要的空间就是“参数个数乘以每个参数的字节数”。在传统的 32 位浮点（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 对较小值是有损的，为了稳定性，一些库就会把梯度也保存在 FP32。可参考 DeepSpeed issue 了解更多信息。

📝 注意

在文献和代码中，有时会把这份 FP32 权重副本称作 “master weights”。

有意思的是，混合精度本身并不会减少总体显存需求，因为虽然参数本身变少了，但又多了一份 FP32 副本，甚至如果把梯度也保存在 FP32，整体开销还会上涨。但它能带来很大好处：用 BF16 进行前向/反向计算可用到 GPU 的低精度优化运算单元，速度更快，同时前向传播中激活所需的显存也能减少，这在大 batch 或长序列时尤为重要。

来看几个常见模型规模下，计算或 BF16 混合精度时，这些存储需求的大致量级：

模型参数量 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 的 80GB）。

目前我们只考虑了模型可以单卡放得下的情况，接着让我们看看另一个显存主要来源：激活值。

激活值占用

激活值的占用比权重、梯度和优化器状态更为复杂，因为它会受输入有关。若你不太确定为什么反向传播需要存储激活值，可参考这篇参考文档。仔细分析 Transformer 反向传播的计算过程后，可估算训练时在混合精度下的总激活显存，结论可参见原版 NVIDIA 的重计算论文，给出的公式大致如下：

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

这里 L 是层数，seq 是序列长度，bs 是 batch size（以样本计），h 是模型隐层维度，n_{heads} 是注意力头数。

激活值需要在反向传播时用来计算梯度。与权重、梯度和优化器状态相比，激活值在计算图中随时在变化，因而需要在整个 forward + backward 周期进行分配和释放。

一个非常重要的观察是，激活值的占用会随序列长度 seq 与 batch size bs 线性增长。这意味着当我们增大全局 batch size 或序列长度时，激活值占用会急剧飙升。让我们看看在不同序列长度下，对于 Llama 系列（bs=1）的激活值显存走势：

图中可以看到：对于短序列（或者 equivalently 小 batch size）时，激活值几乎可以忽略不计。但一旦序列长度达到 2-4k，它就会变成一个相当庞大的开销，而此时参数/梯度/优化器状态的开销也不再是主要矛盾。

对于大输入 tokens（即大 batch size/长序列），激活值会成为主要的显存负担。

有没有办法抑制这种“激活值膨胀”？你问得好！

下面来介绍我们的第一个技巧——激活重计算（Activation Recomputation）。它可以帮助我们将激活值的占用限制在合理范围。它是当下大型模型训练中不可或缺的重要技术。

激活重计算（Activation recomputation）

激活重计算（也叫 gradient checkpointing 或 rematerialization）的核心思想是：在前向传播时丢弃部分激活值，从而省显存；需要它们做反向传播时，再运行一次子前向过程把它们计算回来，换取多一些计算量来节省显存。若不开启重计算，我们会在每个可学习操作（比如 feed-forward，layernorm 等）之间都保存激活值。启用重计算后，我们只保存少量关键位置的激活值，而丢弃其它激活值；当反向传播需要时，再用已保存的激活值进行部分前向运算重算出所需内容。可通过下图直观理解：

鼠标悬停于网络元素可查看详情

在实际实施中常见有几种策略来决定哪些激活值需要保存：

让我们用实践测量来看看各种重计算策略能带来多大程度的显存优化，以及 selective 重计算如何在内存节省和计算开销之间取得平衡：

图中还能看到，越小的模型（h 较小）在长序列（seq 较大）下，激活值占比就越明显，重计算带来的收益也更大。

📝 注意

在测量训练对 GPU/TPU/加速器的利用率时，一般要把重计算计算量纳入总 FLOPs，再和理论峰值 FLOPs 进行对比，以得到实际硬件 FLOPS 利用率 (Hardware FLOPS Utilization, HFU)。因为重计算会增加实际运算量。

但最终我们往往更关心的是，从头到尾完成同样数据量训练所需要的总时间，因此评估不同加速器时，若某加速器拥有足够的显存以完全跳过重计算，它所做的实际运算就会变少，这会导致它的 HFU（硬件利用率）看似降低，但训练速度可能反而更快。为此，有人提出仅统计模型本身前+后向步骤所需的 FLOPs（不含重算）来计算 Model FLOPS Utilization (MFU)。这在比较硬件时有时更有意义。

在当今的训练框架中，FlashAttention（后面会介绍）已成为注意力优化的标配，它在反向传播中就会自动重算注意力得分和中间矩阵，而不存储它们，这本质上就是 selective 重计算的一种。所以只要用上 FlashAttention，你其实已经在用 selective 重计算了。

小结：激活重计算增加了些许 FLOPs（约多 2~30%），却显著减少了显存访问带来的延迟和内存需求。从而带来显存大幅节省。

因此它对当前内存不大但运算能力强的 GPU 来说特别有用，即使多做了一些计算，整体上也会因为减少了更多的内存访问而变得更快。

不过即便如此，激活值的开销依然随 batch size 线性增长，那如果想要用很大的 global batch size，该怎么办？这就要看看我们的另一个法宝——梯度累加（gradient accumulation）！

梯度累加（Gradient accumulation）

梯度累加是个非常直观的方法，用于避免因过大 batch size 带来的显存爆炸。它做法是：把原本的一次大 batch，拆成多个更小的 micro-batch，每次只做一次前后向并计算出梯度，然后把这些梯度累加起来，最后再做一次优化器更新（optimizer.step）。这样，就能在保持相同 global batch size 的前提下减少每次前向时的激活显存占用。

我们把每次前向使用的 batch size 称作 micro batch size (mbs)，把整个全局 batch size（在每次 optimizer step 之间）称为 global batch size (gbs)。如果我们在一次优化步骤中处理了 8 个 micro-batch，那么 global batch size 就是 mbs × 8。

因此，之前我们文中统称的 batch size 其实就是 global batch size。用符号表示就是：

bs = gbs = mbs \times grad\_acc

不过，梯度累加也有个明显缺点：每个优化步骤需要重复多次前后向，因此增加了计算量，从而放慢训练速度。没有免费的午餐啊！

梯度累积使我们可以通过仅计算部分微批次来降低激活内存的使用，因为激活内存在批次大小增大时会线性增长。

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

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

在此之前，让我们快速了解如何通过分布式训练工具箱中最有用的工具之一——性能分析器，来可视化计算和通信情况。这个工具对于理解和验证GPU之间以及计算之间的通信方式以及瓶颈所在将非常有用。

分析 GPU 计算与通信

PyTorch 自带的 profiler 能够跟踪并可视化 CPU 和 GPU 在训练过程中的行为。它可以很方便地查看：

下面是一个简单示例：

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()

这会生成一份 trace，可在 TensorBoard 或 Chrome tracing viewer 中查看，展示：

profile_trace_annotated.png

示例 trace，显示 CPU 线程异步向 GPU 提交 kernel，多条 CUDA stream 并行进行计算和通信

通过 trace 我们能发现很多瓶颈，比如：

这些信息对优化分布式训练性能至关重要。比如，你可以清晰看到梯度同步是否成功和反向计算重叠到一起等。

好了，现在我们可以正式进入多 GPU 训练世界，看看第一个扩展技巧 —— 数据并行！

数据并行（Data Parallelism）

如果想让阅读体验更像播客，可以播放此音频，收听 NotebookLM 主播对以下章节的讨论。

数据并行（DP）的核心思想是：在多块 GPU 上复制相同的模型副本，只是每张卡处理不同的微批数据。如此一来，称为“数据并行”。想必很多人已在简单训练示例中用过 DP，但我们这里会深入探讨更多实现细节。

image.png

在每个GPU使用不同的微批次意味着每个GPU上会有不同的梯度，为了保持各GPU上模型实例的同步，模型实例之间的梯度将在反向传播过程中、优化步骤之前，通过一种称为“all-reduce”的操作进行平均。

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

image.png

一个简单的DP实现会等待反向传播完成以获得所有梯度，然后触发一次all-reduce操作，对所有DP级别的梯度进行同步。但这种计算与通信依次进行的方式绝对不可取！因为我们不希望在通信过程中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实现（点击展开）

这是我们首次展示“计算与通信的重叠”，在本文中我们会多次讨论这一关键技术，它是实现最大扩展效率的必要手段。但我们仍可以进一步提高效率！

优化二：对梯度进行 bucket 化

GPU 的效率在处理大张量时通常更高；若我们对每个参数单独执行 all-reduce，会有很多小通信操作，不够高效。更好的做法是把多个梯度合并到一个大的 bucket，一次 all-reduce，即可显著减少通信开销，提升效率。

示意图：

dp_overlap3.svg

类似打包快递，把许多小包裹装进大箱子里，一次性发走。具体实现上，你会预先分配一个大的 buffer，把多个梯度放进来，然后做一次性 all-reduce。例如：

👉 在 Picotron 中的 bucket DP 实现（点击展开）

优化三：梯度累加的结合

最后，前面我们提到的梯度累加与数据并行结合时，需要留意在累加阶段是否还要进行梯度同步。朴素做法可能每次反向传播都发起 all-reduce，但其实只有在完成了所有 micro-batch 的反向传播后再进行一次 all-reduce 就足够了，能减少通信开销。

在 PyTorch 中通常可用 model.no_sync() decorator 来禁止某些 backward 阶段的梯度同步。

📝 注意

进行通信操作时，张量必须在内存中是连续的才能避免额外的拷贝。因此常常会预先分配与激活值或参数大小相同的连续内存块，专门用于通信。不过这样会带来额外的峰值内存占用。

接下来我们再回到 global batch size 的公式。

重新审视 global batch size

带有数据并行和梯度累加之后，全局 batch size 变为：

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 数量，也可以扩大 global batch size，或试着减小 mbs 看是否速度更快（增大数据并行度会带来通信开销，但减少 mbs 意味着可能浪费一部分硬件算力，需要实验来平衡）

举个例子：假设我们想要训练一个新模型，global batch size 设为 4M tokens，序列长度设为 4k，所以 batch size（以样本数计）约为 1024。我们发现单卡仅能放下 mbs=2（再大就 OOM ），若我们有 128 块 GPU，就可以在每个 GPU 上做 4 次梯度累加来得到全局的 1024 样本，这样就能到达目标 4M tokens。接着如果我们拥有 512 块 GPU，就可以把梯度累加减少到 1 并获得更快的训练速度。

📝 注意

当并行度到达 512 块 GPU 乃至以上时，网络瓶颈（ring latency 等）会使得 DP 的通信无法被完全隐藏，GPU 利用率会下降，吞吐量可能开始恶化。此时就要考虑其他并行维度。

虽然数据并行能够很好地将all-reduce梯度同步与反向计算重叠以节省时间，但当规模增大时，这一优势开始显现不足。原因何在？因为随着我们添加越来越多的GPU（数百甚至数千个），它们之间协调的开销显著增加，网络需求也变得过于庞大，从而抵消了并行带来的好处。结果是，每增加一块GPU，我们的设置效率就会逐渐降低。

让我们通过一些基准测试来看看这一现象在实际中的表现：

可以看到，当 GPU 数量超过某个范围后，吞吐量开始明显下降，而每卡所需内存并不会随着 GPU 数的增加而减少。

数据并行是扩展到更多 GPU 的第一种（也是相对简单的）并行策略。它和梯度累加的原理很像，但通过并行加速微批处理，来提高训练吞吐量！

然而，我们之前也提到，DP 默认要求单卡至少能放得下完整的一次前向（mbs≥1），对于更大的模型（如 70B+）即使激活重计算打开后，也还是放不下。比如：

与此同时，我们也发现当 DP 扩展到数百上千块 GPU 时，通信已经成为巨大瓶颈，还有没有其他策略？其实可以让某些张量放到 CPU ，或者把这些权重/梯度/优化器状态在 GPU 之间做分块切分。接下来我们将探讨两类思路：分片并行（tensor/context/pipeline 并行）和 共享（DeepSpeed Zero/FSDP）。它们彼此独立，也可以组合。

因为 ZeRO 方法和 DP 密切相关，所以先从它开始。

ZeRO（Zero Redundancy Optimizer）

本节介绍 DeepSpeed ZeRO，专门用来减少 LLM 训练中的冗余显存占用。

数据并行可以提高吞吐量，但也带来了额外的内存浪费：所有 DP 副本都要存储同样的优化器状态、梯度和参数，形成重复。ZeRO 的思路正是通过沿数据并行维度对上述对象（优化器状态、梯度、参数）进行切分，来消除冗余。这会要求一定的额外通信操作，以便在需要时重建完整的参数，但相比内存收益常常很值得。

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

你可能会注意到，我们没有把激活值包含在可分片的对象中。因为模型的每个DP副本接收到的微批次不同，每个DP节点上的激活值也各不相同，所以它们不会被复制，也就无法进行分片！

让我们从内存需求角度看看不同阶段的 ZeRO 到底能省多少。

再看内存需求

你可能还记得在上一节中关于标准训练过程中优化器状态、梯度和参数的内存使用情况。我们用\Psi来表示模型参数的数量（之前用的是N，但这里我们采用ZeRO论文中的原始符号）。在混合精度训练（更多细节见后续章节）中，使用Adam优化器时，每项所需存储的内存为：

如果我们只关心不带 FP32 梯度累加的混合精度训练，其总占用就是 2\Psi + 2\Psi + 12\Psi （对应权重、梯度、fp32 权重 + 动量 + 方差），若再算上 FP32 梯度累加，会再加 4\Psi。

如果我们不以fp32累积梯度，总内存消耗为 2\Psi + 2\Psi + 12\Psi；如果以fp32累积梯度，则为 2\Psi + 6\Psi + 12\Psi。为了简单起见，我们现在先关注不进行fp32梯度累积的情况，但你只需将额外的字节加到受ZeRO-2和ZeRO-3影响的梯度项上即可。

ZeRO的理念是将这些对象分片存储到各个数据并行（DP）节点上，每个节点只存储部分数据，在需要时再重构这些数据，从而使内存使用量降低为数据并行度 N_d 的1/N_d：

zero_memory.svg

这里 \Psi 表示参数数量，k 表示优化器状态的内存倍增因子（正如我们刚刚看到的，对于Adam而言 k=12），而 N_d 表示DP的规模。

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

ZeRO-1：分区优化器状态

在传统的分布式训练（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

从实际通信的角度来看，与传统DP相比，ZeRO-1将我们的“all-reduce”梯度通信改为“reduce-scatter”操作，并在优化器步骤后添加了一个all-gather操作，涉及所有参数。它看起来是这样的：

dp_zero1_overlap.svg

如果你一直在跟随，我们会记得在传统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-2的过程与ZeRO-1相同，唯一的区别是它在动态地进行通信和释放。总的来说，ZeRO-2在通信方面与传统DP训练相当。

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

dp_zero2_overlap.svg

现在我们已经对梯度进行了分片，我们完成了吗？还是我们还能继续优化？嗯，差不多。接下来是ZeRO-3！

ZeRO-3：添加参数分区

对于阶段3，我们扩展了之前在DP副本中分片优化器状态和梯度的方法，直到对模型的参数进行分片。

📝 注意

这一阶段在PyTorch的原生实现中也叫做FSDP（Fully Shared Data Parallelism，完全共享数据并行）。在这篇博客中我们将称其为ZeRO-3，但你可以在看到FSDP时理解为ZeRO-3。

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

dp_zero3_fwd.svg

因此，在执行前向传播并依次通过各层时，我们按需获取必要的参数，并在不再需要它们时立即从内存中清除。反向传播的工作方式与此类似，只是流程相反，我们会生成梯度分片：

dp_zero3_bwd.svg

另一个问题是，我们需要在前向和反向步骤中持续进行这些all-gather操作，相当于在一个训练步骤中，比ZeRO-2多进行了 2\cdot \text{num\_layers} -1 次all-gather操作，每次都会带来一些小的 基础延迟 开销，正如我们在下图中看到的：

从通信角度看：相较于 ZeRO-2，ZeRO-3 每一层在 forward、backward 里多做了 2 次 all-gather（分别是正向和反向），以及一次 reduce-scatter 用于梯度，故总共 3\Psi。比 ZeRO-2 的 2\Psi 多了一份。

在前向传播过程中，我们在需要时对参数进行all-gather操作，因此有一个 \Psi 的通信开销。由于我们在前向传播后立即丢弃这些参数，因此在反向传播过程中还需要进行一次all-gather，从而再增加一个 \Psi 的通信开销。最后，我们需要与ZeRO-2中相同的 reduce-scatter 操作来处理梯度，这也需要 \Psi 的通信开销，因此我们最终的总通信成本为 3\Psi，而ZeRO-2的通信成本是 2\Psi。

这听起来似乎有很大的通信开销，但实际上还算可以，因为我们可以通过在前向传播中进行所谓的预取，将下一层参数的通信与当前层的前向传播重叠。在预取过程中，我们会在进行当前层（Layer n）的前向传播时“all-gather” *Layer n+1* 的权重，而在进行当前层（Layer n）的反向传播时，我们则会“all-gather” *Layer n-1* 的权重。当然，只有在DP规模没有过大时，这种重叠才有效。（作为经验法则，DP的规模不应超过512）

从内存的角度来看，我们可以看到我们的公式现在已经达到了最终的形式：\frac{2\Psi +2\Psi+k\Psi}{N_d}，这意味着如果我们能够增加DP的节点数，至少对于模型相关的参数，内存使用量可以无限降低。请注意，这对于中间激活并不起作用，对于这一部分，我们可以使用激活检查点和梯度累积，正如我们在前面章节中看到的。

让我们总结一下迄今为止在DP和ZeRO中的历程：我们已经看到，使用DP可以显著提高训练的吞吐量，仅仅通过增加更多的模型副本来扩展训练。通过ZeRO，我们可以训练通常无法放入单个GPU的模型，通过在DP中分片参数、梯度和优化器状态，同时只产生少量的通信开销。

然而，这里也有一个限制，DP仅在模型的某一层可以放入单个GPU时有效，而ZeRO只能分区参数、梯度和优化器状态，但不能分区激活内存！我们回想一下在激活内存讨论中提到的，激活内存会随着序列长度和批大小的增加而增加。当然，我们可以限制这些参数，但实际上我们并不希望硬件成为训练时仅使用短序列长度的限制因素。

为了克服这些问题，现在是探索一种新的、正交的并行轴 - 张量并行（Tensor Parallelism，TP）。与ZeRO-3依赖于大量参数通信不同，TP提出将参数、梯度、优化器状态和激活分布到不同的设备上，而不需要在GPU之间进行任何模型参数的通信。

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

张量并行（Tensor Parallelism）

若想让阅读体验更像播客，可播放此音频，收听 NotebookLM 主播对以下章节的讨论。

因此，我们已经使用ZeRO对模型的参数、梯度和优化器状态进行了分片，但一旦激活内存超出了我们的内存预算，我们就遇到了限制。欢迎张量并行（Tensor Parallelism，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的基本构建模块，让我们看看如何在transformer层内有效地将它们结合起来！

Transformer块中的张量并行

为了制定遵循的策略，让我们从一个玩具示例转向一个真实的模型构建模块。Transformer模型由两个主要的构建模块组成：前馈层（MLP）和多头注意力（MHA）。我们可以对这两者都应用tensor并行。

前馈部分可以通过“列线性”后跟“行线性”来并行化，其中在前向传播中涉及广播以复制输入，并在后续进行全局归约。需要注意的是，在实际训练中，我们可以确保输入已经在TP等级之间同步，因此不需要广播。这种设置比起从“行线性”开始然后是“列线性”更有效，因为我们可以跳过两个分裂操作之间的中间全局归约。

image.png

现在，我们已经为transformer的前馈部分找到了一个高效的模式，让我们来看一下多头注意力块（MHA）。

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

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

image.png

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

Forward pass in Tensor Parallelism

观察tensor并行MLP（同样适用于Attention）的操作时间线，我们可以更好地理解涉及的权衡。在每个解码器层的前向传播中，我们遇到了一个同步点，即无法与计算重叠的AllReduce操作。这种暴露的通信开销是必要的，以便在最终应用LayerNorm之前组合部分结果跨tensor并行等级。

Tensor并行ism确实有助于减少矩阵乘法的激活内存，因为中间激活被分片到不同的GPU上。然而，对于诸如LayerNorm之类的操作，我们仍然需要收集完整的激活，这意味着我们并没有得到完全可能的内存节省。此外，TP引入了依赖于网络基础设施的显著通信需求。无法完全隐藏这种特定的AllReduce在计算背后的事实，意味着它直接增加了前向传播的关键路径。

让我们更好地查看随着TP度量的扩展而产生的权衡：

增加TP会导致每个GPU的吞吐量减少（左侧），但使得处理更大批次大小成为可能（右侧），展示了在分布式训练中计算效率和内存可用性之间的权衡。

实际应用中，如上图左侧所示，当我们扩展到超过8个GPU时，tensor并行ism的通信开销特别显著。在单个节点内，tensor并行ism可以利用快速的NVLink互连，但跨节点则需要较慢的网络连接。当从TP=8移动到TP=16时，我们观察到显著的下降，而从TP=16到TP=32的下降更为陡峭。在更高的并行度下，通信开销变得如此之高，以至于很快就会主导计算时间。

话虽如此，tensor并行ism通过在GPU之间分布模型参数、梯度、优化器状态和（在一定程度上）激活，提供了重要的内存使用优势。让我们看一下对于一个70B参数模型的影响：

增加tensor并行度减少了每个GPU所需的模型参数、梯度和优化器状态的内存，使得我们能够将一个大型模型适配到单个8 GPU的节点中。

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

📝 备注

关于tensor并行训练中的层归一化有一个有趣的点——由于每个TP等级在all-gather之后看到相同的激活，层归一化权重实际上不需要all-reduce来同步其梯度。它们会自然地在各个等级中保持同步。然而，对于dropout操作，我们必须确保在TP等级间同步随机种子，以维持确定性行为。

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

序列并行（Sequence Parallelism）

序列并行（SP）涉及在输入序列维度上分割模型中由张量并行（TP）未处理的部分（如Dropout和LayerNorm）的激活值和计算，而不是在隐藏维度上进行分割。

📝 注释

“序列并行”这一术语有些过载：本节中的序列并行与张量并行紧密相关，适用于dropout和层归一化操作。然而，当我们转到更长的序列时，注意力计算将成为瓶颈，这就需要诸如Ring-Attention等技术，这些技术有时也称为“序列并行（Sequence Parallelism）”，但为了区分这两种方法，我们将其称为“上下文并行（Context Parallelism）”。所以每当你看到序列并行时，请记住，它是与张量并行一起使用的（与上下文并行不同，后者可以独立使用）。

之所以需要序列并行，是因为这些操作需要访问完整的隐藏维度才能正确计算。例如，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”）在张量并行（TP）和序列并行（SP）区域之间进行过渡。关键挑战是如何高效地管理这些过渡，同时保持低内存使用并确保正确性。

在前向传播中：

在反向传播中：

这些操作“f”和“f*”被称为共轭（conjugate）对，因为它们是互补的——当前向传播中一个是无操作时，反向传播中的另一个就是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: 分片

对于嵌入层：

区域 原始TP TP与SP
嵌入层（行线性按词汇表分片） h: 完整（weight_out是完整的 + all-reduce 为了正确性）
s: 完整		 h: 完整（weight_out是完整的 + reduce-scatter 为了正确性）
s: reduce-scatter 到分片

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

正如我们所见，通过引入序列并行，我们再次显著减少了每个GPU的最大内存使用，从而使得在TP/SP=16的情况下能够适应16k tokens的序列长度，比起传统的TP情况有了更大的提升！（虽然TP=16仍然有点大，正如我们在上一节所见，但我们将在下一节看到如何改进这一点）。

你可能会问，使用TP+SP是否比传统的TP有更多的通信开销？答案是，有也没有。在传统TP的前向传播中，每个Transformer块有两个all-reduce，而在SP中，每个Transformer块有两个all-gather和两个reduce-scatter。因此，SP的通信操作数量是TP的两倍。但由于all-reduce操作可以分解为all-gather + reduce-scatter（参见附录中的Ring AllReduce简介部分），因此它们在通信上的效果是等价的。反向传播也是如此，因为我们只是使用每个操作的共轭操作（no-op ↔ all-reduce 和 all-gather ↔ reduce-scatter）。

如果你细心观察，你会注意到我们在每一层中提到了4次通信操作（2次用于Attention，2次用于MLP）。以下是使用Tensor + Sequence Parallelism时MLP性能剖面的展示：

tp_sp_overlap.svg

就像传统的TP一样，TP+SP也不能轻易与计算操作重叠，这使得吞吐量在很大程度上依赖于通信带宽。这里，像传统TP一样，TP+SP通常只在单个节点内进行（保持TP度数不超过每个节点的GPU数量，例如TP≤8）。

我们可以基准测试在我们扩大张量并行度时，通信开销是如何逐渐成为一个问题的。让我们来测量在4096序列长度的3B模型下，随着TP与SP的扩大，吞吐量和内存利用率的变化：

在这里，我们再次看到计算效率（左侧）和内存容量（右侧）之间的权衡。虽然更高的并行度通过减少激活内存使得处理更大的批次成为可能，但它们也会减少每个GPU的吞吐量，特别是当并行度超过节点内GPU数量时。

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

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

📝 注释

由于SP区域中的LayerNorm操作在序列的不同部分进行，因此它们的梯度将在TP的不同rank之间有所不同。为了确保权重保持同步，我们需要在反向传播过程中对它们的梯度进行all-reduce操作，这类似于数据并行（DP）如何确保权重同步。然而，由于LayerNorm的参数相对较少，这个通信开销较小。

然而，TP和SP有两个限制：1）如果我们增加序列长度，激活内存在TP区域仍然会膨胀；2）如果模型太大，无法适应TP=8，那么我们将由于跨节点连接性问题，遇到巨大的性能下降。

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

上下文并行（Context Parallelism）

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

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

上下文并行的核心思想是将序列并行的方法（也就是沿序列长度进行拆分）的思路应用到已经采用张量并行的模块上。我们将对这些模块沿两个维度进行拆分，从而也减少序列长度带来的影响。经过前面所讨论的内容，你会发现这种方法非常直观，但……这里有一个技巧，所以请保持警惕！

对于上下文并行，就像序列并行一样，我们将沿序列维度拆分输入，但这次我们对整个模型进行拆分，而不仅仅是对之前Tensor+Sequence并行中涉及的部分模型。

拆分序列不会影响大多数模块，如MLP和LayerNorm，因为它们对每个token的处理是独立的。它也不像TP那样需要昂贵的通信，因为只拆分了输入而非权重矩阵。就像数据并行一样，在计算梯度后，会启动一次all-reduce操作以在上下文并行组内同步梯度。

不过，有一个重要例外需要特别注意，那就是注意力模块（呵呵……双关语来啦 :D）。在注意力模块中，每个token需要访问来自所有其他序列token的键/值对，或者在因果注意力的情况下，至少需要关注每个前面的token。

由于上下文并行是沿序列维度将输入分布到各个GPU上，注意力模块将需要各个GPU之间进行充分通信，以交换必要的键/值数据。

如果我们采用简单的方法，这听起来会非常昂贵。但有没有办法能更高效、更快速地完成这一操作呢？幸运的是，有一种核心技术可以高效地处理键/值对的通信，叫做环形注意力。

📝 注释

上下文并行与Flash Attention在概念上有一些相似之处（更多细节稍后会提到）——两者都依赖在线Softmax计算以减少内存使用。虽然Flash Attention侧重于在单个GPU上优化注意力计算本身，但上下文并行通过将序列分布到多个GPU上来实现内存降低。

探索环形注意力（Ring Attention）

在这种注意力机制的实现中，每个GPU首先发起一个异步通信操作，将其键/值对发送给其他GPU。在等待其他GPU数据的同时，它会计算当前已在内存中的那部分数据的注意力得分。理想情况下，在这次计算结束前，下一个来自其他GPU的键/值对就已经接收完毕，使得该GPU在完成第一轮计算后能够立即开始下一轮计算。

让我们来说明这一点。假设我们有4个GPU和4个token的输入。最初，输入序列沿序列维度均匀拆分，因此每个GPU仅拥有一个token及其对应的Q/K/V值。假设Q1、K1和V1分别表示第一个token的查询、键和值，并且它们位于第1个GPU上。注意力计算需要4个时间步来完成。在每个时间步中，每个GPU依次执行以下三个操作：

  1. 以非阻塞的方式将“当前的键和值”发送给下一台机器（在非阻塞模式下的最后一个时间步除外），以便在此步骤尚未完成时即可开始下一步骤
  2. 在本地对已拥有的“当前键和值”计算注意力得分，这通常涉及执行 Softmax(\frac{QK^T}{\sqrt{d}}) * V。
  3. 等待接收来自上一台GPU的键和值，然后返回到步骤1，此时“当前的键和值”即为刚刚从上一台GPU接收到的键/值对。

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

整个过程在4个GPU上的表现如下面的动画所示：

ring-attention.gif

从这个动画中，你应该能明显看出作者为何选择将这种方法称为环形注意力。

不过有一个大问题，那就是环形注意力的简单实现会导致因果注意力矩阵形状造成的GPU间工作不平衡。让我们通过考虑带有因果注意力掩码的注意力得分矩阵来观察Softmax计算：

cp_attnmask.svg

Softmax是按行计算的，这意味着每当一个GPU收到一整行的所有token时，就可以进行计算。我们看到GPU1可以立即计算，因为它一开始就拥有token 1-4，而GPU1实际上不需要从其他GPU接收任何信息。然而，GPU2需要等待第二轮，才能收到token 1-4，从而获得token 1-8的所有值。同时，GPU1的工作量明显比其他GPU要少得多。

让我们看看是否能更好地平衡计算负载：

之字形环形注意力（Zig-Zag Ring Attention） – 平衡计算的实现

我们需要一种更好的方式来分配输入序列。这可以通过不将token纯粹顺序地分配给各个GPU，而是稍微混合一下顺序，从而使每个GPU上都有较早和较晚的token。这种方法被称为之字形注意力，在这种新排列中，注意力掩码将显示出计算分布较为均匀；如果你数一数彩色方块的数量，会发现各个GPU间的计算负载现已平衡。

cp_zigzagmask.svg

与此同时，我们也会看到，为了完成所有行的计算，每个GPU都需要来自其他所有GPU的信息。

我们有两种常见方式来重叠计算和通信：一种是通过执行一次通用的all-gather操作，同时在每个GPU上重新组合所有KV（类似于Zero-3的方式）；另一种是根据需要从每个GPU逐个收集KV对：

cp_overlap_allgather.svg

cp_overlap_all2all.svg

这两种实现方式的关键区别在于它们的通信模式和内存使用：

1. AllGather实现：

2. 全对全（环形）实现：

全对全方法通常在内存效率上更优，但其通信模式稍显复杂；而AllGather方法则更简单，但在注意力计算过程中需要更多的临时内存。

到目前为止，我们已经看到如何通过TP在单个节点上拆分模型以驯服大模型，以及如何利用CP应对长序列带来的激活值爆炸问题。

然而，我们也知道TP在跨节点扩展时并不理想，那么如果模型权重难以容纳在单个节点上，我们该怎么办？这时，另一种并行度——流水线并行，将派上用场！

流水线并行（Pipeline Parallelism）

如想让阅读体验更像播客，可播放此音频，收听 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处理模型的后续部分，以此类推。这样，前向传播过程就变成了依次将数据批次沿着模型传递，并依次使用每个计算设备。

这种方法带来的第一个直接优势是：所需的互连带宽保持较低，因为我们只在模型的少数位置传输中等大小的激活值。与张量并行不同，张量并行需要在每层内部进行多次通信，而这里的通信次数要少得多。

但你可能已经开始隐约察觉到即将出现的问题：“依次” 和 “顺序执行”？！？在并行计算的世界里，这听起来似乎效率不高，特别是在我们刚刚讨论了计算与通信重叠的重要性之后。

确实如此！流水线并行的主要挑战在于如何有效地绕过这种顺序执行的限制，确保GPU始终保持忙碌，避免一个GPU在计算时，其他GPU处于等待状态。下面是一个简单的前向和反向传播示例，展示了GPU的利用情况（数字表示模型的层编号）：

image.png

一个16层模型的流水线并行示例，该模型分布在4块GPU上。数字表示层编号。

图中灰色部分表示剩余的空闲时间，通常称为“气泡（bubble）”。看到这些空闲时间，你可能会感到沮丧，毕竟我们已经花费了大量时间来优化吞吐量。我们可以通过计算“气泡”导致的额外时间来衡量流水线并行的效率。假设 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利用率下降。可以看出，在一个简单的实现中，流水线气泡可能会非常大！

幸运的是，已经有多种流水线并行方案被设计出来，以减少气泡的大小。

我们可以使用的第一个优化方法是，将批次拆分成更小的微批次（microbatches），使它们可以并行或近乎并行地处理，就像在数据并行中做的那样。例如，当第二块GPU在处理微批次1时，第一块GPU可以开始处理微批次2。以下是一个使用8个微批次的调度方案：

pp_afab2.svg

上述调度方式被称为全前向-全反向（AFAB, All-Forward-All-Backward）调度，因为它先执行所有前向传播，然后再执行所有反向传播。其优势在于前向和反向传播仍然是严格顺序的，因此可以保持模型训练代码的整体组织，使这种流水线并行实现方式成为最容易实现的一种。

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

👉 Picotron 中的 AFAB 流水线实现（点击展开）

现在我们来估算这种方法的流水线气泡时间。在第一个示例中，理想情况下处理 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 的比例减少气泡的大小。

然而，除了气泡问题，还有另一个令人头疼的问题：存储所有激活值所需的内存。我们需要将所有的激活值保留在内存中，直到反向传播阶段开始，这会导致内存使用量迅速膨胀，从而使这些流水线并行实现变得不可行。那么，我们能否找到一种方法，避免这种内存膨胀呢？

既然内存膨胀是由反向传播阶段所需的激活存储导致的，我们可以尝试在仍然执行部分前向传播时就开始执行反向传播，这样可以尽早释放部分激活，减少内存占用。

一前一后（1F1B）与 Llama 3.1 方案

该调度方案称为 一前一后（1F1B），因为在中间/稳定状态下，交替执行一个前向传播和一个反向传播。其基本思想是尽早开始反向传播。该调度如下所示：

image.png

如果仔细计算，你会发现气泡的大小仍然相同，因此训练效率并未显著提升。然而，我们仅需存储 p 个微批次的激活（其中 p 为流水线并行度），而不是 m（其中 m 是微批次数量）。这减少了在 AFAB 调度中遇到的激活内存爆炸问题。因此，我们可以增加更多微批次，从而实际减少气泡的影响。

这种设置的主要复杂性（如上图所示）在于前向和反向传播不再是完全顺序执行的，而是在设备之间并行交错执行。这意味着，我们需要在每个设备上独立调度从前向传播到反向传播的切换，而不是像往常那样在一个简单的中央训练循环中统一调度。

这也是流水线并行通常需要对训练代码和建模代码进行大幅修改的原因之一。

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

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

让我们看看 1F1B 流水线并行调度在实践中的扩展情况，并查看我们集群上的一些基准测试结果：

Pipeline Parallelism 吞吐量扩展（不同微批次大小）

左侧图表中，当微批次数量等于或小于流水线并行度减一（m = p - 1）时，我们可以看到流水线气泡的负面影响——性能较低，并且随着流水线并行度的增加甚至下降。而右侧图表显示，当微批次数量远大于流水线并行度（m = 32 \gg p - 1）时，可以改善低并行度时的性能，但在较大并行度时仍然受到限制。实际上，我们无法无限增加微批次数量以维持 m \gg p - 1，因为我们最终会受限于目标全局批量大小。当流水线并行度增加到最大可用微批次数时，我们将不得不按照 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}

因此，我们可以通过增加微批次（microbatches）和交错阶段（interleaved stages）来减少管道气泡（bubble），但需要注意的是，从数量上来看，通信量也会随之增加 v，这实际上是一个权衡。在下图中，你可以看到针对 p=8 的管道并行（PP）设置的几种不同配置，其中 m=1, v=1 是最基础的管道并行方式，而 v=1 代表 AFAB 或 1F1B 方案，而 v \neq 1 则是交错配置。

这里的调度变得更加复杂，因为我们需要决定在特定 GPU 上的特定时刻，应该优先处理较早的微批次通过较晚的层（即尽快完成前向和后向传播循环，这被称为“深度优先”（depth-first），即优先让批次尽快通过整个模型），还是优先处理较晚的微批次通过较早的层（即尽可能填充整个管道，这被称为“广度优先”（breadth-first））。关于这个选择，"Breadth-First Pipeline" 论文 中有详细的解释。

现在，你已经掌握了 Llama 3.1 的管道并行方法的所有关键要素。它采用了一种“一前一后”（1F1B）设置，并结合了交错阶段，同时优先级可调，可在深度优先和广度优先之间调整。

pp_llama3.1_schedule.png

然而，我们尚未探索所有可能的管道调度方法，最近，一些新方法已经被提出，可以将气泡减少到几乎为零！例如，DeepSeek V3/R1 实现中就使用了这些技术。是不是很好奇？让我们在离开管道并行的世界之前，最后快速看一下这些神奇的调度方法吧！

零气泡（Zero Bubble）与 DualPipe

最近，一些更复杂的气泡优化方法被提出，并达到了接近“零气泡”的状态。秘诀在于对涉及的操作进行更加精细的拆分，以实现最高效的交错。例如，DeepSeek V3/R1 的管道实现方法——DualPipe——就几乎达到了零气泡状态。

让我们简要了解一下 ZeroBubble 研究，它是 DualPipe 方法的前身。ZeroBubble 的核心观察点是：矩阵乘法的反向传播实际上涉及两个独立的操作——输入的反向传播（B）和权重的反向传播（W）：

其中，B（输入的反向传播）的输出对于执行更低层的反向传播是必需的，而 W（权重的反向传播）并不是必须立即执行的，它通常只需要在优化器步骤之前完成。如下图所示：

image.png

这意味着 W 可以在同一阶段的 B 之后的任何位置灵活调度。这种灵活性使得我们可以巧妙地安排 W，以填补管道中的气泡。右上角的 ZB-H2 调度就是利用这种精细拆分实现零气泡的示例（理论上的）。

image.png

上图（来自 ZeroBubble 论文）：第一张图（Figure 2）展示的是经典的 1F1B 调度，其中前向传播和后向传播交错进行，但仍然保持了较粗粒度的后向传播。而下方两张图（Figure 3）分别展示了 ZeroBubble 调度的两种变体，它们将后向传播进一步拆分为 B 和 W 这两个更细粒度的操作。最后一种调度方式——ZB-H2——是利用这种细粒度分解，实现零气泡的（理论）示例。

DeepSeek 在其 V3 技术报告 中提出的 DualPipe 方法，对这种分解策略进行了扩展，它引入了两个沿管道并行（PP）维度传播的独立数据流，并通过交错执行来最大限度减少 GPU 的空闲时间。其调度方式如下图所示，比之前的方法更为复杂：

image.png

通常，要完全优化如此复杂的调度方式，需要精确测量各个细粒度操作的执行时间，并利用整数线性规划（ILP）来最小化最终的气泡时间。ZeroBubble 论文 讨论了用于实现此类调度的启发式方法和算法。因此，ZeroBubble 和 DualPipe 调度方式过于复杂，无法在这里提供代码示例，但你应该已经对其中涉及的概念有了大致了解。

至此，我们对管道调度和气泡优化的探索就告一段落了！希望你喜欢这次深入的学习之旅！

现在，是时候转向我们要详细介绍的最后一种并行化方法了——专家并行（Expert Parallelism）。

专家并行（Expert parallelism）

这是我们要讨论的最后一种并行方法。在深入探讨之前，如果你对专家混合（Mixture-of-Experts，MoE）还不熟悉，可以阅读我们之前发布的 这篇简短的博客文章，它能帮助你更好地理解 MoE 体系结构。

近年来，专家混合模型受到了越来越多的关注，例如 GPT-4、Mixtral，以及最近的 DeepSeek-V3/R1。这类模型的基本思想是，每一层不是单独使用一个前馈模块（feedforward module），而是可以并行使用多个模块，并通过不同的路径处理 token。

ep_moe.png

MoE 层的示意图，摘自 Switch Transformers 论文

MoE 层的设计使其能够在专家（expert）维度上轻松实现并行计算，我们称之为 专家并行（Expert Parallelism, EP）。由于前馈层（feedforward layers）完全独立，我们可以将每个专家的前馈层放置在不同的计算节点上。相比于张量并行（TP），EP 更加轻量级，因为它不需要拆分矩阵乘法，只需要将 token 的隐藏状态正确路由到相应的专家即可。

在实际应用中，EP 通常会与其他并行方式结合使用，例如数据并行（Data Parallelism, DP）。这是因为 EP 仅影响 MoE 层，并不会像上下文并行（Context Parallelism）那样在序列长度维度上对 token 进行分片。如果仅使用 EP，GPU 仍然会对所有非 MoE 模块执行冗余计算。通过将 EP 与 DP 结合，我们可以有效地在 GPU 之间分片专家模块和输入批次，如下图所示：

ep_schema.png

来源：专家混合综述（A Survey on Mixture of Experts）

不过，先别着急——在下一节中，我们会专门讨论不同并行策略之间的交互作用，所以如果你还不理解这个图也不用担心。

在实践中，有一些技巧可以提高 EP 的效率，这些技巧与模型设计密切相关。例如，DeepSeek-V3 在路由器中施加了一个约束，确保每个 token 最多被发送到 M 个计算节点（在其设计中为 4 个），从而尽可能让 token 保持在单个节点上，并减少通信开销。虽然专家并行已经存在了一段时间，但随着 MoE 体系结构的兴起，它正在重新受到关注。

我们计划很快在 picotron/nanotron 中添加一个更完整的 EP 示例，敬请期待！

5D 并行概览

恭喜你！你已经了解了用于扩展模型训练的五种并行策略：

  1. 数据并行（DP）——按批次维度并行
  2. 张量并行（TP）——按隐藏维度并行
  3. 序列并行 & 上下文并行（SP/CP）——按序列维度并行
  4. 流水线并行（PP）——按模型层并行
  5. 专家并行（EP）——按模型专家并行

此外，还有三种 ZeRO 策略可以与数据并行结合，以减少内存占用：

  1. ZeRO-1 —— 在 DP 复制之间分片优化器状态
  2. ZeRO-2 —— 在 DP 复制之间分片优化器状态和梯度
  3. ZeRO-3 —— 在 DP 复制之间分片优化器状态、梯度和模型参数

到目前为止，你可能会好奇这些并行和 ZeRO 策略如何相互比较和交互。换句话说，我们应该选择哪些策略进行组合，而哪些应该避免混用？

接下来，我们将分析它们之间的相似性和相互作用。首先，我们将并排比较流水线并行（PP）和 ZeRO-3，它们在某些方面非常相似，但也存在重要的区别。

流水线并行 vs. ZeRO-3 —— PP 和 ZeRO-3 都是通过将模型权重分布在多个 GPU 上，并在模型深度轴上进行计算和通信（例如，在 ZeRO-3 中，我们在计算时预取下一层数据）。在这两种方法中，每个设备都需要完整地计算层操作，而不像 TP 或 EP 那样在子层级别进行计算。

然而，PP 和 ZeRO-3 之间存在几个主要区别：

ZeRO-3 流水线并行（PP）
每个计算单元存储 仅存储部分层参数 存储完整层参数
通信用于传输 模型权重 激活值
调度方式 与模型无关 与模型无关
实现挑战 模型分片与通信复杂 流水线调度复杂
扩展性 偏好较大 mbs 和 seq\_len 以隐藏通信 偏好较大 \text{grad\_acc} 以隐藏计算空档

正如你所看到的，ZeRO-3 和 PP 解决了相同的挑战，但采用了不同的方法，选择哪种方式取决于你是更关注权重的通信，还是激活的通信。虽然它们可以结合使用，但在实践中不常见，因为这样做需要显著增加全局 batch size 以摊销通信成本，从而在全局 batch size、模型大小、网络带宽和训练效率之间形成权衡。如果你决定结合使用它们，ZeRO-3 应该被配置为在一系列 PP 微批次期间将权重保留在内存中，以尽可能减少不必要的通信开销。

另一方面，ZeRO-1 和 ZeRO-2 关注优化器状态和梯度，它们可以轻松与流水线并行（Pipeline Parallelism）结合，并且是互补的。结合使用它们不会带来额外的新挑战。例如，DeepSeek-v3 的训练使用了 PP 结合 ZeRO-1。

张量并行（Tensor Parallelism）（与序列并行（Sequence Parallelism）一起）是天然互补的，并且可以与流水线并行和 ZeRO-3 结合使用，因为它依赖矩阵乘法的分布性质，使得权重和激活可以被分片并独立计算后再合并。

TP & SP 结构图

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

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

上下文并行（Context Parallelism） 和 专家并行（Expert Parallelism） 也可以帮助分片激活，并且可以视为 TP 的补充。前者处理长序列，而后者用于分布式 MoE 训练，它们可以无缝结合使用。

上下文并行（CP） 主要用于解决超长序列训练的挑战，它通过在 GPU 之间沿序列维度分片激活来实现。大多数操作（如 MLP 和 LayerNorm）可以独立处理这些分片的序列，而注意力层需要通信，因为每个 token 需要访问整个序列的 key/value。正如我们在 CP 章节 中所看到的，这通过环形注意力模式（ring attention patterns）高效地处理，实现了计算和通信的重叠。当扩展到极端长的序列（128k+ tokens）时，即使使用完整的激活重计算，单个 GPU 也无法满足注意力计算的内存需求，此时 CP 就尤为重要。

CP 结构图

专家并行（EP） 主要用于训练 MoE（Mixture of Experts）模型，它通过在 GPU 之间分片专门的“专家”（experts），并在计算过程中动态地将 token 路由到相关的专家。EP 中的关键通信操作是 `all-to-all` 操作，它负责将 token 发送到相应的专家，并收集返回的计算结果。虽然这种操作引入了一定的通信开销，但它使得模型容量可以大规模扩展，因为每个 token 在推理（或训练）期间仅由一小部分参数处理。对于大规模专家模型，将专家在 GPU 之间分片变得尤为重要。

EP 结构图

📝 备注

EP 在输入处理方面与数据并行（DP）有相似之处，因此某些实现将专家并行视为数据并行的一个子类别，主要区别在于 EP 使用专门的专家路由，而不是让所有 GPU 处理相同的模型副本。

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

张量 + 序列并行 上下文并行 专家并行
沿隐藏/序列维度分片权重和激活 沿序列维度分片激活 分片专家权重和激活
用于矩阵乘操作的通信（列/行线性） 注意力键/值的通信 用于专家路由的通信
需要特定于模型的实现 除了注意力外都是通用的 除了MoE层外都是通用的
偏好高带宽的节点内通信 偏好大序列长度 需要MoE

总结一切—— 现在，让我们尝试将我们在一个单一图表中看到的所有技术聚合和组合起来。是的，我们迎接这个挑战！

在这个总结图表中，您将找到单个transformers层的激活和模块的插图，以其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 仓库中找到几个脚本，可以用来运行我们上述讨论的所有实验，并能够在现实生活中基准测试您自己的模型和集群。

事实上，我们自己在 数千个分布式配置 上进行了基准测试，涵盖了我们在本书中到目前为止讨论过的每个模型大小以及大量集群配置（即每个节点8xH100s的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实现中的通信重叠，预计它将重新获得性能优势。

基准测试中的经验教训

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

这听起来足够简单：我们可以在我们的集群上轻松启动大量作业。然而，一旦我们启动了第一批实验，问题就开始出现：

在有限的时间内运行所有实验需要额外的工程设计，我们最终花费了大量时间处理诸如：

这些挑战值得一书，但它们教会了我们有关分布式训练基础设施复杂性的宝贵教训。理论上看起来简单的东西，在实践中往往需要对许多运作部件进行仔细关注。

在实践中复现理论结果是具有挑战性的，特别是由于生产训练代码的有限可用性。通过像 nanotron 和 picotron 这样的开源项目，我们希望能够帮助使分布式训练技术更加可访问，并在简单高效的代码库上进行合作，帮助研究人员和从业者充分利用他们的硬件资源。

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

回顾我们迄今为止的讨论，我们的许多讨论都依赖于一个关键假设 - 即可以在GPU上有效地重叠计算和通信，而不会对计算吞吐量产生影响。现实情况更加微妙。当使用像NCCL send/recv这样的常见通信原语时，我们面临计算资源和通信资源之间的隐藏竞争，因为通信核心通常会使用相同的GPU流处理器（SM），这些SM用于计算，导致在通信与计算重叠时吞吐量降低。要真正优化我们的分布式训练，我们需要更深入地了解GPU架构本身。

现在是时候关掉灯光并激活CUDA模式了！

深入GPU - 融合、线程化、混合

如果您希望在阅读过程中增加播客体验，请随时收听 NotebookLM 主持人讨论本书的以下部分。

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

但这忽略了我们可以在更低层次上通过仔细理解我们的模型操作如何在每个GPU上调度和执行来做的所有优化。

本节将深入介绍GPU架构的更多细节，特别是NVIDIA的GPU架构，但通常的想法可以在类似的加速器单元上重复使用。

在覆盖Flash-Attention革命如何有效调度GPU工作负载之前，我们将简要解释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的广告80GB），但访问和查询速度最慢。

image.png

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

GPU的目标是通过利用计算/内存的这种分层组织，尽可能并行地在GPU核心上运行尽可能多的工作负载。

在GPU核心上运行的代码片段称为内核。它可以在高级别上用CUDA或Triton等语言编写，然后编译为NVIDIA GPU使用的低级汇编Parallel Thread Execution（PTX）。

要运行内核，您还需要一个特定的代码部分，称为主机代码，它在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); // 拷贝数据到 GPU cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice); cudaMemcpy(d_B, h_B, size, cudaMemcpyHostToDevice); // 启动 kernel int threadsPerBlock = 256; int blocksPerGrid = (N + threadsPerBlock - 1) / threadsPerBlock; VecAdd<<>>(d_A, d_B, d_C, N); // 将结果拷回 CPU cudaMemcpy(h_C, d_C, size, cudaMemcpyDeviceToHost); // 释放 GPU 显存 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

// 设备端（GPU）代码 __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

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

从这些细节中最重要的是记住，有各种大小和分配约束（各种内存的大小，每个线程束和块中的线程数），需要考虑使用GPU架构的最有效方式。

大多数情况下，您不需要达到这种精度级别，幸运的是，您可以重用社区其他成员准备的内核和代码。但无论如何，我们希望为您提供有关如何开始使用内核的入门指南！

如何使用内核提高性能？

如果您想要添加一个缺少优化内核的新操作，或者加速现有的 PyTorch 函数，从零开始编写内核可能看起来是最直接的方法。然而，创建高性能的 CUDA 内核需要丰富的经验，并且学习曲线陡峭。通常，更好的方法是利用 torch.compile，它通过捕获您的操作并在 Triton 中生成更低级的高性能内核，动态优化 PyTorch 代码。

假设您想要为一个名为指数线性单元（Exponential Linear Unit）的激活函数编写一个内核：

\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 中，有时也无法完全达到设备的峰值性能，因为该语言在处理共享内存和流多处理器（SMs）内的调度等低级细节方面存在限制。Triton 的能力仅限于块及其在 SMs 之间的调度。如果想要获得更深层次的控制，您需要直接在 CUDA 中实现内核，这样才能访问所有底层细节。

深入 CUDA 之后，可以使用各种技术来提高内核的效率。我们将在这里介绍其中几种方法：优化内存访问模式以减少延迟、使用共享内存存储频繁访问的数据，以及管理线程工作负载以最小化空闲时间。

在深入探讨 CUDA 示例之前，让我们总结一下我们所见过的用于编写内核代码以在 GPU 上执行指令的工具：

  1. PyTorch：简单但速度较慢
  2. torch.compile：简单且速度快，但灵活性不足
  3. Triton：更难，但更快、更灵活
  4. CUDA：最难，但最快、最灵活（如果掌握得当）

让我们讨论 CUDA 中最常见的优化技术之一：优化内存访问。GPU 的全局内存（在前面的图表中是最大的内存）相比缓存来说，延迟较高，带宽较低，这通常是大多数应用程序的主要瓶颈。高效地访问全局内存的数据可以极大地提高性能。

内存合并

为了有效利用全局内存的带宽，理解其架构至关重要。在CUDA设备中，全局内存是使用DRAM实现的。

内存合并利用DRAM以突发方式提供数据的方式，或者说是连续内存位置的范围，每当访问一个内存地址时，都会并行读取该地址附近的一系列连续位置。一旦读取完成，这些数据可以快速传输到处理器作为一个突发。在CUDA中，合并利用这种突发行为来最大化内存访问效率，通过确保warp中的线程（即以SIMD方式执行同一指令的32个线程）访问连续的内存位置。例如，如果线程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

问题出在于，在这个内核中，同一个warp中的两个线程，例如Thread ID为(0, 0)和(1, 0)（它们最终会进入同一个warp），会从矩阵B的同一列但是矩阵A的不同行加载数据。由于矩阵元素是按行主序存储的（意味着行元素在连续的内存地址中，如下图所示），线程(0, 0)将加载A_{0,0}，线程(1, 0)将加载A_{1,0}在第一次迭代i = 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的不同列。因此，对于行主序矩阵，内存访问可以被合并。

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

image.png

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

现在让我们来介绍另一种你经常在文献中看到的技术：分块处理（tiling）。

分块处理（Tiling）

分块处理是一种利用共享内存优化内存访问模式的技术。正如我们前面提到的，共享内存是一种小而快速的存储，块内的所有线程都可以访问它。这使得数据可以被多个线程重复使用，从而减少了从较慢的全局内存中重复加载数据的需求。

以矩阵乘法为例，块中的每个线程可能需要从两个矩阵（如 A 和 B）中获取元素。如果每个线程独立地从全局内存加载所需的行和列，就会出现大量冗余加载，因为块中的多个线程会访问重叠的数据。相反，我们可以使用分块处理，将 A 和 B 的一个块（或瓦片）一次性加载到共享内存中，让该块中的所有线程重复使用相同的共享数据。

在分块处理的方法中，每次迭代时，块内的所有线程协同工作，将两个瓦片（一个来自矩阵 A，另一个来自矩阵 B）加载到共享内存中。具体来说，线程加载矩阵 A 的一个瓦片（大小为 BLOCK_SIZE_M × BLOCK_SIZE_K）以及矩阵 B 的一个瓦片（大小为 BLOCK_SIZE_K × BLOCK_SIZE_N）。一旦这些瓦片存入共享内存，线程就可以在这些瓦片上执行矩阵乘法，从而实现高效计算，因为所有必要的数据都可以被快速访问。瓦片乘法的结果存储在一个累积矩阵中，该矩阵保存中间结果。在每次迭代后，当前瓦片乘法的结果都会累加到该矩阵中，直到两个矩阵的所有瓦片都被处理完毕。

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中加载一个元素到共享内存。在这种情况下，实现合并内存访问（coalesced memory access）非常直观：通过将 threadIdx.x 作为局部列索引（localCol），同一个 warp 中的线程可以访问相邻的矩阵元素。块内所有线程完成数据加载后（通过调用 __syncthreads() 确保同步），它们就会计算这两个瓦片的点积。当所有瓦片遍历完成——矩阵 A 在水平方向移动，矩阵 B 在垂直方向移动——最终计算出的结果存入矩阵 C的对应位置。

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

线程粗化（Thread Coarsening）

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

image.png

这些神秘状态名称的含义可以在NVidia的性能指南中找到，在Warp Stall Reasons部分可以阅读到：

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

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

在写入或改进自定义内核时，让我们简要提到一个最重要的考虑因素：最小化控制分歧。

最小化控制分歧

流多处理器（SM）被设计为使用单指令多数据（SIMD）模型执行warp中的所有线程。这意味着在任何给定时刻，一个指令同时为warp中的所有线程获取和执行。当执行warp时，其中的线程在数据的不同段上操作，但遵循相同的指令，因此得名单指令多数据。SIMD的主要优势在于其效率；负责指令获取和调度的控制硬件在多个执行单元之间共享。这种设计最小化了与控制功能相关的硬件开销，使得更大比例的硬件专注于提高算术吞吐量。

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

我们已经介绍了写入自定义内核和改进GPU操作性能和内存占用的一些主要考虑因素。但在转向实际示例之前，还有一个重要的概念需要讨论：“融合内核”。

融合内核（Fused Kernels）

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

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

这个想法在 Horace He 的这些图表中得到了美妙的诠释：

image.png

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

image.png

不要把我们的三角形发送回全局内存，只是再次读取它，我们可以直接一次性完成所有操作。

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

融合内核特别适用于简单点操作的连续操作，这些操作在每个输入标记上都是独立执行的。在这种情况下，没有必要在将计算值带回全局内存之前将其保留在本地。在新的内核中，保留所有值直到执行完连续计算的一系列操作更为高效。

在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矩阵，而是完全避免存储庞大的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的底层优化，而不是对注意力机制本身的改动。具体来说：

Flash Attention 是一个典型案例，展示了当我们深入考虑当前GPU加速器的内存/计算设计时，所能带来的突破性改进。

到目前为止，我们讨论的算子融合技术要求对模型代码进行改动，并为特定操作编写自定义内核，以加速训练。

在计算操作的底层优化的最后部分，我们将探索一系列与模型代码无关的方法，这些方法适用于任何模型，并且已经成为业界标准：混合精度训练（Mixed Precision Training）！

混合精度训练（Mixed Precision Training）

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

顾名思义，混合精度训练涉及在训练过程中混合使用不同的数值精度。PyTorch张量的默认数值精度是单精度浮点格式，即FP32（float32），这意味着每个存储的数值占用32位（4字节）。这些位被分为三个部分：

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通过牺牲更多精度来维持float32的范围，但这是有代价的。在float8的情况下，情况更为严峻，因为e4m3在区间1-2内只能表示7个数字，而e5m2只能表示3个数字。

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

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

事实证明，我们不能完全放弃float32，并且通常需要保持一些部分以全精度进行训练。这就是为什么较低精度训练通常被称为混合精度训练。

现在让我们来看看使用16位进行模型训练，然后看看能否进一步降至8位。

FP16和BF16训练

简单地将所有张量和操作切换到float16通常不起作用，结果通常是发散的损失。然而，原始的混合精度训练论文提出了三种技巧来匹配float32训练：

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

通过这些技术，我们可以实现稳定的训练，同时由于更快的低精度算术运算，获得更高的吞吐量。当然，作为一个充满好奇心的读者——现在可能已经稍微沉迷于最大化吞吐量——您可能会问：我们是否可以比16位精度更进一步、更快？

也许可以！

FP8预训练

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

最近的研究，包括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进一步引入了一种特定的量化方案，其中范围按块（tile）归一化：输入/激活使用1×128，权重和缩放因子使用128×128。这种方法使归一化过程不易受到激活值中异常值的影响。此外，他们还提出了一些额外的技巧，以进一步减少内存和通信开销，具体内容可以在DeepSeek-V3技术报告的第3.3节中找到。

以下是一些已知的FP8训练方法的总结：

GEMM计算精度 主模型权重 累积梯度 模型权重 梯度 优化器状态 总内存占用
bfloat16 + fp32混合精度基线 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 Engine 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%减少）

总体而言，在2025年初，FP8仍然是一种实验性技术，相关方法仍在不断发展。鉴于其明显的优势，它很可能很快成为标准，并取代bf16混合精度训练。想要了解FP8训练技术的开源实现，可以查看nanotron的实现：此PR。

展望未来，下一代NVIDIA Blackwell芯片已宣布将支持FP4训练，这将进一步加速训练，但无疑也会带来新的训练稳定性挑战。

这一部分标志着我们在大规模、高速模型训练领域的长途探索告一段落。是时候让我们的GPU集群慢慢休息，并回顾一路走来的学习成果了。

结论

恭喜你，亲爱的读者，你坚持到了最后！我们完成了一次精彩的旅程：从理解如何在单个GPU上训练简单模型，到掌握在数千个GPU上高效训练Llama-405B和DeepSeek-V3等大规模语言模型的复杂技术。现在，你应该能够相对轻松地理解Llama-3的4D并行架构图：

image.png

在GPU集群上高效训练大型LLM并非易事。我们学习了如何优化计算和GPU间通信，以确保它们始终处于最大化利用率。这涉及为特定模型和集群规模选择合适的并行策略，在可能的情况下重叠通信和计算，并编写自定义核函数，以充分利用硬件架构，使运算尽可能快地执行。

你可能会认为这些知识相对小众，仅适用于少数从事LLM预训练的研究人员。历史上确实如此，但随着AI开发者社区和模型规模的迅速增长，越来越多的人在推理、微调和训练中使用分布式技术，使分布式训练变得越来越普遍。因此，深入学习分布式计算正当其时。

这不仅是你的学习旅程，也是我们的学习之旅！在GPU集群上运行数千次基准测试比我们预想的更具挑战性，我们也希望与你分享我们的学习经验。

那么，接下来呢？

你现在对主要的分布式训练概念有了很好的理解，但同时，我们也只是触及了许多工具和技术的表面。以下是我们推荐的深入学习步骤：

我们希望这本书能帮助你入门分布式训练，并希望你能伴随GPU集群的轰鸣声训练出下一代优秀的模型！

最后的话： 我们对这篇文章感到非常自豪，因此决定为首批读者限量赠送实体印刷版。

如果你是前50名填写邮箱地址的读者之一，我们将在今年晚些时候联系你，并寄送正式印刷的实体书。

如果你希望获得实体书，请填写此表单。

无论你是我们的第一批读者，还是后来者，我们都很高兴你能享受这次知识分享。愿开源与开放科学的精神与你同在！

致谢

我们感谢 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

微软的深度学习优化库，包含 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 的流水线并行实现指南，用于大规模模型训练。

广度优先流水线并行

广泛讨论流水线调度。

All-reduce

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

Ring-flash-attention

结合 Flash Attention 进行高效训练的环形注意力机制实现。

环形注意力教程

讲解环形注意力的概念和实现。

ZeRO 和 3D

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

混合精度训练

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

可视化 6D 网格并行

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

硬件

Fire-Flyer - 一款由 10,000 个 PCI 芯片组成的集群

DeepSeek 关于设计一个包含 10,000 个 PCI GPU 集群的报告。

Meta 的 24k H100 Pods

Meta 详细介绍了他们使用 NVIDIA H100 GPU 构建的大规模 AI 基础设施。

Semianalysis - 10万 H100 集群

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

Modal GPU 术语表

为人类准备的 CUDA 文档

其他

Stas Bekman 的手册

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

Bloom 训练纪实

详细记录了 BLOOM 模型训练过程和面临的挑战。

OPT 日志

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

Harm's 法则：训练小模型更久

研究模型大小与训练开销之间的关系。

Harm's 博客：长上下文训练

研究长上下文训练的数据和训练成本。

GPU Mode

一个 GPU 阅读小组和社区。

EleutherAI YouTube 频道

机器学习可扩展性与性能阅读小组

Google Jax 扩展书

如何扩展您的模型

@fvsmassa 和 @TimDarcet FSDP

约 500 行代码的独立 FSDP 实现

thonking.ai

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

Aleksa 的 ELI5 Flash Attention

Flash Attention 的简单解释

TunibAI 的 3D 并行教程

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

附录

A0：并行编程速成课程

在整个博文中，我们将LLM训练从单个GPU扩展到数百个GPU。这将需要在所有机器之间进行权重、梯度和数据的通信与同步。有一组分布式模式可以实现这一点，称为集体操作。在本节中，我们将进行一个小型的速成课程，涵盖诸如广播、全局归约、分散等操作。让我们深入探讨！

一般的设置是，我们有许多独立的节点，可以是CPU核心、GPU或计算节点。每个节点执行一些计算，然后我们希望将结果或其部分传输到其他节点，用于下一个计算步骤（t+1）。

image.png

也许我们需要将一个节点的结果发送到所有其他节点，或者我们需要汇总每个节点的所有中间结果以报告总体结果。通常情况下，有一个具有显著地位的节点在操作中起到核心作用，在这里用root表示，它是某些操作的目标或源。让我们从最简单的原语之一开始：广播操作。

广播（Broadcast）

一个非常常见的模式是，您在节点1上有一些数据，并希望与所有其他节点共享数据，以便它们可以使用数据进行一些计算。广播操作正是做到了这一点：

image.png

PyTorch原生提供了集体操作，因此我们可以很容易地编写一个小例子来演示广播是如何工作的。我们首先需要使用dist.init_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"在排名 {dist.get_rank()} 上广播前：{tensor}") dist.broadcast(tensor, src=0) print(f"在排名 {dist.get_rank()} 上广播后：{tensor}") init_process() example_broadcast()

您可以使用torchrun --nproc_per_node=3 dist_op.py运行上述脚本（您需要3个GPU，或者根据需要更改nproc_per_node），您应该看到以下输出：

在排名 0 上广播前：tensor([1., 2., 3., 4., 5.], device='cuda:0') 在排名 1 上广播前：tensor([0., 0., 0., 0., 0.], device='cuda:1') 在排名 2 上广播前：tensor([0., 0., 0., 0., 0.], device='cuda:2') 在排名 0 上广播后：tensor([1., 2., 3., 4., 5.], device='cuda:0') 在排名 1 上广播后：tensor([1., 2., 3., 4., 5.], device='cuda:1') 在排名 2 上广播后：tensor([1., 2., 3., 4., 5.], device='cuda:2')

很好，看起来正如预期的那样工作。请注意，排名消息可能会以无序的方式打印出来，因为我们无法控制哪个打印语句首先执行（这里我们为可读性排序了它们）。现在让我们继续进行归约和全局归约模式！

归约 & 全局归约（Reduce & AllReduce）

归约模式是分布式数据处理中最基本的模式之一。其思想是通过一个函数f()（例如求和或平均）来组合每个节点上的数据。在归约范例中，结果仅发送到根节点，而在全局归约情况下，结果广播到所有节点：

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"在排名 {dist.get_rank()} 上归约前：{tensor}") dist.reduce(tensor, dst=0, op=dist.ReduceOp.SUM) print(f"在排名 {dist.get_rank()} 上归约后：{tensor}") init_process() example_reduce()

请注意，在Reduce操作中，仅更新了dst节点上的张量：

在排名 0 上归约前：tensor([1., 1., 1., 1., 1.], device='cuda:0') 在排名 1 上归约前：tensor([2., 2., 2., 2., 2.], device='cuda:1') 在排名 2 上归约前：tensor([3., 3., 3., 3., 3.], device='cuda:2') 在排名 0 上归约后：tensor([6., 6., 6., 6., 6.], device='cuda:0') 在排名 1 上归约后：tensor([2., 2., 2., 2., 2.], device='cuda:1') 在排名 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"在排名 {dist.get_rank()} 上all_reduce之前：{tensor}") dist.all_reduce(tensor, op=dist.ReduceOp.SUM) print(f"在排名 {dist.get_rank()} 上all_reduce之后：{tensor}") init_process() example_all_reduce()

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

在排名 0 上all_reduce之前：tensor([1., 1., 1., 1., 1.], device='cuda:0') 在排名 1 上all_reduce之前：tensor([2., 2., 2., 2., 2.], device='cuda:1') 在排名 2 上all_reduce之前：tensor([3., 3., 3., 3., 3.], device='cuda:2') 在排名 0 上all_reduce之后：tensor([6., 6., 6., 6., 6.], device='cuda:0') 在排名 1 上all_reduce之后：tensor([6., 6., 6., 6., 6.], device='cuda:1') 在排名 2 上all_reduce之后：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"在排名 {dist.get_rank()} 上gather之前：{tensor}") dist.gather(tensor, gather_list, dst=0) if dist.get_rank() == 0: print(f"在排名 0 上gather之后：{gather_list}") init_process() example_gather()

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

在排名 0 上gather之前：tensor([1., 1., 1., 1., 1.], device='cuda:0') 在排名 1 上gather之前：tensor([2., 2., 2., 2., 2.], device='cuda:1') 在排名 2 上gather之前：tensor([3., 3., 3., 3., 3.], device='cuda:2') 在排名 0 上gather之后：[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"在排名 {dist.get_rank()} 上all_gather之前：{tensor}") dist.all_gather(gather_list, tensor) print(f"在排名 {dist.get_rank()} 上all_gather之后：{gather_list}") init_process() example_all_gather()

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

在排名 0 上all_gather之前：tensor([1., 1., 1., 1., 1.], device='cuda:0') 在排名 1 上all_gather之前：tensor([2., 2., 2., 2., 2.], device='cuda:1') 在排名 2 上all_gather之前：tensor([3., 3., 3., 3., 3.], device='cuda:2') 在排名 0 上all_gather之后：[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')] 在排名 1 上all_gather之后：[tensor([1., 1., 1., 1., 1.], device='cuda:1'), tensor([2., 2., 2., 2., 2.], device='cuda:0'), tensor([3., 3., 3., 3., 3.], device='cuda:0')] 在排名 2 上all_gather之后：[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，或者在操作之间使用ReduceScatter模式：

Scatter & ReduceScatter

正如名称所暗示的，Scatter操作的目标是将数据从一个节点分发到所有其他节点。因此，它与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"排名 0：待分发的张量：{scatter_list}") else: scatter_list = None tensor = torch.zeros(5, dtype=torch.float32).cuda() print(f"在排名 {dist.get_rank()} 上scatter之前：{tensor}") dist.scatter(tensor, scatter_list, src=0) print(f"在排名 {dist.get_rank()} 上scatter之后：{tensor}") init_process() example_scatter()

结果显示，空张量已被scatter_list的内容填充：

排名 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')] 在排名 0 上scatter之前：tensor([0., 0., 0., 0., 0.], device='cuda:0') 在排名 1 上scatter之前：tensor([0., 0., 0., 0., 0.], device='cuda:1') 在排名 2 上scatter之前：tensor([0., 0., 0., 0., 0.], device='cuda:2') 在排名 0 上scatter之后：tensor([1., 1., 1., 1., 1.], device='cuda:0') 在排名 1 上scatter之后：tensor([2., 2., 2., 2., 2.], device='cuda:1') 在排名 2 上scatter之后：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"在排名 {rank} 上ReduceScatter之前：{input_tensor}") dist.reduce_scatter(output_tensor, input_tensor, op=dist.ReduceOp.SUM) print(f"在排名 {rank} 上ReduceScatter之后：{output_tensor}") init_process() example_reduce_scatter()

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

在排名 0 上ReduceScatter之前：[tensor([1., 2.], device='cuda:0'), tensor([1., 4.], device='cuda:0'), tensor([1., 8.], device='cuda:0')] 在排名 1 上ReduceScatter之前：[tensor([2., 4.], device='cuda:1'), tensor([ 4., 16.], device='cuda:1'), tensor([ 8., 64.], device='cuda:1')] 在排名 2 上ReduceScatter之前：[tensor([3., 6.], device='cuda:2'), tensor([ 9., 36.], device='cuda:2'), tensor([ 27., 216.], device='cuda:2')] 在排名 0 上ReduceScatter之后：tensor([ 6., 12.], device='cuda:0') 在排名 1 上ReduceScatter之后：tensor([14., 56.], device='cuda:1') 在排名 2 上ReduceScatter之后：tensor([ 36., 288.], device='cuda:2')

让我们简要地看一下一个常见的使用ReduceScatter和AllGather的AllReduce实现：环形AllReduce。

快速关注环形AllReduce

环形AllReduce是AllReduce的一种特定实现，经过优化以实现可伸缩性。与所有设备直接相互通信不同，这可能会造成通信瓶颈，环形All-Reduce可以分解为两个关键步骤：ReduceScatter和AllGather。它的工作原理如下：

  1. ReduceScatter
  2. AllGather

让我们通过以下动画来说明，我们有5个GPU，每个GPU都有长度为5的张量。第一个动画显示了ReduceScatter步骤，最终每个GPU都接收到了特定数据块的减少结果（橙色矩形）：

image.png

接下来的动画展示了AllGather步骤，在此过程结束时，每个GPU获取了AllReduce操作的完整结果：

image.png

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

对于AllReduce，有两个关键点需要记住：

  1. 当N（GPU的数量）较大时，AllReduce的通信成本约为2 \times K。
  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} 秒.") time.sleep(rank) # 模拟不同的处理时间 dist.barrier() print(f"排名 {rank} 在Barrier之后的时间差: {time.time()-t_start:.4f}") init_process() example_barrier()

我们可以看到，尽管第一个排名没有睡眠，但它也需要2秒才能通过Barrier：

排名 0 睡眠 0 秒. 排名 1 睡眠 1 秒. 排名 2 睡眠 2 秒. 排名 0 在Barrier之后的时间差: 2.0025 排名 1 在Barrier之后的时间差: 2.0025 排名 2 在Barrier之后的时间差: 2.0024

我们需要小心地进行这种方式的所有节点同步操作，因为这会打败并行独立操作的目的，可能会减慢整个处理速度。在许多情况下，如果快速节点已经开始处理下一个作业，这可能是可以接受的，因为快速节点在下一个迭代中可能会变慢，从而平衡整个过程中的延迟。

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

NCCL：NVIDIA Collective Communications Library

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

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

在上述引用的PyTorch指南中，还有一些决策树中的细微点可以让读者深入探讨。

现在我们已经涵盖了分布式训练的基本操作和您应该准备好轻松跟随博客文章。

A1：分布式训练性能分析

内核

让我们从假设内核已经集成到PyTorch中开始。作为一个简单的例子，我们可以查看在PyTorch中实现的Layer Normalization函数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 Profiler。例如，考虑以下代码：

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时间排序的汇总分析结果表，限制显示前8个条目 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，调用 vectorized_layer_norm_kernel 内核。

📝 注意

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

虽然PyTorch Profiler提供了快速的性能概述，但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）进行通信。每个桶进行全局归约的通信时间为：

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 参数的模型的通信模式：

全局归约操作的通信时间为：

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上。让我们分析通信模式：

让我们分析是否可以在下一层线性的计算中重叠一个层的全聚合通信。全聚合操作的通信时间为：

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

这个比率告诉我们是否可以成功地将全聚合通信隐藏在下一个线性层的计算背后。有趣的是，这个比率只依赖于隐藏大小 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}

而点对点传输的通信时间为：

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、下一个流水线阶段中层的数量，以及硬件的计算与点对点传输带宽能力之比。

引用

如果在学术场景引用本书：

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

BibTeX citation：

@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},
}