数据并行 DDP:allreduce 之上的 200 行 wrapper

数据并行 DDP:allreduce 之上的 200 行 wrapper

本节摘要:DistributedDataParallel 是 allreduce 之上的一个钩子。包一层模型,从 rank 0 广播初始参数让每 rank 起点一致,给每个参数装反向钩子在反向后 allreduce 梯度,剩下的就是梯度下降——整套模式 200 行。本节实现 DistributedDataParallel 形 wrapper,用 torch.multiprocessing.spawn 在 gloo 后端上起 N 个 CPU rank(文件 rendezvous),配套单进程参考回路证明逐步参数字节级等价。读完本节,你能说清为什么用均值不用和(均值对 world_size 不变,学习率不用重调)、为什么装桶(数千张量摊延迟)、为什么种子要分两个用途(init 用共享种子、shuffle 用 rank 专属种子),以及把可用 DDP 变成生产 DDP 的两大改动(装桶 + 通信重叠)。

对应原课程:Phase 19 · Lesson 77 · data-parallel-ddp(原英文 phases/19-capstone-projects/77-data-parallel-ddp/docs/en.md)。本节属第 20 章「毕业项目」的分布式训练赛道。

学习目标

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

1. 接出一个广播初始参数、反向后 allreduce 梯度的 DistributedDataParallel 形 wrapper。
2. 用 torch.multiprocessing.spawn 在 gloo 后端上以文件 rendezvous 起 N 个 CPU rank。
3. 在同数据上顺序训练同一模型,以逐步参数等价证明梯度同步正确。
4. 辩护**装桶(梯度融合)与重叠(反向时通信)**是把可用 DDP 变成生产 DDP 的两大改动。

一、问题与直觉

一个 10 亿参数、激活 12GB 的模型塞不进一张消费级 GPU;即便塞得下,训练也要数周。数据并行把 batch 切到 N 个 rank,每 rank 在自己的分片上算 forward 与 backward,每步把每 rank 的梯度求和,使 N 份副本保持一致——求和后的梯度是优化器步进的依据。

不做梯度同步,N 份副本到第 2 步就发散,它不再是「一个模型在更多数据上训练」,而是 N 个恰巧共享初始权重的独立模型。梯度同步做坏(每参数一次 allreduce、无重叠、不装桶),网络是瓶颈,GPU 闲等线缆。DDP 的手艺是让梯度同步相对计算近乎免费:典型 PyTorch DDP 靠装桶、重叠反向、NVLink 上的 NCCL 达到;CPU 上用 gloo 也能学到同样的课。

二、从零实现

DDP 需要的三个操作

为什么用均值不用和

Allreduce-SUM 除以 world_size 给出均值梯度。均值对 world_size 不变:在一个 rank 上调好的学习率,在四个 rank 上照常工作,因为每步梯度幅度不变。Allreduce-SUM 不除,你每次改集群规模都得重调学习率。DDP 包装 SUM 并除之,本节同样。

为什么装桶

一个 transformer 有数千个参数张量,每张量一次 allreduce 付数千次 gloo 延迟下限。DDP 把梯度装进 ~25MB 桶,每桶一次 allreduce,同样总字节过线缆但延迟在桶上摊销。本节的微型模型把所有梯度装一个桶,结构照搬。

为什么钉种子

每 rank 必须 torch.manual_seed(seed + rank) 用于 shuffle,但 torch.manual_seed(seed) 用于参数初始化。单一共享种子意味着每 rank 看到同一 batch 顺序(挫败数据并行);rank 专属种子用于参数意味着初始参数差 float epsilon,梯度同步不再让副本一致。种子模式搞错,参数等价测试在第 1 步就失败。

实现骨架

class DistributedDataParallel:
    def __init__(self, model, world_size):
        self.model, self.world_size = model, world_size
        broadcast_params(model, src=0)        # 从 rank 0 广播初始参数

    def sync_grads(self):
        for p in self.model.parameters():
            if p.grad is not None:
                dist.all_reduce(p.grad, op=dist.ReduceOp.SUM)
                p.grad /= self.world_size     # 均值,对 world_size 不变

