← GPU Fundamentals：从 SM 到 Roofline 🏠 回到交互版首页训练显存、混合精度与 LoRA →

第 16 章 · 第三部分 · 扩展 (Scale)

Distributed Training：从 DDP 到 3D 并行

AllReduce/Ring、ZeRO 三阶段、FSDP、张量并行、流水线并行、通信计算重叠与选型指南。

对应课件：LLM Systems（新增）

本章定位：Scale 主线的核心缺口。回答三个问题：① 一张卡为什么不够；② 多卡为什么还会慢；③ DP / TP / PP / ZeRO / FSDP 分别在什么场景用。本章自带 Ring AllReduce 动画与 ZeRO 显存对比两个交互演示。

📘 代码类型约定（Build / Systems 章节统一）

【Runnable】可直接运行（关键逻辑已在本课程验证脚本中实测）
【Skeleton】工程骨架：逻辑完整，需要自备数据/环境/权重
【Pseudo-code】算法示意：用于解释思路，不保证可直接运行

16.1 一张卡为什么不够

先懂直觉

两个瓶颈同时出现：显存（7B 模型训练要 100GB+，单卡 80GB 放不下）和时间（8.4e22 FLOPs 在单卡上要几年）。分布式训练就是「把模型和数据切开，摊到多卡上，同时让通信尽量不拖后腿」。

回顾第 17 章的显存账：混合精度 + Adam 的训练状态约 16~18 bytes/参数（两章口径差异见 16.5 节的说明框）。

模型	仅训练状态	单卡 80GB 能否放下
1B	~18 GB	✅
7B	~126 GB	❌
70B	~1.26 TB	❌（要 16 卡分片起步）

两条路：切数据（DP/ZeRO）或切模型（TP/PP）。真实训练是组合使用（3D 并行）。

16.2 Data Parallel 与 DDP

先懂直觉

最简单的并行：每张卡存一份完整模型，各喂不同的数据，算完各自的梯度后「求和平均」，再各自更新——这样所有卡永远保持同一份参数。同步梯度的操作就是 AllReduce。

GPU0: 模型副本 + batch_0 → 梯度 g0 ┐
GPU1: 模型副本 + batch_1 → 梯度 g1 ├─ AllReduce(平均) → ḡ → 每张卡各自 optimizer.step()
GPU2: 模型副本 + batch_2 → 梯度 g2 ┘

# PyTorch DDP：三行代码的变化
model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])
# 其余训练循环与单卡完全一致——DDP 在 backward 时自动触发梯度 AllReduce

数学上等价于大 batch

g_{DP} = \frac{1}{N}\sum_{i=1}^{N} g_i \quad\Longleftrightarrow\quad \text{用 batch size } N \times B \text{ 做单卡训练}

特点	说明
显存	每张卡都是完整模型（未省显存）
通信	每步一次梯度 AllReduce，通信量 ≈ 2×参数量
优点	实现简单、计算通信可重叠、扩展性好
缺点	模型必须能放进单卡

⚠️ 学习率要调整

DP 等价于 batch 变大 N 倍。经验规则：lr 随 batch 增大而适度放大（linear scaling 或 sqrt scaling），并配合 warmup（第 21 章）。

16.3 AllReduce 与 Ring AllReduce

先懂直觉

AllReduce = 所有卡把各自的梯度加起来，结果发回给所有卡。笨办法是「都发给 GPU0 让它算」（GPU0 带宽成瓶颈）；Ring AllReduce 让数据沿环流转，每张卡的通信量相同，与卡数几乎无关。

两步算法

① Reduce-Scatter（N−1 步）：每卡把自己的数据切成 N 块，
   沿环传递并累加，结束时每卡恰好拥有一块的完整和。
② All-Gather（N−1 步）：把完整块沿环传播，
   结束时所有卡拥有所有块的完整和。

通信量公式（先定义清楚再谈「≈2S」）

设要归约的数据总量为 $S$ 字节（如全部梯度），卡数 $N$ 。Ring AllReduce 分两阶段，每阶段 $N-1$ 步：

\text{每卡发送量} = \underbrace{\frac{N-1}{N}S}_{\text{Reduce-Scatter}} + \underbrace{\frac{N-1}{N}S}_{\text{All-Gather}} = 2\cdot\frac{N-1}{N}S \approx 2S

【假设条件】 $S$ 在两个阶段都被均分成 $N$ 块；忽略协议头与启动开销。

