9.2 数据并行：从 DP 到 FSDP

9.1 的账算出一个尴尬的结论：混合精度 Adam 下，每张卡都要为每个参数准备 16 字节。7B 模型就是 112 GB，一张 80 GB 的卡连「把模型状态放进去」都做不到。

这一节从最自然的切法讲起：不改模型，只把数据摊开。数据并行（data parallelism, DP）是所有框架的默认选项，也是理解后面一切分片技术的起点——因为它恰好包含了一份可以被彻底消除的冗余。

本节回答四个问题：

数据并行每一步到底在通信什么？通信量是多少、和什么有关？
既然每卡都存一份一模一样的优化器状态，为什么不把这部分摊开？——这就是 ZeRO。
摊开之后通信量怎么变？什么时候反而更慢？
工程上怎么用？DistributedDataParallel、ZeRO 各阶段与 FSDP2 分别解决哪一段？

9.2.1 DDP：一步一次梯度同步

定义 9.2.1 · 数据并行（DP / DDP）

把 $k$ 张卡组成一个副本组。每张卡保存一份完整的模型状态（参数、梯度、优化器状态、主权重），只处理全局批的一个分片。第 $t$ 步的流程是：

每张卡用自己的数据分片做前向与反向，得到局部梯度 $\boldsymbol g_i=\nabla_\theta \mathcal L_i$ （ $i=1,\dots,k$ ）；
通过 All-Reduce 求出跨卡平均： $\bar{\boldsymbol g}=\frac1k\sum_{i=1}^{k}\boldsymbol g_i ;$
每张卡用同一个 $\bar{\boldsymbol g}$ 更新自己那份参数。

由于每步用同一份梯度做同一个更新，各副本的参数在每步末完全一致。因此数据并行在数学上等价于「批大小为 $kB$ 的单卡训练」， $k$ 张卡只用来加速，不改变模型。

「简单」与「昂贵」各自在哪

简单在于：不改模型定义、不改前向、不改损失——只改数据加载与梯度同步。所以它的工程风险最低，是任何多卡训练的第一选择。

昂贵在于两点：

显存冗余：每张卡都有一份 $16N$ 的模型状态。 $k$ 张卡总共浪费了 $(k-1)\times16N$ 字节的显存。
每步同步：梯度必须跨卡归约，通信量约等于「两份模型」（9.2.2），而且与批大小无关——批越大，同样的通信摊到越多的计算上，通信占比反而越低。

这两条直接导出后面的一切：冗余 → ZeRO 分片；同步 → 通信重叠与更大批。

为什么「等价于大批量」这件事很重要

因为一旦等价，所有关于大批量训练的结论都可以直接搬过来使用：

大批量训练的已知事实	在数据并行里的对应
全局批 $=kB$ ，梯度方差随批减小	每步梯度是 $kB$ 个样本的平均，噪声更低
批放大 $k$ 倍时学习率常按线性缩放规则放大 $k$ 倍	卡数增加时若不调学习率，有效步长相对变小，收敛变慢
大批量常需学习率预热（warmup）	训练初期与「并行度变更」时都应保留 warmup 段
BatchNorm 等依赖批统计的层行为随批变化	Transformer 用 LayerNorm，无此顾虑——这也是 Transformer 恰好适合 DP 的原因之一

注意线性缩放是启发式：它来自「批越大、梯度噪声越小、可以迈更大的步」的直觉，在大批下有明确的反例（批太大时即使放大学习率也收益递减），所以工程上更稳妥的做法是先用小批验证、再逐步放大并观察（见 1.6）。

定义 9.2.2 · 梯度累积与等效批大小

把一次更新所需的大 batch 拆成 $K$ 个 micro-batch 依次前向 / 反向，梯度累加在缓冲区而不立即更新；攒满 $K$ 个后做一次优化器更新并清零。记单卡 micro-batch 大小为 $B_{\text{micro}}$ 、数据并行度为 $k$ ，则每张卡在一次更新中处理的样本数为 $KB_{\text{micro}}$ ，全局等效批大小为

$B_{\text{global}} = K \times B_{\text{micro}} \times k .$

