MXNet 结合 kubeflow 进行分布式训练

MXnet 是灵活且高效的深度学习库，它可以以数据并行的方式进行单机和多机的多卡训练，本文在介绍上述两种训练方式的同时，也会介绍 MXNet 结合 kubeflow 进行分布式训练的方法.

MXNet 以数据并行的方式进行单机多卡训练

MXNet 支持在多 CPU 和 GPU 上进行训练。其中，这些 CPU 和 GPU 可以分布在不同的服务器上。

一台 GPU 机器上的每块 GPU 都有自己的编号(编号从 0 开始计数)。如果想使用某一块特定的 GPU 卡，即可以在代码中直接指定 context(ctx); 也可以在命令行中传递参数–gpus。

例如: 如果想在 python 中使用 GPU 0 和 GPU 2,可以使用下面的代码创建网络模型。

如果多个 GPU 的计算能力不同，那么可以根据它们的计算性能来划分工作负载。比如，如果 GPU 0 是 GPU 2 性能的 3 倍，那么可以提供一个额外的负载选项 work_load_list=[3, 1]。

如果所有其它超参都相同，在多个 GPU 上训练结果应该和单 GPU 上的训练结果相同。但在实际应用中，由于随机存取（随机顺序或其它 augmentations），使用不同的种子初始化权重和 CuDNN，结果可能不同。

我们可以控制梯度聚合和模型更新(如执行，训练过程)的位置，通过创建不同的 KVStore（数据通信模块）。即可以使用 mx.kvstore.create(type)来创建一个实例，也可以使用程序的参数–kv-store type 来实现功能。

两种常用功能

• local: 所有的梯度都拷贝到 CPU 内存完成聚合，同时在 CPU 内存上完成权值的更新并拷贝回每个 GPUworker。这种方式主要在于 CPU 与 GPU,主要的性能负载在于 CPU 拷贝的负载。

• device: 梯度聚合和权值更新都在 GPU 上完成。GPU 之间的如果支持 Peer to Peer 通信（PCIe or NVLink），将避免 CPU 拷贝的负载，可以大大减轻 CPU 的负担，仅受限于通信带宽。PCIe 与 NVLink 通信带宽不同，NVLink 具备告诉的 Peer to Peer 通信带宽。

注意: 如果有大量的 GPU(比如: >=4),建议使用 device 能获取更好的性能。

MXNet 以数据并行的方式进行多机多卡训练

在介绍 MXNet 分布式训练之前，先介绍几个关键性的概念方便理解 MXNet 的分布式训练:

三种进程类型

• Worker: worker 进程会对每一批次(batch_size)的数据样本进行训练。对批数据进行处理之前，workers 会从 servers 服务器 pull 权重。对批处理数据处理完毕之后 workers 会聚合梯度数据发送给 servers。(如果训练模型的工作负载比较高，建议最好不要把 worker 和 server 运行在相同的机器上)。

• Server: 可以运行多个 server 服务进程，用于存储模型参数并与 worker 进行通讯。

• Scheduler: 一个调度器,在集群中负责调度的角色。主要包含: 等待各个 node 节点数据的上报数据并让各个 node 节点知道彼此的存在并互相通讯。

进程之间的关系如下图所示:

工作流程：

1.worker, server 向 scheduler 注册，获取相关的信息。

2.worker 从 server 端 pull 参数 w。

3.worker 基于参数 w 和数据计算梯度，然后 push 梯度到 server。

4.server 更新参数 w。

5.反复执行 2-4 过程。

KVStore

KVStore 是 MXNet 提供的一个分布式 key-value 存储，用来进行数据交换。KVStore 本质的实现是基于参数服务器。

• 通过引擎来管理数据一致性,这使得参数服务器的实现变得简单，同时使得 KVStore 的运算可以无缝的与其他部分结合在一起。

• 使用两层的通讯结构，原理如下图所示。第一层的服务器管理单机内部的多个设备之前的通讯。第二层服务器则管理机器之间通过网络的通讯。第一层的服务器在与第二层通讯前可能合并设备之间的数据来降低网络带宽消费。同时考虑到机器内和外通讯带宽和延时的不同性，可以对其使用不同的一致性模型。例如第一层用强的一致性模型，而第二层则使用弱的一致性模型来减少同步开销。

通过调用 mxnet.kvstore.create()函数并传递 dist 关键字参数来开启分布式训练的 KVStore 模式:

kv = mxnet.kvstore.create(‘dist_sync’)

分布式训练模式

当 KVStore 被创建并且包含 dist 关键参数就会开启分布式训练模式。通过使用不同类型的 KVStore，可以启用不同的分布式培训模式。具体如下:

• dist_sync: 已同步的方式进行分布式训练,在处理每批次(batch)的数据时,所有的 workers 需要使用相同的模型参数集合。这意味着 servers 参数服务需要接收来自所有 workers 模型参数之后，才能进行下一个批次数据的处理。因此在使用这种分布式训练方式时，server 参数服务需要等到所有的 worker 处理完毕之后，并且如果其中的某一个 worker 异常，会导致整个训练的过程 halts。

• dist_async: 已异步的方式进行分布式训练，server 参数服务只要收到 worker 的计算梯度就会立即更新存储。这意味着哪个 worker 处理完当前的批次数据，就可以继续下一批次数据的处理。因此该种方式的分布式训练方式比 dist_sync 要快,但是需要花费更多的 epochs 去收敛。

• dist_sync_device: 该分布式训练模式与 dist_sync 训练模式相同，只是 dist_sync_device 模式会在多 GPUs 上进行梯度聚合和更新权重，而 dist_sync 是在 CPU 上进行这些操作。这种模式比 dist_sync 要快，因为 GPU 和 CPU 之前的通信，但是会占用更多的 GPU 显存。

• dist_async_device: 该模式和 dist_sync_device 类似，但是是已异步的方式进行的。

开启分布式训练

MXNet 为了用户方便的进行分布式训练提供了一个 tools/launch.py 脚本。并且支持对多种类型的集群资源进行管理,如: ssh,mpirun,yarn,sge。

如上图所示，将代码 clone 到本地后，进入 gluon 目录，使用 image_classification.py 和 CIFAR10 数据集来对 VGG11 模型进行分布式训练。

虽然 MXNet 实现了多机多卡的分布式训练，但是在资源隔离，资源调度，资源限制以及大规模训练时数据共享都是不能满足需求的，所以接下来介绍下 MXNet 基于 Kubeflow 的大规模分布式训练。

MXNet 结合 kubeflow 进行分布式训练

在将 MXNet 结合 kubeflow 进行分布式训练之前，首先需要安装 kubeflow 环境之前已经介绍了，这里就不在详细说明了。当 kubeflow 基础环境部署完成之后，需要针对 MXNet 安装 mxnet-operator。

安装 mxnet-operator

安装完成后，执行以下命令，检验 mxnet 是否安装成功：

kubectl get crd

输出如下内容代表 mxnet-operator 安装成功:

基于 kubeflow 测试 MXNet 分布式训练

1 准备测试的训练镜像

示例代码: https://github.com/deepinsight/insightface

Dockerfile 文件内容:

2 创建分布式网络文件系统数据卷(cephfs)

由于我们是基于 kubernetes 的 pv 和 pvc 的方式使用数据卷，所有集群中需要事先安装好 storage-class install，这样当用户创建 pvc 时，会通过 storage-class 自动的创建 pv。

当创建好 pv 之后，用户可以将该数据卷 mount 到自己的开发机上，并将需要训练的数据集移到该数据卷。用于之后创建训练 worker pod 的时候，挂载到 worker 容器中，供训练模型使用。

3 创建 mxnet 分布式训练任务

4 创建训练任务

kubectl create -f insightface-train.yaml

5 查看任务运行情况

6 查看训练日志的信息

登录到具体的 node 计算节点通过 docker logs 命令查看训练的日志：

docker logs -f fc3d73161b27

总结

虽然已经完成了 mxnet 结合 kubeflow 实现大规模的分布式训练，但是除了功能上的基本跑通，还存在很多因素影响分布式训练的性能，如: GPU 服务器的网络带宽，普通的我们使用的以太网因为通信延迟的原因，会大大影响多机扩展性。InfiniBand（IB）网络和 RoCE 网络因为支持 RDMA,大大降低了通信延迟，相比之下，20G 的以太网格延迟会大大提升。当然，对于现有的普通以太网络，也可以通过别的方法优化通信带宽的减少，比方说梯度压缩。通过梯度压缩，减少通信带宽消耗的同时，保证收敛速度和精度不会有明显下降。MXNet 官方提供了梯度压缩算法，按照官方数据，最佳的时候可以达到两倍的训练速度提升，同时收敛速度和精度的下降不会超过百分之一。还有如果使用分布式网络文件系统进行数据集的存储，如果解决吞吐量和网络延迟的问题。以及本地磁盘是否是 SSD，还是在训练时是否需要对大文件的数据集进行 record.io 文件格式的处理及训练前数据集的切分等等问题，都需要更进一步的处理。

本文转载自公众号 360 云计算（ID：hulktalk）。

原文链接：

https://mp.weixin.qq.com/s/r06G7CGUxrVQ2QDJml7FxA