4、PyTorch DDP入门：DistributedDataParallel的基本用法

聊到多卡训练，DDP（DistributedDataParallel）是绕不开的核心工具。说实话，我刚开始接触分布式训练时，也踩过不少坑。那时候觉得DDP的API看着挺简单，但真正跑起来，各种奇奇怪怪的问题就冒出来了。今天咱们就把它彻底讲透。

4.1 为什么是DDP，而不是DP？

PyTorch里其实有两个并行工具：DataParallel（DP）和DistributedDataParallel（DDP）。很多新手会问，DP用起来更简单，为什么非要学DDP？

我个人的经验是：DP只适合单机小规模实验，生产环境千万别用。原因有三：

DP是单进程多线程，受限于Python的GIL锁，多卡扩展性很差
DP每次前向传播都要把模型复制到各卡，通信开销巨大
DP的梯度同步是串行的，GPU利用率上不去

而DDP是多进程的，每个进程独立维护一个模型副本，只在反向传播时做梯度同步。说白了，DDP才是真正为大规模分布式训练设计的方案。

4.2 初始化进程组：NCCL vs GLOO

使用DDP的第一步，就是初始化进程组。这一步很多人容易搞混，我详细说说。

进程组初始化通过 torch.distributed.init_process_group() 完成。核心参数就两个：后端（backend）和通信方式。

后端	适用场景	特点
NCCL	GPU训练（推荐）	NVIDIA官方出品，GPU间通信效率最高
GLOO	CPU训练或调试	通用性强，但GPU通信效率不如NCCL
MPI	高性能计算集群	需要额外安装MPI库，不常用

我在项目中遇到过一个问题：用GLOO后端跑GPU训练，结果发现训练速度比单卡还慢。后来排查发现，GLOO在GPU间的通信走的是CPU内存中转，效率极低。所以记住：GPU训练永远选NCCL。

核心代码：初始化进程组

import torch
import torch.distributed as dist

# 初始化进程组
dist.init_process_group(
    backend='nccl',          # GPU训练用nccl
    init_method='env://',    # 从环境变量读取rank和world_size
    rank=rank,               # 当前进程的编号
    world_size=world_size    # 总进程数（=GPU数量）
)

4.3 环境变量与rank机制

DDP的多进程通信依赖一套rank机制。每个进程都有一个唯一的rank编号，从0到world_size-1。rank 0通常被当作主进程，负责保存模型、打印日志等。

进程间通信需要知道彼此的存在。常用的方式有两种：

环境变量方式（env://）：通过设置 MASTER_ADDR 和 MASTER_PORT 环境变量，让所有进程连接到同一个主节点
共享文件方式（file://）：通过一个共享文件来同步进程信息，适合没有网络的环境

我个人习惯用环境变量方式，配置起来更灵活。举个例子：

# 在rank 0的机器上
export MASTER_ADDR=192.168.1.100
export MASTER_PORT=29500
export WORLD_SIZE=4
export RANK=0

# 在其他机器上
export MASTER_ADDR=192.168.1.100
export MASTER_PORT=29500
export WORLD_SIZE=4
export RANK=1  # 或2、3

小技巧：MASTER_PORT建议选一个不常用的端口，避免冲突。我一般用29500，这是PyTorch的默认端口。

4.4 使用DDP包装模型

初始化完进程组，下一步就是用DDP包装模型。这一步其实很简单：

from torch.nn.parallel import DistributedDataParallel as DDP

# 创建模型
model = MyModel().to(device)

# 用DDP包装
model = DDP(model, device_ids=[local_rank])

这里有个关键点：device_ids 参数要传入当前进程对应的GPU编号。比如你有4张卡，local_rank就是0、1、2、3。

嗯，这里要注意：DDP只负责梯度同步，不负责数据分发。数据加载需要自己用 DistributedSampler 来切分。

4.5 数据加载：DistributedSampler

DDP模式下，每个进程只处理一部分数据。如果不用DistributedSampler，所有进程都会读到相同的数据，那就白训练了。

from torch.utils.data.distributed import DistributedSampler

