PyTorch DDP单机多卡实战：从零到高效训练的完整指南

1. 为什么你需要DDP？从单卡到多卡的效率跃迁

如果你正在用PyTorch训练模型，尤其是那些参数动辄上亿的大模型，或者处理海量图像、文本数据，那你一定对“训练慢”这件事深有体会。我刚开始做AI项目那会儿，用一张显卡跑一个ResNet-50在ImageNet上的训练，一个epoch就要好几个小时，调一次参数等一天是家常便饭。后来显卡多了，很自然地就想：能不能把几张卡都用上，让训练快起来？

PyTorch确实提供了两种主流的多卡训练方案：DataParallel (DP)和DistributedDataParallel (DDP)。很多新手，包括当年的我，会先尝试DP，因为它太简单了，真的就是一行代码：model = nn.DataParallel(model)。把模型往上一套，代码几乎不用改，看起来数据就在多张卡上跑起来了。但用不了多久，你就会发现不对劲。我印象最深的一次是训练一个BERT-Large模型，明明用了4张卡，但其中一张卡的显存占用比其他三张高出快4个G，风扇狂转，而其他卡却有点“悠闲”。这就是DP著名的“负载不均衡”问题。

DP采用的是**Parameter Server（参数服务器）**模式。想象一下，有一个“总指挥”（默认是第0号GPU），它负责把模型复制到各个“工人”（其他GPU）上。每个“工人”处理一部分数据，算完梯度后，都要把梯度汇报给“总指挥”。“总指挥”收集所有梯度，平均一下，更新自己的模型参数，再把新参数分发给各个“工人”。这个过程中，“总指挥”的通信和计算压力巨大，成了瓶颈。所以你会看到第0张卡特别忙，显存也吃得最多，而其他卡的算力没有被充分利用。当模型很大或者数据批次（batch size）很大时，这个瓶颈会非常明显，有时甚至比单卡训练还慢，因为你多卡通信的开销可能已经抵消了并行计算带来的收益。

而DDP采用的是All-Reduce（全规约）算法。这个模式里，没有唯一的“总指挥”，每个GPU都是平等的“工人”。每个工人都有完整的模型副本，处理自己分到的那份数据。计算梯度后，所有工人聚在一起开个会（通过高速的NCCL后端），高效地互相通信，把所有梯度汇总、平均。这个平均操作是大家一起完成的，通信压力均匀分摊到了每张卡上。平均后的梯度，每个工人自己独立更新自己的模型参数。由于大家的起点（模型参数）一样，更新的梯度也一样，所以更新后的参数仍然保持一致。

这种去中心化的方式，让每张卡的负载非常均衡，通信效率也高得多。实测下来，在单机多卡（比如4卡、8卡）的场景下，DDP的训练速度几乎可以做到接近理想的线性加速。也就是说，4张卡训练，速度差不多是单卡的4倍。这才是我们压榨多卡硬件性能该有的样子。所以，尽管DDP的配置比DP多几步，但为了真正的效率，投入这点学习成本是完全值得的。官方也早已推荐使用DDP替代DP，即使是单机多卡。

2. 动手之前：彻底搞懂DDP的三个核心概念

在写代码之前，我们得先把DDP里几个绕不开的概念掰扯清楚。很多教程直接甩代码，但如果不明白这些参数的含义，一旦出错，debug起来会非常痛苦。我自己就曾经因为local_rank没设置对，导致所有进程都挤到一张卡上，闹了笑话。

第一个是rank。你可以把它理解为一个进程的“全局身份证号”。在分布式训练的世界里，每一个进程都有一个唯一的rank编号。如果是多机多卡，rank可能代表“第几号机器”；但在我们今天的主题——单机多卡里，rank就直接代表“第几号GPU进程”。假设我们用4张卡，那么就会启动4个进程，它们的rank分别是0, 1, 2, 3。

第二个是world_size。这个就简单了，它代表“全世界”总共有多少个进程在并肩作战。在单机多卡下，world_size就等于你使用的GPU数量。你告诉程序world_size=4，它就明白要拉起4个兄弟一起干活。

