首发于机器不学习

2020 Rethinking GPU 集群上的分布式训练

一、引言

受疫情影响，转眼间在家待了有小半年了。在一些机缘巧合下，第一次用到了 48 张 V100 的 GPU 集群，也参加了业界和学界的一些前沿会议（线上会议大幅降低了参会成本，估计会成为新常态）。在交流过程中，可以越来越明显地感受到算法从业人员工作内容的转变。工业界组件库日渐成熟，新组件的复现层出不穷，在屎山般的 pipline 上对细节进行调优对算力的要求越来越高；学术界 SOTA 饱和，学术重心偏向新思路的设计，但为了赶超 SOTA 也只能无奈 NAS 新思路基础上的排列组合。可以预见，在不远的未来，GPU 的军备竞赛会随着内卷日趋严重，对 GPU 集群的理解和使用将成为 2021 年研究生的必修课。这几天正好要调优模型（没超 SOTA 多少，没法跟老板交差呀），在 GPU 集群真香后，打算将手头几台零散的 V100 也整合起来，在此将整个安装、编码的过程记录下来，方便还没有入坑的同学从 0 到 1 入门～

二、安装与部署

其实，作为著名的资源调度和任务管理工具，Slurm 已经被广泛应用在各高校的集群中。物理所用 Slurm 管理 CPU 集群已经有很长的历史，深度学习兴起后计算所用 Slurm 管理 GPU 集群也逐渐成为实际标准。

很多成立不久的研究所和研究组的 GPU 设备增长迅速（拿我们组来说，我就用了 3 台 GPU 机器），但是普遍来看管理水平还有待积淀因此尚未组成集群。如果各位同学也用到了 3 台及以上的 GPU 节点，但仍处于野蛮生长的状态无法充分利用 GPU 涨点，向老板汇报并部署一套 Slurm 是一个不错的选择。

近年来 Slurm、FNS 技术迭代较快，一些 tutorial 可能不再适用或绕了弯路。因此这里还是为那些需要从 0 到 1 部署的同学提供了我的部署方案，以便大家在 23 分钟内拥有一个 Slurm 管理的 GPU 集群（实测）。

1. 安装 Slurm

slurm 依赖于 munge，先用 apt 安装好：

sudo apt install munge

至于 Slurm 本身，推荐直接用 apt 安装 stable 的版本，省时省力，当然也可以从官网下载最新版自己编译。安装的时候只需要安装 slurm-wlm 就可以了，apt 会自动帮我们安装其依赖的 slurmd 等组件。

sudo apt install slurm-wlm slurm-wlm-doc -y

查看安装是否成功：

slurmd -V
# slurm 15.08.7
slurmd -C

2. 开启 Slurm 服务

首先删除默认配置文件，新建一个（如果重装 slurm-wlm 也需要先删除配置）

sudo rm /etc/slurm-llnl/slurm.conf 
sudo vi /etc/slurm-llnl/slurm.conf

配置文件可以参考下面内容，我的服务器 hostname 为 v10032，请改成你自己的 hostname。当然，头铁的话也可以自己去官网找到相应版本的配置生成工具，生成配置：

