【深度学习框架】pytorch之分布式数据并行化DDP

---

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时，包含了对预训练模型的导入，本意只是用到其中的一个函数，但它导致了这些参数既非叶节点还需要梯度，因此报错。最后只截取所需的函数，去掉了预训练模型的导入即可。

参考文献

https://zhuanlan.zhihu.com/p/178402798