数据并行、ZeRO 与 FSDP

数据并行（data parallelism，DP）是多卡训练里最先学到、也最常用的一维并行：每张卡放一份完整模型，各自吃不同的数据，反向之后把梯度平均一下。听起来一句话就讲完了，但真要把它用好，有一串问题要回答：梯度同步能不能和反向计算同时进行？同步应该按参数一个个发还是攒成一批？梯度累积的时候要不要每次都同步？卡多到什么程度会被网络卡住？每张卡都存一份完整的优化器状态是不是太浪费？

最后一个问题的答案是 ZeRO 和 FSDP：它们在不增加多少通信的前提下，把每张卡上的冗余状态去掉。用到的 AllReduce、ReduceScatter、AllGather 及其代价公式，推导见 集合通信与分片矩阵乘法。

1. 先把一步训练拆开

单卡上，一个训练步（step）分三段：

1. 前向：输入一批样本，逐层计算，得到 loss，同时把反向需要的中间结果（激活）存下来；
2. 反向：从最后一层往前，逐层算出 loss 对激活和对参数的梯度；
3. 优化器步：用梯度更新参数，例如 Adam 要先更新一阶、二阶动量，再更新参数。

记模型参数量为 \(\Psi\)。在常见的 bf16 混合精度 + Adam 训练里，每个参数要存：bf16 参数 2 字节、bf16 梯度 2 字节、fp32 主参数 4 字节、fp32 一阶动量 4 字节、fp32 二阶动量 4 字节，一共 16 字节；如果梯度也在 fp32 里累积，再加 4 字节。这些和激活一起就是训练时的显存大头，细节见 训练的账本。

另一个要用到的概念是梯度累积（gradient accumulation）。一次前向+反向处理的样本数叫 micro-batch size（mbs），两次优化器步之间总共处理的样本数叫 global batch size（gbs）。如果一张卡一次只能放下 mbs 个样本，但我们想要更大的 gbs，就连续做 grad_acc 次前向+反向，把梯度累加起来（实践中取平均，好让结果与累积步数无关），再做一次优化器步：

梯度累积用时间换显存：激活只需要放一个 micro-batch 的量，但要多跑几次前向反向，而且需要一块在整个训练步内常驻的梯度累加缓冲。注意这几次前向反向之间互不依赖，只是输入样本不同——既然如此，为什么不把它们放到不同的卡上同时算？这就是数据并行。

2. 朴素数据并行

2.1 做法

有 \(N_d\) 张卡（DP degree，也直接叫 DP size）。每张卡上放一份完全相同的模型副本（"model instance"），每个训练步：

1. 把这一步的 batch 切成 \(N_d\) 份，每张卡拿一份 micro-batch；
2. 每张卡独立地做前向和反向，得到各自不同的梯度 \(g_r\)；
3. 对梯度做一次 AllReduce（求和再除以 \(N_d\)，即求平均），之后每张卡上的梯度都相同；
4. 每张卡用相同的梯度做相同的优化器步，模型副本保持一致。

2.2 为什么结果和单卡大 batch 一样

设整个 batch 有 \(B\) 个样本，loss 是平均值 \(L = \frac{1}{B}\sum_{i=1}^{B}\ell_i\)。把样本均分到 \(N_d\) 张卡，第 \(r\) 张卡上的局部 loss 是 \(L_r = \frac{N_d}{B}\sum_{i\in S_r}\ell_i\)，于是

每张卡算自己的 \(\nabla L_r\)，平均一下，就精确等于大 batch 的梯度。这里假设每张卡上的样本数相同；如果 loss 按 token 平均而各卡 token 数不同，需要按 token 数加权，否则会有偏差。所以数据并行在数学上就是"并行版的梯度累积"，不改变优化过程，只改变墙钟时间。

2.3 用分片记号写出来

为了把通信算清楚，我们把 Transformer 近似成一串 MLP 块（大模型里注意力只占 FLOPs 的一小部分；门控矩阵的行为和 \(W_\text{in}\) 一样，先忽略），每层两个矩阵 \(W_\text{in}[D, F]\) 和 \(W_\text{out}[F, D]\)，输入 \(\text{In}[B, D]\)，这里 \(B\) 是这一步的总 token 数。数据并行就是沿 mesh 轴 X（参与数据并行的那组卡）切 batch 维、权重完全复制：

前向里收缩维 \(D\)、\(F\) 都没被切，是上一篇里的"情形 1"，完全不需要通信。反向时：