为保持梯度尺度不变，损失要按整个等效批归一化：每个 micro-batch 的损失先乘 $1/K$ 再累加（或等价地，对累积的梯度最后除以 $K$ ）。

梯度累积是「用时间买批大小」

它把「显存」与「批大小」解耦了：显存只决定 $B_{\text{micro}}$ （单次有多少 token 的激活在场），而 $K$ 想设多大就设多大，代价是同一份计算要串行跑 $K$ 遍。所以：

省的是激活：任意时刻只有一份 micro-batch 的激活驻留；
不省模型状态：参数、梯度、优化器状态与 $K$ 无关；
不省通信（在 DDP 里）：不需要每个 micro-batch 都同步梯度，整个等效批只同步一次——这一点比很多人以为的更好。在标准实现里，梯度在本地累加、只在最后做一次 All-Reduce，于是通信量不变（仍是两份模型），但随着计算量增大，通信与计算之比下降了 $K$ 倍。

例 9.2.1 · 用梯度累积把批凑够

目标：全局批 $B_{\text{global}}=256$ 条序列，每条 $T=1024$ 。硬件：4 张卡，每卡在 $17$ GB 激活预算下最多一次放下 $B_{\text{micro}}=2$ 条序列。求梯度累积步数 $K$ 与每步通信次数。

解

约束检查。 全局批 $=K\times B_{\text{micro}}\times k$ ： $256 = K\times 2\times 4 \Longrightarrow K=32 .$

每步通信次数。 数据并行在每个优化器步骤同步一次梯度，因此 K 个 micro-batch 只触发 1 次 All-Reduce；通信量仍为 $\approx2P$ 字节（ $P$ 为模型大小的字节数），与 $K$ 无关。

对照：如果不做累积会怎样。 为了凑够 256 的批，必须把 $B_{\text{micro}}$ 提到 64（4 卡 × 64 = 256），激活会涨到现在的 32 倍 → 必然 OOM。结论：累积把「批大小」这个优化变量从显存约束里解救出来，代价是单步耗时长 $K$ 倍、且并行度（能同时在场的 token 数）下降，因此算力利用率通常会降低——这就是「累积不是免费的批」的含义。 $\blacksquare$

∎

9.2.2 通信量：为什么每次要传「两份模型」

结论 9.2.1 · 环状 All-Reduce 的通信量

把 $k$ 张卡组成一个环，对大小为 $S$ 字节的张量做 All-Reduce（先求和、再让所有卡拿到结果）。环状算法把张量切成 $k$ 块，先做 $k-1$ 步 reduce-scatter（每步传 $S/k$ ），再做 $k-1$ 步 all-gather（每步传 $S/k$ ）。于是每张卡的总通信量为

$V_{\text{ring}}(S)=2\times\frac{k-1}{k}\,S\ \xrightarrow{\ k\ \text{较大}\ }\ 2S .$

用于梯度同步时 $S\approx$ 模型字节数 $P$ （bf16 下 $P=2N$ ），所以每步每卡通信量约等于两份模型，且与批大小、micro-batch 数完全无关。

「两份模型」是怎么来的，为什么和批无关

为啥是两份：All-Reduce 可以拆成「先归约、再广播」。第一遍（reduce-scatter）把各卡的梯度块互相传一圈求和，每卡传走的量是 $S(k-1)/k$ ；第二遍（all-gather）把求和结果传回来，又是同样的量。两遍加起来 $\approx2S$ 。

为啥和批无关：通信的对象是梯度张量，它的形状只由模型结构决定。批大小决定的是「算出这个梯度花了多少矩阵乘」，所以它影响计算 / 通信比，不影响通信量本身。

这条推论非常有用：加卡会让通信量略微增大（ $2S(k-1)/k$ 逼近 $2S$ ）但计算量按卡数分摊，所以「卡越多、每卡的计算越少、通信占比越高」——这正是 9.1.4 里通信墙的定量形态。

例 9.2.2 · 7B 模型的每步通信时间

在 8 张卡上对 7B 模型做数据并行，bf16 梯度。集群为节点内 NVLink（每卡约 300 GB/s 单向量级）＋节点间 InfiniBand 400 Gb/s（≈ 50 GB/s）。分别估算环状 All-Reduce 的每卡通信量，以及「全部跨机」时的时间下界。

