Deep Learning:PyTorch 基于docker 容器的分布式训练实践

引言

PyTorch distributed currently only supports Linux.
 
这句话是来自 pytorch 官网 的 torch.distributed 部分，说明 pytorch 支持分布式训练，而且只在linux 上支持。
 
torch.distributed supports three backends, each with different capabilities.
 
同样来自上述 页面，pytorch 在分布式训练中，支持三种后端（backend）进行集群管理或者通信。

那么，在什么情况下选用什么样的backend？

• 经验法则
  • 使用NCCL后端进行分布式GPU培训
  • 使用Gloo后端进行分布式CPU培训。
• 具有InfiniBand互连的GPU主机
  • 使用NCCL，因为它是目前唯一支持InfiniBand和GPUDirect的后端。
• GPU主机与以太网互连
  • 使用NCCL，因为它目前提供最佳的分布式GPU训练性能，特别是对于多进程单节点或多节点分布式训练。如果您遇到NCCL的任何问题，请使用Gloo作为后备选项。（请注意，Gloo目前运行速度比GPU的NCCL慢。）
• 具有InfiniBand互连的CPU主机
  • 如果您的InfiniBand已启用IP over IB，请使用Gloo，否则请使用MPI。我们计划在即将发布的版本中为Gloo添加InfiniBand支持。
• 具有以太网互连的CPU主机
  • 除非您有特殊原因要使用MPI，否则请使用Gloo。

实践部分

• 首先，假设，我们要起一个分布式训练，起 2 个 rank。
• 把他们起来两台机器上，机器 1 的 ip 为 192.168.61.55，机器 2 的 ip 为 192.168.61.56
• 机器 1 上起 1 个 rank，机器 2 上起 1 个 rank

我这里使用的示例代码来自于 github 上的一段开源代码， 地址是 https://github.com/pytorch/examples/tree/master/mnist 。我对这段代码做了些更改，使其支持起 rank 的操作，测试的时候一台机器上起一个 rank 。你可以在 网上搜索到类似的修改。
 
但是，这不是一个最佳的测试例子，原因是因为这段代码缺少 类似 all_reduce 这样的操作，缺少运算过程中所有 rank 的权重进行修改、因而引起的 PCI 通信或者 InfiniBand 网络通信（非系统通信）等等。我不明白的是，网上有很多人把这段代码作为一个 非常有效的例子 用来 测试分布式运算，估计都是没有深究，互相抄抄改改。
 
我找到了一段最佳实验代码， 地址是 https://github.com/seba-1511/dist_tuto.pth/blob/gh-pages/train_dist.py 。奈何不是做 深度学习 出身，修改了几次也没能成功修改成需要符合需求的测试代码（有高手的话，麻烦能修改好，联系我测试测试）。这段代码是自带 创建 rank 的，并且是单机运行的。我想要的是，自己创建每一个 rank，并且支持多机分布式。

多机分布式训练

• 需要安装 python3 环境，并且安装 torch 和 torchvision 。两台机器的 python 环境应当保持一致。
• /mnt/pytorch/ 是一个分布式存储挂载目录。
• 执行程序需要下载数据，如果你需要数据，可以在 https://download.csdn.net/download/github_37320188/11670737 查找 （由于没有积分了，小收几个积分用用，不能老做公益呀）。当然也可以自己下载，有合适的网络环境的话，相对来说就是花点时间。

# 机器 1 上输入命令
python3 /mnt/pytorch/mnist.py --rank 0 --world-size 2 --init-method tcp://192.168.61.55:9090 --backend gloo

# 机器 2 上输入命令
python3 /mnt/pytorch/mnist.py --rank 1 --world-size 2 --init-method tcp://192.168.61.55:9090 --backend gloo

其实这已经是一个多机训练的例子了。但是我建议在操作之前，先单节点 跑 2 个 rank 试试。
并且 run 一下 官方 示例 的一些代码，玩一下。
 
这个操作中，以 gloo 作 backend ，以 rank0 作为 master 通信，9090 机器 1 的未被占用通信端口。相关概念可以仔细阅读官方说明。

docker 分布式训练

• 首先你需要一个 支持 python 、torch、torchvision 的 docker image。如果选择省事，我建议去docker hub 下载一个官方的 tag 为 pytorch/pytorch:latest 的 image 。
• 然后，这个 image 自带了 conda 、python3、torch，但是你需要自己安装 torchvision ，打包一个新的 image。建议通过 pip 安装，并且受限于 国内环境问题，你可能一次成功不了，可能需要多尝试几次。
• 最后，如果实在 操作 docker 容器的耐心都没有了，你可以从我的 docker hub 里下载一个。tag 是 docker.io/loongc/pytorch:pytorch

操作一波

# 机器 1 
docker run -it --rm  -v /mnt/pytorch/:/pytorch -w /pytorch --net=host loongc/pytorch:pytorch  python mnist.py --rank 0 --world-size 2 --init-method tcp://192.168.61.55:9090 --backend gloo

# 机器 2
docker run -it --rm  -v /mnt/pytorch/:/pytorch -w /pytorch --net=host loongc/pytorch:pytorch  python mnist.py --rank 1 --world-size 2 --init-method tcp://192.168.61.55:9090 --backend gloo

讨论

以上部分，测试都是在 多节点CPU 上运行的。如果你需要在GPU 进行运算：