# slurm.conf file generated by configurator easy.html.
# Put this file on all nodes of your cluster.
# See the slurm.conf man page for more information.
#
ControlMachine=v10032
#ControlAddr=
#
#MailProg=/bin/mail
MpiDefault=none
#MpiParams=ports=#-#
ProctrackType=proctrack/pgid
ReturnToService=1
SlurmctldPidFile=/var/run/slurm-llnl/slurmctld.pid
#SlurmctldPort=6817
SlurmdPidFile=/var/run/slurm-llnl/slurmd.pid
#SlurmdPort=6818
SlurmdSpoolDir=/var/spool/slurmd
SlurmUser=slurm
#SlurmdUser=root
StateSaveLocation=/var/spool/slurm-llnl
SwitchType=switch/none
TaskPlugin=task/none
#
#
# TIMERS
#KillWait=30
#MinJobAge=300
#SlurmctldTimeout=120
#SlurmdTimeout=300
#
#
# SCHEDULING
FastSchedule=1
SchedulerType=sched/backfill
SelectType=select/linear
#SelectTypeParameters=
#
#
# LOGGING AND ACCOUNTING
AccountingStorageType=accounting_storage/none
ClusterName=cluster
#JobAcctGatherFrequency=30
JobAcctGatherType=jobacct_gather/none
#SlurmctldDebug=3
#SlurmctldLogFile=
#SlurmdDebug=3
#SlurmdLogFile=
#
#
# COMPUTE NODES
GresTypes=gpu
NodeName=v10032 CPUs=32 Gres=gpu:1 State=idle
PartitionName=debug Nodes=v10032 Default=YES MaxTime=INFINITE State=UP

可以注意到，在 slurm.conf 中我声明了 GPU 相关资源 Gres=gpu:1，即 v10032 有一个可用 GPU。这个可用的 GPU 资源还需要单独在 gres.conf 中注册：

sudo vi /etc/slurm-llnl/gres.conf

可以参考如下配置：

Name=gpu File=/dev/nvidia7 # 只使用第 8 个 GPU
# 查询可用的 GPU 挂载，可以适用 ls /dev | grep nvidia

根据配置文件创建目录并设置权限给 slurm 用户（不要乱改）：

sudo rm -rf  /var/spool/slurm-llnl
sudo mkdir /var/spool/slurm-llnl
sudo chown -R slurm.slurm /var/spool/slurm-llnl
sudo rm -rf /var/run/slurm-llnl/
sudo mkdir /var/run/slurm-llnl/
sudo chown -R slurm.slurm /var/run/slurm-llnl/

按顺序启动程序（一个 start 完再 start 下一个）：

sudo systemctl start slurmd
sudo systemctl enable slurmd
sudo systemctl start slurmctld
sudo systemctl enable slurmctld

如果需要修改配置，记得（按顺序）重新启动：

sudo systemctl restart slurmd
sudo systemctl restart slurmctld

测试能否成功运行：

sinfo

执行一个简单任务：

salloc
srun nvidia-smi
exit

3. 从单机到集群上的 Slurm

在每台机器上执行上述步骤，确保 slurm 能够在单机正常运行。之后我们需要额外配置 munge 做机器间的通信，首先选定一台机器作为控制节点，其他为计算节点。在控制机器上生成 munge.key：

sudo create-munge-key

将 key 拷贝出来，并使用 scp 复制到其他计算机器上（复制出来时需要更改权限）：

sudo chmod 400 /etc/munge/munge.key
sudo chown munge.munge /etc/munge/munge.key
sudo cp /etc/munge/munge.key ~/key
sudo chmod 777 ~/key
sudo chown zhangzhi.zhangzhi ~/key
scp ~/key v10016:~/key

在计算机器上恢复 key 的权限和所有者信息，放到与控制机器相同的位置，并重启 munge 服务（如果不重启会报错 unmunge: Error: Invalid credential）：

sudo chmod 400 ~/key
sudo chown munge.munge ~/key
sudo cp ~/key /etc/munge/munge.key
sudo systemctl restart munge

修改所有机器上的 slurm.conf（注意所有机器保持一致），重新定义好控制节点和计算节点：

