pytorch多卡训练nn.DataParallel和nn.DistributedDataParallel比较

本文来源于：Pytorch Distributed - 知乎 Pytorch Distributed

现在的模型越来越大，并行显得越来越重要，而众所周知，pytorch 的并行文档写的非常不清楚，不仅影响使用，甚至我们都不知道他的工作原理。一次偶然的机会，我发现了几篇在这方面写的很好的文章，因此也准备参考别人的（参考的文章在Reference部分列出）再结合自己的使用经验总结一下。

nn.DataParallel

Pytorch的数据并行方式，是经常使用的单机多卡的并行方式。

工作原理

这种方式只有一个进程，假设我们有一个8卡的机器，然后设置的batchsize是128，那么单卡的batchsize就是32，在训练的时候，他把我们的模型分发到每个GPU上，然后在每个GPU上面对着32张图进行前向和反向传播得到loss，之后将所有的loss汇总到0卡上，对0卡的参数进行更新，然后，非常重要的一点：他再把更新后的模型传给其他7张卡，这个 *** 作是在每个batch跑完都会执行一次的。
可以想象一下，每个batch，GPU之间都要互相传模型，GPU的通信显然成为了一个极大的瓶颈，直接导致GPU在一段时间是闲着的，也就导致了GPU的利用率低的问题。

缺点

1. GPU利用率低
2. 不支持多机多卡
3. 不支持混合精度

使用方法

nn.DataParallel的使用方法非常的简单，首先你需要保障可用的GPU数目，这个可以通过参数设置CUDA_VISIBLE_DEVICES，你可以直接在你的bash输入

export CUDA_VISIBLE_DEVICES=1

代表程序将只能看到你的编号为1的显卡（这里默认从0开始编号，所以其实就是只能看到你的第二张卡），你也可以设置多张：

export CUDA_VISIBLE_DEVICES=0,1,2,3

或者你也可以在你的python代码里面进行设置：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

注意，pytorch会把它看到的GPU从0排序，也就是假设你有八张卡，你设置的后四张可见，但是对于pytorch，他以为你就这4张，所以在他内部程序使用的时候，是0123来指代这些卡。
还有一个

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"

是让设备ID=物理ID，在CUDA_VISIBLE_DEVICES之前用，不过一般不需要这个，默认的就是正常的顺序。

其次你需要在你的模型这里设置：

model = nn.DataParallel(model).cuda()

剩余的就和单卡的设置是一样的了。

nn.DistributedDataParallel

分布式的数据并行在很大程度上解决了上面数据并行的问题，支持多机多卡，支持混合精度，是现在非常实用的一种方式。

工作原理

nn.DistributedDataParallel 和 nn.DataParallel 最大的区别是，不再来回同步模型了，太麻烦了，我们需要同步的其实只有梯度，梯度同步后，每个GPU做自己的参数更新就行，那么在每个GPU有独立工作的需求后，多进程就被拿了出来，每个进程控制一块GPU。
总结原理就是：每个进程控制一个GPU，然后每个GPU计算自己分得的数据的梯度后，通过进程之间的通信进行综合所有梯度的 *** 作(all-reduce)，每个GPU独立更新参数(实际上就是同样的东西算了好多遍)，但是这里更新完每个GPU的模型参数是一样的。(插一句话：这里其实体现了计算机科学现在的一个瓶颈，就是IO瓶颈超越了计算速度瓶颈，多计算比使用IO同步要更快)

带来的改变

从原理上看其实就只是添加了一个多进程，但是其实多进程的添加，给整个训练添加了很多的改变。
1. 设置seed，如果每个进程的模型使用不同的seed，那就没办法同步，整个训练就没有意义了 2. 每个进程需要知道自己是第几个进程(rank)，以及一共有多少个进程(word_size) 3. 每个进程需要知道自己load哪部分数据，因为每个进程load的数据是不一样的。这里有个坑，估计很少有人会提，laoddata的时候pytorch是有一个num_worker，这个代表使用多进程加载，而我们现在已经使用了多进程，那么其实一共跑在机器上的是n*num_worker，如果这个太多，也会出一些问题，也就是在已经多进程的基础上就不用再搞很多进程了。