解

通信量。 梯度字节数 $S=2N=2\times7\times10^{9}=14\times10^{9}$ 字节 $=14$ GB。 $V_{\text{ring}}=2\times\frac{k-1}{k}S=2\times\frac{7}{8}\times14\ \text{GB}=24.5\ \text{GB/卡}.$

跨机时间下界。 $\beta\approx50$ GB/s： $t_{\text{comm}}\approx\frac{24.5}{50}\approx0.49\ \text{s}.$ （没有减去重叠，也没有计入延迟与小消息效率损失，因此这是乐观下界。）

与计算对比。 单步前向 + 反向计算量 $\approx6NT_{\text{tok}}$ 。若每卡 1024 条序列 × 1024 token $=1.05\times10^{6}$ token，则 $C\approx6\times7\times10^{9}\times1.05\times10^{6}\approx4.4\times10^{16}$ FLOPs；按单卡 $10^{14}$ FLOP/s 有效算力（bf16 且利用率中等）估 $t_{\text{compute}}\approx440$ s。相比之下通信是零头。

换一个场景就不同了：若每卡只有 64 条序列（token 数降 16 倍）， $t_{\text{compute}}\approx27$ s，通信 0.49 s 仍可接受；但若模型再大 10 倍、token 数不变，通信涨到 4.9 s、计算涨到 4400 s，通信占比反而下降——因为通信随 $N$ 线性、计算也随 $N$ 线性。真正让通信占比失控的是「token 数太少」（小批、短序列），这也解释了为什么大规模训练偏爱「小卡数 × 大批 × 长序列」。 $\blacksquare$

∎

9.2.3 批变大之后：学习率、预热与「大批量陷阱」

数据并行把全局批从 $B$ 变成 $kB$ 。批变了，优化过程也会变，必须同步处理三件事。

线性缩放规则与它的边界

规则（启发式）：批大小放大 $k$ 倍时，学习率也乘 $k$ ： $\eta_{\text{new}} \approx k\,\eta_{\text{base}},\qquad\text{并配合 warmup} .$

为什么这样猜：批越大，梯度估计的方差越小（噪声 $\propto 1/\sqrt{B}$ ），单步方向更可靠，于是可以迈更大的步。写成「步长 × 批」的形式就是 $\eta B$ 保持不变的直觉。

为什么它只是启发式：噪声减小带来的收益随批增大而递减，而大步长的风险（跑到损失曲面的坏区域、发散）是递增的。所以在大批区域，正确做法往往是把缩放指数从 $1$ 降到 $1/2$ （即 $\eta\propto\sqrt{k}$ ），或干脆不放大、改为多跑几步。工程判据：放大后若训练早期 loss 出现尖峰或 NaN，说明缩放过猛。

warmup 为什么必要：训练初期参数离最优点远、梯度尺度不确定，而 Adam 的二阶矩 $v$ 还在用很短的窗口估计（1.6），此时按自适应步长走大步容易失控。让学习率在前若干步（常见 1%–5% 的总步数）线性升到目标值，是并行训练里几乎必开的保险。

并行改动会**同时**改变三个量

新手最容易忽略的一点：把卡数从 $k$ 改到 $k'$ 时，以下三件事同时变了——

全局批 $K B_{\text{micro}} k$ 变了（若保持每卡 micro-batch 不变）；
于是每步的样本数变了，学习率的最优值随之变化；
同时墙钟时间与通信占比也变了（9.2.2）。

所以「加卡」不是纯粹的加速：它是一次优化超参数的迁移。做对照组时要把「同样的全局批」（改 $K$ 或 $B_{\text{micro}}$ 抵消）与「同样的学习率调度」写清楚，否则两次实验的差异无法归因。这也是工具箱要求记录「配置 + 种子 + 数据版本」的原因。

9.2.4 ZeRO-1/2/3：把冗余摊开

定义 9.2.3 · ZeRO 的三级分片

ZeRO（Zero Redundancy Optimizer）观察到： $k$ 张卡各存一份完全相同的模型状态，是纯冗余。它把可摊开的部分切成 $k$ 份，每卡只存 $1/k$ ，按粒度分三级递进（记 $N$ 为参数量，假设数据并行度为 $k$ ）：

