【分布式训练】多gpu训练出现负载不均衡,尝试DistributedDataParallel分布式训练

某次训练深度学习模型时,使用roberta-large模型作为基础模块,起初使用DataParallel 的方式,进行单机多卡训练,卡数为2,每张卡显存为10G。

训练期间发现,无法使用较大的batch_size,batch_size最大为4。同时,训练时间增加到3个小时一个epoch,时间开销太大。

观察GPU利用显存率,0卡占用显存明显比1卡占用的要多,这也是DataParallel 这种模式的弊端。
DataParallel 数据传输过程包括:

1.前向传播得到的输出结果gather到主cuda计算loss

2.scatter上述loss到各个cuda

3.各个cuda反向传播计算得到梯度后gather到主cuda后 主cuda的模型参数被更新

4.主cuda将模型参数broadcast到其它cuda设备上,至此完成权重参数值的同步。
综上,dp大概是有4次输出传输。

因此大量的计算就在0卡上进行使用,所以导致负载不均衡。

而DistributedDataParallel分布式训练方式,数据传输过程:

1.前向传播的输出和loss的计算都是在每个cuda独立计算的,梯度all-reduce到所有的cuda(这里可以认为是在‘传输‘梯度)这样初始参数相同,para.grad也相同,反向传播后参数就还是保持一致的。

还有一大区别是,DDP是利用多进程,每个GPU上运行一个进程,而DP则是单进程多线程,共同维护一个optimizer。DDP采用多进程的方法,从而避免了Python解释器GIL带来的性能开销。也就避免了多GPU训练时出现负载不均衡的问题。

代码:

import argparse
import torch.distributed as dist

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=0, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()
dist.init_process_group(backend='nccl')
print("args.local_rank:", args.local_rank)
torch.cuda.set_device(args.local_rank)

dl_module.cuda()
dl_module = torch.nn.parallel.DistributedDataParallel(dl_module, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True)


# data

1. 使用Dataset加载数据,无论时官方还是自定义的Dataset
train_data = Dataset(data)
2. 在加入DataLoader之前,进行
train_sampler = DistributedSampler(train_data)
train_generator = DataLoader(
            train_data,
            batch_size=batch_size,
            # shuffle=True,
            sampler=train_sampler,
            num_workers=num_workers,
            collate_fn=self._train_collate_fn
        )

3. 在训练时,每次epoch后面加上set_epoch,保证sampler样本是随机的
for epoch in range(epochs):
      train_sampler.set_epoch(epoch)

调整完成后,每轮epoch的时间,由3h降为1h,提升了2/3的训练效率。

参考:
https://blog.csdn.net/searobbers_duck/article/details/115299691
https://zhuanlan.zhihu.com/p/206467852

你可能感兴趣的:(深度学习,GPU,分布式训练,torch,深度学习,负载均衡)