dOut[B_X, D]      = ...
dW_out[F, D]{U_X} = Tmp[B_X, F] ·_B dOut[B_X, D]    # 沿被切开的 B 收缩：部分和
dW_out[F, D]      = AllReduce(dW_out[F, D]{U_X})    # 不在关键路径上，可异步
dTmp[B_X, F]      = dOut[B_X, D] ·_D W_out[F, D]
dW_in[D, F]{U_X}  = In[B_X, D] ·_B dTmp[B_X, F]
dW_in[D, F]       = AllReduce(dW_in[D, F]{U_X})     # 同上
dIn[B_X, D]       = dTmp[B_X, F] ·_F W_in[D, F]     # 传给前一层

权重梯度是沿 batch 维收缩出来的，而 batch 维被切开了，所以每张卡只得到一个"沿 X 未归约"的部分和 \(\{U_X\}\)，需要 AllReduce，这就是上一篇的"情形 3"。关键的一点：这些 AllReduce 不在关键路径上。算前一层的 \(d\text{In}\) 只需要本地的 \(d\text{Tmp}\) 和 \(W\)，不需要等 \(dW\) 同步完成。只要总通信时间不超过总计算时间，就可以完全藏在计算后面。张量并行就没有这个好性质。

2.4 一次 AllReduce 要多久

在 \(N\) 个 rank 上用环形算法 AllReduce 一个 \(V\) 字节的数组，每个 rank 收发 \(2\frac{N-1}{N}V\) 字节，带宽受限时

其中 \(W\) 是每张卡（或每个节点）的单向出口带宽。H100 节点内 NVLink 是每卡 450 GB/s，跨节点每个节点出口是 8×400 Gbps = 400 GB/s。拿一个 8B 参数的模型来说，bf16 梯度 \(V = 16\) GB：节点内 8 卡约 \(2 \times \frac78 \times 16\text{e}9 / 450\text{e}9 \approx 62\) ms，跨节点约 \(2 \times 16\text{e}9/400\text{e}9 = 80\) ms（实际带宽通常只能跑到标称值的七八成）。藏不藏得住，取决于它和反向计算时间的比（第 7 节）。

2.5 朴素实现的问题

最朴素的写法是等反向全部结束，拿到所有梯度，再发起一次 AllReduce，结束后做优化器步。这样计算和通信是串行的：AllReduce 的那几十毫秒里 GPU 的计算单元完全闲着，而且每一步都要白等这么久。DDP 用三个办法解决这件事：重叠、分桶、与梯度累积配合。

3. 优化一：梯度同步与反向重叠

反向是从最后一层往前算的。最后一层的梯度一算完，它就已经是最终值了，后面再怎么算前面的层也不会改变它。所以没必要等所有层算完：某一层的梯度一就绪，就可以立刻开始对它做 AllReduce，同时 GPU 继续算更前面的层。

PyTorch 里的实现方式是给每个参数挂一个 hook，梯度累加完就触发：

def register_backward_hook(self, hook):
    """给所有需要梯度的参数注册 hook。"""
    for p in self.module.parameters():
        if p.requires_grad is True:
            p.register_post_accumulate_grad_hook(hook)

# hook 里对 p.grad 发起异步 AllReduce（async_op=True），
# 把返回的 work handle 存起来；优化器步之前统一 wait()。

通信 kernel 跑在单独的 CUDA stream 上，与计算 kernel 并发执行。"计算与通信重叠"这个思路，在张量并行、流水线并行、上下文并行里还会反复出现。

4. 优化二：梯度分桶

按参数逐个 AllReduce 有个问题：一个模型有成百上千个参数张量，很多很小（LayerNorm 的权重、偏置），每次集合通信都有固定开销——kernel 启动、同步、环上每一跳的延迟。上一篇的 α-β 模型里，环形 AllReduce 的时间是

消息越小，延迟项占比越大；两项相等时 \(V^\ast = N\alpha W\)。量级上，8×H100 节点内 AllReduce 一个很小的张量，实测要 16 µs 左右；这段时间够 NVLink 传 7 MB 的数据。一个 4 kB 的 LayerNorm 权重单独发一次，几乎全是固定开销。再加上 GPU 上无论计算还是通信，大张量都比很多小张量高效。所以做法是：把若干个参数的梯度攒到一个"桶"（bucket）里，桶满了才发起一次 AllReduce。就像寄东西，几个大箱子比一堆小包裹省事。PyTorch DDP 默认的桶大小是 25 MB（bucket_cap_mb）。

