分布式训练：如何加速你的模型训练？

一、为什么要使用分布式训练

当模型比较大或训练数据比较多时，单 GPU 训练速度会非常慢。GPU 适合深度学习中大量的矩阵或向量计算（流式处理过程），但单个 GPU 的计算能力是有限的，因此需要多个 GPU 协同工作，这就是分布式训练。[pdf_17]

分布式训练的两大核心问题：

二、单机单卡基础操作（GPU 使用三步）

第一步：判断 GPU 可用性

import torch

# 判断当前机器是否有可用 GPU
torch.cuda.is_available()    # 返回 True 或 False

# 得到目前可用 GPU 数量
torch.cuda.device_count()    # 返回 GPU 数量

第二步：获得 GPU 实例

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

• torch.device 代表将 torch.Tensor 分配到的设备，是一个设备对象实例
• cuda:0 表示使用第一块 GPU，不声明 :0 默认从第一块开始

第三步：将数据或模型推到 GPU

data = torch.ones((3, 3))
print(data.device)   # 输出：cpu

device = torch.device("cuda:0")
data_gpu = data.to(device)
print(data_gpu.device)   # 输出：cuda:0

# 模型迁移同样使用 to() 函数
net = nn.Sequential(nn.Linear(3, 3))
net.to(device)

三、DataParallel（DP）——单进程控制多 GPU

1. 工作原理

DP 是单进程控制多 GPU 的方式，具体流程如下：[pdf_17]

2. 主要问题

• 主 GPU 负载过高：主 GPU 负载和使用率比其它 GPU 高，导致 GPU 负载不均衡
• 效率问题：需要频繁复制模型和汇总梯度

3. 使用方式

定义：

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

完整示例：

class ASimpleNet(nn.Module):
    def __init__(self, layers=3):
        super(ASimpleNet, self).__init__()
        self.linears = nn.ModuleList([nn.Linear(3, 3, bias=False) for i in range(layers)])
    def forward(self, x):
        print("forward batchsize is: {}".format(x.size()[0]))
        x = self.linears(x)
        x = torch.relu(x)
        return x

batch_size = 16
inputs = torch.randn(batch_size, 3)
labels = torch.randn(batch_size, 3)
inputs, labels = inputs.to(device), labels.to(device)
net = ASimpleNet()
net = nn.DataParallel(net)
net.to(device)

for epoch in range(1):
    outputs = net(inputs)

★

上述示例中，GPU 数量为 4，batch size 为 16，每个 GPU 获得的数据量都是 4 个。DataParallel 会自动帮我们将数据切分、加载到相应 GPU，将模型复制到相应 GPU，进行正向传播计算梯度并汇总。[pdf_17]

四、DistributedDataParallel（DDP）——多进程控制多 GPU

1. 工作原理

DDP 是多进程控制多 GPU 的方式。系统会为每个 GPU 创建一个进程，不再有主 GPU，每个 GPU 执行相同的任务。[pdf_17]

★

官方推荐：使用 DistributedDataParallel 进行分布式训练。[pdf_17]

2. 分布式基本概念

3. DDP 使用三步骤

第一步：初始化进程组

torch.distributed.init_process_group(backend, init_method=None, world_size=-1, rank, group_name)

# 使用 nccl 后端初始化
torch.distributed.init_process_group(backend="nccl")

第二步：模型并行化

torch.nn.parallel.DistributedDataParallel(module, device_ids=None, output_device=...)

用法与 DataParallel 基本相同，只需将 DataParallel 替换为 DistributedDataParallel：

net = torch.nn.parallel.DistributedDataParallel(net)

第三步：创建分布式数据采样器

使用 DDP 时，不再是从主 GPU 分发数据到其他 GPU 上，而是各 GPU 从自己的硬盘上读取属于自己的那份数据。

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
data_loader = DataLoader(train_dataset, batch_size=batch_size, sampler=train_sampler)

★

注意：在建立 DataLoader 的过程中，如果 sampler 参数不为 None，那么 shuffle 参数不应该被设置。[pdf_17]

五、DP 与 DDP 对比总结

★

DP 并不是完整的分布式计算，只是将一部分计算放到了多张 GPU 卡上，在计算梯度的时候，仍有主 GPU 负载过重的现象。而 DDP 刚好能解决 DP 的上述问题。[pdf_17]

六、实战：PyTorch 官方 ImageNet 分布式训练示例

1. DDP 初始化

if args.distributed:
    if args.dist_url == "env://" and args.rank == -1:
        args.rank = int(os.environ["RANK"])
    if args.multiprocessing_distributed:
        args.rank = args.rank * ngpus_per_node + gpu
    dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                            world_size=args.world_size, rank=args.rank)

• args.distributed 为 True 表示使用 DDP，反之表示使用 DP
• ngpus_per_node 表示每个节点的 GPU 数量

2. 模型并行化逻辑（CPU/DP/DDP 选择）

if not torch.cuda.is_available():
    print('using CPU, this will be slow')
elif args.distributed:
    # 使用 DDP
    if args.gpu is not None:
        torch.cuda.set_device(args.gpu)
        model.cuda(args.gpu)
        args.batch_size = int(args.batch_size / ngpus_per_node)
        args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
    else:
        model.cuda()
        model = torch.nn.parallel.DistributedDataParallel(model)
elif args.gpu is not None:
    torch.cuda.set_device(args.gpu)
    model = model.cuda(args.gpu)
else:
    # 使用 DP
    model = torch.nn.DataParallel(model).cuda()

3. 创建分布式数据采样器

if args.distributed:
    train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
else:
    train_sampler = None

train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
    num_workers=args.workers, pin_memory=True, sampler=train_sampler)

4. 为每个 GPU 启动进程

ngpus_per_node = torch.cuda.device_count()
if args.multiprocessing_distributed:
    args.world_size = ngpus_per_node * args.world_size
    mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
else:
    main_worker(args.gpu, ngpus_per_node, args)

5. 模型保存——只使用主进程中的一个 GPU 保存

if not args.multiprocessing_distributed or (args.multiprocessing_distributed
        and args.rank % ngpus_per_node == 0):
    save_checkpoint({
        'epoch': epoch + 1,
        'arch': args.arch,
        'state_dict': model.state_dict(),
        'best_acc1': best_acc1,
        'optimizer': optimizer.state_dict(),
    }, is_best)

★

注意：使用 DDP 意味着使用多进程，如果直接保存模型，每个进程都会执行一次保存操作。因此，只使用主进程中的一个 GPU 来保存即可。[pdf_17]

七、每课一练

问题：在 torch.distributed.init_process_group(backend="nccl") 函数中，backend 参数可选哪些后端，它们分别用于什么场景？[pdf_17]

解答：

八、核心要点总结

九、小结

1. 分布式训练的目的：当模型较大或训练数据较多时，通过多 GPU 协同工作加速训练。
2. 单机单卡 GPU 三步：torch.cuda.is_available() → torch.device("cuda:0") → data.to(device)。
3. DataParallel（DP）：单进程控制多 GPU，有主 GPU，负载不均衡，效率较低。
4. DistributedDataParallel（DDP）：多进程控制多 GPU，无主 GPU，负载均衡，速度更快，官方推荐使用。
5. DDP 三步骤：初始化进程组 → 模型并行化 → 创建分布式数据采样器（DistributedSampler）。
6. 模型保存注意：DDP 多进程时，只使用主进程中的一个 GPU 保存模型，避免重复保存。
7. backend 选择：nccl 用于 GPU 训练，gloo 用于 CPU 训练。