阶段	额外分片的对象	每卡模型状态（混合精度 Adam）	该阶段引入的通信
DP（基线）	—	$16N$	1 次梯度 All-Reduce
ZeRO-1	优化器状态（ $m$ 、 $v$ 、fp32 主权重 = $12N$ ）	$4N+12N/k$	仍是 All-Reduce，通信量不变
ZeRO-2	再分梯度（ $2N$ ）	$2N+2N/k+12N/k$	用 reduce-scatter 替代 All-Reduce，通信量约减半
ZeRO-3	再分参数（ $2N$ ）	$16N/k$	计算前 all-gather 参数 + 梯度 reduce-scatter，通信量升至约 1.5 倍

读法：每一级都把更多东西切成 $k$ 份，每卡最终只剩 $16N/k$ ；但通信的「代价曲线」不是单调的——阶段 2 通信下降（因为只需归约、无需广播），阶段 3 通信上升（因为每次用到参数都要临时聚齐）。

ZeRO 不是「新并行」，而是「给数据并行去冗余」

它保留了数据并行的两个核心优点——不改模型结构、每卡都参与算数据——只把副本里冗余的那部分挪走。所以可以这样理解三级：

ZeRO-1 最划算：省 $12N-12N/k$ 的显存，通信一分不多。工程上「显存差一点」时第一反应就该是它。
ZeRO-2 再省梯度那 $2N$ ，同时通信还降了一半：看起来是「白拿」，但代价是归约后的梯度只在拥有它的卡上是完整的，实现复杂度上升（更新时要让各卡读到属于自己分片的那部分状态）。
ZeRO-3 把参数也分了，于是「每卡都有一份完整模型」这个前提被打破——每次计算某一层之前必须把该层参数聚合起来，算完立刻释放。这就是它与 FSDP 的接口形态（9.2.4）。省得最多，通信也最多。

一句话判据：能放得下就用 DP（+ ZeRO-1/2），放不下才上 ZeRO-3 / FSDP。

例 9.2.3 · 8 卡上三级 ZeRO 的每卡模型状态

$N=7\times10^{9}$ ， $k=8$ ，混合精度 Adam。分别算出 DP、ZeRO-1、ZeRO-2、ZeRO-3 的每卡模型状态，并说明 80 GB 卡在哪一级开始「不紧张」。

解

按定义逐项代入（单位 GB， $1\,\text{GB}=10^{9}$ 字节）：

配置	参数	梯度	优化器状态	主权重	合计 / 卡
DP	$14$	$14$	$28+28$	$28$	$112$
ZeRO-1	$14$	$14$	$3.5+3.5$	$3.5$	$38.5$
ZeRO-2	$14$	$1.75$	$3.5+3.5$	$3.5$	$26.25$
ZeRO-3	$1.75$	$1.75$	$3.5+3.5$	$3.5$	$14$

核对公式：ZeRO-1 应为 $4N+12N/k=28+84/8=28+10.5=38.5$ GB ✔；ZeRO-2 为 $2N+2N/k+12N/k=14+1.75+10.5=26.25$ GB ✔；ZeRO-3 为 $16N/k=14$ GB ✔。

结论：DP 与 ZeRO-1 在 80 GB 卡上剩不下多少给激活（尤其 $38.5$ GB 之外还有通信缓冲与碎片）；ZeRO-2 开始舒适，ZeRO-3 留出约 60 GB 给激活。这也解释了为什么「7B 微调选 ZeRO-2/FSDP」是社区常见默认。 $\blacksquare$

∎

什么时候分片**反而更慢**

分片买来显存，付出的是通信与实现开销。以下四种情形要警惕：

情形	为什么变慢	怎么判断
$k$ 很小（2–4 卡）	省下的显存有限，但引入归约路径与每层 all-gather 的固定延迟	对比开 / 关 ZeRO-3 的每步耗时
模型本来放得下	平白多出通信；ZeRO-3 每层要 all-gather 参数，小模型上通信占比极高	先算 $16N \lt M_{\text{cap}}$ 就别开
跨机且消息小	每层参数分片后会「碎成很多小消息」，小消息受延迟而非带宽支配	看 profiler 里通信内核的次数是否暴涨
序列短、批小	计算量本身就小，通信无法被摊薄（例 9.2.2 的反面）	提高每卡 token 数后再比较