桶也不是越大越好：桶越大，第一个桶要等越久才凑满，而且最后一个桶是在反向完全结束后才凑满的，它的通信一定露在外面。极端情况下桶 = 整个模型，就退化成了朴素的串行实现。

拖一拖可以看到几种状态：桶很小时（几十 kB），桶的个数成千上万，固定开销累积起来让通信总时长远超带宽项；桶很大时（几 GB），只有一两个桶，最后一个桶整个露在外面；中间有一段宽的"甜区"，通信几乎全部藏进反向里，只剩最后一个小桶的尾巴。再把每卡 token 数调小，反向计算变短而通信量不变（梯度大小只取决于参数量），无论桶怎么分都藏不住了——这就是通信瓶颈（第 7 节）。

5. 优化三：与梯度累积配合

数据并行和梯度累积可以叠加使用：每张卡连续做几次 micro-batch 的前向反向，累加梯度，最后做一次同步和优化器步。问题是第 3 节的 hook 会在每一次反向时都触发 AllReduce。中间那几次同步完全是浪费：最终需要的只是累加完之后的梯度和，在最后一次反向时同步一次效果相同，开销却只有 1/grad_acc。

PyTorch 的做法是用 model.no_sync() 上下文管理器在不需要同步的反向里关掉 hook：

for i, batch in enumerate(micro_batches):
    last = (i == len(micro_batches) - 1)
    ctx = contextlib.nullcontext() if last else model.no_sync()
    with ctx:
        loss = model(batch).loss / len(micro_batches)
        loss.backward()        # 只有最后一次触发梯度 AllReduce
optimizer.step()
optimizer.zero_grad()

6. 重新计算全局 batch size

加上数据并行，全局 batch size 变成三个因子的乘积：

在 LLM 里我们更常用 token 数来计 batch：\(\text{gbs}_\text{tokens} = \text{gbs} \times \text{seq\_len}\)。给定一个目标 gbs，grad_acc 和 dp 可以互换：dp 翻倍、grad_acc 减半，gbs 不变。两者的区别在于数据并行是真并行，梯度累积是串行的，所以实践中尽量把 dp 开大，只有在卡不够用、dp 达不到目标 gbs 时才用梯度累积补齐。

6.1 一个配置流程

实际搭一个只用数据并行的训练，可以按这个顺序定参数：

1. 先确定最优的全局 batch size（以 token 计）：查文献，或者跑实验看收敛；
2. 再确定训练序列长度：同样查文献或做实验。主预训练阶段一般取 4k–8k token（DeepSeek-V3 用 4K，Llama 3 405B 绝大部分时间用 8K），长上下文通常在训练末期再扩：网上比这更长的文档本来就很少，而注意力的开销随序列长度增长；
3. 于是知道了 gbs（样本数）。在单卡上不断增大 mbs 直到显存放不下，得到最大 mbs；
4. 根据可用卡数确定 dp，剩下的用 \(\text{grad\_acc} = \text{gbs}/(\text{mbs}\times\text{dp})\) 补齐。

举个例子：想要 gbs = 4M token、序列长度 4k，那么 gbs ≈ 1024 个样本（取最近的 2 的幂）。假设单卡只放得下 mbs = 2，手上有 128 张卡，那么 \(\text{grad\_acc} = 1024 / (2 \times 128) = 4\)。如果突然有了 512 张卡，就可以 mbs = 2、grad_acc = 1，同样的 gbs，一步快 4 倍。

如果算出来 grad_acc 小于 1，说明卡太多了（GPU-rich）：可以不用完所有卡、试更大的 gbs，或者调小 mbs——后者以单卡效率为代价换总吞吐。

6.2 数据并行的天花板

dp 也不能无限加。环形 AllReduce 的延迟项随卡数线性增长，到了几百张卡（视网络而定，常见的量级是 512 张左右），信号绕环一圈的时间已经不可忽略，DP 通信不再能被完全重叠，吞吐开始下降。实测也是这样：超过某个规模后，再加 DP rank，每卡吞吐明显下降，而每卡显存一点不变。显存不变，是因为每张卡仍然存着完整的模型、梯度和优化器状态。

7. 什么时候被通信卡住

7.1 推导