第三个是local_rank。这是最容易混淆，也最关键的一个。local_rank指的是在当前机器内部，这个进程使用的是哪一块本地GPU。在单机多卡环境下，rank和local_rank的值通常是一一对应的。比如rank为2的进程，它的local_rank通常也是2，意味着它使用本机的第2号物理GPU（注意，这里的编号可能受CUDA_VISIBLE_DEVICES环境变量影响）。

为什么需要local_rank？因为我们需要在每个进程内部，明确地告诉PyTorch：“嘿，你这个进程就去用那块叫local_rank的显卡！” 这是通过torch.cuda.set_device(local_rank)来实现的。这样，多个进程才不会争抢同一块GPU的资源。

我画个简单的表格帮你理清关系（以单机4卡为例）：

这三个概念是DDP通信和协作的基石。DDP的启动器（如torchrun）会自动帮你管理这些rank信息，并通过命令行参数--local-rank（或环境变量）传递给每个进程，我们只需要在代码里接住它并用起来就行。

3. 从零开始：将一个单卡训练脚本改造为DDP

光讲理论有点枯燥，我们直接拿一个最普通的PyTorch单卡训练脚本开刀，一步步把它改成DDP多卡版本。我会把每一步的意图和注意事项都讲明白。假设我们有一个简单的图像分类任务。

第一步：导入必要的模块并解析参数。首先，除了常规的torch、nn、DataLoader，我们必须导入分布式相关的核心模块：torch.distributed和DistributedSampler。同时，我们需要使用argparse来接收一个由启动器自动传递的关键参数：--local_rank。

import torch import torch.nn as nn import torch.distributed as dist from torch.utils.data import Dataset, DataLoader from torch.utils.data.distributed import DistributedSampler import argparse # 解析参数 parser = argparse.ArgumentParser() parser.add_argument('--batch_size', type=int, default=32) parser.add_argument('--epochs', type=int, default=10) # 这个参数非常重要！它由 torch.distributed.launch 或 torchrun 自动传入，不需要手动指定。 parser.add_argument('--local_rank', type=int, default=0) args = parser.parse_args()

注意，--local_rank这个参数我们不需要也不应该在命令行手动指定。当你用DDP的专用方式启动脚本时，启动器会为每个进程自动赋予一个不同的local_rank值。

第二步：初始化进程组并设置当前GPU。这是DDP设置的核心步骤，必须在创建模型和数据加载器之前完成。

# 初始化进程组。backend指定通信后端，单机多卡一定要用'ncc'，这是NVIDIA GPU间最高效的通信库。 dist.init_process_group(backend='nccl') # 获取当前进程的本地rank，并设置当前进程使用的GPU设备。 # 这行代码确保了不同进程使用不同的GPU。 torch.cuda.set_device(args.local_rank) device = torch.device('cuda', args.local_rank)

init_process_group就像是让所有进程加入了一个微信群，他们之后可以通过这个“群”来同步梯度。nccl后端是针对NVIDIA GPU优化的，性能远好于gloo（后者通常在CPU分布式训练或多机场景下作为备选）。

第三步：用DistributedSampler包装数据集。在单卡训练中，我们通常直接在DataLoader里设置shuffle=True来打乱数据。但在DDP中，每个进程只处理数据的一部分，我们需要一个“分配器”来确保：

1. 不同进程拿到不同的数据子集，避免数据重复。
2. 每个epoch数据都能被充分打乱。 这就是DistributedSampler的工作。

# 假设你已经定义好了你的 dataset # dataset = YourDataset(...) # 创建DistributedSampler，它会根据总进程数（world_size）和当前进程rank来分配数据索引。 sampler = DistributedSampler(dataset, shuffle=True) # 创建DataLoader。注意：这里传入了sampler，就不要再设置shuffle=True了。 # sampler自己会处理打乱逻辑。batch_size是每个进程的批次大小。 dataloader = DataLoader(dataset, batch_size=args.batch_size, sampler=sampler)

这里有个关键点：args.batch_size现在指的是每个GPU进程的批次大小。如果你设置batch_size=32，并且用了4张卡，那么全局的有效批次大小就是32 * 4 = 128。这是DDP训练中调整学习率等重要超参数时需要牢记的。

第四步：将模型移至GPU并用DDP包装。这一步和DP很像，但本质不同。

# 1. 创建模型实例 model = YourModel().to(device) # 2. 用DDP包装模型 model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)