dataset = MyDataset()
sampler = DistributedSampler(
    dataset,
    num_replicas=world_size,
    rank=rank
)
dataloader = DataLoader(
    dataset,
    batch_size=32,
    sampler=sampler,
    shuffle=False  # sampler自带shuffle功能
)

我曾经犯过一个低级错误：在DDP模式下还设置 shuffle=True，结果每个epoch数据打乱方式不一致，导致模型收敛变慢。后来才发现DistributedSampler内部已经处理了shuffle，不需要再额外设置。

重要提醒：每个epoch开始时，记得调用 sampler.set_epoch(epoch)，否则每个epoch的数据顺序都一样，模型容易过拟合。

4.6 launch命令详解

写好了训练脚本，怎么启动呢？PyTorch提供了两种启动方式。

方式一：torchrun（推荐）

从PyTorch 1.10开始，官方推荐用 torchrun 替代老旧的 torch.distributed.launch。用法很简单：

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

# 多机训练（每机4卡，共2机）
torchrun --nnodes=2 --nproc_per_node=4 --rdzv_endpoint=192.168.1.100:29500 train.py

torchrun会自动设置环境变量：LOCAL_RANK、RANK、WORLD_SIZE。你只需要在脚本里读取就行：

local_rank = int(os.environ['LOCAL_RANK'])
rank = int(os.environ['RANK'])
world_size = int(os.environ['WORLD_SIZE'])

方式二：torch.distributed.launch（旧版）

如果你还在用老版本PyTorch，可能会遇到这个：

python -m torch.distributed.launch --nproc_per_node=4 train.py

不过我个人建议尽快迁移到torchrun。它解决了旧版的一些痛点，比如自动处理错误重启、更好的日志输出等。

4.7 完整训练模板

最后，我贴一个完整的DDP训练模板，这是我项目里一直在用的结构：

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

def main():
    # 1. 初始化进程组
    dist.init_process_group(backend='nccl')
    local_rank = int(os.environ['LOCAL_RANK'])
    torch.cuda.set_device(local_rank)
    
    # 2. 创建模型并包装
    model = MyModel().cuda(local_rank)
    model = DDP(model, device_ids=[local_rank])
    
    # 3. 数据加载
    dataset = MyDataset()
    sampler = DistributedSampler(dataset)
    dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)
    
    # 4. 优化器
    optimizer = torch.optim.Adam(model.parameters())
    
    # 5. 训练循环
    for epoch in range(10):
        sampler.set_epoch(epoch)
        for data, target in dataloader:
            data, target = data.cuda(local_rank), target.cuda(local_rank)
            optimizer.zero_grad()
            output = model(data)
            loss = F.cross_entropy(output, target)
            loss.backward()
            optimizer.step()
        
        # 只在rank 0上保存模型
        if dist.get_rank() == 0:
            torch.save(model.module.state_dict(), f'checkpoint_{epoch}.pt')

if __name__ == '__main__':
    main()

避坑指南：保存模型时要用 model.module.state_dict() 而不是 model.state_dict()。因为DDP包装后，模型多了一层module属性。我曾经因为这个bug，加载模型时参数名对不上，排查了半天。

4.8 知识体系总览

为了让你更直观地理解DDP的整体流程，我画了一张图：

这张图把DDP的核心流程串起来了。你想想看，从初始化到模型包装，再到数据切分和训练循环，每一步都有它的道理。刚开始接触时可能会觉得步骤多，但跑通一次之后，后面就顺了。

好了，DDP的基本用法就讲到这里。记住：初始化进程组、DDP包装模型、DistributedSampler切分数据、torchrun启动，这四个步骤缺一不可。下次遇到多卡训练，按这个流程走，基本不会出大问题。

想要更多资料，联系可可

📚 课程目录

4、PyTorch DDP入门：DistributedDataParallel的基本用法

4.1 为什么是DDP，而不是DP？

4.2 初始化进程组：NCCL vs GLOO

4.3 环境变量与rank机制

4.4 使用DDP包装模型

4.5 数据加载：DistributedSampler

4.6 launch命令详解

方式一：torchrun（推荐）

方式二：torch.distributed.launch（旧版）

4.7 完整训练模板

4.8 知识体系总览