继续用 2.3 节的 MLP 模型。设每张卡算力 \(C\)（FLOP/s），集合通信带宽 \(W\)，DP 轴上 \(X\) 张卡，这一步全局 \(B\) 个 token。只看反向（前向没有通信）。

计算时间：每层反向有 4 个矩阵乘（每个矩阵的权重梯度和输入梯度各一个），每个 \(2 \cdot \frac{B}{X} \cdot D \cdot F\) FLOPs：

通信时间：每层两个权重矩阵，各 \(DF\) 个参数、bf16 下 \(2DF\) 字节，各做一次 AllReduce（代价是 2 × 字节数 / 带宽）：

两者能重叠时，每层耗时 \(\approx \max(T_\text{math}, T_\text{comms}) = 8DF\cdot\max\!\big(\frac{B}{XC}, \frac{1}{W}\big)\)。计算受限的条件是 \(T_\text{math} \gt T_\text{comms}\)：

左边是每张卡的 token 数，右边是硬件的"网络算术强度"——每通信一个字节，卡能算多少 FLOPs。\(D\)、\(F\) 都约掉了，模型多大都一样。直观上：计算量随每卡 batch 线性增长，通信量只取决于参数量、与 batch 无关；每卡 batch 太小，计算就盖不住通信。这和单卡 roofline 里"batch 要大于某个阈值才能计算受限"是同一个道理（见 Roofline 一篇）。

7.2 代入硬件

跨节点时，AllReduce 分两段：先在节点内走 NVLink，再在节点间走网卡，耗时取两段中慢的那段。H100 节点内每卡 450 GB/s，节点间每节点 400 GB/s，瓶颈在节点间。H800 的 NVLink 只有每卡单向 200 GB/s，节点内那段（\(V\cdot\frac78 / 200\text{e}9 = V/229\text{e}9\)）反而比节点间（\(V/400\text{e}9\)）慢，所以表里用 229 GB/s。拿几次公开的训练来对照：

• Llama 3 405B 用了约 16K 张 H100，按 2475 token/卡，batch 至少要 40M token 才不被 DP 通信卡住，而它最大的 batch 是 16M——在这么多卡上只靠纯数据并行，做不到计算受限。
• DeepSeek-V3 在 2048 张 H800 上训练，batch 从 3072 条 4K 序列逐步加到 15360 条（约 12.6M → 63M token）。按 4300 token/卡算，2048 张卡需要约 8.8M token；如果它是稠密模型，这个 batch 已经够大。但它是 MoE，门槛还要再乘一个因子，见 7.3 节。
• TPU 上的门槛更低：TPU v5p 三个轴全做 DP，每卡 850 token，一个 8960 芯片的 pod 只要 7.6M token 的 batch。

7.3 两个修正

DP 很小时：环形 AllReduce 的精确代价带一个 \(\frac{X-1}{X}\) 因子，\(X=2\) 时只有一半。把它带进去，门槛变成 \(\frac{B}{X} \gt \frac{C}{W}\cdot\frac{X-1}{X}\)。跨节点时这个因子按节点数算：只跨 2 个节点，节点间那段只要 \(V/800\text{e}9\)，瓶颈变回节点内的 NVLink（\(V\cdot\frac78/450\text{e}9 = V/514\text{e}9\)），H100 的门槛从 2475 降到约 \(9.9\text{e}14/5.14\text{e}11 \approx 1900\) token/卡。

MoE：每个 token 只激活 \(E\) 个 expert 里的 \(k\) 个，计算按 \(k\) 算，但每层要同步全部 \(E\) 个 expert 的梯度：

门槛放大 \(E/k\) 倍。gpt-oss-120b 是 \(E=128,\ k=4\)，在 H100 上跨节点需要 \(32 \times 2475 \approx 79{,}000\) token/卡；DeepSeek-V3 是 256 个 routed expert 选 8 个，同样是 32 倍，在 H800 上要 \(32 \times 4300 \approx 138{,}000\) token/卡，而它实际每卡只有约 31K token。所以 MoE 不能只靠数据并行同步 expert 的梯度，这是它要配合专家并行的原因之一。

7.4 换成整个模型来算

同一件事也可以不拆到每层，直接用整个模型来写：梯度总量 ≈ 参数量，反向计算 \(t_\text{compute} = 4\cdot\text{tokens}\cdot\Psi / \text{peak\_flops}\)，按 bf16 每参数 2 字节，AllReduce 全部梯度 \(t_\text{comm} = 2\Psi \cdot 2\frac{N_d-1}{N_d} / \text{peak\_bw}\)，于是