device_ids和output_device参数告诉DDP模型的主设备是当前local_rank对应的GPU。经过DDP包装后，你的model在正向传播、反向传播时，内部已经自动包含了跨进程的梯度同步（All-Reduce）逻辑，对你来说是透明的。

第五步：在训练循环中设置sampler。为了让每个epoch的数据划分不同，需要在每个epoch开始时调用sampler.set_epoch(epoch)。

for epoch in range(args.epochs): # 关键！设置当前epoch，保证每个进程每个epoch的数据划分都不同 dataloader.sampler.set_epoch(epoch) for batch_idx, (data, target) in enumerate(dataloader): data, target = data.to(device), target.to(device) # ... 正常的训练步骤 ... output = model(data) loss = criterion(output, target) optimizer.zero_grad() loss.backward() # 在这里，DDP已经自动完成了所有进程的梯度同步！ optimizer.step()

第六步：处理日志和模型保存。由于多个进程同时在跑，如果你不做任何处理，每个进程都会打印日志，你的终端会被刷屏。通常我们只让rank 0的主进程来打印日志和保存模型，避免重复操作。

if dist.get_rank() == 0: # 判断是否是主进程（rank 0） print(f'Epoch [{epoch+1}/{args.epochs}], Loss: {loss.item():.4f}') torch.save(model.module.state_dict(), 'model.pth') # 注意是 model.module

这里又一个重点：保存模型时，用的是model.module.state_dict()。因为model现在是一个DDP包装后的对象，其内部的原始模型可以通过.module属性访问。

4. 如何正确启动DDP训练？两种主流方式详解

代码写好了，怎么运行它？你不能再用python train.py了。PyTorch提供了两种启动多进程的方式，我强烈推荐第二种。

方式一：使用torch.distributed.launch（旧版，但仍可用）这是比较传统的启动方式，你需要指定总进程数（即GPU数）。

python -m torch.distributed.launch --nproc_per_node=4 train.py --batch_size 32 --epochs 50

• --nproc_per_node=4：指定在当前节点（机器）上启动4个进程。
• 你的训练脚本train.py里必须能解析--local_rank参数（就像我们第一步做的那样），启动器会自动为每个进程传入不同的值。

方式二：使用torchrun（新版推荐）从PyTorch 1.9+开始，官方推荐使用torchrun，它是torch.distributed.launch的升级版，API更简洁，功能也更强大。

torchrun --nproc_per_node=4 train.py --batch_size 32 --epochs 50

命令几乎一样，只是把python -m torch.distributed.launch换成了torchrun。torchrun会自动设置一些必要的环境变量，包括LOCAL_RANK，你在代码里可以通过os.environ['LOCAL_RANK']来获取，或者依然用argparse解析--local-rank（注意这里是带中划线的）。我个人的习惯是在代码里同时兼容两种获取方式：

local_rank = int(os.environ.get('LOCAL_RANK', args.local_rank))

启动命令执行后，你会看到程序瞬间启动了4个进程，它们各自占用一块GPU，开始协同训练。第一次看到终端里飞速滚动的、来自不同进程的日志（如果你还没做只让rank 0打印的处理），可能会觉得有点乱，但这正是分布式训练在工作的标志。

5. 避坑指南：DDP实战中常见的“坑”与解决方案

踩过坑才能成长。下面是我和同事们在实际项目中用DDP时，遇到过的一些典型问题及其解决办法。

坑一：CUDA error: out of memory (OOM)这可能是最常遇到的错误。在DDP中，OOM可能原因更复杂。

• 每个进程的batch size过大：记住，DDP中设置的batch_size是每个GPU的。如果你习惯单卡设batch_size=64，4卡DDP时还设64，那全局batch size就是256，显存占用直接x4。解决方案是适当调小每个进程的batch_size。
• 没有设置torch.cuda.set_device：如果忘记设置，默认所有进程都可能尝试使用第0块GPU，导致0号卡OOM。务必确保每个进程都正确设置了设备。
• 模型或数据没有.to(device)：如果数据或模型还在CPU上，但计算图涉及到了分布在GPU上的DDP模型，可能会引发意想不到的显存问题。确保所有计算都在正确的设备上进行。

