pytorch单机多卡DistributedDataParallel (DDP)分布式训练

本文内容参考链接：
当代人应当掌握的5种Pytorch并行训练方法（单机多卡）【CSDN】
pytorch单机多卡训练【知乎】
PyTorch分布式训练【简书】
pytorch基于DistributedDataParallel进行单机多卡的分布式训练【CSND】

原作者写的详细易懂，本人进行粗略总结，并自己具体实施，详细内容请移步原文。

pytorch分布式训练

笔者所知道的常见分布式训练方式有两种，第一种是nn.DataParallel (DP)，第二种是nn.parallel.DistributedDataParallel (DDP)。

DP：（使用单进程控）将模型和数据加载到多个 GPU 中，控制数据在 GPU 之间的流动，协同不同 GPU 上的模型进行并行训练。DataParallel的整个并行训练过程利用python多线程实现。

• 优点： 部署起来简单方便，不需要过多的修改训练代码。
• 缺点：分布式效果较差，训练速度提升不明显。负载不均衡问题。gpu0所承担的任务明显要重于其他gpu。

DDP：

• 优点： 多线程训练，torch 就会自动将代码分配给 n 个进程，分别在 n 个 GPU 上运行。通过 NCCL 实现 GPU 通信，解决了 DataParallel 速度慢，GPU负载不均衡的问题。
• 缺点：部署起来较麻烦，且pytorch对于分布式训练有多次更新，训练方式一直在更新迭代。

本文主要对DDP分布式训练进行实践。

分布式参数初始化

第一步引入一些支持分布式通讯的一些代码。

import os
local_rank = int(os.environ["LOCAL_RANK"])

import torch.distributed as dist
dist.init_process_group(backend="nccl")

local rank是每个进程的标签，对于八个进程，local_rank这个变量最后会被分配0-7的整数。pytorch 1.10后local_rank参数可通过os得到（代替argparse方法）。使用 init_process_group 设置GPU 之间通信使用的后端和端口。

数据集分布式划分

使用 DistributedSampler 对数据集进行划分，让dataset被分布式地sample到不同的进程上。将每个 batch 划分成几个 partition，在当前进程中只需要获取和 rank 对应的那个 partition 进行训练：

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., shuffle=False, sampler=train_sampler)

train_dataset进行分布式处理，加快模型训练过程中参数更新。
val_dataset不进行分布式处理，验证时模型参数不更新。
这里DataLoader中batch_size参数为单张卡上batch大小。shuffle设置为False，在dataloader中,sampler和 shuffle不能同时设置，否则报错：

ValueError: sampler option is mutually exclusive with shuffle

模型分布式包装

使用 DistributedDataParallel 包装模型，它能帮助我们为不同 GPU 上求得的梯度进行 all reduce（即汇总不同 GPU 计算所得的梯度，并同步计算结果）。all reduce 后不同 GPU 中模型的梯度均为 all reduce 之前各 GPU 梯度的均值：

torch.cuda.set_device(local_rank)
model = your_model()
model = torch.nn.parallel.DistributedDataParallel(model.cuda(), device_ids=[local_rank], output_device=local_rank)

记得先将model放到cuda上model.cuda()，再利用DistributedDataParallel 包装，否则报错：

ValueError: DistributedDataParallel device_ids and output_device arguments only work with single-device/multiple-device GPU modules or CPU modules, but got device_ids [1], output_device 1, and module parameters {device(type=‘cpu’)}.

模型保存与加载

evaluation时，加上对于local rank的判断，保证只有主线程在指定的时间会做evaluation，保存模型等。

if local_rank == 0:
    with torch.no_grad():
        model.eval()
        ......
        torch.save(...)

model被distributed处理之后，torch会为model所有层的参数加一个module.的前缀。模型重新加载时与初始化模型模块名字对不上，可将保存模型内的所有module去掉：

checkpoint = torch.load(pretrain_weights)
model.load_state_dict(
    {k.replace('module.', ''): v for k, v in checkpoint['model'].items()},
    strict=True)

整体训练大致框架

import os
local_rank = int(os.environ["LOCAL_RANK"])

import torch.distributed as dist
dist.init_process_group(backend="nccl")

if __name__ == '__main__':
	# 训练数据集分布式
	train_dataset = ...
	train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
	train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., shuffle=False, sampler=train_sampler)
	val_dataset = ...
	val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=..., shuffle=False)
	
 	# 模型分布式
	torch.cuda.set_device(local_rank)
	model = your_model()
	model = torch.nn.parallel.DistributedDataParallel(model.cuda(), device_ids=[local_rank], output_device=local_rank)
 
	optimizer = ...
	criterion= ...
 
	for epoch in range(epochs):
		for batch_idx, (data, target) in enumerate(train_loader):
		    images = images.cuda(non_blocking=True)
		    target = target.cuda(non_blocking=True)
		    ...
		    output = model(images)
		    loss = criterion(output, target)
		    ...
		    optimizer.zero_grad()
		    loss.backward()
		    optimizer.step()
		# 模型测试与保存
		if local_rank == 0:
            with torch.no_grad():
                model.eval()
				......
			torch.save(...)

模型训练

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train_SID_distributed.py

CUDA_VISIBLE_DEVICES=0,1指定训练时可见GPU，--nproc_per_node=2每个物理节点（就是一台机器，节点内部可以有多个GPU）上面进程的数量，与可见GPU数量保持一致。