这里 tokens 指每卡的 token 数，和 7.1 节（带上 7.3 节的 \(\frac{N_d-1}{N_d}\) 修正）完全一致：参数量约掉了，只剩"每卡 token 数要大于 \(C/W\)"。

7.5 这和"临界 batch size"有什么关系

上面的条件有一个反直觉的推论：卡数固定时，总 batch 越小，越容易被通信卡住。反过来，只要能不断加大 batch，DP / FSDP 就能扩到任意多的卡。问题在于 batch 不能无限加大：batch 大到一定程度，梯度噪声已经很小，再加大收益递减（这个拐点常被称为临界 batch size），有时还会不稳定。

所以在算力充足时，真实的问题往往是反过来的：batch size 由 scaling law 定死了，卡数也定了（很多），要找一种切分方式，让这么"小"的 batch 也能铺满这么多卡。举两个例子：DeepSeek-V2 的 batch 最大是 9216 条 4K 序列，约 38M token，在 TPU v5p 上三个轴都做 DP（850 token/卡），大约能铺到 44,000 张芯片；一个 16M token 的 batch（Llama 3 405B 最后阶段的大小）可以铺到约 18,800 张。再往上就需要换一种并行维度（张量并行、流水线并行），它们不靠 batch 维来分摊。

8. 数据并行的显存账

数据并行把激活按 batch 分摊到了各卡（每卡只存自己 micro-batch 的激活），这对激活占大头的训练很有帮助。但它对参数、梯度、优化器状态完全没有帮助：每张卡都存一整份。按第 1 节的计数，bf16 混合精度 + Adam，每参数 16 字节（不算 fp32 梯度累积）：

记优化器相关的那部分为 \(k\Psi\)，Adam 下 \(k = 12\)。一个 7.5B 的模型就是 \(16 \times 7.5 = 120\) GB，单张 80 GB 的卡放不下，这还没算任何激活。一个好记的速算：光是 bf16 参数就是参数量 × 2 字节，70B → 140 GB。

就算省掉 fp32 主参数和梯度缓冲，只存 bf16 参数加 fp32 的 Adam 一阶、二阶动量（每参数 10 字节），纯数据并行能训练的最大模型也只有 HBM / 10：96 GB 的 TPU v5p 大约 9B。而这还是 batch 只有 1 个 token、不存任何激活时的绝对上界，实际远到不了。结论是：纯数据并行只适合几 B 以下的模型。要训更大的模型，至少要把参数或优化器状态切开。

切开有两条正交的路：一是"并行"——张量、上下文、流水线并行，把计算本身拆开；二是"分片"（sharding）——DeepSpeed ZeRO 或 PyTorch FSDP，计算仍按数据并行做，只是把各种状态分片存放、用时再取。两者可以组合。

9. ZeRO：把冗余去掉

ZeRO（Zero Redundancy Optimizer）的观察是：数据并行里每个 rank 存着完全相同的参数、梯度和优化器状态，这是纯粹的冗余。把它们沿 DP 维切成 \(N_d\) 份，每个 rank 只存 \(1/N_d\)，需要完整数据时再通过通信临时拼回来。按切的范围分三级：

• ZeRO-1：切优化器状态；
• ZeRO-2：再切梯度；
• ZeRO-3：再切参数（PyTorch 的原生实现叫 FSDP）。

激活为什么不在列表里？每个 DP rank 吃不同的 micro-batch，激活本来就各不相同，不存在冗余，自然没什么可切的。

每张卡的显存（不含激活）：

如果在 ZeRO-2 / ZeRO-3 里用 fp32 做梯度累积，梯度项再加 \(\frac{4\Psi}{N_d}\)（在 DP 和 ZeRO-1 里是完整的 \(4\Psi\)）。通信量这一列按每卡收发的元素数计：一次 ReduceScatter 或 AllGather 算 \(\Psi\)，一次 AllReduce 算 \(2\Psi\)（都省略了 \(\frac{N-1}{N}\) 因子）。

9.1 ZeRO-1：切优化器状态

普通 DP 里，所有 rank 在梯度 AllReduce 之后拿着相同的梯度，各自做一模一样的优化器步——同样的计算做了 \(N_d\) 遍。ZeRO-1 让第 \(r\) 个 rank 只负责 \(1/N_d\) 的参数：只存这部分参数的 fp32 主参数和 Adam 动量，也只更新这部分。一个训练步变成：

