前两天为了多占几个机器多占几块卡,让程序跑得更快一点,踩了一下分布式训练的坑。
写在前面
为什么要分布式训练?
- 可以用多张卡,总体跑得更快
- 可以得到更大的 BatchSize
- 有些分布式会取得更好的效果
分布式训练可以分为:
- 单机多卡,DataParallel(最常用,最简单)
- 单机多卡,DistributedDataParallel(较高级)
- 多机多卡,DistributedDataParallel(最高级)
TL;DR
- DataParallel 最简单,但是效率不高
- DDP 的话,Horovod 最简单,但是依然有小坑
单机多卡,DataParallell
怎么写
from torch.nn import DataParallel
device = torch.device("cuda")
model = MyModel()
model = model.to(device)
model = DataParallel(model)
注意之后如果需要调用 model 里面的函数,需要用 model.module.XXX
怎么跑
哦,正常跑。
坑
没啥坑
单机多卡,DistributedDataParallel
前面的 DataParallel 并不是完整的分布式计算,只是将一些部分的计算(例如,前向和反向)放到了多张卡上,某些东西(例如,梯度)计算的时候仍然是「一卡有难,八方围观」,可能会将第一张卡撑爆,并且和 DDP 对比的话效率实在不高。
首先增加几个概念:
world_size
:总共有几个 Workerrank
:这个 Worker 是全局第几个 Workerlocal_rank
:这个 Worker 是这台机器上的第几个 Worker- 每个 Worker 可以用一张或者多张卡,但习惯于一个 Worker 只用一张卡
PyTorch 原生 DDP
PyTorch 大约是在 0.4 版本的时候加上了对 DDP 的支持,目前大家的评价是「已经基本成熟」。但是需要改的代码比较多,并且用起来坑比较多。好处是……可以少用几个第三方库,并且有 PyTorch 官方背书。
首先,打开 PyTorch 的官方 Tutorial,放旁边,以备不时之需。
怎么改
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from torch.utils.data import DataLoader
gpu_list = [0,1,2,3] # 假设这个 Worker 用 0-3 号共 4 张卡
device = torch.device(f"cuda:{gpu_list[0]}") # 这里注意这个 0
# 首先看模型
model = MyModel()
model = model.to(device)
model = DDP(model, device_ids=device_ids) # 前面只能放到 list 中的第 0 个 GPU 上,然后这里再分发到其他 device 上
# 大坑注意:如果 model 没梯度(例如 GAN 里面用 VGG 提取 StyleLoss),那么不能放 DDP 里
optimizer = optim.SGD(model.parameters(), lr=0.001)
# 然后看数据集
dataset = MyDataset()
sampler = DistributedSampler(dataset)
loader = DataLoader(dataset, sampler=sampler, batch_size=...)
前面说了,一般来说,一个卡上跑一个 Worker,所以 gpu_list 里面应该只有一个值,这个值既是 local_rank,又是本机上使用第几块卡。
OK,以上都是坑比较少的,真正坑比较多的是如何做环境初始化:
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'
dist.init_process_group("gloo", rank=rank, world_size=world_size)
上面的这段代码复制自 PyTorch 的文档。这个 init_process_group 东西支持很多种初始化方式:
- env 的方式:就像上面代码一样,好像常用于单机多卡
- host+port 的方式跟 env 其实没有本质区别,但好像要每台机器上都配置一下
- 共享文件的方式:比如我有个多个机器都能访问到的文件夹,那么可以在这里创建个文件,里面好像会自动给你写进去各种配置。——但是我一直没成功。
具体文档在这里: https://pytorch.org/docs/stable/distributed.html
怎么跑
第二步是怎么将这个程序跑起来。这里也有两种方法:
- 用
torch.distributed.launch
:
python -m torch.distributed.launch --nproc_per_node=4 main.py
- 用 torch.multiprocessing:
import torch.multiprocessing as mp
def main(rank, your_custom_arg_1, your_custom_arg_2): # 这里第一个 rank 变量会被 mp.spawn 函数自动填充,可以充当 local_rank 来用
pass # 将前面那一堆东西包装成一个 main 函数
mp.spawn(main, nprocs=how_many_process, args=(your_custom_arg_1, your_custom_arg_2))
我这里比较迷茫的是, global rank 怎么获取,global size 怎么传进去,并且为啥文件的方式不生效……折腾了大概半天没有结果,放弃。
坑
- 真的被各种 Rank 之类的东西搞晕了
- 如果模型没梯度,那么不能放到 DDP 里
- 如果使用 LMDB,并且在 init 里面写了类似 self.dataset = lmdb.open(“the_file”) 之类的话,那么大概率会报「不能 pickle Environment/Curser」之类的错误。
- 需要想办法将 LMDB 延迟初始化。
- 例如在 get_item 里面每次判断:if self.lmdb_file is None: self.lmdb_file = lmdb.open(“the_file.lmdb”) 之类。
Horovod【推荐】
Horovod 是一个第三方包,可用于多种框架的加速。优点是代码改动比较少。缺点是有小坑。
首先下包:
pip install horovod --no-cache-dir
怎么改
import horovod.torch as hvd
from torch.utils.data. distributed import DistributedSampler
from torch.utils.data import DataLoader
# 初始化一些环境
hvd.init()
torch.cuda.set_device(hvd.local_rank())
device = torch.device(f"cuda:{hvd.local_rank()}") # 这一句,如果喜欢写 to_device 的话可能会比较坑,但不是说推荐用 to_device 么?
# 数据集部分:要包一层
dataset = MyDataset()
sampler = DistributedSampler(dataset, num_replicas=hvd.size(), rank=hvd.rank())
loader = DataLoader(dataset, sampler=train_sampler, batch_size=...)
# 模型部分:要包一层
model = MyModel()
model = model.to(device)
optimizer = optim.SGD(model.parameters())
optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters()) # 关键:将 optimizer 包装一下
hvd.broadcast_parameters(model.state_dict(), root_rank=0) # 关键:将 model 的参数分发到其他 worker 上
# 其他东西:
# 收集上来的 Loss,可能会被变成了一个一维 Tensor,需要做一下 average 之类的操作
那代码也就差不多改成这样了。
怎么跑
第一种方式:用 horovodrun 来跑(其实背后是调用了 mpirun):
horovodrun -np 4 -H localhost:4 python train.py
第二种方式:自己用 mpirun 来跑。这个我没有尝试。
坑
- Horovod 好像只支持全图完全反向,不支持「只反向一半的图」。(例如,GAN,G 和 D 全部前向后,只想优化 G 或者只想优化 D,这个好像不支持)
- GPU Memory 使用量小增加(某 GAN,BS=1 的时候 MEM 刚好溜边,换用 Horovod 刚好超显存,囧)
写在最后
- DP 是「0 号卡 Memory 占得多(因为负责计算梯度)」,DDP 好像 0 号卡反而占得比其他卡少一些
- GPU 使用率的话,比之前有一点点提高
- GPU Memory 方面,Horovod 方式使用量比其他两个方式稍微大一点(如果模型刚好卡在 Memory 边缘,那么 DDP 之后可能反而单卡里面放不下了)
- 最终效果有没有提高?没做对比实验,不太好说
- 如果喜欢做进度条、打 Log,那么一台机器上会同时出现所有 Worker 的进度条或者 Log,非常乱。所以需要设置只有
worker=0
或者global_worker=0
的 Worker 才有权利做下面这几个事情- 跑进度条
- 打 Log
- 保存模型
- 打 TensorBoard
- 保存可视化结果
参考资料
发表回复