坑二：训练速度没有提升，甚至更慢如果发现4卡训练速度和单卡差不多，甚至更慢，需要排查：

• 通信开销过大：对于非常小的模型（比如几万参数），梯度同步的通信时间可能比计算时间还长，导致加速比不佳。DDP更适合大模型。
• 数据加载是瓶颈：如果DataLoader的num_workers设得太小（比如为0），数据预处理可能跟不上GPU计算。可以尝试增大num_workers。但要注意，在DDP中，每个进程都会创建自己的数据加载worker，总worker数会是num_workers * world_size，设置过大可能耗尽系统内存。一般设置为4 * GPU数量是个不错的起点。
• 没有使用DistributedSampler或使用错误：如果每个进程都读取了全部数据，那计算量没变，还增加了通信开销，自然会慢。

坑三：验证或测试时的重复计算问题在训练中，我们通常用DistributedSampler来分配训练数据。但在验证或测试时，我们往往希望每个进程都跑一遍完整的数据集，然后汇总指标。如果你不小心在验证的DataLoader里也用了DistributedSampler，那每个进程只会验证一部分数据，最后你得到的准确率是“部分数据”上的，是错误的。 解决方案是，为验证集创建一个普通的DataLoader（使用默认的SequentialSampler或RandomSampler），并且只在rank 0进程上进行验证和日志记录。或者，你也可以继续用DistributedSampler，但需要确保每个进程拿到全部验证数据（设置shuffle=False），并在计算指标时使用torch.distributed的all_gather等函数来跨进程收集结果。

坑四：随机种子与可复现性为了实验可复现，我们通常会设置随机种子。但在DDP中，如果你只在主进程设置torch.manual_seed，其他进程的随机状态可能不同，导致数据划分、模型初始化等出现不一致。 解决方案是，在初始化进程组之后，为每个进程设置相同的随机种子，并确保DistributedSampler的shuffle逻辑也基于此种子。

def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) set_seed(42 + dist.get_rank()) # 可以给不同进程一个基础偏移，但通常设成一样也行 # 然后在每个epoch: dataloader.sampler.set_epoch(epoch)

坑五：保存和加载检查点保存时，如之前所说，用model.module.state_dict()。加载时，如果你是从一个单卡训练的检查点开始做DDP微调，需要先将权重加载到model.module上。

# 加载检查点 checkpoint = torch.load('checkpoint.pth', map_location=device) model.module.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) # 注意：如果优化器状态是在多卡环境下保存的，加载可能更复杂，需要处理不同设备。

更稳健的做法是，在保存时不仅存模型权重，也存下epoch、loss等信息，并且只由rank 0进程来执行保存操作。

6. 性能调优：让你的DDP训练再快一点

基础功能跑通后，我们可以关注一些优化点，进一步榨干硬件性能。

优化一：梯度累积应对超大Batch Size有时，我们受限于单卡显存，无法设置太大的batch_size，但学术论文或某些任务又要求很大的全局batch size。这时可以用梯度累积。原理很简单：让每个进程连续计算多个小batch，但不立即更新权重（optimizer.step()），而是将多个batch的梯度累加起来，等累积到目标“虚拟batch size”时，再执行一次真正的权重更新和梯度清零。

accumulation_steps = 4 # 累积4步 for batch_idx, (data, target) in enumerate(dataloader): ... loss = criterion(output, target) # 将loss除以累积步数，使得梯度累加的平均效果与大批次一致 loss = loss / accumulation_steps loss.backward() if (batch_idx + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()

在DDP中，梯度累积的代码和单卡几乎一样，因为DDP的梯度同步发生在loss.backward()时。每个进程独立累积梯度，在optimizer.step()时，大家基于同步后的平均梯度进行更新。

优化二：混合精度训练（AMP）混合精度训练即用半精度（float16）做前向和反向传播，用单精度（float32）维护权重。这能大幅减少显存占用，并可能利用Tensor Core加速计算。PyTorch提供了torch.cuda.amp模块，与DDP兼容性很好。

from torch.cuda.amp import autocast, GradScaler scaler = GradScaler() # 梯度缩放，防止半精度下梯度下溢 for data, target in dataloader: optimizer.zero_grad() with autocast(): # 自动混合精度上下文 output = model(data) loss = criterion(output, target) # 用scaler缩放损失，反向传播 scaler.scale(loss).backward() # scaler更新优化器，并自动unscale梯度 scaler.step(optimizer) # 更新scaler的缩放因子 scaler.update()

