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

写在前面

为什么要分布式训练？

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

参考资料

• https://pytorch.org/tutorials/intermediate/ddp_tutorial.html
• https://pytorch.org/docs/stable/distributed.html
• https://zhuanlan.zhihu.com/p/98535650
• https://horovod.readthedocs.io/en/latest/pytorch.html
• https://horovod.readthedocs.io/en/latest/running_include.html