• 首先，你需要节点本身 支持gpu 运算，安装 cuda 环境
• 如果你希望在 docker 中进行 gpu 运算，那么 docker image 需要支持gpu 运算。幸运的是，前文提到的 image 支持gpu 运算

 
如果，环境已经准备好了，就可以做一些 GPU 相关的测试了。
 
使用 基于 Docker 的 GPU 分布式运算，如果不需要 诸如 运算性能 特殊的需求，你可以直接使用 上一节 的 操作方法即可。如果 需要 诸如使用 分布式文件系统通信、环境变量通信、InfiniBand通信，使用 nccl 作为 backend 等等，将在下一节做一个总结。

backend 与通信方式 的一些总结

引言中，已经提到 3 种 backend，分别是：MPI、GLOO、NCCL。这 3 种 backend 支持了 3 种 通信 方式：tcp、环境变量、分布式文件系统（当前使用的 pytorch 版本如此，后续的版本就不知道了）。
 
所以，大致上可以认为不同 rank 之间的通信就是 由 backend 与 通信方式的 组合完成的。当然，在不同的组合方式所带来的 通信速度（影响运算整体性能）是不同的，能实现的功能也是不太一样的（详情见引言中的表格）。
 
本文前面的内容当中，使用的 backend 都是GLOO，使用的通信方式都是 tcp。
 
本节内容，主要围绕 NCCL 的 3 种 通信操作方式展开，以及如何使用 InfiniBand。
MPI 的方式我没有测过，也不打算测试，主要原因是需要 额外的依赖 。

NCCL 与 tcp

nccl 与 gloo 使用 tcp 的方式是一致的，不需要多余的操作。在运行 python 脚本的时候，只需要将传入 backend 的参数 gloo 改为 nccl 即可。

NCCL 与 环境变量

nccl 使用环境变量，相对于 tcp 要复杂一些。

• 首先，需要将传入 backend 的参数 gloo 改为 nccl
• 其次，将传入 init-method 的参数 由 tcp://ip:port 改为 env://
• 另外，容器启动的时候的需要给容器设置 2 个环境变量 MASTER_ADDR ，MASTER_PORT

还是在 192.168.61.55 和 192.168.61.56 两台机器上操作，启动大致如下：

# 机器 1
docker run -it --rm --env-file /mnt/pytorch_env/master -v /mnt/pytorch/:/pytorch -w /pytorch --net=host loongc/pytorch:pytorch python mnist.py --rank 0 --world-size 2 --init-method env:// --backend nccl

# 机器 2
docker run -it --rm --env-file /mnt/pytorch_env/master -v /mnt/pytorch/:/pytorch -w /pytorch --net=host loongc/pytorch:pytorch python mnist.py --rank 1 --world-size 2 --init-method env:// --backend nccl

# 文件 /mnt/pytorch_env/master 内容如下
MASTER_ADDR=192.168.61.55
MASTER_PORT=9090

上述 /mnt/pytorch_env/master 文件的内容中 MASTER_ADDR=192.168.61.55 是 rank0 所在机器的 ip ，MASTER_PORT=9090 是该机器上一个未被征用的 端口

NCCL 与 分布式文件系统

nccl 使用分布式文件系统，那么就需要把分布式文件系统 映射进 容器当中，在这里，我是在 挂载目录下 /mnt/pytorch/ 创建了一个文件 backendnccl ，因为本身这个挂载 就是分布式文件系统。

# 以机器 1 上的操作为例
docker run -it --rm -v /mnt/pytorch/:/pytorch -w /pytorch --net=host loongc/pytorch:pytorch python mnist.py --rank 0 --world-size 2 --init-method file:///pytorch/backendnccl --backend nccl

使用 InfiniBand

原本，pytorch 是支持 IP over IB 的，所以可操作 InfiniBand 的方式主要有以下几种：

• 在使用 tcp方式 进行通信 时，可以不使用 IP 地址，而直接使用 IB (InfiniBand) 卡的地址
• 在使用环境变量过程中，设置的环境变量中 将 MASTER_ADDR 设置成 IB 卡的地址

另外，NCCL和Gloo后端都会尝试找到要使用的正确网络接口，如果自动检测到的接口不正确，您可以使用以下环境变量覆盖它（适用于相应的后端）：
 
在使用 过程中，如果确定 使用 nccl，并且打算用 InfiniBand ，推荐通过环境变量（写进 前文提到的 文件 /mnt/pytorch_env/master ） NCCL_SOCKET_IFNAME=ib0

额外说明，如果你的机器上有多张 InfiniBand 卡，你可以通过 NCCL_SOCKET_IFNAME 指定任意一张，如 ib0 或者 ib1 这种
 
但是，你可以直接使用 NCCL_SOCKET_IFNAME=ib 进行指定，这样所有的 ib 卡都有机会被调用起来用于通信。你可以通过相关工具，如 ibdump 进行监控 ib 卡的发包 和 数据的发送与接受。
 
经过分析，我们认为在 以发送数据包 的方式进行通信的过程中，除了真正用于计算的数据包，还可能包括 机器通信 的数据。

总结

这里，实现了手动启动多个rank 进行 基于 docker 的分布式训练。那么，自动化进行分布式训练就变得比较容易了。
 
可以通过 mesos 或者 k8s 启动多个 容器，进行集群管理，完成分布式训练。通过 设计 api 指定启动 指定数目的 rank，指定硬件资源，以及是否使用 InfiniBand 。
 
如果有更好的实现 分布式训练 PyTorch ，欢迎留言讨论。