第4章：PyTorch DDP实战：DistributedDataParallel使用详解

分布式训练，说白了就是把一块GPU干不完的活，分给多块GPU一起干。PyTorch的DistributedDataParallel（DDP）是目前最主流的方案。我刚开始接触时也踩过不少坑，今天就把这些经验一次性倒出来。

4.1 DDP的核心原理

DDP不是简单的数据复制。它采用了一种叫「All-Reduce」的通信模式。每个GPU独立计算梯度，然后通过NCCL后端把所有梯度加起来求平均，再同步更新参数。

嗯，这里要注意：DDP在每个进程里维护一个完整的模型副本，而不是把模型切分到不同GPU上。这和模型并行有本质区别。

4.2 启动方式：torchrun vs multiprocessing

PyTorch提供了两种启动DDP的方式。我个人强烈推荐torchrun，它更现代、更稳定。

4.2.1 torchrun（推荐）

# 单机多卡
torchrun --nproc_per_node=4 train.py

# 多机多卡
torchrun --nnodes=2 --nproc_per_node=8 --rdzv_endpoint=master:1234 train.py

torchrun会自动处理环境变量设置。你只需要在代码里读取这些变量：

import torch
import torch.distributed as dist

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

4.2.2 multiprocessing（不推荐）

这是老方法，用torch.multiprocessing.spawn启动。我在项目中遇到过一个问题：子进程的异常很难捕获，调试起来特别痛苦。

# 不推荐，除非你有特殊需求
import torch.multiprocessing as mp

def train(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    # ... 训练逻辑

if __name__ == "__main__":
    mp.spawn(train, args=(4,), nprocs=4)

4.3 核心代码模板

这是我多年总结出来的DDP训练模板，直接拿去用：

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

def train():
    # 1. 初始化进程组
    dist.init_process_group("nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    
    # 2. 创建模型并包装
    model = MyModel().to(local_rank)
    ddp_model = DDP(model, device_ids=[local_rank])
    
    # 3. 数据加载（关键！）
    dataset = MyDataset()
    sampler = DistributedSampler(dataset)
    dataloader = DataLoader(dataset, sampler=sampler, batch_size=32)
    
    # 4. 优化器
    optimizer = torch.optim.Adam(ddp_model.parameters())
    
    # 5. 训练循环
    for epoch in range(10):
        sampler.set_epoch(epoch)  # 重要！保证每个epoch数据打乱
        for batch in dataloader:
            optimizer.zero_grad()
            loss = ddp_model(batch)
            loss.backward()
            optimizer.step()
    
    # 6. 清理
    dist.destroy_process_group()

if __name__ == "__main__":
    train()

4.4 性能调优

DDP的性能瓶颈通常不在计算，而在通信。我调优时主要关注这几个方面：

参数	推荐值	说明
batch_size	每个GPU 32-128	太小会导致通信开销占比过高
gradient_accumulation	2-8步	模拟更大batch，减少通信频率
find_unused_parameters	False	设为True会降低性能
bucket_cap_mb	25（默认）	梯度桶大小，影响通信粒度

4.5 常见坑与避坑指南

我曾经在项目里被这些坑折磨过，今天全部分享出来：

1. 数据加载不一致：每个进程必须加载相同的数据集，但用不同的采样器。不要手动切分数据！
2. 模型保存与加载：只保存rank 0的模型。保存时用model.module.state_dict()而不是model.state_dict()。
3. BN层问题：DDP默认同步BN，但小batch时效果反而差。可以用SyncBatchNorm.convert_sync_batchnorm()手动控制。
4. 日志打印：所有进程都会打印日志。我习惯用if rank == 0: print(...)来避免刷屏。
5. 梯度累积：如果用了梯度累积，记得在optimizer.step()前调用model.require_backward_grad_sync = False来跳过不必要的同步。

4.6 性能对比

我用ResNet-50在4卡V100上做过测试，结果如下：

配置	吞吐量（images/sec）	加速比
单卡	280	1x
4卡DDP（默认）	980	3.5x
4卡DDP（调优后）	1080	3.86x

你看，调优后接近线性加速。说白了，DDP的通信开销是固定的，只要计算量够大，加速比就接近理想值。