数据并行与 ZeRO / FSDP

数据并行（DP）是最朴素的并行：每个 GPU 拿一份完整模型，各自处理不同数据，每步同步梯度。它扩展性最好、实现最简单，但有一个致命约束——每张卡必须装下整个模型加优化器状态。ZeRO（Zero Redundancy Optimizer）与 PyTorch FSDP 用"分片"解决这个约束，是今天大模型训练的默认起点。

朴素数据并行的显存账

一个 7B 模型，16 位混合精度训练（AdamW），单卡显存：

组成	大小	计算
参数（BF16）	14 GB	2 B/参数
梯度（BF16）	14 GB	2 B/参数
优化器状态（FP32 参数 + 动量 + 方差）	84 GB	12 B/参数
合计（不含激活）	112 GB	16 B/参数

16 字节/参数是混合精度训练的通用公式。7B 模型 112 GB，一张 80 GB 的卡装不下；70B 模型 1120 GB，八卡也装不下。这就是 DP 的天花板。

ZeRO 的三步分片

ZeRO 的观察：DP 里每张卡都保存着完全相同的优化器状态、梯度和参数——这是纯冗余。按卡数分片即可：

阶段	分片什么	单卡显存（7B，8 卡）	新增通信
DP（基线）	无	112 GB	梯度 all-reduce
ZeRO-1	优化器状态	14 + 14 + 10.5 ≈ 38.5 GB	同 DP（all-reduce 顺手完成分片）
ZeRO-2	优化器 + 梯度	14 + 1.75 ≈ 15.75 GB	同 DP（reduce-scatter）
ZeRO-3	优化器 + 梯度 + 参数	~2 GB（+激活）	前向/反向各一次 all-gather

三阶段的代价递增：

ZeRO-1：免费。梯度的 all-reduce 本来就要做，用 reduce-scatter + 优化器分片更新，通信量与朴素 DP 完全一样
ZeRO-2：仍然免费。把 all-reduce 换成 reduce-scatter（本来就是它的半程），通信量不变，还省了
ZeRO-3：有代价。每层前向和反向都要临时 all-gather 参数（用完即丢），通信量约为朴素 DP 的 1.5 倍，但靠 overlap 可以藏掉大半

经验法则：单卡装得下就用 ZeRO-2（免费省显存）；装不下再上 ZeRO-3（通信贵一点但能装下百亿模型）。

FSDP：PyTorch 的 ZeRO-3

FSDP（Fully Sharded Data Parallel）是 PyTorch 原生的 ZeRO-3 实现，思想相同、工程包装更好：

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(model,                      # 参数/梯度/优化器状态全分片
             sharding_strategy=...,      # FULL_SHARD / SHARD_GRAD_OP / ...
             auto_wrap_policy=...,      # 按层切 unit（关键超参）
             mixed_precision=...)

三个实践要点：

一、wrap 粒度是核心超参。 FSDP 按"unit"分片与 all-gather：unit 太大（整模型）则首层收集要一次性装下全部参数，显存峰值高；unit 太小（每层）则 all-gather 次数多、小通信开销大。按 transformer block 包一层是默认起点（transformer_auto_wrap_policy）。

二、通信重叠要确认开启。 FSDP 默认尝试把 all-gather 与计算 overlap；use_orig_params、forward_prefetch、backward_prefetch 的组合影响实际重叠效果。看 GPU 利用率曲线：锯齿状剧烈波动 = 重叠没生效。

三、与 DP/ZeRO 的关系是演进不是替代。 朴素 DP → ZeRO-1/2 → ZeRO-3/FSDP，是同一条轴上的不同分片深度。FSDP 还支持 HSDP（混合分片：节点内 replicate + 节点间 shard），把同步量进一步摊薄到"每步只同步节点间的分片"。

通信量与重叠：DP 的真实成本

朴素 DP 每步的通信 = 梯度 all-reduce， $2\frac{p-1}{p} \times 14$ GB（7B）。ZeRO 化后通信量不变甚至略降（reduce-scatter 只传一次），关键是能不能藏进计算里：

backward:  ──层N──层N-1──层N-2──...
                │      │      │
grad sync:      └──────┴──────┴──  每层梯度算完立刻 reduce-scatter
                                     （与后面层的反向计算重叠）

梯度按层产出、按层同步，与剩余层的反向计算天然重叠——重叠好的 DP/Zeo-2，梯度同步几乎免费。万卡集群里这正是把 560 ms 的同步（见 RDMA 的代价计算）藏掉的手段。

批量大小的连带影响

DP 的另一个自由度：全局 batch = 单卡 batch × 卡数 × 梯度累积步数。三个约束互相拉扯：

约束	要求	手段
显存装得下激活	单卡 batch 不能太大	激活重计算/卸载
训练稳定	全局 batch 有下限（太小噪声大）	梯度累积
吞吐	每卡 batch 尽量大（打满算力）	显存允许内拉满

经验：先用梯度累积把全局 batch 调到稳定区间，再尽量增大单卡 batch 打满算力，最后用激活技术解决显存溢出。 全局 batch 与学习率要联动调（线性缩放规则）。

什么时候不用 DP

DP/ZeRO 解决"模型太大"，但解决不了"单卡计算太慢"：

单卡 batch 拉满后 MFU 仍低（如极长序列的注意力）→ 需要序列并行
模型单层就超过显存（千亿模型的单层 TP-embedding）→ 需要张量并行
MoE 架构 → 专家并行
这些场景叠加起来就是混合并行——DP 是混合并行里的"基础层"，永远存在，但很少单独存在。

本章要点

混合精度训练显存 ≈ 16 B/参数（参数 2 + 梯度 2 + 优化器 12）；7B = 112 GB。
ZeRO 三步：优化器分片（免费）→ +梯度分片（免费）→ +参数分片（通信 ×1.5，但能装百亿模型）。
FSDP 是 PyTorch 的 ZeRO-3：wrap 粒度是核心超参，按 transformer block 起步。
梯度按层产出按层同步，与反向计算重叠后几乎免费；锯齿状利用率 = 重叠没生效。
全局 batch = 单卡 batch × 卡数 × 累积步数，三个约束按序满足。
DP 解决"装不下"，不解决"算不满"——单卡 MFU 低时看其他并行方式。