multitask训练torch_Pytorch多机多卡分布式训练

被这东西刁难两天了&＃xff0c;终于想办法解决掉了&＃xff0c;来造福下人民群众。

关于Pytorch分布训练的话&＃xff0c;大家一开始接触的往往是DataParallel&＃xff0c;这个wrapper能够很方便的使用多张卡&＃xff0c;而且将进程控制在一个。唯一的问题就在于&＃xff0c;DataParallel只能满足一台机器上gpu的通信&＃xff0c;而一台机器一般只能装8张卡&＃xff0c;对于一些大任务&＃xff0c;8张卡就很吃力了&＃xff0c;这个时候我们就需要面对多机多卡分布式训练这个问题了&＃xff0c;噩梦开始了。

官方pytorch(v1.0.10)在分布式上给出的api有这么几个比较重要的&＃xff1a;

torch.nn.parallel.DistributedDataParallel &＃xff1a;

这个从名字上就能看出来与DataParallel相类似&＃xff0c;也是一个模型wrapper。这个包是实现多机多卡分布训练最核心东西&＃xff0c;它可以帮助我们在不同机器的多个模型拷贝之间平均梯度。

2. torch.utils.data.distributed.DistributedSampler:

在多机多卡情况下分布式训练数据的读取也是一个问题&＃xff0c;不同的卡读取到的数据应该是不同的。dataparallel的做法是直接将batch切分到不同的卡&＃xff0c;这种方法对于多机来说不可取&＃xff0c;因为多机之间直接进行数据传输会严重影响效率。于是有了利用sampler确保dataloader只会load到整个数据集的一个特定子集的做法。DistributedSampler就是做这件事的。它为每一个子进程划分出一部分数据集&＃xff0c;以避免不同进程之间数据重复。

# 分布式训练示例

from torch.utils.data import Dataset, DataLoader

from torch.utils.data.distributed import DistributedSampler

from torch.nn.parallel import DistributedDataParallel

dataset &＃61; your_dataset()

datasampler &＃61; DistributedSampler(dataset, num_replicas&＃61;world_size, rank&＃61;rank)

dataloader &＃61; DataLoader(dataset, batch_size&＃61;batch_size_per_gpu, sampler&＃61;datasampler)

model &＃61; your_model()

model &＃61; DistributedDataPrallel(model, device_ids&＃61;[local_rank], output_device&＃61;local_rank)

其他部分就和正常训练代码无异了。

得提的几个点&＃xff1a;

和dataparallel不同&＃xff0c;dataparallel需要将batchsize设置成n倍的单卡batchsize&＃xff0c;而distributedsampler使用的情况下&＃xff0c;batchsize设置与单卡设置相同。

这里有几个新的参数&＃xff1a;world size, rank, local rank, rank。world size指进程总数&＃xff0c;在这里就是我们使用的卡数&＃xff1b;rank指进程序号&＃xff0c;local_rank指本地序号&＃xff0c;两者的区别在于前者用于进程间通讯&＃xff0c;后者用于本地设备分配。这个时候真正麻烦的地方来了&＃xff1a;

想要使用DistributedDataParallel&＃xff0c;得先完成多进程的初始化&＃xff0c;就是这个&＃xff1a;

torch.distributed.init_process_group()

看官方的说明&＃xff1a;

gloo基本只支持cpu&＃xff0c;不考虑。mpi需要在本地重新编译pytorch&＃xff0c;感兴趣的朋友可以试试。nccl对gpu支持良好还不需要重新编译&＃xff0c;在下和官方都强烈推荐这个作为backend。

pytorch作者推荐的初始化方式&＃xff1a;

我最后的实现也是利用这种方式。但我面临的问题是&＃xff1a;如何在我们的slurm集群上完成这个初始化并进行训练&＃xff0c;那么问题就变成了如何在slurm集群上把你分配到的ip写进程序里。两个办法&＃xff1a;

1.srun指定-n 进程总数以及 –ntasks-per-node 每个节点进程数&＃xff0c;这样就可以通过os.environ获得每个进程的节点ip信息&＃xff0c;全局rank以及local rank&＃xff0c;有了这些就可以很方便很方便的完成初始化。推荐使用该方法(感谢评论区大佬指点)

