前两天为了多占几个机器多占几块卡，让程序跑得更快一点，踩了一下分布式训练的坑。
写在前面
为什么要分布式训练？
- 可以用多张卡，总体跑得更快
- 可以得到更大的 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
- 保存可视化结果
参考资料
发表回复