3、PyTorch DDP入门：DistributedDataParallel原理、初始化进程组、数据加载器改造、模型包装与同步、损失函数与梯度同步

好，咱们直接进入正题。DDP 这个东西，说白了就是 PyTorch 官方给咱们准备的多卡训练利器。我刚开始接触分布式训练时，也踩过不少坑，尤其是 DataParallel 那个老古董，性能差得让人抓狂。后来换成 DDP，才真正体会到什么叫「多卡加速」。这一节，我会把 DDP 的核心原理和实操步骤掰开揉碎讲清楚。

3.1 DDP 原理：它到底干了什么？

先说说原理。DDP 的全称是 DistributedDataParallel。它跟 DataParallel 最大的区别在于：每张卡都有自己的模型副本和优化器，数据被切分成多份，每张卡只处理自己那份。

为什么会这样设计？你想想看，如果所有卡共享一个模型，那通信开销会大得离谱。DDP 的做法是：每张卡独立前向传播、独立计算梯度，然后在反向传播时，通过 AllReduce 操作把梯度同步到所有卡上。这样每张卡拿到的梯度都是一样的，模型更新也就一致了。

我个人习惯把 DDP 的工作流程总结成三步：

1. 数据切分：每张卡拿到不同的 batch 数据
2. 独立计算：各自前向、反向，算出梯度
3. 梯度同步：AllReduce 把梯度平均后分发给所有卡

嗯，这里要注意：DDP 只在反向传播时做通信，前向传播是完全独立的。这也是它比 DataParallel 快很多的原因之一。

3.2 初始化进程组：第一步不能错

用 DDP 之前，必须先初始化进程组。这一步说白了就是告诉 PyTorch：「嘿，我要用几张卡，它们之间怎么通信。」

我见过不少新手在这步翻车。最常见的错误是忘记设置 MASTER_ADDR 和 MASTER_PORT。这两个环境变量指定了主节点的地址和端口，所有进程都要通过它来协调。

import torch
import torch.distributed as dist

def init_process(rank, world_size):
    # rank: 当前进程的编号（0,1,2,...）
    # world_size: 总进程数（等于 GPU 数量）
    dist.init_process_group(
        backend='nccl',          # GPU 训练用 nccl
        init_method='env://',    # 从环境变量读取配置
        rank=rank,
        world_size=world_size
    )
    # 设置当前 GPU
    torch.cuda.set_device(rank)

3.3 数据加载器改造：DistributedSampler 是关键

数据加载这块，核心就一句话：用 DistributedSampler 替换普通的 Sampler。

DistributedSampler 会自动把数据集切分成 world_size 份，每个进程只拿自己那份。而且它还能保证每个 epoch 的数据顺序不同——通过设置 shuffle=True 和调用 set_epoch() 实现。

from torch.utils.data import DataLoader, DistributedSampler

dataset = YourDataset(...)
sampler = DistributedSampler(
    dataset,
    num_replicas=world_size,
    rank=rank,
    shuffle=True
)
dataloader = DataLoader(
    dataset,
    batch_size=32,
    sampler=sampler,   # 用 sampler 替代 shuffle
    num_workers=4,
    pin_memory=True
)

# 每个 epoch 开始前必须调用
for epoch in range(num_epochs):
    sampler.set_epoch(epoch)  # 保证每个 epoch 数据顺序不同
    for batch in dataloader:
        # 训练代码
        pass

3.4 模型包装与同步：DDP 包装器

模型包装是最简单的一步。把模型放到 GPU 上，然后用 DistributedDataParallel 包一层就行。

from torch.nn.parallel import DistributedDataParallel as DDP

model = YourModel().cuda(rank)
model = DDP(model, device_ids=[rank])

# 之后正常使用即可
output = model(input_data)

嗯，这里有个细节：device_ids 参数必须指定当前进程的 GPU 编号。如果不指定，DDP 会默认使用 cuda:0，那就乱套了。

DDP 包装后，模型的参数访问方式会变。比如原来用 model.fc.weight，现在得用 model.module.fc.weight。我个人习惯在包装前保存一份 raw_model = model，方便后续操作。

3.5 损失函数与梯度同步：AllReduce 的幕后工作

损失函数这块，DDP 不会自动帮你做任何特殊处理。你该怎么算还怎么算。但梯度同步是自动发生的——DDP 在反向传播时，会拦截梯度计算，然后执行 AllReduce 操作。

具体来说，当你调用 loss.backward() 时，DDP 会在每个参数的梯度计算完成后，自动触发 AllReduce，把所有卡上的梯度取平均，然后写回每张卡。这样每张卡更新模型时，用的梯度都是一样的。

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

for batch in dataloader:
    optimizer.zero_grad()
    
    output = model(inputs)
    loss = criterion(output, labels)
    
    loss.backward()  # 这里 DDP 会自动同步梯度
    
    optimizer.step()  # 所有卡用相同的梯度更新模型

另外，如果你需要记录 loss 的值，建议只从 rank 0 进程打印，否则控制台会被刷爆。

if rank == 0:
    print(f'Epoch {epoch}, Loss: {loss.item()}')

3.6 完整训练脚本示例

最后，给一个完整的单机多卡训练脚本框架。你直接套用就行。

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

def train(rank, world_size):
    # 1. 初始化进程组
    dist.init_process_group('nccl', rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    
    # 2. 数据加载器
    dataset = YourDataset()
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)
    
    # 3. 模型包装
    model = YourModel().cuda(rank)
    model = DDP(model, device_ids=[rank])
    
    # 4. 优化器和损失函数
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    criterion = torch.nn.CrossEntropyLoss()
    
    # 5. 训练循环
    for epoch in range(10):
        sampler.set_epoch(epoch)
        for inputs, labels in dataloader:
            inputs, labels = inputs.cuda(rank), labels.cuda(rank)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            if rank == 0:
                print(f'Epoch {epoch}, Loss: {loss.item():.4f}')
    
    # 6. 清理
    dist.destroy_process_group()

if __name__ == '__main__':
    world_size = torch.cuda.device_count()
    torch.multiprocessing.spawn(train, args=(world_size,), nprocs=world_size)

好了，DDP 的入门内容就这些。说白了就是四步：初始化进程组、改造数据加载器、包装模型、正常训练。你只要按这个流程走，多卡训练基本不会出大问题。