2.salloc&＃xff0c;这个就相对霸道一些&＃xff0c;直接指定几个节点自己拿来用&＃xff0c;这样就很容易选出来通信用的节点&＃xff0c;再随便给个端口&＃xff0c;我们就能完成初始化。相比1还是麻烦不少。

关于获取节点信息的详细代码&＃xff1a;

import os

os.environ[&＃39;SLURM_NTASKS&＃39;] #可用作world size

os.environ[&＃39;SLURM_NODEID&＃39;] #node id

os.environ[&＃39;SLURM_PROCID&＃39;] #可用作全局rank

os.environ[&＃39;SLURM_LOCALID&＃39;] #local_rank

os.environ[&＃39;SLURM_STEP_NODELIST&＃39;] #从中取得一个ip作为通讯ip

贴段差不多能跑的代码吧&＃xff1a;

import torch

torch.multiprocessing.set_start_method(&＃39;spawn&＃39;)

import torch.nn as nn

from torch.utils.data import Dataset, DataLoader

from torch.utils.data.distributed import DistributedSampler

from torch.nn.parallel import DistributedDataParallel

import os

def dist_init(host_addr, rank, local_rank, world_size, port&＃61;23456):

host_addr_full &＃61; &＃39;tcp://&＃39; &＃43; host_addr &＃43; &＃39;:&＃39; &＃43; str(port)

torch.distributed.init_process_group("nccl", init_method&＃61;host_addr_full,

rank&＃61;rank, world_size&＃61;world_size)

num_gpus &＃61; torch.cuda.device_count()

torch.cuda.set_device(local_rank)

assert torch.distributed.is_initialized()

rank &＃61; int(os.environ[&＃39;SLURM_PROCID&＃39;])

local_rank &＃61; int(os.environ[&＃39;SLURM_LOCALID&＃39;])

world_size &＃61; int(os.environ[&＃39;SLURM_NTASKS&＃39;])

# get_ip函数自己写一下 不同服务器这个字符串形式不一样

# 保证所有task拿到的是同一个ip就成

ip &＃61; get_ip(os.environ[&＃39;SLURM_STEP_NODELIST&＃39;])

dist_init(ip, rank, local_rank, world_size)

# 接下来是写dataset和dataloader&＃xff0c;这个网上有很多教程

# 我这给的也只是个形式&＃xff0c;按自己需求写好就ok

dataset &＃61; your_dataset() #主要是把这写好

datasampler &＃61; DistributedSampler(dataset, num_replicas&＃61;world_size, rank&＃61;rank)

dataloader &＃61; DataLoader(dataset, batch_size&＃61;batch_size_per_gpu, sampler&＃61;source_sampler)

model &＃61; your_model() #也是按自己的模型写

model &＃61; DistributedDataPrallel(model, device_ids&＃61;[local_rank], output_device&＃61;local_rank)

# 此后训练流程与普通模型无异

照上面写好train.py之后(叫啥都行&＃xff0c;这儿就叫train.py吧)&＃xff0c;slrum指令写这样&＃xff1a;

# 这里是3台机器&＃xff0c;每台机器8张卡的样子

srun -n24 --gres&＃61;gpu:8 --ntasks-per-node&＃61;8 python train.py

—————————————分割线—————————————–

以下是关于方法2的补充&＃xff0c;1已经足够用了&＃xff0c;保留这部分只是为了文章完整。

然而&＃xff0c;事情并没有这么简单。

如果一次salloc多个节点&＃xff0c;那么接下来的srun指令默认是每个节点执行一个拷贝&＃xff0c;也就是说&＃xff0c;我们的rank是无法保证两个节点上是不一样的。怎么办&＃xff1f;

参考这个

需要多少个节点&＃xff0c;就开多少个窗口&＃xff0c;每个窗口salloc一个节点&＃xff0c;就能解决上面的问题。

最后&＃xff0c;如何方便的得到多个rank不一样的进程&＃xff0c;其实pytorch已经有模块可以很好的解决这个问题&＃xff1a;

参考&＃xff1a;

采用torch.distributed.launch模块可以很容易为每个进程得到不一样的localrank&＃xff0c;再在不同的节点指定不同的rank初始值&＃xff0c;slurm群组上就可以开心的pytorch多机分布式训练啦。