1. 前向：每个 rank 用同一份完整的 bf16 参数，算不同的 micro-batch；
2. 反向：每个 rank 得到完整的（但各不相同的）梯度；
3. 对梯度做 ReduceScatter：rank \(r\) 只拿到自己负责那 \(1/N_d\) 的梯度之和；
4. 每个 rank 用本地的优化器状态更新自己那 \(1/N_d\) 的 fp32 参数，再转成 bf16；
5. 对 bf16 参数做 AllGather，把别人更新好的那些片段拿回来，每个 rank 又有了完整的新参数。

显存从 \(16\Psi\) 降到 \(4\Psi + 12\Psi/N_d\)。通信呢？普通 DP 是一次 AllReduce；ZeRO-1 是一次 ReduceScatter 加一次 AllGather——而 AllReduce 本来就等于 ReduceScatter + AllGather。所以总通信量不变，只是 AllGather 从"梯度"挪到了"更新后的参数"上。ReduceScatter 本身只有 AllReduce 一半的代价。

新出现的 AllGather 也要想办法藏起来。两种思路：

• 在优化器步里重叠：第一片参数一更新完就开始 AllGather，同时更新后面的片；
• 在前向里重叠：下一步前向用到第 \(\ell\) 层时才需要它的完整参数，所以每层的 AllGather 可以和前面层的前向计算重叠。

这两种都需要精细地使用 hook 和分桶，自己实现不容易。实践中更省事的做法是直接用 PyTorch 原生的 FSDP，把 FSDP 单元设成整个模型（见 10.3 节），就能得到类似的效果。

9.2 ZeRO-2：再切梯度

ZeRO-1 里每个 rank 只用自己那 \(1/N_d\) 的梯度去更新，其他部分的梯度存着也没用。那就别存：反向时不做 AllReduce，直接对梯度做 ReduceScatter，每个 rank 只保留自己那份，其余的发出去后立刻释放。

显存变成 \(2\Psi + \frac{(2+k)\Psi}{N_d}\)。\(N_d\) 很大时趋于 \(2\Psi\)，相对基线 \(16\Psi\) 最多省 8 倍。通信和 ZeRO-1 完全一样（梯度 RS + 参数 AG），也就和普通 DP 一样，区别只是边通信边释放内存。如果用 fp32 做梯度累积，只需要为自己那 \(1/N_d\) 保留 fp32 梯度，用来累加 ReduceScatter 送来的 bf16 梯度。

ZeRO-2 相对 ZeRO-1 基本没有额外开销，只是实现更复杂一些，所以通常 ZeRO-2 是更好的选择。

9.3 ZeRO-3 / FSDP：连参数也切

最后一步是把 bf16 参数也切了，每个 rank 只常驻 \(1/N_d\) 的参数。前向和反向需要某一层的完整参数时，临时 AllGather 过来，用完立刻丢掉：

• 前向：逐层 AllGather 参数 → 计算 → 释放参数；
• 反向：反着来，逐层再 AllGather 一次参数（前向时已经扔掉了）→ 算梯度 → 对这层梯度做 ReduceScatter → 释放参数和非本 rank 的梯度。

通信量：前向 AllGather 全部参数一遍，\(\Psi\)；反向再 AllGather 一遍，\(\Psi\)；梯度 ReduceScatter，\(\Psi\)。合计 \(3\Psi\)，是 DP / ZeRO-2 的 \(2\Psi\) 的 1.5 倍。次数上，ZeRO-3 每一步要做 \(2L - 1\) 次逐层的 AllGather（\(L\) 为层数）：前向每层一次；反向本来也是每层一次，但反向最先用到的是最后一层，它的参数刚在前向里用完，实现上会留着不释放、直接复用，于是反向只需 \(L-1\) 次。每次 AllGather 都带一段固定延迟，层数多、每层小的时候这笔开销不能忽略。

预取：多出来的通信听起来很多，但它可以藏起来。前向算第 \(n\) 层时，就开始 AllGather 第 \(n+1\) 层的参数；反向算第 \(n\) 层时，AllGather 第 \(n-1\) 层。只要单层的计算时间大于单层参数的 AllGather 时间，就藏得住（10.2 节会算这个条件）。和普通 DP 一样，DP 规模到 512 左右以后，延迟开始让重叠失效。

