torch DDP模式并行

torch.nn.DataParallel

只需要一行代码就可以实现DataParallel:

if torch.cuda.device_count() > 1:  #判断是不是有多个GPU    print("Let's use", torch.cuda.device_count(), "GPUs!")    model = nn.DataParallel(model,device_ids=range(torch.cuda.device_count()))

DataParallel方法虽然代码非常简单，但是GPU会出现负载不均衡的问题（"一个gpu干活，其它gpu看戏"）。第一个GPU（12GB）可能占用了10GB，剩余的GPU却只使用了2GB -- 这是极大的浪费。那么，为什么会出现这样的现象呢？我们来仔细研究一下下面这个流程图就明白了：

我们首先需要指定一个GPU(0)作为master GPU，它必须承担最多的负载。

第一步：GPU(0)从锁页内存(分配主机内存时锁定该页，让其不与磁盘交换。必须保证自己内存空间充足的时候才用锁页内存！)中取得一个minibatch。

第二步：GPU(0)把这个minibatch去scatter到其他GPU上，每个GPU拿到1/n的数据，这就是data parallel的定义。

第三步：GPU(0)把模型也分配到其他GPU上，现在模型是同步的

第四步：在每个GPU上并行的做forward pass，得到每个sub-minibatch的输出

第五步：GPU(0)把所有GPU的输出gather到自己这里，计算loss

第六步：GPU(0)把loss scatter到其他GPU，分别进行backward pass计算梯度

第七步：GPU(0)把梯度集中起来进行reduce，计算出梯度的总和

第八步：GPU(0)用这个梯度更新自己的模型。而其他模型没有更新，是stale版本。

我们可以看到，nn.DataParallel做了很多不必要的操作，导致其性能不佳：

第一步的时候，从page-lock memory先取数据到GPU(0),然后由GPU(0)再进行分发。为什么不能直接从page-lock memory中分发到各个GPU呢？
第五步的时候，GPU(0)拿到每个GPU的output然后计算loss，那么为什么不能每个GPU自己计算自己的loss呢？
最后一步参数的更新是在GPU(0)上进行的，此时其他GPU上的模型和GPU（0）不再同步。所以每次都需要broadcast 模型进行显式的同步。

由于loss的计算、梯度的reduce、模型的更新都是在GPU(0)上进行，所以造成了第一个GPU的负载远远大于剩余其他的GPU。

torch.nn.DistributedDataParallel

真正使用了多进程（process），而DataParallel只是使用了多线程（thread）：

DataParallel: 一个host一个process；DistributedDataParallel: 一个GPU一个process

DistributedDataParallel是对每个GPU使用一个进程，适用于单机多卡和多机多卡的场景。它的主要工作流程如下图所示，可以看到相比DataParallel做了很多的优化：

每个GPU都有一个process，分别计算在这个GPU上的loss、gradient，然后用ring-all-reduce的方式进行聚合。得到综合的gradient之后同时更新本GPU上的模型参数，所以模型一直都是同步的，不需要手动同步。

DistributedDataParallel is proven to be significantly faster than torch.nn.DataParallel for single-node multi-GPU data parallel training.

调用过程：

（1）初始化：sync the processes

#初始化使用nccl后端torch.distributed.init_process_group(backend="nccl")

Gloo（用于CPU）和NCCL（用于GPU）都是用于多个node之间通信的库。从下面的图中我们可以看到，torch.nn.parallel.DistributedDataParallel是依赖于torch.distributed的，而torch.distributed提供的是多机通信的一些原语。它依赖于Gloo/NCCL/MPI backend。

而torch.nn.DataParallel根本就是虚假的分布式！因为它只能使用于单个node，而且是单进程(single process)的，它的并行化全部依赖于线程(thread)。

（2）使用DistributedSampler

多gpu训练是，我们希望同一时刻在每个gpu上的数据是不一样的，这样相当于batch size扩大了N倍，因此起到了加速训练的作用。DistributedSampler用于把锁页内存的一个minibatch再进行切分，分配到不同的GPU进程中去。要保证这样的划分是没有overlap的。

每个进程都知道自己的local rank 和 host number,就可以推算出自己的global rank, 然后获得自己应有的那一份sub-minibatch数据，保证每个进程拿到的数据没有重叠。

