数据并行的数学目标很简单:每张卡处理不同 mini-batch,各自算梯度, 然后把梯度平均,让所有 replica 做同一次 optimizer step。 但工程实现有两条路:nn.DataParallel 是单进程多线程, DistributedDataParallel 是多进程,每个进程绑定一张 GPU。 后者绕开 Python GIL、通信路径更清楚、可跨机扩展,也是今天真正要掌握的对象。

DDP 的三个不变量

Day 15 已经跑过 torchrun,今天要把训练细节补齐: local rank 绑定 GPU、初始化进程组、用 DistributedSampler 切数据、 包 DDP、每个 epoch 调 sampler.set_epoch, 最后只让 rank 0 保存 checkpoint。

import os, torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler

def setup():
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl")
    return local_rank, dist.get_rank(), dist.get_world_size()

def main():
    local_rank, rank, world_size = setup()
    device = torch.device("cuda", local_rank)

    model = TinyModel().to(device)
    model = DDP(
        model,
        device_ids=[local_rank],
        bucket_cap_mb=25,
        gradient_as_bucket_view=True,
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

    dataset = MyDataset()
    sampler = DistributedSampler(dataset, shuffle=True)
    loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=32,
        sampler=sampler,
        num_workers=4,
        pin_memory=True,
    )

    for epoch in range(3):
        sampler.set_epoch(epoch)  # 保证各 rank 每个 epoch 的 shuffle 不同但可对齐
        for x, y in loader:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)
            loss = loss_fn(model(x), y)
            loss.backward()       # DDP 在 backward 中触发梯度同步
            optimizer.step()

    if rank == 0:
        torch.save(model.module.state_dict(), "ckpt.pt")
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

启动与日志

# 单机 4 卡
torchrun --standalone --nproc_per_node=4 train_ddp.py

# 打开分布式调试信息
TORCH_DISTRIBUTED_DEBUG=DETAIL NCCL_DEBUG=INFO \
torchrun --standalone --nproc_per_node=4 train_ddp.py

# 多机示意:每台机器都运行同一命令,只改 node_rank
torchrun \
  --nnodes=2 \
  --nproc_per_node=8 \
  --node_rank=0 \
  --master_addr=10.0.0.1 \
  --master_port=29500 \
  train_ddp.py

梯度累积:正确使用 no_sync

accum = 4
for step, (x, y) in enumerate(loader):
    sync_now = ((step + 1) % accum == 0)
    ctx = model.no_sync() if not sync_now else contextlib.nullcontext()
    with ctx:
        loss = loss_fn(model(x), y) / accum
        loss.backward()
    if sync_now:
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

初学 DDP 最常见误解是"backward 完成后统一 AllReduce"。 如果真这么做,通信完全暴露在计算之后,扩展性会很差。 DDP 真正做的是:给每个参数的梯度 accumulator 挂 hook; 当某个参数梯度 ready,Reducer 记录它所在 bucket 的 ready 状态; bucket 里所有梯度都 ready 后,立刻异步发 AllReduce。

同步的精确语义

DDP 的 bucket 是一组参数梯度的扁平 buffer。bucket 太小,AllReduce 次数多、launch overhead 高; bucket 太大,要等很久才 ready,overlap 变差。 默认 bucket_cap_mb 常见值是 25 MiB,但最优值取决于模型层大小、网络带宽、反向计算时间和 batch size。

四组实验记录法

# 只改 bucket_cap_mb,其它都固定
for cap in 6 12 25 50; do
  TORCH_DISTRIBUTED_DEBUG=DETAIL NCCL_DEBUG=WARN \
  torchrun --standalone --nproc_per_node=4 train_ddp.py --bucket-cap-mb $cap
done

ddp_data = model._get_ddp_logging_data()
if dist.get_rank() == 0:
    for key in [
        "bucket_sizes",
        "avg_backward_comm_time",
        "avg_backward_compute_comm_overlap_time",
        "can_set_static_graph",
    ]:
        print(key, ddp_data.get(key))

torch/nn/parallel/distributed.py 是 Python 外壳, 真正高性能同步逻辑在 C++ Reducer 和 c10d/NCCL 后端里。 读源码时不要从第 1 行硬啃,按"构造时做了什么、forward 前后做了什么、backward hook 何时触发、调试信息怎么拿"四条线读。

源码心智伪代码

# 这不是 PyTorch 原源码,是读 distributed.py 时的心智压缩版
class DistributedDataParallel:
    def __init__(self, module, bucket_cap_mb=25, ...):
        self.module = module
        self.process_group = default_pg
        self.parameters = list(module.parameters())
        self.reducer = Reducer(
            parameters=self.parameters,
            bucket_bytes_cap=bucket_cap_mb * 1024 * 1024,
            process_group=self.process_group,
        )
        self._register_autograd_hooks()
        self._sync_module_states()

    def forward(self, *args, **kwargs):
        self.reducer.prepare_for_forward()
        out = self.module(*args, **kwargs)
        if self.require_backward_grad_sync:
            self.reducer.prepare_for_backward(out)
        return out

    # autograd 在每个参数 grad ready 时调用
    def on_grad_ready(param):
        bucket = self.reducer.bucket_for(param)
        bucket.mark_ready(param)
        if bucket.all_ready():
            bucket.work = dist.all_reduce(bucket.buffer, async_op=True)

    def finalize_backward(self):
        for bucket in self.buckets:
            bucket.work.wait()
            bucket.buffer.div_(world_size)
            bucket.copy_back_to_param_grads()

1. 用一张时序图解释 DDP backward 中的梯度同步:参数 grad ready、bucket ready、异步 AllReduce、copy back、optimizer.step 分别在什么时候发生。
2. 对比 nn.DataParallel 和 DistributedDataParallel:进程模型、通信对象、性能瓶颈、扩展性分别有什么不同。
3. 解释 bucket_cap_mb 太小和太大分别会带来什么问题,并设计一个 profiler 实验找到更好的值。
4. 说明 no_sync 的正确使用方式,为什么 forward 必须放进上下文,以及 loss 为什么要除以 accumulation steps。
5. 按源码阅读路线列出 distributed.py 的 6 个入口,说明每个入口回答 DDP 的哪个核心问题。