口径	数值	说明
每卡发送量（常用口径）	$\approx 2S$	与 N 无关（ $N\to\infty$ 时趋近 2S）
每卡发送 + 接收总量	$\approx 4S$	双向链路视角
全网总流量	$\approx 2(N-1)S$	所有卡都在发，随 N 线性增长

关键结论：单卡的通信量不随卡数增长（≈2S），所以 DDP 能扩展到成千上万卡；但整个网络的总流量随卡数线性增长——这是大规模训练集群网络的瓶颈所在。

🎮 交互：Ring AllReduce 分步动画交互演示 · 静态阅读版

点「下一步」看 4 张卡如何沿环传递并累加（Reduce-Scatter），再把完整结果传播回去（All-Gather）。方块颜色表示该块包含了几张卡的数据。

在交互版中打开这个演示 →

方案	每卡发送量	特点
朴素（全发 GPU0 汇总再广播）	O(N·S)	GPU0 带宽成为瓶颈
Ring AllReduce	$\approx 2S$ （与 N 无关）	带宽最优；延迟 $O(N)$ 步
Tree AllReduce	$\approx 2S$ （同量级）	延迟 $O(\log N)$ ，适合小消息/低延迟场景

小消息看延迟（步数），大消息看带宽（每卡流量）。这就是 NCCL 会按消息大小自动选择 Ring / Tree 等算法的原因。

📘 NCCL 与拓扑

实际使用的是 NVIDIA 的 NCCL 库，它会自动利用 NVLink / NVSwitch（卡间专用高速互联）和 InfiniBand（跨机）。拓扑意识：同一台机器内用 NVLink（数百 GB/s）远快于跨机网络（几十 GB/s）——这直接决定了并行策略的选择。

16.4 通信与计算重叠

先懂直觉

通信很慢，但如果让通信和计算同时进行，它就可以「免费」隐藏掉。DDP 的做法是把梯度按桶（bucket）分组：某层梯度算完就立即开始 AllReduce，同时反向传播继续算后面的层。

时间线（理想情况）：
反向:   [layer5] [layer4] [layer3] ...
通信:            [allreduce5] [allreduce4] ...
                 ↑ 与前面的计算重叠，几乎不占额外时间

技术	作用
Gradient Bucketing	小梯度合并成大消息（提高带宽利用率）
通信流（comm stream）	通信走独立 CUDA stream，与计算并行
`bucket_cap_mb` 参数	控制桶大小（越大越省启动开销，越小越早开始通信）

⚠️ 反例：为什么「多卡还会更慢」

如果通信无法重叠（如每步都做参数 AllGather 的 naive ZeRO-3 实现）、或跨机带宽太小，扩展效率会低于 1（8 卡只快 6 倍甚至更少）。判断标准是 扩展效率 = 实际加速比 / 卡数。

16.5 ZeRO：把训练状态切开

先懂直觉

DDP 的问题：每张卡都存了一份完整的「显存三件套」（参数、梯度、优化器状态）。但这些数据在每张卡上是重复的。ZeRO 的思想：既然重复，那就切开分给各卡，用的时候再 AllGather 回来。

三阶段（逐级把重复数据分片）

显存口径说明（重要）：下表按 16 bytes/参数（bf16 参数 2 + bf16 梯度 2 + fp32 主权重 4 + Adam m/v 8）计算，只含 model states（参数/梯度/优化器状态），不含 activations 与临时缓冲。第 17 章按「梯度也保留 FP32」的保守口径算成 18 bytes/参数。两者都对，取决于实现——关键是口径一致，不要混用。

阶段	分片什么	每卡 model states 显存（16 B/参数口径）	通信量（相对基线）
ZeRO-0（DDP）	无	16Ψ	1×（梯度 AllReduce）
ZeRO-1	优化器状态（m/v + fp32 主权重）	4Ψ + 12Ψ/N	≈1.5×
ZeRO-2	+ 梯度	2Ψ + 14Ψ/N	≈1×（与基线相当）
ZeRO-3	+ 参数	16Ψ/N	≈1.5×

通信量口径来自 ZeRO 论文的对照分析：ZeRO-2 用「梯度 Reduce-Scatter + 参数 All-Gather」替代「梯度 AllReduce」，总量与基线相同；ZeRO-1/3 因额外的参数 All-Gather 约为 1.5×。

🎮 交互：ZeRO 显存对比交互演示 · 静态阅读版

拖动参数量和 GPU 数，看四个阶段的单卡显存如何下降，以及能否放进 80GB 卡。

在交互版中打开这个演示 →

🧮 数字感受（7B 模型，8 卡）

