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

文章目录

    • 1. 引言
    • 2. Quick Start
    • 3. 基本概念
    • 4. DDP使用流程
      • 4.1 launch启动
      • 4.2 spawn启动
    • 5. 不是很相关的一些bug
    • 参考文献


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

你可能感兴趣的:(pytorch,分布式,python)