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

引言

MXNet supports distributed training enabling us to leverage multiple machines for faster training.
 
MXNet支持分布式培训，使我们能够利用多台机器进行更快速的培训。这段话来自于 MXNet官网 ，说明了MXNet 支持跨越设备运行。
 
How to Start Distributed Training?
那么MXNet是怎样实现分布式训练的？
 
仔细阅读官方文档，官方文档给出了一个示例并写了这样一段话：
For distributed training of this example, we would do the following:
If the mxnet directory which contains the script image_classification.py is accessible to all machines in the cluster (for example if they are on a network file system), we can run:

../../tools/launch.py -n 3 -H hosts --launcher ssh python image_classification.py --dataset cifar10 --model vgg11 --epochs 1 --kvstore dist_sync

这是个使用 ssh 方式进行分布式训练的例子。

其中使用的测试代码，来源于 https://github.com/apache/incubator-mxnet

注意，这里的 launch.py 是一个非常重要的工具，如果仔细阅读 python 源码，会发现正是由于执行它才能实现分布式训练，它目前支持了5种分布式或并发训练方式：

launch 方式

–launcher denotes the mode of communication. The options are:

• ssh if machines can communicate through ssh without passwords. This
  is the default launcher mode.
• mpi if Open MPI is available
• sge for Sun Grid Engine
• yarn for Apache Yarn
• local for launching all processes on the same local machine. This can be used for debugging purposes.

官网介绍了5种方式进行分布式或并行训练。这些方式都是以类似于集群管理的方式进行分布式训练。如果阅读整个测试代码，你可能还会发现其他几种集群管理方式，如mesos，不知道什么原因，官网没有介绍，也没有明确说支持 mesos 调度。

实际上，跟 tensorflow 或者 pytorch 进行分布式训练稍有不同。
在tensorflow 或者 pytorch 进行分布式训练，可能需要自己手动 或者 通过 mpi 工具起 不同角色，如 tensorflow 中的 ps 和 worker ，pytorch 中的 rank 。
而 MXNet 起不同的角色，全部都交给 launch.py 做完了。相当于对普通用户进行了一定程度上的屏蔽。
 
MXNet 这样做有一定的好处，普通用户只需要关注 训练脚本 的编写，而不需要关注 分布式计算 集群 如何运作。相对应的，当用户想抛开 launch.py 进行多机器多节点分布式训练 时，这也会成为弊病。因为MXNet 官网并没有仔细介绍如何手动启动 分布式训练。

官方单机并发训练

export COMMAND='python example/gluon/image_classification.py --dataset cifar10 --model vgg11 --epochs 1 --kvstore dist_sync'
DMLC_ROLE=server DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=server DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=scheduler DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=worker DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=worker DMLC_PS_ROOT_URI=127.0.0.1 DMLC_PS_ROOT_PORT=9092 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND

这是MXNet官方给出的在单节点机器，起一个 scheduler ，两个 worker， 两个 server 进行分布式训练。但是实际情况中，我们希望将 上述三种角色， 以在不同主机上进行启动，进行大规模硬件资源进行分布式训练。
 
另外，我们还希望通过仅仅提供硬件资源，以云计算的方式提供服务，用户之间具备隔离。也就是我们希望所有计算服务，以 docker 容器化方式启动。通过指定 worker 和 server 的数量以及要使用的资源，来启动容器。
 
如此一来，我们不能使用官方的 推荐 的 集群管理方式进行分布式训练。我们真正需要实现的是，launch.py 实现的功能，并以自动化的方式在分布式集群上 启动 docker 容器， 进行分布式训练。
 
所有的用户可以创建自己的 docker image，push 到分布式 image 仓库，并挂载不同的存储，达到用户隔离的目的。

这个 例子中用的 cifar10 数据 在执行过程中下载非常慢，我把它 放在https://download.csdn.net/download/github_37320188/11584523 ， （由于没有积分了，小收几个积分用用，不能老做公益呀）

实践部分

如果已经尝试过绕过 launch.py 进行分布式训练，最直接的方法是修改官方给出的 单机 并发训练 脚本。往上翻 可以看到。
 
实际上，该脚本通过设置环境变量的方式，设置 不同 角色的参数，然后执行 python 脚本。
官方对相关参数已经有部分解释。

• DMLC_ROLE: Specifies the role of the process. This can be server,
  worker or scheduler. Note that there should only be one scheduler.
  When DMLC_ROLE is set to server or scheduler, these processes start
  when mxnet is imported.
