DDP 与 FSDP

DDP 与 FSDP

本节摘要:多 rank 训练是两个集合通信加一条规则。启动时广播(broadcast)参数,反向后平均(all-reduce)梯度,各 rank 永不对当前在步几上产生分歧。数据并行(DDP):每 rank 在批的不同切片上跑同一模型,优化器步前平均梯度。FSDP:模型装不下单设备,每 rank 持每参数的一份,前向时逐层 all-gather 重组全张量、用后丢弃。痛在簿记:参数跨 rank 漂移则运行静默损坏;平均梯度不平均损失则仪表盘撒谎;集合后端对拓扑达不成一致则运行永远挂。修复是手写一遍集合通信,再不信你复现不出的包装器。本节跑在 CPU(gloo 后端),CUDA 不假设;同代码切 nccl 即上多 GPU。

对应原课程:Phase 19 · Lesson 48 · distributed-fsdp-ddp(原英文 phases/19-capstone-projects/48-distributed-fsdp-ddp/docs/en.md)。本节属「预训练/分布式」赛道第七节。

学习目标

阅读完本节,你应当能够:

1. 用 gloo 后端跨 N rank 起进程组,无需特殊硬件。
2. 实现最小 DDP 包装:构造时广播、反向后 all-reduce 梯度。
3. 证明每 rank 梯度的 all-reduce 匹配单进程在拼接输入上的梯度。
4. 草绘 FSDP 参数分片:每 rank 持一片,前向 all-gather 全张量、用后丢。

一、问题与直觉

模型装得下单设备,数据集装不下。优化预算要每墙秒看 N 倍样本。第一杠杆是数据并行:每 rank 在批不同切片上跑同模型,优化器步前平均梯度。第二杠杆是 FSDP:模型也装不下,每 rank 持每参数的一份,前向时逐层重组全张量。

两个要紧的集合通信

DDP 契约是构造时 broadcast、反向后 all_reduce。FSDP 草图在每层前向前加 all_gather。

梯度平均匹配单进程梯度

在 N rank 上跨批 B 例训的模型,必产出与单进程在 N*B 批上训相同的梯度。诀窍是每 rank 梯度求和再除 N 给平均损失梯度,正是交叉熵(mean 归约)在全批上会产出的。本节代码用 max-abs-diff < 1e-3 断言手写 all-reduce 梯度与参考单进程梯度一致。

FSDP 草图

内存赢是精确的:每 rank 参数内存降到 1/N。成本是 gather,每次前向付。生产 FSDP 把 gather 与上一层计算重叠,墙时间成本远小于朴素会计。本节对每参数做往返,断言重组与原张量位等。

CPU 与 gloo 后端

CUDA 是生产目标,但同代码路径在 CPU 上存在。gloo 是 CPU 集合后端,比 GPU 上 nccl 慢几个数量级,但 API 面相同。本节进程组用 backend="gloo",rank 用 torch.multiprocessing spawn(而非 torchrun);两者最终到同一 torch.distributed 调用。多 GPU 节点上,改动只是 backend="nccl"、设备张量、torchrun 启动。

二、从零实现

code/main.py 五步。

起进程组:

os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = str(port)                  # bind-and-close 取空闲端口
dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)
​

MASTER_ADDR/MASTER_PORT 是 rendezvous:每 rank 拨同主机同端口。

构造时广播:MinimalDDP.__init__ 遍历每参数与缓冲,调 dist.broadcast(tensor, src=0)。rank 0 的值成规范初始化——无此,每 rank 用自己种子初始化,rank 从步一就分歧。

反向后 all-reduce 梯度:

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)                    # 求和后除 N -> 均梯度
​

每 rank 得相同均梯度,优化器步是同输入的函数,故参数跨运行保持同步。

证等价:manual_all_reduce_matches_single_process 在 rank 0 建同模型,比 all-reduce 后梯度与单进程在拼接输入上会算的梯度,max-abs-diff 约 1e-8。

FSDP 往返:fsdp_round_trip_sketch 展平每参数、填到 world_size 倍数、切片、all-gather、去填,每 rank 重组等于原张量。

设计要点:broadcast 在构造时是防 rank 分歧的根——无它各 rank 从步一就不同。all-reduce 后除 world_size 给均梯度,匹配交叉熵 mean 归约。FSDP 的 gather 一次前向付一次,生产上靠与上层计算重叠隐藏成本。本节每步都断言位等,因分布式静默损坏最难诊断。

三、框架对比

PyTorch 的 DistributedDataParallel 加:反向后梯度钩子使 all-reduce 与反向重叠、分桶 all-reduce 把几个小梯度合并成一个集合、第 44 节用的 no_sync 上下文。PyTorch FSDP 加:每层扁平参数视图使每 rank 持一块连续缓冲、下层 unshard 与本层计算重叠、可选 CPU offload 分片。形状不变:启动 broadcast、反向后 reduce、装不下时分片。本节手写让你看清两个集合通信的裸调用、为什么 broadcast 防 rank 分歧、为什么除 N 给均梯度。生产栈(DeepSpeed ZeRO、Megatron-LM)在 FSDP 基础上还分片优化器状态与梯度(ZeRO-2/3),但根契约仍是 broadcast + all-reduce。

四、可复用产物

code/main.py + outputs/skill-distributed-fsdp-ddp.md(新训练脚本食谱:CPU 用 gloo、GPU 用 nccl 起进程组,模型包 DDP 壳(broadcast 构造、reduce 反向后),可选 all_gather 模式分片参数)。demo 默认 world_size=2,两 CPU 进程 spawn、经 gloo 互谈、退零。outputs/ddp-demo.json 捕获每 rank 参数和、all-reduce 后梯度范数、FSDP 往返结果、手写 vs 参考梯度 diff。

五、练习

1. 4 rank:--world-size 4 跑,确认参数跨运行散度 < 1e-3。
2. AVG 算子:把手写均换成 dist.all_reduce(op=dist.ReduceOp.AVG),计时差异。
3. 重叠 all-reduce:给 DDP 加反向后钩子使 all-reduce 与反向其余重叠,量墙时间改善。
4. FSDP 分桶:把 FSDP 草图改成按层分桶 gather,与按参数 gather 比墙时间。
5. rank 故障注入:中途杀一 rank,讨论进程组如何检测、是否挂死、如何优雅退出。

本节要点回顾

1. 两个集合通信:broadcast(一对多)+ all-reduce(多对多求和),加各 rank 步调一致。
2. broadcast 防分歧:构造时从 rank 0 广播,否则各 rank 从步一就不同。
3. all-reduce 后除 N:给均梯度,匹配交叉熵 mean 归约,保证参数同步。
4. 梯度等价可证:all-reduce 梯度与单进程拼接输入梯度一致,max-abs-diff < 1e-8。
5. FSDP 1/N 内存:每 rank 持一片,前向 all-gather 重组、用后丢,生产靠重叠隐藏成本。
6. gloo/nccl 同 API:CPU 用 gloo、GPU 用 nccl,代码结构不变。

下一节,我们做「LM 评估框架」——把任务定义、指标、运行器、排行榜包进一个短、可换的形状。