AI 原生开放教材

外观

AI 工程从零到一（中文）源码研读·中文原文 ↗原件 ↓仓库 ↗

从零实现分布式数据并行与 FSDP ​

多 rank 训练就是两个 collective 加一条铁律。启动时广播参数，backward 后平均梯度，永远不让各 rank 对当前是第几步产生分歧。

类型： Build 语言： Python 前置要求： 第19阶段第42-45课 预计时间： ~90 分钟

学习目标 ​

  • 用 gloo 后端在 N 个 rank 之间拉起一个 process group，不需要特殊硬件。
  • 实现一个最小 DDP wrapper：构造时广播参数，backward 后 all-reduce 梯度。
  • 证明各 rank 梯度的 all-reduce 等价于单进程在拼接输入上的梯度。
  • 概述 FSDP 参数分片：每个 rank 持有一个切片，forward 时 gather 出完整 tensor，用完即丢。

问题背景 ​

模型放得进一张卡。数据放不进。优化预算要求每秒看到 N 倍的样本。第一个杠杆是数据并行：每个 rank 跑同一个模型、看不同的 batch 切片，然后在 optimizer step 之前平均梯度。第二个杠杆是 FSDP：模型也放不进一张卡了，于是每个 rank 只持有每个参数的一部分，forward 时逐层重建完整 tensor。

痛点是记账。如果参数在各 rank 之间发生漂移，run 就被静默损坏了。如果平均了梯度但没平均 loss，仪表盘就在骗人。如果 collective 后端对拓扑达不成一致，run 就永远挂住。解决办法是亲手写一遍 collective，别再信任你无法复现的 wrapper。

本课在 CPU 上运行。不假设有 CUDA。gloo 后端随每个 PyTorch 构建一起发布，接受 torch.multiprocessing worker；同样的代码在多 GPU 节点上切到 nccl 也不需要改结构。

核心概念 ​

mermaid
flowchart TB
  init[rank 0 进程] --> seed[在 rank 0 上初始化模型]
  init --> spawn[启动 rank 1..N-1]
  spawn --> pg[init_process_group: backend, world_size, master_addr, master_port]
  pg --> bcast[从 rank 0 广播模型参数]
  bcast --> loop[每个 rank 的训练循环]
  loop --> shard[每个 rank: 自己的 batch 切片]
  shard --> fwd[本地前向 + 反向]
  fwd --> ar[all_reduce 梯度, 除以 world_size]
  ar --> step[每个 rank 用相同梯度执行 optimizer.step]
  step --> loop

两个最重要的 collective ​

Collective做什么什么时候用
broadcast把一个 tensor 从一个 rank 复制到所有其他 rank参数初始化、scheduler 状态、任何一对多同步
all_reduce对一个 tensor 在所有 rank 间做 sum（或 mean、max），每个 rank 都拿到结果backward 后的梯度平均
all_gather每个 rank 贡献一个 tensor，每个 rank 拿到拼接结果logits 收集、FSDP 参数 unshard

DDP 的合约是：构造时 broadcast，backward 后 all_reduce。FSDP 的概述额外加了每层 forward 前的 all_gather。

梯度平均等价于单进程梯度 ​

一个模型在 N 个 rank 上用 B 个样本的 batch 训练，产生的梯度必须与单进程在 N*B 的 batch 上训练一致。关键是：把各 rank 梯度求和再除以 N 就得到平均 loss 梯度，也就是 cross entropy 用 mean reduction 在完整 batch 上会产生的东西。课程代码用 max-abs-diff < 1e-3（手动 all-reduce 梯度 vs 参考单进程梯度）来断言这一点。

FSDP 概述 ​

mermaid
flowchart LR
  param[完整参数] --> split[拆分为 N 个等长扁平分片]
  split --> r0[rank 0 持有分片 0]
  split --> r1[rank 1 持有分片 1]
  split --> rN[rank N-1 持有分片 N-1]
  r0 --> gather[前向前 all_gather]
  r1 --> gather
  rN --> gather
  gather --> full[每个 rank 上的完整张量]
  full --> fwd[该层前向传播]
  fwd --> drop[丢弃完整张量, 仅保留分片]

内存节省是精确的：每个 rank 的参数内存降到 1/N。代价是 gather——每次 forward 都要付。生产级 FSDP 把 gather 和上一层的计算重叠，所以实际挂钟开销比朴素计算预测的小得多。本课对每个参数做往返并断言重建结果与原始值 bit-equal。

CPU 与 gloo 后端 ​

CUDA 是生产目标，但同样的代码路径在 CPU 上也存在。gloo 是 CPU collective 后端。它比 GPU 上的 nccl 慢几个数量级，但 API 面完全一样。本课的 process group 用 backend="gloo" 初始化，rank 用 torch.multiprocessing 而非 torchrun 启动；两者最终调的都是相同的 torch.distributed 接口。在多 GPU 节点上，唯一的变化是 backend="nccl"、tensor 放到设备上，以及用 torchrun 来启动。

