【问题解决】DDP | 如何使用 DDP 模式来训练模型

在训练 pytorch 模型时，多卡并行训练能够很大程度上提升模型的训练效率

一般来说，有两种多卡并行的训练方式：DP 和 DDP

一、DP 和 DDP 的区别

DP（Data Parallelism）和DDP（Distributed Data Parallelism）是深度学习中用于训练模型的两种并行计算策略。

1、Data Parallelism (DP):

DP通常指在单台机器的多个GPU上实施的数据并行处理。在这种设置中，模型的每个副本都放在不同的GPU上，每个GPU都会处理输入数据的不同批次。
所有GPU完成自己的前向和反向传播后，它们会将梯度发送到主GPU，主GPU会对这些梯度进行平均，然后更新模型的权重。
然后这些更新的权重会被复制到其他所有GPU上，以保持模型的一致性。
DP的缺点是它受限于单个机器的资源，因此当模型或数据集非常大时，可能会受到内存或带宽的限制。

2、Distributed Data Parallelism (DDP):

DDP是一种更加高级的并行计算策略，它允许跨多台机器的多个GPU进行模型训练。
在DDP中，每个节点（机器）都可能有一个或多个GPU，每个节点都在其GPU上运行模型的一个副本，并处理数据的不同部分。
DDP的关键在于它使用了一种更加高效的梯度聚合策略。每个节点都独立地完成前向和反向传播，并计算出梯度。梯度不是首先发送到主节点，而是在所有节点之间直接同步，通常是通过一种称为“All-Reduce”的操作。
这种方法减少了通信瓶颈，因为它不需要所有数据都通过单个主节点，并且可以更有效地扩展到大规模的计算资源。

总结来说，DP是单机多GPU的并行策略，而DDP是跨多台机器的多GPU并行策略。DDP通常在大规模分布式训练场景中更为有效，因为它可以更好地利用分布式系统的计算和存储资源。

二、如何将模型包装进 DDP

import torch
import torch.distributed as dist
import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
import datetime
import os# 初始化进程组
dist.init_process_group(backend='nccl', init_method='env://')# 设置本地设备
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
device = torch.device("cuda")# 创建模型
model = ...  # 替换为你的模型
device = torch.device("cuda:{}".format(rank))
model.to(device)# 重点一：包装数据集
# load training data
train_data = Dataset(...)
train_sampler = DistributedSampler(train_data)
training_data_loader = DataLoader(dataset=train_data, batch_size=batch_size, sampler=train_sampler, drop_last=True, num_workers=opt.num_workers)# 重点二：包装模型
ddp_model = DDP(model, device_ids=[rank])# 定义损失函数和优化器
criterion = ...
optimizer = ...# 训练循环
for epoch in range(num_epochs):# 重点三：设置 sampler 的 epoch，DistributedSampler 需要这个来维持各个进程之间的随机种子，也就是保证所有进程在数据洗牌时使用的随机种子是一致的，这样每个进程就会得到不同的数据子集，但整个训练集上的采样是一致的，train_sampler.sampler.set_epoch(epoch)for iteration, data in enumerate(training_data_loader):inputs, labels = datainputs, labels = inputs.to(device), labels.to(device)optimizer.zero_grad()outputs = ddp_model(inputs)loss = criterion(outputs, labels)loss.backward()  # DDP将在这里同步梯度optimizer.step()# 日志和其他操作if rank == 0:# 打印日志或写入日志文件...# 清理
dist.destroy_process_group()

如果你在使用DDP训练时没有调用 train_sampler.set_epoch(epoch)，那么在每个epoch中所有的进程都将以相同的方式从数据集中采样数据。这意味着每个进程将获得相同的数据子集，导致以下几个问题：

数据重复：所有的模型副本将在每个epoch中学习相同的数据，这样就减少了模型训练的总体数据多样性。
并行效率降低：数据的重复使用降低了并行训练的效率，因为模型的不同副本不再是在不同的数据子集上训练，而是在复制的数据上训练。
收敛问题：由于数据的多样性降低，模型可能更难收敛到一个好的解，或者可能导致过拟合，因为模型只看到了数据的一个子集。
泛化能力下降：模型的泛化能力可能会受到影响，因为它没有在整个数据集的不同采样上进行训练。
因此，为了确保有效的分布式训练和数据的多样性，在每个epoch开始时调用 train_sampler.set_epoch(epoch) 是很重要的。这将为每个进程提供一个不同的数据视图，从而使整个模型能够从整个数据集中学习，并提高训练的效率和最终模型的性能。
如果不使用这个 train_sampler.set_epoch(epoch)，也可能会导致 8 卡训练结果没有单卡训练结果好，因为不使用 train_sampler.set_epoch(epoch) 的话，即使有多个进程，每个进程在每个 epoch 中采样的数据将会是一样的，这意味着所有 gpu 卡在每个 epoch 都在训练相同的数据

三、如何训练

# 方法一：老方法
python -m torch.distributed.launch --nproc_per_node=8 --nnodes=1 --use_env --node_rank=0 --master_port=12348 train.py --aug xxx

python -m torch.distributed.launch：这部分是告诉Python运行torch.distributed.launch模块。这个模块是PyTorch的一部分，用于帮助用户启动多个工作进程进行分布式训练。
--nproc_per_node=8：这个参数指定每个节点（node）上要启动的进程数。在这个例子中，你将在单个节点上启动8个进程。
--nnodes=1：这个参数指定了参与分布式训练的节点总数。在这个例子中，你只使用了一个节点。
--node_rank=0：这个参数指定了当前节点的排名。在分布式训练中，每个节点都有一个唯一的排名，用于在节点之间通信。在只有一个节点的情况下，这个排名通常是0。
--master_port=12348：这个参数指定了主节点上用于分布式训练的通信端口。所有的节点都会连接到这个端口来进行通信。

# 方法二：新方法
OMP_NUM_THREADS=8 torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 --master_port=11347 train.py --aug xxx

OMP_NUM_THREADS=8: 这是一个环境变量，用于设置OpenMP使用的线程数。OpenMP是一个支持多平台共享内存并行编程的API。在这里，设置OMP_NUM_THREADS=8意味着每个进程将尝试使用8个线程进行计算，这有助于优化CPU上的并行计算性能。
torchrun: 这是PyTorch 1.9版本引入的一个新工具，用来替代python -m torch.distributed.launch。torchrun是一个简化的命令行工具，用于启动分布式训练。
--nproc_per_node=8: 这个参数指定了每个节点（node）上要启动的进程数。在这里，它设置为8，意味着在当前节点上会启动8个训练进程。
--nnodes=1: 这个参数指定了参与分布式训练的节点数。这里设置为1，表示只有一个节点参与训练。
--node_rank=0: 这个参数指定了当前节点的排名。因为只有一个节点，所以排名是0。在多节点训练中，每个节点会有一个唯一的排名。
--master_port=11347: 这个参数指定了主节点用于通信的端口。在分布式训练中，各个节点需要通过网络进行通信，这个端口就是用于这种通信的。
train.py: 这是你的训练脚本，torchrun会运行这个脚本作为分布式训练的一部分。