# slurm.conf file generated by configurator easy.html.
# Put this file on all nodes of your cluster.
# See the slurm.conf man page for more information.
#
ControlMachine=v10032
#ControlAddr=
#
#MailProg=/bin/mail
MpiDefault=none
#MpiParams=ports=#-#
ProctrackType=proctrack/pgid
ReturnToService=1
SlurmctldPidFile=/var/run/slurm-llnl/slurmctld.pid
#SlurmctldPort=6817
SlurmdPidFile=/var/run/slurm-llnl/slurmd.pid
#SlurmdPort=6818
SlurmdSpoolDir=/var/spool/slurmd
SlurmUser=slurm
#SlurmdUser=root
StateSaveLocation=/var/spool/slurm-llnl
SwitchType=switch/none
TaskPlugin=task/none
#
#
# TIMERS
#KillWait=30
#MinJobAge=300
#SlurmctldTimeout=120
#SlurmdTimeout=300
#
#
# SCHEDULING
FastSchedule=1
SchedulerType=sched/backfill
SelectType=select/linear
#SelectTypeParameters=
#
#
# LOGGING AND ACCOUNTING
AccountingStorageType=accounting_storage/none
ClusterName=cluster
#JobAcctGatherFrequency=30
JobAcctGatherType=jobacct_gather/none
#SlurmctldDebug=3
#SlurmctldLogFile=
#SlurmdDebug=3
#SlurmdLogFile=
#
#
# COMPUTE NODES
GresTypes=gpu
NodeName=v10032 CPUs=32 Gres=gpu:1 State=idle
NodeName=v10016 CPUs=32 Gres=gpu:1 State=idle
PartitionName=debug Nodes=v10016,v10032 Default=YES MaxTime=INFINITE State=UP

修改所有机器上的 gres.conf 定义本机的 GPU 资源：

# v10016 节点上的 gres.conf
Name=gpu File=/dev/nvidia7 # v10016 也只使用第 8 个 GPU

按顺序重启所有 slurm 服务：

sudo systemctl restart slurmd
sudo systemctl restart slurmctld

执行一个简单的任务进行测试：

salloc -N2
srun nvidia-smi
exit

4. 使用文件共享系统

为了存放代码及数据，供计算节点使用，我们需要在存储节点上安装文件共享系统，为计算节点提供统一的代码和数据的访问地址。如果没有存储节点，可以在控制节点安装。（这里以控制节点 172.31.224.79 为例）用如下命令，安装 NFS server 用于共享文件夹，以备控制节点和计算节点（172.31.234.192）访问：

sudo apt-get install -y nfs-kernel-server

创建一个用于共享的文件夹：

mkdir /home/zhangzhi/Data/exports

添加配置信息，指定共享的文件夹、可以访问该文件夹的主机（需要写 IP 地址）、权限等信息：

/home/zhangzhi/Data/exports 172.31.224.79(rw,sync,no_root_squash,no_subtree_check) 172.31.234.192(rw,sync,no_root_squash,no_subtree_check)

重启 NFS server 使配置生效：

sudo /etc/init.d/rpcbind restart
sudo /etc/init.d/nfs-kernel-server restart

在所有需要访问该共享文件夹的控制节点和计算节点安装客户端，用于将共享的文件夹挂载到其路径上：

sudo apt-get install -y nfs-common

查看共享的文件夹：

showmount -e 172.31.224.79

挂载到（所有需要访问该共享文件夹）的同一位置上：

# 这里 controller 既做控制节点，又做存储节点，因此我们只需要再在计算节点 172.31.234.192 上挂载，所有控制节点和计算节点就都可以访问到这个位置了。
mkdir /home/zhangzhi/Data/exports
sudo mount -t nfs 172.31.224.79:/home/zhangzhi/Data/exports /home/zhangzhi/Data/exports

如果挂载到不想要的地方了，就 umount 掉：

sudo unmount /home/zhangzhi/Data/exports

三、多机多卡训练脚本

做完上述准备工作，我们就可以在 /home/zhangzhi/Data/exports 文件夹内编写分布式训练代码进行多机多卡的分布式训练了。slurm 会将写好的 python 程序在每个节点上分别执行，调用节点上定义的 GPU 资源进行运算。我们要做的是告诉每一个执行的任务，要用哪些训练哪一部分数据，反向传播的结果如何合并。

完整的 ImageNet 训练脚本请见：github ！

完整的 ImageNet 训练脚本请见：github ！

完整的 ImageNet 训练脚本请见：github ！