figure
cg-allreduce-ring

动手构建 ​

code/main.py 是可运行的产物。

第一步：拉起 process group ​

python
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = str(port)
dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)

MASTER_ADDR 和 MASTER_PORT 是 rendezvous 点：每个 rank 拨同一个主机的同一个端口。课程用 bind-and-close 技巧找空闲端口，避免多个 run 共享同一台机器时的冲突。

第二步：构造时广播 ​

MinimalDDP.__init__ 遍历每个参数和 buffer，调用 dist.broadcast(tensor, src=0)。rank 0 的值成为规范初始化。没有这一步，每个 rank 用自己的 seed 初始化，从第一步就开始分叉。

第三步：backward 后 all-reduce 梯度 ​

python
def all_reduce_grads_(module, world_size):
    for p in module.parameters():
        if p.grad is None:
            p.grad = torch.zeros_like(p.data)
        dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
        p.grad.data.div_(world_size)

每个 rank 最终拿到相同的平均梯度。optimizer step 现在是对每个 rank 上相同输入的函数，这就是参数在整个 run 中保持同步的原因。

第四步：证明等价性 ​

manual_all_reduce_matches_single_process 在 rank 0 上构建同一个模型，把 all-reduce 后的梯度与单进程在拼接输入上计算的梯度做比较。max-abs-diff 大约在 1e-8。

第五步：FSDP 往返 ​

fsdp_round_trip_sketch 展平每个参数，pad 到 world_size 的倍数，切片，all-gather，再去 pad。每个 rank 的重建结果等于原始值。这是 unshard 步骤；反向操作（forward 后 re-shard）就是从 gather 出的 tensor 上切一片。

运行：

bash
python3 code/main.py

默认 world size 为 2。两个 CPU 进程启动，通过 gloo 互相通信，退出码为零。输出 outputs/ddp-demo.json 记录每个 rank 的参数和、all-reduce 后的梯度范数、FSDP 往返结果，以及手动梯度 vs 参考梯度的 diff。

实际应用 ​

生产训练栈调用的是同样的原语。PyTorch 的 DistributedDataParallel 额外加了：backward 后的梯度 hook 把 all-reduce 和 backward 重叠、bucket 化的 all-reduce 把多个小梯度合成一个 collective，以及第46课用过的 no_sync context。

PyTorch 的 FSDP 额外加了：每层一个 flat 参数视图让每个 rank 持有一块连续 buffer、下一层的 unshard 与当前层的计算重叠，以及可选的 CPU offload。

形状不变：启动时广播，backward 后 reduce，参数放不下就分片。

交付产物 ​

outputs/skill-distributed-fsdp-ddp.md 是新训练脚本的 recipe：用 gloo 起 CPU 的 process group、用 nccl 起 GPU 的，把模型套进一个 DDP shell（构造时广播、backward 后 reduce），可选地用 FSDP 概述中的 all_gather 模式做参数分片。

练习 ​

  1. 用 --world-size 4 运行，确认整个 run 中参数 spread 保持在 1e-3 以下。
  2. 把手动平均替换为 dist.all_reduce(op=dist.ReduceOp.AVG)，计时对比差异。
  3. 给 DDP wrapper 加一个 post-backward hook，让 all-reduce 和剩余 backward 重叠；测量挂钟改善。
  4. 实现 FSDP re-shard 步骤：forward 之后把完整 tensor 换回本地 shard。确认每个 rank 的内存降低了。
  5. 在 CUDA 机器上切换到 nccl 后端。记录哪些环境变量变了、哪些没变。

关键术语 ​

术语日常说法实际含义
Backend"gloo 或 nccl"实现 collective ops 的库；gloo 用于 CPU，nccl 用于 GPU
World size"总 rank 数"组里的进程数量；组是 collective 操作的单位
Rank"Worker id"组内的进程标识符，从零开始
All-reduce"把梯度加起来"对一个 tensor 在所有 rank 间求和，每个 rank 最终拿到相同结果
Unshard"把参数拼回来"通过 all_gather 从各 rank 的切片重建完整 tensor

延伸阅读 ​

  • PyTorch torch.distributed 文档，本课所依赖的 collective 语义。
  • gloo 库的 collective 列表，形状与 CUDA 版 nccl 原语一致。
  • 第19阶段第46课的梯度累积模式，把 DDP all-reduce 包在 no_sync 中。
  • 第19阶段第47课的 checkpoint 布局，能在 DDP 和 FSDP 的 run 中存活。
  • PyTorch FSDP 文档，本课概述的参数分片的生产级实现。