阶段	单卡显存
ZeRO-0	112 GB（放不下）
ZeRO-1	38.5 GB
ZeRO-2	26.3 GB
ZeRO-3	14 GB

从「8 卡也放不下」到「8 卡轻松放下」——代价是通信量依次增加。

选择口诀

ZeRO-1 几乎总是划算（省显存、通信增加最少）；ZeRO-2 更省；ZeRO-3 最省但通信最多，要配合 prefetch 掩盖。

16.6 FSDP：ZeRO-3 思想的 PyTorch 实现

先懂直觉

FSDP（Fully Sharded Data Parallel）的思想与 ZeRO-3 相同（参数分片存、用的时候 AllGather、用完释放），是 PyTorch 原生实现，工程细节（分片方式、prefetch 调度、混合分片策略）与 DeepSpeed ZeRO-3 略有差异但目标一致：让你用「写 DDP 的方式」获得接近 ZeRO-3 的显存收益。

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(model, sharding_strategy=ShardingStrategy.FULL_SHARD)  # ≈ ZeRO-3
# 另有 SHARD_GRAD_OP（≈ZeRO-2）、HYBRID_SHARD（机内全切、跨机复制）等策略

FSDP 策略	概念上对应	场景
`FULL_SHARD`	ZeRO-3（全部切）	显存最紧
`SHARD_GRAD_OP`	ZeRO-2（切梯度与优化器状态）	通信与显存折中
`HYBRID_SHARD`	机内全切 + 机间复制	多机训练常用（省跨机通信）

16.7 Tensor Parallel：把单层切开

先懂直觉

DP 切数据、ZeRO 切状态，但单个矩阵乘还是太大或太慢怎么办？把矩阵本身切开：一张卡算左半列、一张卡算右半列，再拼起来。这就是张量并行（Megatron 的核心）。

两种切法（Megatron 的经典设计）

列并行（用于 QKV 投影、FFN 第一层）：

Y = XW, \quad W = [W_1, W_2] \implies Y_1 = XW_1, \ Y_2 = XW_2

输入 $X$ 每张卡都要（复制），输出按列切分；
不需要通信（各卡算自己的部分列）。

行并行（用于 Attention 输出投影、FFN 第二层）：

Y = XW, \quad X = [X_1, X_2], \ W = \begin{pmatrix}W_1\\W_2\end{pmatrix} \implies Y = X_1W_1 + X_2W_2

每张卡算部分和，最后 AllReduce 一次。

一层 Transformer 的通信次数（Megatron 经典实现）

QKV 投影（列并行，无通信）
→ Attention（本地计算）
→ 输出投影（行并行，1 次 AllReduce）
→ FFN 第一层（列并行，无通信）
→ FFN 第二层（行并行，1 次 AllReduce）
合计：每层前向 2 次 AllReduce，反向再 2 次

【假设条件】以上为 Megatron 经典（未开 Sequence Parallel） 的切分方式：列并行输出「切分的结果」，行并行通过 AllReduce 汇总。若开启 Sequence Parallel，通信会改为 All-Gather + Reduce-Scatter 的组合（总量相近，但更省激活显存）。

特点	说明
显存	参数/激活都按卡数切分
通信	每层都要通信（频繁），所以必须在机内 NVLink 上做
典型配置	TP ≤ 8（不超过单机 GPU 数）

⚠️ TP 不宜跨机

TP 每层通信 2 次，对带宽极其敏感。跨机 InfiniBand 只有 NVLink 的 1/10 带宽 → TP 通常限制在机内，跨机用 DP/PP。

16.8 Pipeline Parallel：把层切开

先懂直觉

按深度切：GPU0 放第 1-8 层、GPU1 放第 9-16 层……数据像流水线一样从前往后流。代价是「气泡」：流水线没填满时有的卡在空转。

micro-batch 1: GPU0[fwd] → GPU1[fwd] → GPU2[fwd] → GPU3[fwd]
micro-batch 2:            GPU0[fwd] → GPU1[fwd] → ...
（1F1B 调度：前向和反向交错，让气泡最小化）

气泡率 ≈ (P - 1) / (M + P - 1)

【假设条件】经典 1F1B 调度（非 interleaved）、每个 micro-batch 的前向/反向耗时近似相等、忽略阶段间通信时间。interleaved 调度（virtual pipeline）可以进一步减小气泡。
P = 流水段数，M = micro-batch 数量

特点	说明
显存	参数按层切分，激活可优化（只存阶段边界）
通信	只传阶段边界的激活（量小、可重叠）
关键	micro-batch 越多，气泡越小

