跳到主要内容
分布式训练

第3章：数据并行(DP、DDP、FSDP)

掌握数据并行三代方案的演进逻辑、核心机制和工程实践，理解 AllReduce 梯度同步与参数分片的本质区别

数据并行 DDP FSDP AllReduce ReduceScatter

📖 本章概述

数据并行是最基础也最常用的并行策略。本章按”为什么演进”的逻辑线，讲清三代方案各自解决什么问题、引入什么代价。

📑 章节结构

1. DataParallel（DP）——单进程多线程

  • 工作原理：单进程持有模型，每次前向时将模型 Scatter 到多卡，输出 Gather 回主卡
  • 致命缺陷：Python GIL 限制、主卡显存/计算负载不均、通信经过 PCIe 而非 NVLink
  • 定位：仅用于快速原型验证，生产环境不使用

2. DistributedDataParallel（DDP）——生产标配

  • 核心机制：
    • 每 GPU 一个独立进程，各自持有完整模型副本
    • init_process_group 建立通信组
    • DistributedSampler 保证各进程数据不重叠
    • 反向传播中通过 AllReduce 同步梯度
  • Bucket 机制与计算通信重叠：
    • 参数梯度按反向计算顺序分组为 Bucket
    • 每个 Bucket 就绪后立即发起 AllReduce，不等全部反向完成
    • 通信与后续 Bucket 的反向计算重叠（overlap）
  • 通信量分析：每步 AllReduce 通信量 = 2Ψ2\Psi（Ring AllReduce 下每卡收发总量）
  • 使用要点：DistributedSampler + set_epoch()、仅 rank 0 日志/保存、model.module 访问原始参数
  • 局限：每卡存完整模型 → 参数 + 梯度 + 优化器状态必须 ≤ 单卡显存

3. FullyShardedDataParallel（FSDP）——显存优化的数据并行

  • 核心思想：不让每卡冗余存完整状态，而是分片存储，计算时按需 AllGather 组装
  • 工作流程：
    1. 初始状态：每卡只存 1N\frac{1}{N} 的参数/梯度/优化器状态
    2. 前向：AllGather 重组完整参数 → 计算 → 释放非本地分片
    3. 反向：AllGather 参数 → 计算梯度 → ReduceScatter 分片梯度 → 释放完整参数
    4. 更新：每卡只更新自己负责的分片
  • 分片策略选择：
    • FULL_SHARD（ZeRO-3）：参数+梯度+优化器全部分片，显存最省
    • SHARD_GRAD_OP（ZeRO-2）：梯度+优化器分片，参数不分片
    • HYBRID_SHARD：机内全分片 + 机间数据并行，多机训练推荐
    • NO_SHARD：等价于 DDP，用于调试
  • 通信量分析：每步 = AllGather（前向）+ AllGather + ReduceScatter（反向）= 3Ψ3\Psi，比 DDP 多 50%
  • 工程配置：auto_wrap_policy（按模块大小或类型自动分片）、MixedPrecision 配置、DCP 保存/加载
  • FSDP2 简介：per-parameter 分片、基于 DTensor、与 TP 更好组合

4. DDP vs FSDP 选型指南

  • 对比表：显存占用、通信量、代码复杂度、适用规模
  • 决策规则：单卡装得下 → DDP（简单高效）；装不下 → FSDP；多机大模型 → FSDP + HYBRID_SHARD

5. 动手实验

  • 实验 A：将单卡脚本改造为 DDP 版本（30 分钟）
  • 实验 B：在 2-4 卡上对比 DDP vs FSDP 的峰值显存与训练吞吐
  • 实验 C：测试不同 sharding_strategy 对显存和速度的影响

🎯 本章学习目标

  • 能解释 DP → DDP → FSDP 每一步演进解决了什么问题、引入了什么代价
  • 能用 PyTorch DDP 将单卡脚本改为多卡训练（含 Sampler、set_epoch、模型保存）
  • 能配置 FSDP 的分片策略和自动包装策略
  • 能计算 DDP 和 FSDP 每步的通信量并解释差异