• DMLC_PS_ROOT_URI: Specifies the IP of the scheduler
• DMLC_PS_ROOT_PORT: Specifies the port that the scheduler listens to
• DMLC_NUM_SERVER: Specifies how many server nodes are in the cluster
• DMLC_NUM_WORKER: Specifies how many worker nodes are in the cluster

 
但是尝试修改参数过后，如果成功找到方向的话，会发现，官方给出的环境变量是不够的。

解决思路

这里提供 3 个解决思路去进行解决这个问题。

1. 仔细阅读 launch.py 的源码，这是官方给的示例，虽然不清楚什么原因官方没有仔细解释这部分功能、并将说明写的再详细些。但是 这部分确实实现了分布式 训练，代码中对所有必须的环境变量进行了export，具有非常重要的参考的价值。
2. 如果没有 完整 阅读 launch.py 的耐心或者准备，你可以想办法打印出以 ssh 方式分布式训练时 export 的环境变量，但不保证能找到所有的。
3. 最后，网络上有一些开源的产品，做了类似自己实现 MXNet 分布式训练，不依赖 launch.py ，并以云计算的方式提供服务。但是不保证他们在文档方面说的很详细，不过能实现说明研究过 手动进行分布式训练 。这里推荐一个 https://github.com/bytedance/byteps，来源于 Bytedance

实际操作

一开始我打印出以 ssh 方式分布式训练时 export 的环境变量，用来确定了确实是忽略了一些环境变量。
 
然后，我才开始阅读 launch.py 的代码，整个官方源码包其他的代码，找到了一些官方说明没有提到的变量。
 
操作一波

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

# 机器 1 输入
export COMMAND='python3 /mnt/mxnet-test/incubator-mxnet/example/gluon/image_classification.py --dataset cifar10 --model vgg11 --epochs 1 --kvstore dist_sync'
DMLC_ROLE=scheduler DMLC_PS_ROOT_URI=192.168.61.55 DMLC_PS_ROOT_PORT=9091 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 $COMMAND &
DMLC_ROLE=server DMLC_PS_ROOT_URI=192.168.61.55 DMLC_NODE_HOST=192.168.61.55 DMLC_PS_ROOT_PORT=9091 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 DMLC_SERVER_ID=0 $COMMAND &
DMLC_ROLE=worker DMLC_PS_ROOT_URI=192.168.61.55 DMLC_NODE_HOST=192.168.61.55 DMLC_PS_ROOT_PORT=9091 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 DMLC_WORKER_ID=0 $COMMAND

# 机器 2 输入
export COMMAND='python3 /mnt/mxnet-test/incubator-mxnet/example/gluon/image_classification.py --dataset cifar10 --model vgg11 --epochs 1 --kvstore dist_sync'
DMLC_ROLE=server DMLC_PS_ROOT_URI=192.168.61.55 DMLC_NODE_HOST=192.168.61.56 DMLC_PS_ROOT_PORT=9091 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 DMLC_SERVER_ID=1 $COMMAND &
DMLC_ROLE=worker DMLC_PS_ROOT_URI=192.168.61.55 DMLC_NODE_HOST=192.168.61.56 DMLC_PS_ROOT_PORT=9091 DMLC_NUM_SERVER=2 DMLC_NUM_WORKER=2 DMLC_WORKER_ID=1 $COMMAND

python 版本为 3.6.8
scheluder 不需要跟 任何其他 角色 绑定到 同一个 节点，也能独立运行。
甚至 可以把 这 5 个服务起在 5台机器上
DMLC_PS_ROOT_PORT 指定的是 起 scheduler 的机器上未被分配的可用端口
/mnt/mxnet-test/ 是一个分布式存储挂载目录

这里看到的是，相对于 官网 给到的介绍 多了几个环境变量：

• DMLC_NODE_HOST 启动 worker 或者 server 标示自己所在主机的地址，可以是 ip 也可以是 ib 卡的地址
• DMLC_SERVER_ID 启动 server 是标记 index ，从 0 开始 计数，不可重复
• DMLC_WORKER_ID 启动 worker 是标记 index ，从 0 开始 计数，不可重复

如果你使用了ib 卡，除了直接指定 DMLC_NODE_HOST=ib卡的地址 ，还可通过 DMLC_INTERFACE 环境变量指定
除此之外，指定 DMLC_PS_ROOT_URI 和 DMLC_NODE_HOST 不能使用主机名，必须使用 ip 地址 或者 ib 卡的地址。原因是 MXNet 的通信依赖于 zmq ，其不支持 主机名。具体原因或者代码可以网上查阅