🧮 气泡算例

4 段流水线（P=4），micro-batch M=32：气泡率 = 3/(32+3) ≈ 8.6%。 M=4 时：3/7 ≈ 43%（严重浪费）。所以 PP 必须配合较大的 M 使用。

16.9 Sequence / Context / Expert Parallel（了解即可）

并行方式	切什么	场景
Sequence Parallel	LayerNorm/Dropout 等逐元素算子的序列维	与 TP 配合（Megatron），省激活显存
Context Parallel	长序列的注意力上下文	超长上下文训练（如 128K+）
Expert Parallel	MoE 的不同专家放不同卡	MoE 模型（如 Mixtral）

16.10 3D 并行与选型指南

真实的大规模训练是「三维组合」：

总 GPU 数 = DP × TP × PP
例：LLaMA-3 405B ≈ 训练用 16K GPU（TP=8 机内 × PP=16 × DP=... 组合）

决策表（收藏）

症状	方案	为什么
单卡放得下模型，只想更快	DDP	最简单，通信可重叠
单卡放不下（显存不够）	ZeRO-1/2 或 FSDP	先切优化器状态与梯度
ZeRO-3 通信太重	TP=8（机内）+ FSDP	机内 NVLink 扛得住频繁通信
模型层数多、跨机想少通信	PP	只传边界激活，配大 M 减少气泡
想要极限规模	DP × TP × PP 组合	各取所长：TP 机内、PP 跨机少通信、DP 兜底

📘 从「能用」到「高效」的检查顺序

显存够不够（先解决能不能训）
扩展效率（8 卡是否接近 8 倍）
通信占比（profiler 里看 allreduce 时间占比，经验目标 < 15%，非硬标准）
MFU（第 14 章：目标 35%~50%）

🔑 本章必须记住

概念	一句话
DP/DDP	每卡完整模型 + 梯度 AllReduce，等价大 batch
Ring AllReduce	每卡通信 ≈ 2S，与卡数几乎无关
通信重叠	梯度分桶 + 独立通信流
ZeRO-1/2/3	依次切优化器状态 / 梯度 / 参数
FSDP	PyTorch 版 ZeRO-3，策略可选
TP	切矩阵（列并行无通信+行并行 AllReduce），限机内
PP	切层，气泡 ≈ (P−1)/(M+P−1)，靠大 M 缓解
3D	DP × TP × PP 组合，按瓶颈选型

📝 小测验

DDP 训练中每步的梯度 AllReduce，其作用是？

A同步数据加载

B对所有卡的梯度求平均，使各卡参数保持一致（等价于大 batch 训练）

C压缩参数

D加速前向计算

显示答案与解析

正确答案：B

各卡用不同数据算出梯度后 AllReduce 求平均，再各自更新，因此所有副本始终保持相同参数，数学上等价于用 N×B 的大 batch 训练单卡。

📝 小测验

ZeRO-3 相比 ZeRO-2 多切分了什么？代价是什么？

A多切梯度，代价是显存

B多切参数，代价是前向/反向都要 AllGather 参数，通信量增加

C多切优化器状态，代价是精度

D无区别

显示答案与解析

正确答案：B

ZeRO-3（FSDP 的 FULL_SHARD）连参数都分片，单卡显存降到 16Ψ/N，但每次前向/反向用到哪层就要把该层参数 AllGather 回来，通信量最大。

📝 小测验

为什么张量并行（TP）通常限制在单机内（≤8 卡）？

A因为只有 8 张卡的驱动支持

BTP 每层都要通信（前向 2 次 AllReduce），跨机带宽远低于机内 NVLink，通信会成为瓶颈

C因为参数量太小

D因为跨机不支持张量运算

显示答案与解析

正确答案：B

TP 的通信频率高（每层），适合高带宽的 NVLink 机内互联；跨机要做就用流水线并行（通信少）或 DP。

📝 小测验

流水线并行的「气泡」如何计算和缓解？

A气泡 = P，无法缓解

B气泡率 ≈ (P−1)/(M+P−1)，通过增加 micro-batch 数量 M 来减小

C气泡 = 0，因为前向反向并行

D气泡与 micro-batch 无关

显示答案与解析

正确答案：B

P 个流水段需要 P−1 步填充/排空，micro-batch 越多，填充成本占比越小。4 段 + 32 microbatch 时气泡约 8.6%。

← GPU Fundamentals：从 SM 到 Roofline 🏠 回到交互版首页训练显存、混合精度与 LoRA →