torch.distributed collectives原语

dist.broadcast(tensor,src,group): 对于group中的所有GPU编号，把tensor从src GPU分发到其他的GPU process中
dist.reduce(tensor,dst,op,group): Applies op to all tensor in group and store the result in dst.
dist.all_reduce(tensor,op,group): Same as reduce, but the result is stored in all processes.
dist.scatter(tensor,src,scatter_list,group): copy the ith tensor scatter_list[i] to the ith process.
dist.gather(tensor, dst, gather_list, group): copy tensor from all processes in group in dst
dist.all_gather(tensor_list,tensor,group): copy tensor from all processes to tensor_list, on all processes.

下面我们来看一个例子：

import torchimport argparsefrom torch import distributed as dist
print(torch.cuda.device_count())  # 打印gpu数量torch.distributed.init_process_group(backend="nccl")  # 并行训练初始化，'nccl'模式print('world_size', torch.distributed.get_world_size()) # 打印当前进程数
# 下面这个参数需要加上，torch内部调用多进程时，会使用该参数，对每个gpu进程而言，其local_rank都是不同的；parser.add_argument('--local_rank', default=-1, type=int)  args = parser.parse_args()torch.cuda.set_device(args.local_rank)  # 设置gpu编号为local_rank;此句也可能看出local_rank的值是什么


'''多卡训练加载数据:注意shuffle与sampler是冲突的，并行训练需要设置sampler，此时务必# 要把shuffle设为False。但是这里shuffle=False并不意味着数据就不会乱序了，而是乱序的方式交给# sampler来控制，实质上数据仍是乱序的。'''train_sampler = torch.utils.data.distributed.DistributedSampler(My_Dataset)dataloader = torch.utils.data.DataLoader(ds,                                         batch_size=batch_size,                                         shuffle=False,                                         num_workers=16,                                         pin_memory=True, ##锁页内存                                         drop_last=True,                                         sampler=self.train_sampler)
'''多卡训练的模型设置：'''def average_gradients(model):  ##每个gpu上的梯度求平均    size = float(dist.get_world_size())    for param in model.parameters:        dist.all_reduce(param.grad.data,op = dist.reduce_op.SUM)        param.grad.data /= size
My_model = My_model.cuda(args.local_rank)  # 将模型拷贝到每个gpu上.直接.cuda()也行，因为多进程时每个进程的device号是不一样的My_model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(My_model) # 设置多个gpu的BN同步My_model = torch.nn.parallel.DistributedDataParallel(My_model,                                                      device_ids=[args.local_rank],                                                      output_device=args.local_rank,                                                      find_unused_parameters=False,                                                      broadcast_buffers=False)
'''开始多卡训练：'''for epoch in range(200):    train_sampler.set_epoch(epoch)  # 这句莫忘，否则相当于没有shuffle数据    My_model.train()    for idx, sample in enumerate(dataloader):        inputs, targets = sample[0].cuda(local_rank, non_blocking=True), sample[1].cuda(local_rank, non_blocking=True)        opt.zero_grad()        output = My_model(inputs)        loss = My_loss(output, targets)          loss.backward()        average_gradient(My_model) ##计算梯度的平均        opt.step() ##根据梯度更新模型

'''多卡测试(evaluation)：'''if local_rank == 0:    My_model.eval()    with torch.no_grad():        acc = My_eval(My_model)    torch.save(My_model.module.state_dict(), model_save_path)dist.barrier() # 这一句作用是：所有进程(gpu)上的代码都执行到这，才会执行该句下面的代码
'''其它代码'''

在分布式的GPU上进行SGD，每个GPU在计算完梯度之后，需要把不同GPU上的梯度进行all-reduce，即计算所有GPU上梯度的平均值：dist.all_reduce(param.grad.data, op=dist.reduce_op.SUM), 意为我们对每个模型参数的梯度进行all_reduce, 采用求和的形式计算all_reduce.

最后一步，运行上述代码的方式（8卡为例）：

python3 -m torch.distributed.launch --nproc_per_node=8 DDP.py

参考资料：

熊勒个猫：和nn.DataParallel说再见

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

格子不太方：pytorch 并行训练之DistributedDataParallel（代码样例和解释）

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

Columbia University COMS 6998: Deep Learning System 课件

link

猜你喜欢