下面我们对重点部分进行分步讲解：

  1. 任务与进程的管理

我们首先需要获得每个任务（对应每个节点）的基本信息，以便针对任务的基本信息处理其应当负责的数据。在使用 slurm 执行 srun python 代码时，python 可以从环境变量 os.environ 中获取当前 python 进程的基本信息，这里只介绍需要用到的部分：

import os
local_rank = os.environ['SLURM_PROCID'] # 当前任务的编号（比如节点 1 执行 1 号任务，节点 2 执行 2 号任务）
world_size = os.environ['SLURM_NPROCS'] # 共开启的任务的总数（共有 2 个节点执行了 2 个任务）
job_id = os.environ['SLURM_JOBID'] # 当前作业的编号（这是第 1 次执行 srun，编号为 1）

接下来，在每个任务（节点）中，我们需要为节点中的每个 GPU 资源分配一个进程，管理该 GPU 应当处理的数据：

当前节点的 GPU 的数量可以由 torch.cuda 查询得到：

ngpus_per_node = torch.cuda.device_count()

其后，我们使用 torch.multiprocessing 创建 ngpus_per_node 个进程，其中，每个进程执行的函数为 main_worker ，该函数调用所需要的由 args 传入：

mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))

spawn 创建每个进程时，会为 main_worker 注入当前进程编号和输入参数，可以按顺序取用：

def main_worker(gpu, ngpus_per_node, args):
  gpu -> 当前的进程编号，也就是 gpu 编号
  ngpus_per_node -> 我们传入的当前节点的 gpu 的数量
  args -> 我们传入的其他参数
  ...

在编写 main_worker 时，我们首先需要解决的问题是：不同节点、或者同一节点间的不同进程之间需要通信来实现数据的分割、参数的合并。我们可以使用 pytorch 的 dist 库在共享文件系统上创建一个文件进行通信。

例如，在控制节点和计算节点都可以访问的 /home/zhangzhi/Data/exports 目录下，slurm 同时在所有节点上开始执行任务，编号为1的任务会创建一个文件，其他任务（根据用户给定的地址）轮训该文件是否存在，直到创建完成任务之间的通信开始。

import torch.distributed as dist

def main_worker(gpu, ngpus_per_node, args):
  dist_url = "file://dist_file.{}".format(job_id)
  rank = local_rank * ngpus_per_node + gpu
  dist.init_process_group(backend='nccl', init_method=dist_url, world_size=world_size, rank=rank)
  ...

2. 每个进程中 pipline 的管理

完成进程创建和通信后，下一步就是实现我们常用的 pipline 了：加载模型、加载数据、正向传播、反向传播。

这里，我们把模型加载进当前进程所对应的 GPU 中。和单机训练的不同之处仅限于 device 不是任意定义的，而是分配给 main_worker 的 GPU 编号。此外，model 外还需要使用 DistributedDataParallel 进行一下封装，便于 pytorch 帮助我们同步训练参数：

def main_worker(gpu, ngpus_per_node, args):
  dist_url = "file://dist_file.{}".format(job_id)
  rank = local_rank * ngpus_per_node + gpu
  dist.init_process_group(backend='nccl', init_method=dist_url, world_size=world_size, rank=rank)
  ...
  torch.cuda.set_device(gpu)
  model.cuda(gpu)
  model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu])

接着，我们把当前进程对应的数据段采样出来，也加载到对应的 GPU 中。我们也可以使用 pytorch 的 dist 库轻松实现这个采样过程：

def main_worker(gpu, ngpus_per_node, args):
  dist_url = "file://dist_file.{}".format(job_id)
  rank = local_rank * ngpus_per_node + gpu
  dist.init_process_group(backend='nccl', init_method=dist_url, world_size=world_size, rank=rank)
  ...
  torch.cuda.set_device(gpu)
  model.cuda(gpu)
  model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
  ...
  train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
  train_loader = torch.utils.data.DataLoader(train_dataset,
                                             batch_size=args.batch_size,
                                             num_workers=2,
                                             pin_memory=True,
                                             sampler=train_sampler)
  for i, (images, target) in enumerate(train_loader):
    images = images.cuda(gpu, non_blocking=True)
    target = target.cuda(gpu, non_blocking=True)