9.2.5 FSDP 与 FSDP2：把 ZeRO-3 做进框架

定义 9.2.4 · FSDP（全分片数据并行）

FSDP（Fully Sharded Data Parallel）是 ZeRO-3 思想在 PyTorch 原生实现中的落地。它的核心循环是「按需聚齐、用完即散」：

分片持有：参数、梯度、优化器状态都以分片形式常驻（每卡 $16N/k$ ）；
前向：进入第 $l$ 层前，对该层参数做 all-gather，得到完整参数 → 计算 → 立即释放（只保留分片）；
反向：再次 all-gather 该层参数 → 算梯度 → reduce-scatter（每卡只留自己那 $1/k$ 份梯度，同时在拥有分片的卡上更新优化器状态）；
通信重叠：现代实现会在算第 $l$ 层时预取第 $l+1$ 层的参数，把 all-gather 藏进计算里。

FSDP2 是 PyTorch 后续版本中对 FSDP 的重写（截至 2025 年以官方文档为主要依据，无对应论文）：把「按层展平成一维大张量」的做法改为逐参数的 DTensor 分片。带来三点变化：① 每个参数的 Shard 布局显式可见，可与张量并行等其他分片方式组合；② 不必再维护「扁平参数 + 视图」的映射，checkpoint 与状态字典更直观；③ 混合精度与分片粒度的控制更细（如只把某些层分片）。

「按需聚齐」为什么是这一切的核心

把这个过程想成图书馆借书：

DP 相当于「每人都买全套书放在自己桌上」——便宜好用，但桌子（显存）必须够大；
ZeRO-1/2 相当于「只有书（参数与梯度）各人一套，但读书笔记（优化器状态）共享，按需查」；
ZeRO-3 / FSDP 相当于「连书都只有一套，谁要读哪本就把哪本调过来，读完立刻还回去」——桌子瞬间够用，但调书的人（通信）变得非常忙。

所以 FSDP 的调优重点不在「分片多少」，而在调书能否与读书重叠：预取深度（prefetch）与分片粒度（每层一块还是每几个参数一块）决定了 all-gather 能不能被计算盖住。这也是 FSDP2 把分片粒度改成「逐参数」的动机之一：粒度更细 → 更容易做流水与组合。

# ── Pathway A：数据并行（最省心，模型状态放得下时首选）──────────────
import torch, torch.nn as nn, torch.distributed as dist

dist.init_process_group("nccl")                      # 初始化进程组
model = MyTransformer().cuda()
model = nn.parallel.DistributedDataParallel(model)   # 每步自动 All-Reduce 梯度
# 注意：DDP 不切分任何状态，每卡都是完整 16N

# ── Pathway B：FSDP2（ZeRO-3 语义，逐参数分片）────────────────────────
from torch.distributed.fsdp import fully_shard, MixedPrecisionPolicy

policy = MixedPrecisionPolicy(param_dtype=torch.bfloat16,   # 参数聚合时用 bf16
                              reduce_dtype=torch.bfloat16)  # 梯度归约时用 bf16
fully_shard(model, mp_policy=policy)      # 也可逐层调用：fully_shard(model.layer_i)
# 语义：前向 all-gather 该层参数 → 算 → 释放；反向再来一次 → reduce-scatter 梯度
# 每卡模型状态 ≈ 16N / k（含优化器状态分片）

# ── Pathway C：ZeRO 各阶段（DeepSpeed 一类框架的配置片段）─────────────
# 概念上等价于下面这段"伪配置"，字段名以你所用框架版本为准（需核实）：
#   zero_optimization:
#     stage: 2                 # 1=分片优化器状态 2=+梯度 3=+参数
#     offload_optimizer: none  # 打开会引入主机内存往返，见 9.5.4
#     stage3_prefetch_bucket_size: 5e7

三套接口的选择判据