显存：\(\frac{(2+2+k)\Psi}{N_d}\)，原则上 \(N_d\) 越大越省，可以无限降下去——但只限于模型相关的状态。外加两层临时的完整参数（正在用的和正在预取的）。激活一点没减少，那要靠激活重计算和梯度累积。

9.4 用计算器验证

ZeRO 论文给过一个表：\(\Psi = 7.5\text{B}\)、\(k = 12\)、\(N_d = 64\) 时，每卡显存基线 120 GB，ZeRO-1 31.4 GB，ZeRO-2 16.6 GB，ZeRO-3 1.88 GB。计算器的默认值就是这组参数。

拖动时注意几件事：ZeRO-1、ZeRO-2 在 \(N_d\) 增大时很快饱和在 \(4\Psi\)、\(2\Psi\)，因为它们没切的部分不随 \(N_d\) 下降；只有 ZeRO-3 一直按 \(1/N_d\) 下降。70B 模型在 64 卡上 ZeRO-2 也要 140 GB 以上，必须上 ZeRO-3 或者别的并行。

10. 用分片记号看 FSDP

用上一篇的分片记号，FSDP 就是：权重沿同一个 DP 轴 X 切开（切的是 \(D\) 维），激活照旧沿 batch 维切开：

这时 \(\text{In}[B_X, D]\cdot W_\text{in}[D_X, F]\) 是"一边的收缩维被切"（情形 2），而且 \(B_X\) 和 \(D_X\) 用的是同一个轴，所以必须先把权重 AllGather 回来。完整算法：

# 前向
W_in[D, F]   = AllGather(W_in[D_X, F])     # 不在关键路径：可在上一层计算时进行
Tmp[B_X, F]  = In[B_X, D] ·_D W_in[D, F]    # 之后可以扔掉 W_in[D, F]
W_out[F, D]  = AllGather(W_out[F, D_X])    # 同上
Out[B_X, D]  = Tmp[B_X, F] ·_F W_out[F, D]

# 反向
dW_out[F, D]{U_X} = Tmp[B_X, F] ·_B dOut[B_X, D]
dW_out[F, D_X]    = ReduceScatter(dW_out[F, D]{U_X})   # 可异步
W_out[F, D]       = AllGather(W_out[F, D_X])           # 可提前做
dTmp[B_X, F]      = dOut[B_X, D] ·_D W_out[F, D]        # 之后扔掉 W_out[F, D]
dW_in[D, F]{U_X}  = In[B_X, D] ·_B dTmp[B_X, F]
dW_in[D_X, F]     = ReduceScatter(dW_in[D, F]{U_X})    # 可异步
W_in[D, F]        = AllGather(W_in[D_X, F])            # 可提前做
dIn[B_X, D]       = dTmp[B_X, F] ·_F W_in[D, F]         # 传给前一层

和 2.3 节的纯 DP 对比，唯一的变化是：反向里的 AllReduce 被拆成了 ReduceScatter（立即做，得到切开的梯度）和 AllGather（推迟到下一次用权重之前做）。上一篇证明过 AllGather 的转置就是 ReduceScatter，所以前向 AllGather 权重、反向 ReduceScatter 梯度，是求导的必然结果。

ZeRO 名字的来历也在这里：不做任何多余的计算或存储。普通 DP 里每张卡 AllReduce 全部梯度、更新全部优化器状态、更新全部参数，都是重复劳动；ZeRO 分片后，ReduceScatter 梯度，只更新自己那片优化器状态和参数，用时再 AllGather 参数。

10.1 FSDP 的通信瓶颈和 DP 一样

前向和反向的计算 / 通信比例相同，看前向即可（前向 2 个矩阵乘，各一次 AllGather，AllGather 代价是 AllReduce 的一半）：

门槛与纯 DP 完全相同。也就是说，只要每卡 batch 大到纯 DP 能计算受限，就可以无代价地升级到 FSDP，省下大量参数和优化器状态的显存。前向确实多了通信，但它与前向的计算按同样比例增长，不改变 roofline。总结一句：DP 和 FSDP 都在每卡 batch 小于 \(C/W\) 时变成带宽受限（TPU v5p 单轴 2550，用几个 mesh 轴做 DP 就除以几；H100 跨节点约 2475）。

10.2 逐层看预取藏不藏得住

