MxNet系列——how_to——multi_devices

2023-05-16

博客新址: http://blog.xuezhisd.top
邮箱：xuezhisd@126.com

在多个CPU/GPUs上以数据并行方式运行MXNet

MXNet 支持在多个CPUs和GPUs上进行训练。其中，这些CPUs和GPUs可能位于不同的物理机上。

数据并行 vs 模型并行

MXNet模式使用数据并行的方式将工作负载划分到多个设备上。假如有 n 个设备，每一个设备都将获得完整的模型，并使用 1/n 的数据进行训练。结果（比如，梯度）和更新后的模型在不同设备之间进行通信。

MXNet也支持模型并行。这时，每个设备上维护模型的一部分。当模型非常大以至于单个设备无法存储时，模型并行非常有用。有一个关于如何在多层LSTM上使用模型并行方法的教程 a tutorial 。本教程主要关注 数据并行。

单机上的多个GPUs

划分工作负载

MXNet 默认会将一个数据批均匀地划分到每一个GPU上。加入批大小是 b，GPUs的数目是 k，因此在每一次迭代过程中，每一个GPU将会在 b/k 个样本上执行前向和后向传播。在更新模型之前，会将所有GPUs上的梯度都加起来。

如何使用

为了使用 GPUs，需要编译支持GPU的MXNet。
比如，在配置文件 config.mk 中设置 USE_CUDA=1，然后再 make。
(更多选项，请查看 MXNet installation guide)。

如果一台主机安装了一个或多个GPU，每个GPU都会有一个编号（编号从0开始计数）。如果想使用某个特定的显卡，既可以指定在代码中指定环境（context） ctx；也可以在命令行中传递参数 --gpus。例如，如果想在Python中使用GPU 0和GPU 2，可以使用下面的代码创建网络模型。

import mxnet as mx
#列表包含多个GPU
model = mx.model.FeedForward(ctx=[mx.gpu(0), mx.gpu(2)], ...)

如果程序接受参数 --gpus ，比如 example/image-classification，那么可以尝试下面的代码。

python train_mnist.py --gpus 0,2 ...

高级用法

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

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

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

有两个常用的类型，

local: 所有的梯度都复制到CPU内存，并且权重也在那里进行更新。
device: 梯度聚集和权重更新都在GPU上进行。它也会尝试使用GPU的P2P通信（它可以加速通信）。但该选项会使用更多的GPU内存。

如果有大量的GPU（比如，>=4），我们建议使用 device，以获得更佳性能。

使用多个机器进行分布式训练

我们可以通过简单地修改 KVStore ，实现在多个机器上运行MXNet。

dist_sync。它的行为和 local相似。但一个主要区别是 batch-size 此处表示每个机器上的批大小。机器数量为 n，批大小为 b，dist_sync的行为等价于批大小为 n/b local类型。
dist_device_sync。它和 dist_sync之间的区别，与 device 和 local之间的区别相同。即梯度聚集和权重更新都在GPU上进行。
dist_async 。它执行一步更新。一旦从任何机器上接收到梯度，立即更新权重。这种更新是具有原子性（不可分），即同一权重在相同的时间不可能出现两个更新。然而却无法保证顺序。

如何启动一个作业

如果使用分布式训练，需要先设置 USE_DIST_KVSTORE=1，再进行编译。
(更多选项，请查看 MXNet installation guide)。

启动一个分布式作业和单机运行稍有不同。MXNet 提供了 tools/launch.py ，它利用 ssh, mpi, sge, 或 yarn 启动作业。

假定我们位于目录 mxnet/example/image-classification 下。希望在数据集mnist上使用脚本 train_mnist.py 训练LeNet。

单机运行时的命令：

python train_mnist.py --network lenet

现在有两个可以使用SSH进行通信的机器，我们希望在这两个机器上训练网络。首先，将这两个主机的IPs或主机名保存在 hosts 文件中。比如，

$ cat hosts
172.30.0.172
172.30.0.171

接下来，两个机器需要都可分访问mxnet文件夹，比如网络文件系统。

运行以下命令，将在多机上启动MXNet。

../../tools/launch.py -n 2 --launcher ssh -H hosts python train_mnist.py --network lenet --kv-store dist_sync

注意：除了包括单机运行参数，此处还需：

使用 launch.py 来提交作业。
--launcher 提供启动器。如果所有的机器可以互相ssh，那么使用 ssh；如果 mpirun 可用，那么使用 mpi；如果Sun公司的网格引擎（Sun Grid Engine）可用，使用 sge；如果Apache Yarn可用，使用 yarn。
-n 计算节点（worker nodes）的数量
-H 指定host文件（ ssh 和 mpi 使用）
--kv-store 使用 dist_sync 或 dist_async

同步目录

现在开始考虑 MXNet 文件夹不可访问的情况。我们可以先MXNet库复制到当前目录。

cp -r ../../python/mxnet .
cp -r ../../lib/libmxnet.so mxnet

然后，使用 launch.py 和参数 --sync-dst-dir 将当前目录同步到所有机器的 /tmp/mxnet 目录下。

../../tools/launch.py -n 2 -H hosts --sync-dst-dir /tmp/mxnet \
   python train_mnist.py --network lenet --kv-store dist_sync

使用一个特定的网络接口

MXNet 一般会选择第一个可用的网络接口。但对于有多个接口的机器，我们通过环境变量 DMLC_INTERFACE 可以指定使用哪个网络接口进行数据通信。例如，下面的代码中使用了网络接口 eth0。

# 前面的命令设置网络接口（网卡）
# 后面的命令和上面相同
export DMLC_INTERFACE=eth0; ../../tools/launch.py ...

调试连接

通过设置 PS_VERBOSE=1，可以查看调试日志。比如，

export PS_VERBOSE=1; ../../tools/launch.py ...