最后，进行正常的正向和反向传播就可以了，dist 库会帮我们自动解决不同进程之间的梯度/参数合并：

def main_worker(gpu, ngpus_per_node, args):
  dist_url = "file://dist_file.{}".format(job_id)
  rank = local_rank * ngpus_per_node + gpu
  dist.init_process_group(backend='nccl', init_method=dist_url, world_size=world_size, rank=rank)
  ...
  torch.cuda.set_device(gpu)
  model.cuda(gpu)
  model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[gpu])
  ...
  train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
  train_loader = torch.utils.data.DataLoader(train_dataset,
                                             batch_size=args.batch_size,
                                             num_workers=2,
                                             pin_memory=True,
                                             sampler=train_sampler)
  for i, (images, target) in enumerate(train_loader):
    images = images.cuda(gpu, non_blocking=True)
    target = target.cuda(gpu, non_blocking=True)
    ...
    output = model(images)
    loss = criterion(output, target)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

3. 跑起来！

完整的 pytorch + slurm 在 ImageNet 上的训练代码可以从 github 上下载：

git clone https://github.com/tczhangzhi/pytorch-distributed

下载完成后，进入 pytorch-distributed-slurm 文件夹，并使用 srun 执行：

cd pytorch-distributed
srun -N2 --gres gpu:1 python distributed_slurm_main.py --dist-file dist_file

当 batch 很大时，可能会遇到这个错误：

RuntimeError: cuDNN error: CUDNN_STATUS_NOT_SUPPORTED. This error may appear if you passed in a non-contiguous input.

这是由于矩阵过大时 cudnn 的加速导致的 tensor 不连续，如果非要保持 batch size 可以禁用 cudnn 的加速（当然，最好降低 batch size 否则速度较慢）：

torch.backends.cudnn.enabled = False

四、后记

在交流过程中，笔者发现工业界 horovod + NFS 很受青睐。但是大部分高校仍采用上述的 Slurm + NFS + DDP 或 Slurm + NFS + DDP2（在多机间使用 DDP，单机多卡间使用 DP）的方案。还有一些细分的应用领域，比如联邦学习，基于 k8s 的 FATE 也火了一阵。但是限于其 task-special 的特点，恐怕不适用于日常魔改的科研。

从笔者的个人体验来讲，由于 Slurm 使用的时间较长，其性能和稳定性都高于 horovod。由于目前个人实现的 horovod 版性能较差（实际 ring all-reduce 应该性能更好），因此还有待总结相关最佳实践，后续会考虑添加与 horovod 的横向对比。

此外，最近笔者也试用了 pytorch-lightning 开箱即用的多机多卡功能，感觉较为满意，正在整体迁移重构。后续也会考虑添加与 pytorch-lightning 的横向对比。

最后，在完成相关分布式实践的开源后，笔者计划在此基础上继续探讨在分布式训练中的 CNN trick 库和 gridsearch 脚本，希望对同学们充分利用 GPU 涨点有所帮助。目前笔者用到较多的功能主要包括：试验驱动的git、分布式的log、分布式的超参封装等。在实现过程中也参照了一些优秀的项目，例如 hpman、fitlog、hyperopt 和 ray 但是感觉仍很不满意。比起臃肿的架构，小而美才是魔改的灵魂伴侣。因此相关部分的开源计划也只能留到下一个 blog 啦～

注：由于最近试验排期较紧，所以行文比较潦草，也请大家多多提 issue 和 pr 喔～

编辑于 2020-06-21 12:27
分布式计算
PyTorch
深度学习（Deep Learning）

文章被以下专栏收录