【分布式训练】多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