def worker(rank, world_size, ...):
    init_process(rank, world_size, backend="gloo", rendezvous="file")
    torch.manual_seed(seed)                   # 参数 init:共享种子
    model = MiniMLP()
    ddp = DistributedDataParallel(model, world_size)
    torch.manual_seed(seed + rank)            # shuffle:rank 专属种子
    for batch in shard(loader, rank, world_size):
        loss = model(batch).loss()
        loss.backward()
        ddp.sync_grads()
        optimizer.step()
​

code/main.py 实现:MiniMLP(3 层 MLP,秒级收敛、暴露接线)、DistributedDataParallel、worker(完整 gloo 训练回路)、_reference_single_process_loop(单 rank 顺序训练同数据,供测试逐字节参数等价)。运行 python3 code/main.py 输出逐步训练表,对比单进程 loss/参数校验和与 4 rank DDP 运行,两条路径到 float epsilon 产出同 loss 曲线,证明梯度同步正确。

三、框架对比

业界对比:PyTorch DDP 是典型实现,torch.nn.parallel.DistributedDataParallel(model) 接好装桶、重叠、no_sync 上下文;HuggingFace Accelerate 加 launcher 处理 torchrun 环境变量与模型 wrap,底下还是 DDP;Megatron-LM 数据并行 把 DDP 与张量并行组合给大模型,数据并行那块仍是「反向后 allreduce」。它们的共识与本节一致:广播参数 + 反向后 allreduce 梯度 + 除以 world_size 取均值,生产再加装桶、重叠、NCCL。

四、生产中的硬化模式

三种模式把 DDP 硬化到可上线。

找未用参数。 某 forward 路径有条件跳过参数(早退、MoE 路由),被跳参数无梯度,但 DDP 的桶就绪钩子仍等它,allreduce 死锁。find_unused_parameters=True 让 DDP 在归约前看哪些参数有梯度,代价是每步一次图遍历——除非 forward 有分支否则关掉。

静态图优化。 forward 跨步稳定时,static_graph=True 让 DDP 预算桶调度,省的几 ms 在 10000 步上复利。

梯度累积要小心。 跨 K 个微批累梯度而不每次同步是 10 倍吞吐赢。DDP 暴露 no_sync() 上下文管理器暂停反向后 allreduce;忘用管理器,你白 allreduce K 次,吞吐跌到地板。

五、可复用产物

• DistributedDataParallel wrapper + worker 回路:本节是 78 节 ZeRO 的对照基线,也是 81 节端到端 demo 的数据并行那一半。
• 种子模式(init 共享、shuffle rank 专属):任何数据并行训练复用,否则参数等价测试第 1 步就失败。
• _reference_single_process_loop:逐字节参数等价的回归网,任何重写都可比对。

六、练习

1. 可配桶:加可配大小的梯度桶,在更深模型上测相对「每参数一次 allreduce」的加速。
2. 实现 no_sync():作上下文管理器,验证梯度累积跨 K 微批与单进程基线一致。
3. 加 find_unused_parameters:让 forward 有时跳过某 MLP 层,不加标志应死锁。
4. 换 barrier 同步:用 torch.distributed.barrier() 替代 allreduce,体会 allreduce 与 barrier 同步的差别。
5. 测同步开销:对 batch 1、16、256 测梯度同步占步时间比例,解释缩放。

本节要点回顾

1. DDP = allreduce 之上的钩子:广播初始参数 + 反向后 allreduce 梯度 + 除以 world_size,200 行。
2. 用均值不用和:均值对 world_size 不变,学习率不用重调;SUM 后除以 world_size。
3. 装桶:数千张量装 ~25MB 桶,延迟在桶上摊销;同样字节过线缆。
4. 种子分用途:init 用共享种子(参数一致),shuffle 用 rank 专属种子(数据并行生效)。
5. 对单进程参考逐字节等价:证明梯度同步正确,两条路径到 float epsilon 同 loss 曲线。
6. 三大硬化:找未用参数(分支 forward)、静态图优化(稳定 forward 预算桶)、no_sync(梯度累积 10 倍吞吐)。
7. 生产再加:NCCL、重叠反向、find_unused_parameters,本节 gloo 学同样课。

下一节,我们将进入「ZeRO 参数分片」——把每参数 allreduce 换成 reduce_scatter,让每 rank 只存优化器状态的一片,把显存占用从 O(M) 降到 O(M/N)。