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
    • 保存可视化结果

参考资料


评论

  1. 林曦 的头像
    林曦

    DataParallell就是辣鸡,DDP快太多啦

    1. 小金鱼儿 的头像
      小金鱼儿

      但是DataParallel最好写呀~

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注