PyTorch 分布式训练引擎拆解：DDP 到 FSDP 的设计演进
实验目标
用 torchrun 跑起 DDP/FSDP 很容易，但"为什么 DDP 显存翻倍""为什么 FSDP 卡间通信这么多""梯度什么时候同步的"这类问题，决定了你能否在千卡规模把引擎调到最优。本篇从框架引擎机制视角拆解 PyTorch 分布式训练：
- DDP 的梯度同步到底怎么做的？通信与计算如何重叠？
- FSDP 相对 DDP 改了什么？参数分片与通信层级如何设计？
- 什么场景选 DDP、什么场景选 FSDP？
本篇为机制解析 + 配置实践路线，命令与配置为通用实践示例，不编造性能基准；具体实现以所用 PyTorch 版本为准。
第一段：DDP——梯度同步的工程化设计
DDP（DistributedDataParallel）的逻辑很简单：每卡一份完整模型 + 各自数据，反向传播后 AllReduce 梯度，各卡用平均后的梯度更新。工程难点在同步的时机与开销。
1. 梯度桶（Gradient Bucket）：把 AllReduce 的次数压下来
朴素做法是每个参数的梯度各 AllReduce 一次——千亿参数就有上万个梯度张量、上万次通信，开销不可接受。DDP 的做法是梯度桶：
python
# 概念示意：DDP 内部把连续参数的梯度打包成桶，按桶做 AllReduce
# bucket_cap_mb：默认 25MB，桶大小影响通信粒度与重叠效果
ddp_model = torch.nn.parallel.DistributedDataParallel(
model,
bucket_cap_mb=25, # 桶容量，小模型可调小（如 5MB）
gradient_as_bucket_view=True, # 复用桶内存，减少拷贝
)- 桶 = 通信粒度：桶越大通信次数越少但单个桶等待时间越长；桶越小重叠更细但通信次数多。25MB 是默认平衡点，小模型可调小；
gradient_as_bucket_view=True：让参数梯度直接以桶内存视图存在，避免额外拷贝（显存与耗时双收益）。
2. 通信与计算重叠：反向传播的"边算边通"
DDP 的核心优化是梯度就绪即通信（gradient ready → bucket ready → AllReduce 异步启动），与后续层的反向计算重叠：
backward 反向传播
├─ 层 N 计算梯度 → 桶就绪 → AllReduce 启动（异步，不阻塞）
├─ 层 N-1 计算梯度（与 N 的通信重叠）→ ...
└─ 全部完成 → 最后一次 AllReduce 完成 → 更新- 重叠是否有效取决于桶切分与层顺序：桶按参数注册顺序切分，若某桶跨越多个层，就绪时间被最长层拖累——这就是为什么模型注册顺序影响 DDP 性能；
- 检查重叠效果：用 PyTorch Profiler 看 backward 阶段是否存在明显的通信空洞（AllReduce 在等待梯度）。
3. DDP 的显存账
- 每卡持有完整模型 → 显存 ≈ 单卡模型 + 激活 + 梯度 + 优化器状态；
- 梯度 AllReduce 的通信量与模型大小成正比（每卡都要发自己的梯度），千卡时通信是瓶颈而非显存；
- 工程观察：DDP 适合"模型能塞进单卡、但数据/算力需要多卡"的场景（微调、中小模型）；模型大到单卡放不下时，进入 FSDP。
第二段：FSDP——把"每卡一份模型"改成"每卡一份分片"
FSDP（Fully Sharded Data Parallel）的核心思想：参数、梯度、优化器状态都按数据并行维度分片，计算时再 All-Gather 出需要的参数。
1. 三级分片：参数/梯度/优化器状态
| 维度 | DDP | FSDP |
|---|---|---|
| 参数 | 每卡完整副本 | 每卡 1/DP 分片，用时 All-Gather |
| 梯度 | 每卡完整 | 每卡分片（Scatter-Reduce 后只保留本卡分片） |
| 优化器状态 | 每卡完整 | 每卡分片（显存大头省在这里） |
| 通信 | 每步 AllReduce 梯度 | All-Gather 参数 + Scatter-Reduce 梯度（通信量约 2-3 倍于 DDP） |
python
# FSDP 基础用法
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
fsdp_model = FSDP(
model,
sharding_strategy=torch.distributed.fsdp.ShardingStrategy.FULL_SHARD,
mixed_precision=torch.distributed.fsdp.MixedPrecision(
param_dtype=torch.bfloat16,
reduce_dtype=torch.bfloat16,
buffer_dtype=torch.bfloat16,
),
)2. 通信层级：Flatten 与 All-Gather/Scatter-Reduce
- 参数 flatten：FSDP 把（受管）参数展平成连续存储，一次 All-Gather 取回一整块，而不是逐个参数通信——降低通信次数；
- 前向：每层/每块计算前 All-Gather 该块参数，计算后释放（显存按需取用）；
- 反向：各卡计算本卡分片的梯度 → Scatter-Reduce 聚合 → 只保留本卡分片梯度 → 更新本卡分片优化器状态；
- 通信量账：每步 All-Gather（参数）+ Scatter-Reduce（梯度）≈ 2 倍参数量的通信，比 DDP 的 1 倍多——FSDP 用更多通信换显存，网络是瓶颈时未必更快。
3. FSDP 的调优旋钮
| 旋钮 | 作用 | 工程建议 |
|---|---|---|
sharding_strategy | FULL_SHARD / SHARD_GRAD_OP（只分片梯度，省通信）/ NO_SHARD | 显存够则 SHARD_GRAD_OP 减少通信 |
| 分片粒度（auto_wrap_policy） | 按层/块分片 vs 整模型分片 | 按层分片 → 更细粒度 All-Gather，重叠更好 |
cpu_offload | 参数/梯度卸载到 CPU | 显存极限场景，注意 CPU↔GPU 传输开销 |
forward_prefetch / backward_prefetch | 提前预取下一块参数 | 与计算重叠通信，默认开启更优 |
第三段：选择矩阵——什么场景用什么
| 场景 | 推荐 | 理由 |
|---|---|---|
| 中小模型（单卡可放）+ 多卡扩展数据 | DDP | 通信最少、实现简单、显存不是瓶颈 |
| 大模型（单卡放不下）训练 | FSDP（+ 必要时 CPU offload） | 显存分片是硬约束，通信换显存 |
| 大模型微调（LoRA 等） | FSDP | 分片后显存给激活/可训练参数留空间 |
| 超大规模 + 极端显存 | FSDP + 张量并行组合 | 单靠 FSDP 无法拆分单层超大权重（如超大 Embedding/FFN），需 TP 叠加 |
工程观察：选择不是"越先进越好"——DDP 在中小模型上是通信最省的方案；FSDP 的价值在"显存分片解锁更大模型"。混合场景（EP×TP×DP×FSDP，10 章 MoE 篇）是另一层复杂度，先理解 DDP/FSDP 的通信-显存权衡，再叠加并行维度。
分布式引擎检查清单
- [ ] DDP：桶大小与模型规模匹配（小模型调小 bucket_cap_mb）；Profiler 确认反向通信无空洞
- [ ] DDP：
gradient_as_bucket_view=True已开启 - [ ] FSDP：分片粒度与模型结构匹配（按层 auto_wrap），混合精度位置正确
- [ ] FSDP：网络带宽评估——通信量约 2 倍参数，确认 03 章网络能支撑
- [ ] 大模型场景：FSDP 与 TP 叠加方案明确（谁分片、谁切层）
- [ ] 训练日志/Profiler 纳入通信耗时占比观测
工程启示
- 引擎机制决定调优方向：DDP 的桶、FSDP 的分片粒度，都是"通信次数 × 等待时间"的权衡——理解机制，才知道旋钮往哪拧。
- 显存与通信是同一枚硬币：FSDP 用 2 倍通信换显存，网络弱时反而更慢；选型先算账。
- Profiler 是引擎的听诊器：通信空洞、All-Gather 等待都体现在火焰图里，比猜有效。
下一篇预告：训练引擎就绪，推理侧同样需要"测量先行"——下一篇《LLM 推理性能基准方法：吞吐、TTFT 与成本实测》拆解推理服务的统一测法。