将AMP和DDP结合，是当前训练大模型的标配，能显著提升训练速度和扩大模型容量。

优化三：DataLoader的配置DataLoader的num_workers和pin_memory对数据加载速度影响巨大。

• num_workers: 如前所述，设置过小是瓶颈，设置过大会占用大量内存。建议从4开始，根据CPU核心数和内存情况调整。监控GPU利用率，如果发现GPU经常空闲等待数据，就适当增加num_workers。
• pin_memory=True: 将数据锁页内存中，可以加速从CPU到GPU的数据传输。在数据预处理不复杂的情况下，开启它通常能带来一点性能提升。

7. 完整实战代码与进阶思考

最后，我把前面所有步骤整合成一个完整的、带有验证和模型保存的DDP训练脚本框架。你可以以此为基础，替换成自己的模型和数据集。

import os import torch import torch.nn as nn import torch.distributed as dist import torch.optim as optim from torch.utils.data import Dataset, DataLoader from torch.utils.data.distributed import DistributedSampler from torch.cuda.amp import autocast, GradScaler import argparse def main(): # 1. 解析参数 parser = argparse.ArgumentParser() parser.add_argument('--local_rank', type=int, default=0, help='local rank for distributed training') parser.add_argument('--batch_size', type=int, default=32) parser.add_argument('--epochs', type=int, default=100) parser.add_argument('--lr', type=float, default=0.001) args = parser.parse_args() # 2. 初始化分布式环境 dist.init_process_group(backend='nccl') torch.cuda.set_device(args.local_rank) device = torch.device('cuda', args.local_rank) # 3. 固定随机种子（可选，为了可复现） seed = 42 torch.manual_seed(seed + args.local_rank) # 4. 准备数据 # train_dataset = YourDataset(...) # val_dataset = YourDataset(...) train_sampler = DistributedSampler(train_dataset, shuffle=True) train_loader = DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler, num_workers=4, pin_memory=True) # 验证集不用DistributedSampler，或者用但要注意指标收集 val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=4, pin_memory=True) # 5. 构建模型 model = YourModel().to(device) model = nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank) # 6. 定义损失函数和优化器 criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=args.lr) scaler = GradScaler() # 混合精度训练 # 7. 训练循环 for epoch in range(args.epochs): # 设置epoch，保证每个epoch数据划分不同 train_loader.sampler.set_epoch(epoch) model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() with autocast(): output = model(data) loss = criterion(output, target) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() # 只在主进程打印日志 if args.local_rank == 0 and batch_idx % 100 == 0: print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] Loss: {loss.item():.6f}') # 8. 验证（只在主进程进行） if args.local_rank == 0: model.eval() val_loss = 0 correct = 0 with torch.no_grad(): for data, target in val_loader: data, target = data.to(device), target.to(device) output = model(data) val_loss += criterion(output, target).item() pred = output.argmax(dim=1, keepdim=True) correct += pred.eq(target.view_as(pred)).sum().item() val_loss /= len(val_loader.dataset) accuracy = 100. * correct / len(val_loader.dataset) print(f'\nValidation set: Average loss: {val_loss:.4f}, Accuracy: {correct}/{len(val_loader.dataset)} ({accuracy:.2f}%)\n') # 9. 保存模型（只在主进程） torch.save({ 'epoch': epoch, 'model_state_dict': model.module.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': val_loss, }, f'checkpoint_epoch_{epoch}.pth') # 10. 清理分布式进程组 dist.destroy_process_group() if __name__ == '__main__': main()

这个脚本已经是一个功能相对完整的生产级DDP训练框架雏形。当你熟练掌握了单机多卡DDP后，你会发现它的思想是通用的。未来如果你需要扩展到多机多卡，主要的改动在于初始化进程组时需要指定init_method（如使用环境变量MASTER_ADDR和MASTER_PORT），而模型、数据并行部分的代码几乎可以保持不变。DDP的设计很好地屏蔽了底层通信的复杂性，让我们能更专注于模型和算法本身。