PyTorch 如何使用分布式训练

前两天为了多占几个机器多占几块卡,让程序跑得更快一点,踩了一下分布式训练的坑。

写在前面

为什么要分布式训练?

  1. 可以用多张卡,总体跑得更快
  2. 可以得到更大的 BatchSize
  3. 有些分布式会取得更好的效果

分布式训练可以分为:

  1. 单机多卡,DataParallel(最常用,最简单)
  2. 单机多卡,DistributedDataParallel(较高级)
  3. 多机多卡,DistributedDataParallel(最高级)

TL;DR

  1. DataParallel 最简单,但是效率不高
  2. 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 对比的话效率实在不高。

首先增加几个概念:

  1. world_size:总共有几个 Worker
  2. rank:这个 Worker 是全局第几个 Worker
  3. local_rank:这个 Worker 是这台机器上的第几个 Worker
  4. 每个 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 东西支持很多种初始化方式:

  1. env 的方式:就像上面代码一样,好像常用于单机多卡
  2. host+port 的方式跟 env 其实没有本质区别,但好像要每台机器上都配置一下
  3. 共享文件的方式:比如我有个多个机器都能访问到的文件夹,那么可以在这里创建个文件,里面好像会自动给你写进去各种配置。——但是我一直没成功。
    具体文档在这里: https://pytorch.org/docs/stable/distributed.html

怎么跑

第二步是怎么将这个程序跑起来。这里也有两种方法:

  1. torch.distributed.launch
python -m torch.distributed.launch --nproc_per_node=4 main.py
  1. 用 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 怎么传进去,并且为啥文件的方式不生效……折腾了大概半天没有结果,放弃。

  1. 真的被各种 Rank 之类的东西搞晕了
  2. 如果模型没梯度,那么不能放到 DDP 里
  3. 如果使用 LMDB,并且在 init 里面写了类似 self.dataset = lmdb.open("the_file") 之类的话,那么大概率会报「不能 pickle Environment/Curser」之类的错误。
  4. 需要想办法将 LMDB 延迟初始化。
  5. 例如在 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 来跑。这个我没有尝试。

  1. Horovod 好像只支持全图完全反向,不支持「只反向一半的图」。(例如,GAN,G 和 D 全部前向后,只想优化 G 或者只想优化 D,这个好像不支持)
  2. GPU Memory 使用量小增加(某 GAN,BS=1 的时候 MEM 刚好溜边,换用 Horovod 刚好超显存,囧)

写在最后

  1. DP 是「0 号卡 Memory 占得多(因为负责计算梯度)」,DDP 好像 0 号卡反而占得比其他卡少一些
  2. GPU 使用率的话,比之前有一点点提高
  3. GPU Memory 方面,Horovod 方式使用量比其他两个方式稍微大一点(如果模型刚好卡在 Memory 边缘,那么 DDP 之后可能反而单卡里面放不下了)
  4. 最终效果有没有提高?没做对比实验,不太好说
  5. 如果喜欢做进度条、打 Log,那么一台机器上会同时出现所有 Worker 的进度条或者 Log,非常乱。所以需要设置只有 worker=0 或者 global_worker=0 的 Worker 才有权利做下面这几个事情
    • 跑进度条
    • 打 Log
    • 保存模型
    • 打 TensorBoard
    • 保存可视化结果

参考资料

《PyTorch 如何使用分布式训练》有2条留言

留下评论