你的处境	选哪套	理由
单卡装得下模型状态，只想加速	DDP ／ `torchrun`	零改造，通信量最小（ $2S$ ）
差一点装不下（ $16N$ 略超预算）	ZeRO-1 / ZeRO-2	省 $12N$ （或再多 $2N$ ），通信不增加甚至减半
差很多，必须把参数也分片	FSDP2 / ZeRO-3	每卡 $16N/k$ ；务必开预取并评估通信占比
需要与张量并行 / 流水线并行组合	FSDP2（DTensor 分片）或 Megatron 系	分片布局显式，便于组合（9.6.5）

数据并行与 ZeRO 的七个高频误解

「数据并行让每张卡装下更小的模型」。恰恰相反：DP 要求每卡都装得下完整模型状态。它只是让计算变快、批变大。
「批越大，梯度通信量越大」。通信量由模型大小决定（ $\approx2S$ ），与批大小无关。批越大，同样的通信摊到更多计算上，通信占比反而更低。
「梯度累积会成倍增加通信」。 $K$ 个 micro-batch 只在最后同步一次；通信量不变，通信占比还下降了约 $K$ 倍。
「ZeRO 阶段越高、训练越快」。阶段 1/2 的收益是显存（1 通信不变、2 通信减半），阶段 3 换来显存却增加通信。阶段 3 只在「否则根本跑不起来」时才是正确选择。
「开了 ZeRO-3 就不需要激活重计算」。两者管的是不同的账：ZeRO-3 分的是模型状态，激活一分不减。
「FSDP 的 each-layer all-gather 是纯开销」。实现会在计算前预取下一层参数，理想情况下大部分 all-gather 被计算掩盖；但预取是有代价的（多留一份参数分片），需要在 bucket 大小与显存之间取平衡。
「梯度累积与数据并行可以相互替代」。二者都扩大等效批，但机制不同：DP 用更多卡同时算（省墙钟时间、不省模型状态），累积用同一批卡串行算（不省任何显存、反而慢）。典型用法是两者叠加（9.7 的五步决策）。

即时自测 · 9.2

Q1 数据并行每步每卡的通信量约为多少（用模型字节数 $P$ 表示）？它与批大小的关系是什么？

Q2 $N=13\times10^{9}$ 、 $k=8$ 卡、混合精度 Adam：手算 DP / ZeRO-1 / ZeRO-3 每卡的模型状态显存。8 张 80 GB 卡上，ZeRO-1 够用吗？如果激活还要 40 GB/卡呢？

Q3 为什么 ZeRO-2 的通信量比基线减半，而 ZeRO-3 反而增加？

Q4 一个训练任务用了 FSDP，GPU 利用率只有 40%，nvidia-smi 显示显存充足。最可能的两个原因是什么？给出对应的检查动作。

参考要点

Q1 约 $2P$ 字节（环状算法， $P$ 为模型状态中「梯度」部分的字节数；bf16 时 $P=2N$ ），且与批大小无关；精确式为 $2P(k-1)/k$ ， $k$ 大时逼近 $2P$ 。

Q2 $16N=16\times13=208$ GB。DP： $208$ GB/卡；ZeRO-1： $4N+12N/k=52+156/8=52+19.5=71.5$ GB/卡；ZeRO-3： $16N/k=26$ GB/卡。ZeRO-1 不够（71.5 GB 之外几乎没有余量，即使不考虑碎片与 40 GB 激活）；加了 40 GB 激活后只有 ZeRO-3（26 GB）能过，总占用 66 GB 量级，仍应留余量。

Q3 基线要把「归约 + 广播」都做一遍才能让每卡拿到完整平均梯度；ZeRO-2 中每卡只需要自己那 $1/k$ 份梯度，因此可以用只做归约的 reduce-scatter 替代 All-Reduce，通信量从 $\approx2S$ 降到 $\approx S$ 。ZeRO-3 额外需要每层参数的 all-gather（前向 + 反向各一次），这部分是全新的通信，净效果约 $1.5\times2S$ 。

Q4 最可能：① 通信未重叠——每层 all-gather 与计算串行，GPU 在等网卡；检查 profiler 中通信内核占比与「计算空隙」。② 分片粒度 / bucket 过小或预取不足，导致小消息多、延迟高；检查 bucket 大小配置与预取开关，试着增大 bucket。其他可接受的答案：每卡 token 数太少（算力利用率低）、CPU 侧数据加载成为瓶颈（num_workers 不足）。