在pytorch中数据并行训练涉及到nn.DataParallel和nn.parallel.DistributedDataParallel两个模块,也就是DP和DDP
数据并行的含义
每个 GPU 复制一份模型,将一批样本分为多份输入各个模型并行计算
当一张 GPU 可以存储一个模型时,可以采用数据并行得到更准确的梯度或者加速训练,因为求导以及加和都是线性的,数据并行在数学上也有效
1.DP
DP使用数据并行的方式只需要将原来单卡的module用DP改成多卡
model = nn.DataParallel(model)
DP 基于单机多卡,所有设备都负责计算和训练网络,除此之外,device[0] (并非 GPU 真实标号而是输入参数 device_ids 首位) 还要负责整合梯度,更新参数,主要有三个过程
1)各卡分别计算损失和梯度
2)所有梯度整合到 device[0]
3)device[0] 进行参数更新,其他卡拉取 device[0] 的参数进行更新
所有卡都并行运算,将梯度收集到device[0]和device[0]分享模型参数给其他GPU这三个主要过程
DP实现的单机训练不能算是严格意义上的分布式训练,但是其原理和分布式训练算法里的Parameter Server架构很相近
DP总结下来有以下几点不足
负载不均衡,device[0] 负载大一些
通信开销,GPU数量增加导致通信时间线性增长
单进程,由于python的GIL全局解释器锁,多线程始终只能使用一个CPU
2.DDP
import torch from torch.nn.parallel import DistributedDataParallel as DDP parser = argparse.ArgumentParser() parser.add_argument("--save_dir", default='') parser.add_argument("--local_rank", default=-1) parser.add_argument("--world_size", default=1) args = parser.parse_args() # 初始化后端 # world_size 指的是总的并行进程数目 # 比如16张卡单卡单进程 就是 16 # 但是如果是8卡单进程 就是 1 # 等到连接的进程数等于world_size,程序才会继续运行 torch.distributed.init_process_group(backend='nccl', world_size=ws, init_method='env://') torch.cuda.set_device(args.local_rank) device = torch.device(f'cuda:{args.local_rank}') model = nn.Linear(2,3).to(device) # train dataset # train_sampler # train_loader # 初始化 DDP,这里我们通过规定 device_id 用了单卡单进程 # 实际上根据我们前面对 parallel_apply 的解读,DDP 也支持一个进程控制多个线程利用多卡 model = DDP(model, device_ids=[args.local_rank], output_device=args.local_rank).to(device) # 保存模型 if torch.distributed.get_rank() == 0: torch.save(model.module.state_dict(), 'results/%s/model.pth' % args.save_dir)
DDP的原理和DP的区别
1)多进程
和 DP 不同, DDP 采用多进程,最推荐的做法是每张卡一个进程从而避免上一节所说单进程带来的影响。前文也提到了 DP 和 DDP 共用一个 parallel_apply 函数,所以 DDP 同样支持单进程多线程多卡操作,自然也支持多进程多线程,不过需要注意一下 world_size。
2)通信效率
DP 的通信成本随着 GPU 数量线性增长,而 DDP 支持 Ring AllReduce,其通信成本是恒定的,与 GPU 数量无关。
3)同步参数
DP 通过收集梯度到 device[0],在device[0] 更新参数,然后其他设备复制 device[0] 的参数实现各个模型同步;
DDP 通过保证初始状态相同并且改变量也相同(指同步梯度) ,保证模型同步。
DDP 也是数据并行,所以每张卡都有模型和输入。我们以多进程多线程为例,每起一个进程,该进程的 device[0] 都会从本地复制模型,如果该进程仍有多线程,就像 DP,模型会从 device[0] 复制到其他设备。
DDP 通过 Reducer 来管理梯度同步。为了提高通讯效率, Reducer 会将梯度归到不同的桶里(按照模型参数的 reverse order, 因为反向传播需要符合这样的顺序),一次归约一个桶。其中桶的大小为参数 bucket_cap_mb 默认为 25,可根据需要调整
DDP 通过在构建时注册 autograd hook 进行梯度同步。反向传播时,当一个梯度计算好后,相应的 hook 会告诉 DDP 可以用来归约。当一个桶里的梯度都可以了,Reducer 就会启动异步 allreduce 去计算所有进程的平均值。allreduce 异步启动使得 DDP 可以边计算边通信,提高效率。当所有桶都可以了,Reducer 会等所有 allreduce 完成,然后将得到的梯度写到 param.grad
参考原文:https://blog.csdn.net/qq_39967751/article/details/123483190