后端的选择

使用了多进程，就意味着需要进行进程之间的通信，pytorch提供给的后端有三种，一种是基于mpi的，一种是gloo，最后一种是基于nccl的。但是其实能够用的只有一种，也就是nccl，因为mpi的需要本地编译pytorch，里面必然有很多坑，估计也没什么人用，基于gloo的对GPU的支持很差，基本只支持CPU，所以现在绝大多数的人都用的是基于nccl的，本博客后面的也是基于nccl。

使用方法

首先需要import一些东西

import torch.distributed as dist

在这种方式下最重要的就是获得几个重要的信息，第一是wordsize，第二个是rank，第三个是使用nccl时使用tcp通信的ip和port。
我们先假设我们有四个节点，每个节点有8块GPU，我们需要在main获取完args后进行下面的 *** 作：

args.world_size = args.gpus * args.nodes                
os.environ['MASTER_ADDR'] = '10.57.23.164'              
os.environ['MASTER_PORT'] = '8888'                      
mp.spawn(train, nprocs=args.gpus, args=(args,))

wordsize是节点数乘以每个节点的GPU数量，addr和port直接设置成为环境变量，你也可以用下面的方式在bash里面设置。

export MASTER_ADDR=10.57.23.164
export MASTER_PORT=8888

最后再开启多进程传入参数即可。 然后是进入def train(gpu, args):里面：

torch.manual_seed(0)
rank = args.nr * args.gpus + gpu                              
dist.init_process_group(                                   
    backend='nccl',                                         
    init_method='env://',                                   
    world_size=args.world_size,                              
    rank=rank                                               
)

先拿到这个进程的rank，rank=第几个节点x8+这个节点的第几个进程，根据这些关键按信息对pytorch分布式进行初始化。

model = nn.parallel.DistributedDataParallel(model, device_ids=[gpu])

模型初始化，以及GPU分配

train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset,
        num_replicas=args.world_size,
        rank=rank
    )

sampler记得使用分布式的sampler，要传入rank，这样就能只加载这个进程需要的数据。train_loader的参数shuffle=False, sampler=train_sampler这里就不能shffle了，为什么留给读者自己思考，欢迎在评论区进行讨论。 最后是跑，因为我们有四个节点，所以需要启动四次，在第0个节点上这样跑

python main.py -n 4 -g 8 -nr 0

在123节点上这样跑

python main.py -n 4 -g 8 -nr i

就是告诉程序一共有几个节点，每个节点有几个GPU，目前终端是第几个节点。 注意，在dist.init_process_group这里是阻塞的，也就是会等待所有参与的进程都准备就绪才会继续开始，所以不用担心四个终端启动时间的问题，但是还是要先0后123。

使用进阶

很显然，上面四个节点分四个终端启动的方式是很蠢的，实际上肯定不会这么做，并且实际上节点都会有其任务管理系统，不会和上面这样这么简单，因此这部分就slurm这个任务调度工具为例，说说在slurm下，pytorch的多进程怎么做。
其实跟上面的核心是一样的，使用slurm的主要区别是获取那些关键信息的方法，其他的并没有大的区别。
我们在启动srun的时候指定-n 进程总数以及 --ntasks-per-node 每个节点进程数，就能在os里面拿到上面需要的一些东西了

rank = int(os.environ['SLURM_PROCID'])
local_rank = int(os.environ['SLURM_LOCALID'])
world_size = int(os.environ['SLURM_NTASKS'])
# get_ip是个字符串处理函数，需要根据自己的集群的名字等等信息调试一下
ip = get_ip(os.environ['SLURM_STEP_NODELIST'])

这里面的rank就是上面的rank的值，wordsize也是上面的那个，ip和我们上面设置的ip是一样的作用，local_rank其实就是前面train传入的gpu参数，也就是在这个节点上的进程值。
这些信息有着落了之后，其他的就和前面的方法一样了，最后只需使用类似的命令开始task即可：

srun -n32 --gres=gpu:8 --ntasks-per-node=8 python main.py