4. PyTorch DDP详解

聊到多卡训练，PyTorch DDP 绝对是绕不开的核心组件。说实话，我见过不少团队在单卡上跑得飞起，一上多卡就各种翻车。问题出在哪？多半是对 DDP 的理解还停留在「调个 API 就行」的层面。

今天我就把 DDP 的底裤扒开，从原理到实战，把那些容易踩坑的地方都捋一遍。嗯，保证你看完能少走半年弯路。

4.1 DDP 到底在干什么？

先问个问题：DDP 和 DP（DataParallel）有啥区别？

DP 是 PyTorch 最早的多卡方案。它的逻辑很简单——一个主卡负责 forward，然后把梯度广播到其他卡。但问题来了：主卡成了瓶颈，而且 Python 的 GIL 锁让多线程并行效率大打折扣。我早期项目用过 DP，8 张卡跑起来，实际加速比连 3 倍都不到，气得我直接重写。

DDP 就不一样了。它采用多进程方案，每个 GPU 对应一个独立进程。每个进程都有自己的模型副本和优化器。前向传播各自算，反向传播各自算梯度，然后通过 AllReduce 通信把梯度同步到所有卡上。

说白了，DDP 的核心思想就是：各自为战，梯度共享。

关键区别一句话总结：

  • DP：单进程多线程，主卡负责梯度聚合，通信开销大
  • DDP：多进程，每个进程独立计算，AllReduce 同步梯度，效率高

下面这张图是我自己画的 DDP 工作流程，你看一眼就明白了：

DDP 多卡训练流程 GPU 0 (进程 0) GPU 1 (进程 1) GPU 2 (进程 2) 数据分片 0 数据分片 1 数据分片 2 前向传播 前向传播 前向传播 反向传播 反向传播 反向传播 AllReduce 梯度同步 模型参数初始化（相同） 每个 GPU 独立计算，梯度通过 AllReduce 同步后统一更新参数

4.2 进程组初始化——别小看这一步

DDP 的第一步就是初始化进程组。很多人直接抄官网代码，结果跑分布式时各种连不上。我踩过这个坑，当时排查了整整一天，最后发现是 MASTER_ADDR 写错了。

来看标准初始化代码：

import torch.distributed as dist
import torch.multiprocessing as mp

def init_process(rank, world_size, master_addr='127.0.0.1', master_port='29500'):
    # 设置环境变量
    os.environ['MASTER_ADDR'] = master_addr
    os.environ['MASTER_PORT'] = master_port
    
    # 初始化进程组
    dist.init_process_group(
        backend='nccl',      # GPU 训练用 nccl
        init_method='env://', # 从环境变量读取配置
        rank=rank,           # 当前进程编号
        world_size=world_size # 总进程数
    )
    
    # 设置当前 GPU
    torch.cuda.set_device(rank)

我的经验之谈：

  • MASTER_PORT 别用默认的 29500，容易跟别人冲突。我习惯用 29501 + 随机偏移
  • 多机训练时，MASTER_ADDR 必须填主节点的 IP，别填 127.0.0.1
  • 如果遇到 Address already in use，八成是端口被占了

进程组初始化完成后，每个进程都有一个 rank 编号。rank 0 通常作为主进程，负责打印日志、保存模型等操作。其他进程只管算。

4.3 数据加载器适配——最容易出幺蛾子的地方

数据加载这块，我见过最经典的错误就是：每个进程都加载了完整的数据集。结果 8 张卡跑的是 8 份一模一样的数据，梯度一模一样，同步了个寂寞。

正确的做法是用 DistributedSampler。它会自动把数据集切分成 world_size 份，每个进程只拿自己那份：

from torch.utils.data.distributed import DistributedSampler

# 创建分布式采样器
sampler = DistributedSampler(
    dataset,
    num_replicas=world_size,
    rank=rank,
    shuffle=True,
    seed=42
)

# 创建 DataLoader
dataloader = DataLoader(
    dataset,
    batch_size=32,
    sampler=sampler,  # 用 sampler 替代 shuffle
    num_workers=4,
    pin_memory=True
)

# 每个 epoch 要调用 sampler.set_epoch(epoch)
# 不然每个 epoch 的数据顺序都一样
for epoch in range(num_epochs):
    sampler.set_epoch(epoch)
    for batch in dataloader:
        # 训练代码
        pass

注意！

  • 用了 DistributedSampler 后，DataLoader 的 shuffle 参数必须设为 False
  • 每个 epoch 必须调用 sampler.set_epoch(epoch)，否则数据顺序永远不变
  • batch_size 是每个 GPU 的 batch size，总 batch size = batch_size × world_size

我曾经有个项目，训练 loss 曲线特别平滑，但验证集效果很差。排查了半天，发现就是忘了调 set_epoch，导致每个 epoch 数据顺序一样，模型学到了数据顺序的「假规律」。

