3. 单机多卡同步训练实战：DDP 原理与代码实战

说实话，很多同学第一次接触多卡训练时，第一反应就是去查「怎么用 DataParallel」。我当年也是这么干的。但用了一段时间后，发现 DataParallel 有个硬伤——它把模型复制到每张卡上，但梯度计算还是靠主卡来汇总。说白了，主卡成了瓶颈，卡一多反而更慢。

后来我换成了 DistributedDataParallel（DDP），才真正体会到什么叫「多卡就该这么玩」。DDP 的核心思想其实很简单：每张卡各自维护一份模型副本，各自算梯度，然后通过 all-reduce 操作把梯度同步一下。没有主卡依赖，没有通信瓶颈。

核心区别一句话总结：

  • DataParallel：主卡汇总梯度，再广播回去。卡越多，主卡越累。
  • DistributedDataParallel：每张卡独立算梯度，通过 all-reduce 同步。各卡平等，效率高。

3.1 DDP 的工作原理

我们先画一张图，把 DDP 的流程理清楚。我个人习惯在讲代码之前，先让脑子里有个「地图」。

DDP 单机多卡同步训练流程 DistributedSampler 数据分片到各卡 GPU 0 GPU 1 GPU 2 前向传播 前向传播 前向传播 反向传播 反向传播 反向传播 All-Reduce 梯度同步 各卡独立更新参数（同步后梯度一致）

你看，整个流程其实就四步：数据分片 → 各卡独立前向/反向 → all-reduce 同步梯度 → 各卡更新参数。这里有个关键点：all-reduce 是在反向传播过程中自动触发的，你不需要手动调用。DDP 在注册了 hook 之后，梯度算完就自动同步了。

3.2 代码实战：从单卡到 DDP

好，我们直接上代码。我假设你已经有一个单卡训练的脚本了，现在要改成 DDP。其实改动量不大，核心就几个步骤。

3.2.1 初始化进程组

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

def setup(rank, world_size):
    """初始化分布式环境"""
    # 设置后端为 NCCL（NVIDIA 的通信库，速度最快）
    dist.init_process_group(
        backend='nccl',
        init_method='tcp://127.0.0.1:23456',
        rank=rank,
        world_size=world_size
    )
    # 把当前进程绑定到指定 GPU
    torch.cuda.set_device(rank)

我个人的习惯：init_method 我一般用 'env://' 配合环境变量，这样更灵活。但初学者用 tcp 方式更直观，不容易出错。

3.2.2 分布式采样器

这个环节特别容易踩坑。你想想看，如果每张卡都读取完整的数据集，那每张卡算出来的梯度其实是一样的，多卡训练就失去了意义。所以我们需要 DistributedSampler 来给每张卡分配不同的数据子集。

from torch.utils.data import DataLoader, Dataset
from torch.utils.data.distributed import DistributedSampler

# 假设你有一个自定义数据集
dataset = MyDataset(...)

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

# DataLoader 使用这个采样器
dataloader = DataLoader(
    dataset,
    batch_size=32,      # 注意：这是每张卡的 batch_size
    sampler=sampler,
    num_workers=4,
    pin_memory=True
)

我曾经犯过的错：一开始我把 batch_size 设成了 32 × 卡数，结果每张卡又拿了 32 条数据，总 batch 变成了 32 × 卡数 × 卡数。嗯，那次训练直接 OOM 了。记住：batch_size 是每张卡的 batch_size，总 batch_size = batch_size × 卡数。

3.2.3 包装模型

from torch.nn.parallel import DistributedDataParallel as DDP

def train(rank, world_size):
    setup(rank, world_size)
    
    # 创建模型，移到对应 GPU
    model = MyModel().to(rank)
    
    # 用 DDP 包装
    model = DDP(model, device_ids=[rank])
    
    # 优化器
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    
    # 损失函数
    criterion = torch.nn.CrossEntropyLoss()
    
    for epoch in range(num_epochs):
        # 重要：每个 epoch 要设置 sampler 的 epoch
        sampler.set_epoch(epoch)
        
        for batch_idx, (data, target) in enumerate(dataloader):
            data, target = data.to(rank), target.to(rank)
            
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            
            if batch_idx % 10 == 0 and rank == 0:
                print(f'Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}')
    
    cleanup()

这里有个细节：sampler.set_epoch(epoch) 这行代码很多人会漏掉。如果不设置，每个 epoch 的数据打乱顺序是一样的，模型学到的模式会有偏差。我刚开始也忘了，后来发现验证集准确率波动特别大，排查了半天才找到原因。

3.2.4 启动训练

def main():
    world_size = torch.cuda.device_count()
    print(f"发现 {world_size} 张 GPU")
    
    # 使用 spawn 启动多进程
    mp.spawn(
        train,
        args=(world_size,),
        nprocs=world_size,
        join=True
    )

if __name__ == "__main__":
    main()

3.3 模型保存与加载

模型保存这块，很多同学会困惑：DDP 包装后的模型怎么保存？直接保存 model.state_dict() 行不行？

答案是：可以，但要注意取 model.module。DDP 把原始模型包装了一层，如果你直接保存 model.state_dict()，加载的时候也要用 DDP 包装才能加载。这显然不合理。

3.3.1 保存模型（推荐方式）

def save_checkpoint(model, optimizer, epoch, filepath):
    # 关键：取 model.module
    if isinstance(model, DDP):
        model_state = model.module.state_dict()
    else:
        model_state = model.state_dict()
    
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model_state,
        'optimizer_state_dict': optimizer.state_dict(),
    }
    torch.save(checkpoint, filepath)
    print(f"Checkpoint saved to {filepath}")

3.3.2 加载模型

def load_checkpoint(filepath, model, optimizer=None):
    checkpoint = torch.load(filepath, map_location='cpu')
    
    # 加载模型参数
    if isinstance(model, DDP):
        model.module.load_state_dict(checkpoint['model_state_dict'])
    else:
        model.load_state_dict(checkpoint['model_state_dict'])
    
    if optimizer:
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    
    return checkpoint['epoch']

我个人的做法：保存的时候统一用 model.module.state_dict()，加载的时候统一用 model.load_state_dict()。这样不管训练时用不用 DDP，保存的 checkpoint 都是通用的。你想想看，如果哪天你想在单卡上做推理，直接加载就行，不用改代码。

3.4 避坑指南

最后，我把自己踩过的坑整理一下，希望能帮你省点时间。

常见问题 原因 解决方案
训练 loss 不下降 DistributedSampler 没设置 set_epoch 每个 epoch 前调用 sampler.set_epoch(epoch)
OOM 内存溢出 batch_size 理解错误，总 batch 太大 batch_size 设为每张卡的大小，不要乘以卡数
模型加载报错 key 不匹配 保存了 DDP 包装后的 state_dict 保存时用 model.module.state_dict()
多卡训练比单卡还慢 batch_size 太小，通信开销占比大 适当增大每张卡的 batch_size

嗯，基本上掌握了这些，单机多卡训练这块你就没什么大问题了。DDP 的设计其实很优雅——它把分布式通信的复杂性封装在框架内部，你只需要关注模型和数据本身。说白了，就是「框架帮你搞定通信，你只管算梯度」。

想要更多资料，联系可可

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