预取能把 AllGather 藏起来的条件是单层计算比单层通信长。每层 \(16h^2\) 个参数（bf16 下 \(32h^2\) 字节），AllGather 一层参数需要 \(t_\text{comm} = 32h^2\cdot\frac{N_d-1}{N_d}/\text{peak\_bw}\)；这一层前向计算 \(t_\text{compute} = 2\cdot\text{seq}\cdot\text{mbs}\cdot 16h^2/\text{peak\_flops}\)。比值：

\(\text{seq}\cdot\text{mbs}\) 就是每卡 token 数，结论还是"每卡 token 数要大于 \(C/W\)"。H100 节点内 \(C/W \approx 2200\)：seq = 4096、mbs = 1 就够；但如果为了省激活把 mbs 压到很小、序列又短，FSDP 的预取就藏不住了。

10.3 FSDP 单元

FSDP 不是逐个参数做 AllGather 的，而是把模型分成若干"单元"（通常一个 Transformer 块一个单元），以单元为粒度 gather、计算、释放、预取——和第 4 节的分桶是同一个权衡：单元太小，通信次数多、固定开销大；单元太大，临时的完整参数占显存多，而且第一个单元的 gather 无法与计算重叠。把单元设成整个模型，就退化成在前向开始前一次性 gather 全部参数，参数在前向之前整份 gather 回来、一直用到反向结束，梯度在反向结束后一次性 ReduceScatter；常驻显存的只有分片后的梯度和优化器状态，效果上就是 ZeRO-1 / 2。

11. 小结与局限

要点：

• 数据并行 = 并行版梯度累积，数学上与大 batch 等价；
• 梯度 AllReduce 不在关键路径上，用 hook 与反向重叠，用分桶摊薄固定开销，用 no_sync 避免梯度累积时的重复同步；
• \(\text{gbs} = \text{mbs}\times\text{grad\_acc}\times\text{dp}\)，优先加 dp；
• 每卡 token 数低于 \(C/W\)（H100 约 2200–2500）时被通信卡住；batch 又不能无限加大，这限制了纯 DP 能用的卡数；
• ZeRO 用 RS + AG 代替 AllReduce，通信量不变（ZeRO-1/2）或只多 50%（ZeRO-3），显存可以降到 \(1/N_d\)。

但 ZeRO 有两个绕不过去的局限：

1. 激活不能切。激活随序列长度和 batch 增长，ZeRO 一点也帮不上忙，只能靠激活重计算和梯度累积。我们并不想因为硬件限制就只能训短序列。
2. 每层必须放得下。ZeRO-3 每次仍要把一整层的参数 gather 到一张卡上计算；而且计算仍然沿 batch 维分摊，batch 固定时卡数一多就会通信受限。

张量并行换了一个方向：沿模型的隐藏维切矩阵乘本身，参数、梯度、优化器状态和激活都被切开，而且完全不需要在卡之间传参数。

参考与延伸阅读

• Nanotron, The Ultra-Scale Playbook：Data Parallelism（三个优化、Revisiting global batch size、Our journey up to now）和 Zero Redundancy Optimizer（ZeRO-1/2/3）两节，附录 A2 Typical Scales（每层 \(16h^2\) 的估算）和 A3 Math for Compute/Communication Overlap（7.4、10.2 节的整模型 / 单层写法）。DP 扩展性的基准测试也在这里。配套的 Picotron 仓库有带重叠和分桶的 DP 实现，适合对照代码读。
• Austin et al., How to Scale Your Model：Part 5 · Data Parallelism、FSDP、Takeaways（2.3、7.1、10 节的分片记号推导与 TPU 门槛）；Part 12 · GPU 上的数据并行 roofline（H100 门槛与 MoE 修正）。
• Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models：显存表（9.4 节的对照数字）与 3Ψ 通信量。
• PyTorch 文档：DDP 设计说明（hook、分桶、no_sync）与 FSDP2 fully_shard。
• 训练配置的公开数字：The Llama 3 Herd of Models（batch 4M → 8M → 16M token）、DeepSeek-V3 Technical Report（2048 张 H800，batch 3072 → 15360 条 4K 序列）、DeepSeek-V2（batch 2304 → 9216 条 4K 序列）。
• 小消息 AllReduce 的实测延迟（演示里 α′ 的默认值）：nccl-tests 在 单节点 8×H100 和 2 / 4 节点 H100 上的公开结果。
• 想再看一遍 DP 的工程细节，推荐 Simon Boehm 的 Data-Parallel Distributed Training of Deep Learning Models。