4.4 模型同步机制——AllReduce 到底怎么工作的？

模型同步是 DDP 的核心。每次反向传播结束后，DDP 会自动触发 AllReduce 操作，把各个 GPU 上的梯度加起来求平均，然后广播回去。

具体流程是这样的：

  1. 前向传播：每个 GPU 用自己的数据计算 loss
  2. 反向传播：每个 GPU 计算梯度，存在各自的 .grad 里
  3. 梯度同步：DDP 的 hook 自动触发 AllReduce，把梯度同步到所有 GPU
  4. 参数更新：每个 GPU 用同步后的梯度更新自己的模型参数

你可能会问：那模型参数怎么保证一致？

答案是：初始化时所有 GPU 的模型参数是一样的（通过 torch.manual_seed 保证）。每次更新用的梯度也是一样的（AllReduce 同步过）。所以参数永远保持一致。

来看 DDP 的完整使用示例：

import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def train(rank, world_size):
    # 1. 初始化进程组
    dist.init_process_group('nccl', rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    
    # 2. 创建模型并移到 GPU
    model = MyModel().to(rank)
    
    # 3. 用 DDP 包装模型
    ddp_model = DDP(
        model,
        device_ids=[rank],
        output_device=rank,
        find_unused_parameters=False  # 如果有参数没用到梯度，设为 True
    )
    
    # 4. 创建优化器（用 ddp_model.parameters()）
    optimizer = torch.optim.SGD(ddp_model.parameters(), lr=0.01)
    
    # 5. 创建数据加载器
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)
    
    # 6. 训练循环
    for epoch in range(10):
        sampler.set_epoch(epoch)
        for batch in dataloader:
            inputs, labels = batch
            inputs, labels = inputs.to(rank), labels.to(rank)
            
            outputs = ddp_model(inputs)
            loss = criterion(outputs, labels)
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    
    # 7. 只在 rank 0 上保存模型
    if rank == 0:
        torch.save(ddp_model.module.state_dict(), 'model.pth')
    
    # 8. 清理
    dist.destroy_process_group()

几个容易忽略的细节：

  • 保存模型时用 ddp_model.module.state_dict()，不是 ddp_model.state_dict()
  • find_unused_parameters 默认 False。如果你的模型有参数没参与 loss 计算（比如 BERT 的某些层），要设为 True
  • DDP 会自动处理梯度同步，你不需要手动调用任何通信函数

4.5 避坑指南——我踩过的那些坑

做 DDP 训练这几年，我踩过的坑能写一本书。挑几个典型的说说：

问题 现象 原因 解决方案
训练速度不提升 8 卡比单卡只快 2-3 倍 DataLoader 的 num_workers 太小，GPU 在等数据 增大 num_workers，用 pin_memory=True
loss 不下降 loss 震荡或发散 学习率没调整。总 batch size 变大，lr 要相应增大 lr = base_lr × world_size（线性缩放规则）
显存溢出 OOM 错误 每个 GPU 的 batch_size 太大 减小 batch_size，或用梯度累积
通信卡死 程序卡在某个 epoch 不动 某个进程挂了，AllReduce 等不到它 检查所有进程是否正常启动，用 timeout 参数

我曾经遇到过一个特别诡异的问题：8 卡训练，前 10 个 epoch 正常，第 11 个 epoch 突然卡死。查了三天，最后发现是数据加载器里有个 random.shuffle 操作，导致不同进程的数据顺序不一致，DDP 的 hook 在等梯度同步时死锁了。

从那以后，我养成了一个习惯：所有随机操作都要用 DistributedSampler 控制，不要自己写 shuffle。

4.6 性能调优小技巧

最后分享几个我常用的调优技巧：

  • 梯度累积：如果显存不够，可以用梯度累积模拟大 batch。但要注意，梯度累积时 DDP 的 AllReduce 会在每次 loss.backward() 后触发，所以累积 N 步会触发 N 次通信。我一般累积 4-8 步。
  • 混合精度训练：DDP + AMP 是绝配。用 torch.cuda.amp 可以再提速 30%-50%。
  • 通信压缩：如果网络带宽是瓶颈（比如多机训练），可以试试梯度压缩，比如 torch.distributed.algorithms.ddp_comm_hooks 里的 PowerSGD。
  • 异步通信：DDP 默认是同步通信，所有 GPU 算完梯度才开始 AllReduce。如果 GPU 算力不均衡，可以试试异步模式，但实现起来比较复杂。

嗯，DDP 的内容差不多就这些。说白了，它就是个帮你搞定梯度同步的工具。你只要把进程组初始化、数据加载、模型包装这三件事做对，剩下的交给 DDP 就行。

记住一句话：DDP 不背锅，背锅的都是配置。遇到问题先检查环境变量，再检查数据加载，最后检查模型结构。按这个顺序排查，90% 的问题都能解决。

← 上一章 📖 返回目录 下一章 →