docker 分布式训练

如果已经通过了上述 多机 分布式训练。后面就会变的稍微容易些。
首先，需要一个支持 python3 的 image，并且 支持 mxnet。你可以 在 docker hub 上找到一个

docker pull loongc/mxnet:v1

操作一波

• 假设 要启动 5 个容器，1 个运行 scheduler，2 个 worker， 2 个 server
• 准备两台机器，机器 1 的 ip 为 192.168.61.55， 机器2 的 ip 为 192.168.61.56
• 机器 1 上起 scheduler 容器，1 个 worker 容器，1 个server 容器
• 机器 2 上起 1 个 worker 容器， 1 个 server 容器

# 机器 1 上的 启动 shell 脚本 mxnet-test.sh
docker run -d --env-file /tmp/mxnet_env/worker  \
--name worker_1 \
-v /mnt/mxnet-test/incubator-mxnet:/incubator-mxnet \
-w /incubator-mxnet/example/gluon/ \
--net=host loongc/mxnet:v1 \
python3 image_classification.py --dataset cifar10 --model vgg11 --epochs 10 --kvstore dist_sync

docker run -d --env-file /tmp/mxnet_env/server  \
--name server_1 \
-v /mnt/mxnet-test/incubator-mxnet:/incubator-mxnet \
-w /incubator-mxnet/example/gluon/ \
--net=host loongc/mxnet:v1 \
python3 image_classification.py --dataset cifar10 --model vgg11 --epochs 10 --kvstore dist_sync

docker run -d --env-file /tmp/mxnet_env/scheduler  \
--name scheduler \
-v /mnt/mxnet-test/incubator-mxnet:/incubator-mxnet \
-w /incubator-mxnet/example/gluon/ \
--net=host loongc/mxnet:v1 \
python3 image_classification.py --dataset cifar10 --model vgg11 --epochs 10 --kvstore dist_sync

这里在机器 1 上启动了 3 个容器。
其中 的 /tmp/mxnet_env/worker ， /tmp/mxnet_env/server ， /tmp/mxnet_env/scheduler 文件 中写的是环境变量

这里举个例子

# 机器1 上的 scheduler 配置 /tmp/mxnet_env/scheduler
DMLC_ROLE=scheduler
DMLC_PS_ROOT_URI=192.168.61.55
DMLC_PS_ROOT_PORT=9091
DMLC_NUM_SERVER=2
DMLC_NUM_WORKER=2

# 机器1 上的 worker 配置 /tmp/mxnet_env/worker 
DMLC_ROLE=worker
DMLC_PS_ROOT_URI=192.168.61.55
DMLC_PS_ROOT_PORT=9091
DMLC_WORKER_ID=0
DMLC_NODE_HOST=192.168.61.55
DMLC_NUM_SERVER=2
DMLC_NUM_WORKER=2

# 机器1 上的 server 配置 /tmp/mxnet_env/server
DMLC_ROLE=server
DMLC_PS_ROOT_URI=192.168.61.55
DMLC_PS_ROOT_PORT=9091
DMLC_NODE_HOST=192.168.61.55
DMLC_SERVER_ID=0
DMLC_NUM_SERVER=2
DMLC_NUM_WORKER=2

这些是机器 1 上的配置， 机器 2 上的配置很相似
区别在于：

1. 机器 2 上的启动脚本 mxnet-test.sh 不需要启动 scheduler 容器，直接拿 机器 1 上脚本 改改 删掉 scheduler 部分
2. 机器 2 上不需要 /tmp/mxnet_env/scheduler 配置脚本
3. 机器 2 上 /tmp/mxnet_env/worker 和 /tmp/mxnet_env/server 两个文件中 DMLC_NODE_HOST 修改成机器2 的地址，DMLC_WORKER_ID、DMLC_SERVER_ID 修改为 1 （从 0 开始 计数）

分别启动两台机器上的 启动脚本 mxnet-test.sh ，无严格先后顺序。可以查看 /mnt/mxnet-test/incubator-mxnet/example/gluon/image-classification.log 查看执行情况。

结论

已经实现了手动分布式 训练，那么离 自动化 实现分布式训练就不困难了。
 
可以通过 k8s 或者 mesos 进行集群管理，通过调用 api 启动 指定数目的 worker 和 server 容器，并指定硬件资源，确定是否使用 ib 卡等。
 
如果有更好的实现 分布式训练 MXNet ，欢迎留言讨论。