1. 引言
DistributedDataParallel（DDP）是一个支持多机多卡、分布式训练的深度学习工程方法。它通过Ring-Reduce的数据交换方法提高了通讯效率，并通过启动多个进程的方式减轻Python GIL的限制，从而提高训练速度。即是，将数据并行划分到多个进程（一般一个进程是一张卡），各进程初始化模型并由各自的数据训练，再通过Ring-Reduce进行梯度交换与合并，实现进程数倍数的效率。
2. Quick Start
这是正常的单卡代码：
## main.py文件
import torch
# 构造模型
model = nn.Linear(10, 10).to(local_rank)
# 前向传播
outputs = model(torch.randn(20, 10).to(rank))
labels = torch.randn(20, 10).to(rank)
loss_fn = nn.MSELoss()
loss_fn(outputs, labels).backward()
# 后向传播
optimizer = optim.SGD(model.parameters(), lr=0.001)
optimizer.step()
## Bash运行
python main.py
这是加入DDP的多卡代码：
## main.py文件
import torch
# 新增：
import torch.distributed as dist
# 新增：从外面得到local_rank参数
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1)
args = parser.parse_args()
local_rank = args.local_rank
# 新增：DDP backend初始化
torch.cuda.set_device(local_rank)
dist.init_process_group(backend='nccl') # nccl是GPU设备上最快、最推荐的后端
# 构造模型
device = torch.device("cuda", local_rank)
model = nn.Linear(10, 10).to(device)
# 新增：构造DDP model
model = DDP(model, device_ids=[local_rank], output_device=local_rank)
# 前向传播
outputs = model(torch.randn(20, 10).to(rank))
labels = torch.randn(20, 10).to(rank)
loss_fn = nn.MSELoss()
loss_fn(outputs, labels).backward()
# 后向传播
optimizer = optim.SGD(model.parameters(), lr=0.001)
optimizer.step()
## Bash运行
# 改变：使用torch.distributed.launch启动DDP模式，
# 其会给main.py一个local_rank的参数。这就是之前需要"新增:从外面得到local_rank参数"的原因
python -m torch.distributed.launch --nproc_per_node 4 main.py
3. 基本概念
在两台机器，每台8张显卡，共16张显卡，16的并行数下，DDP会同时启动16个进程。下面介绍一些分布式的概念。
group：即进程组。默认情况下，只有一个组。
world size：表示全局的并行数，即是2x8=16。
# 获取world size，在不同进程里都是一样的，得到16
torch.distributed.get_world_size()
rank：表示当前进程的序号，用于进程间通讯。对于16的world sizel来说，就是0,1,2,…,15。其中，rank=0的进程就是master进程。
# 获取rank，每个进程都有自己的序号，各不相同
torch.distributed.get_rank()
local_rank：每台机子上的进程的序号。机器一上有0,1,2,3,4,5,6,7，机器二上也有0,1,2,3,4,5,6,7
# 获取local_rank
torch.distributed.local_rank()
4. DDP使用流程
4.1 launch启动
## main.py文件
import torch
import argparse
# 新增1:依赖
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
# 新增2：从外面得到local_rank参数，在命令行用launch启动时会提供
parser = argparse.ArgumentParser()
parser.add_argument("--local_rank", default=-1)
FLAGS = parser.parse_args()
local_rank = FLAGS.local_rank
# 新增3：DDP backend初始化
# a.根据local_rank来设定当前使用哪块GPU
torch.cuda.set_device(local_rank)
# b.初始化DDP，使用默认backend(nccl)就行。如果是CPU模型运行，需要选择其他后端。
dist.init_process_group(backend='nccl')
# 新增4：定义并把模型放置到单独的GPU上，需要在调用`model=DDP(model)`前。
device = torch.device("cuda", local_rank)
model = nn.Linear(10, 10).to(device)
# 新增5：之后才是初始化DDP模型
model = DDP(model, device_ids=[local_rank], output_device=local_rank)
## 数据集
my_trainset = torchvision.datasets.CIFAR10(root='./data', train=True)
# 新增1：使用DistributedSampler进行各进程的采样
train_sampler = torch.utils.data.distributed.DistributedSampler(my_trainset)
# 需要注意的是，这里的batch_size指的是每个进程下的batch_size。也就是说，总batch_size是这里的batch_size再乘以并行数(world_size)。
trainloader = torch.utils.data.DataLoader(my_trainset, batch_size=batch_size, sampler=train_sampler)
for epoch in range(num_epochs):
# 新增2：设置sampler的epoch，DistributedSampler需要这个来维持各个进程之间的相同随机数种子
trainloader.sampler.set_epoch(epoch)
for data, label in trainloader:
data = data.to(local_rank)
label = label.to(local_rank)
prediction = model(data)
loss = loss_fn(prediction, label)
loss.backward()
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
optimizer.step()
# 1. save模型的时候，和DP模式一样，有一个需要注意的点：保存的是model.module而不是model。
# 因为model其实是DDP model，参数是被`model=DDP(model)`包起来的。
# 2. 我只需要在进程0上保存一次就行了，避免多次保存重复的东西。
if dist.get_rank() == 0:
torch.save(model.module, "saved_model.ckpt")
调用方法，在单机下：
CUDA_VISIBLE_DEVICES="0,1,2,3" python -m torch.distributed.launch --nproc_per_node 4 --master_port 53453 main.py
在多机下：
## Bash运行
# 假设我们在2台机器上运行，每台可用卡数是8
# 机器1：
python -m torch.distributed.launch --nnodes=2 --node_rank=0 --nproc_per_node 8 \
--master_adderss $my_address --master_port $my_port main.py
# 机器2：
python -m torch.distributed.launch --nnodes=2 --node_rank=1 --nproc_per_node 8 \
--master_adderss $my_address --master_port $my_port main.py
4.2 spawn启动
def run(rank, world_size):
dist.init_process_group(backend='nccl', init_method=args.dist_url, rank=local_rank, world_size=args.world_size)
# training code...
def run_demo(demo_fn, world_size):
parser.add_argument("--gpus", "-G", default='0', type=str)
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
args.world_size = len(args.gpus.split(','))
port_id = 10000 + np.random.randint(0, 1000)
args.dist_url = 'tcp://127.0.0.1:' + str(port_id)
mp.spawn(run,
args=(args,),
nprocs=args.world_size,
join=True)
调用方法:
python main.py --gpus 0,1,2,3
5. 不是很相关的一些bug
在使用mp.spawn训练模型时，报如下错。
RuntimeError: Cowardly refusing to serialize non-leaf tensor which requires_grad, since autograd does not support crossing process boundaries. If you just want to transfer the data, call detach() on the tensor before serializing (e.g., putting it on the queue).
调bug后，发现是在构造dataset时，包含了对预训练模型的导入，本意只是用到其中的一个函数，但它导致了这些参数既非叶节点还需要梯度，因此报错。最后只截取所需的函数，去掉了预训练模型的导入即可。