用华为MindSpore进行分布式训练

技术背景

分布式和并行计算，在计算机领域是非常重要的概念。对于一些行外人来说，总觉得这是一些很简单的工作，但是如果我们纵观计算机的硬件发展史，从CPU到GPU，再到TPU和华为的昇腾（NPU），乃至当下的热点量子计算机（QPU），其实就是一个分布式与并行计算的发展史。从简单的数据并行，到算法并行，到图的并行，最后是量子叠加所带来的物理并行。因此能否做好分布式与并行的技术，很大程度上决定了一个工具的性能上限，本文我们一起来研究一下MindSpore分布式训练的方法。

环境部署

在前面的博客中，我们探讨过用Docker和Singularity容器等方案来安装MindSpore的CPU版本和GPU版本，感兴趣的读者可以翻一翻这些历史博客。而这篇文章中，我们将默认已经在本地安装好一个MindSpore的GPU环境，以此为前提进行探讨在单机多GPU卡的环境下去使用MindSpore的分布式功能。比较完整的介绍可以参考这个官方地址，里面包含了完整的安装部署和使用的介绍。这里我们仅针对本地Ubuntu的环境介绍基本安装和使用方法。

安装openmpi

这里一共需要安装2个软件，我们都是采取了源码安装的方法，首先到这个MindSpore给出的下载链接中下载对应版本的源码：

按照版本配套我们下载这个4.0.3的tar.gz的压缩包：

下载完成后可以大致的按照这里面的操作指导进行源码安装：

安装成功的话运行mpirun --version可以看到版本号。有这里面有两个需要提醒的点是：1. 解压缩不一定用这里面给出的指令，可以用自己的；2. 如果按照这个指引的prefix，后面需要在LD_LIBRARY_PATH这个环境变量中加入/usr/local/lib/这个路径，否则会有一些so库无法被识别。

安装NCCL

MindSpore在GPU上的分布式通信是采用了NCCL这个工具，同样的我们先去这里面给出的链接找源码包下载安装：

这里推荐使用Local的安装方案：

找到适配自己的系统版本的软件包之后，就可以按照下述的指引一步一步安装：

其中如果在apt-key这一步执行一次失败的话，可以多执行两次看看。

环境测试

安装成功后，openmpi和NCCL都安装成功后，可以用如下的初始化示例测试一下环境的部署情况：

# test-init.py

from mindspore import context
from mindspore.communication.management import init

if __name__ == "__main__":
    context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
    init("nccl")

如果执行的结果如下所示，没有任何的报错信息，则说明环境部署成功：

dechin@ubuntu2004:~/projects/mindspore/test$ mindspore test-init.py 
dechin@ubuntu2004:~/projects/mindspore/test$

MindSpore分布式训练

首先可能会有人好奇，上一个章节中最后用的mindspore是个啥指令？这其实是我自己在本地配置的一个mindspore-gpu版本的快捷命令：

dechin@ubuntu2004:~/projects/$ cat ~/.bashrc | grep mindspore
alias mindspore='singularity exec --nv /home/dechin/tools/singularity/mindspore-gpu_1.2.0.sif python'

如果觉得还算方便的话，读者也可以按照自己的喜好去alias一个快捷命令。

接下来我们要看一个实现的案例，这个代码来自于前面的这一篇博客，只是我们在这个代码基础之上，加入了上一个章节中的初始化代码init()，完整代码如下所示：

# test_nonlinear.py

from mindspore import context
from mindspore.communication.management import init
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
init()
import numpy as np
from mindspore import dataset as ds
from mindspore import nn, Tensor, Model
import time
from mindspore.train.callback import Callback, LossMonitor

def get_data(num, a=2.0, b=3.0, c=5.0):
    for _ in range(num):
        x = np.random.uniform(-1.0, 1.0)
        y = np.random.uniform(-1.0, 1.0)
        noise = np.random.normal(0, 0.03)
        z = a * x ** 2 + b * y ** 3 + c + noise
        yield np.array([[x**2], [y**3]],dtype=np.float32).reshape(1,2), np.array([z]).astype(np.float32)

def create_dataset(num_data, batch_size=16, repeat_size=1):
    input_data = ds.GeneratorDataset(list(get_data(num_data)), column_names=['xy','z'])
    input_data = input_data.batch(batch_size)
    input_data = input_data.repeat(repeat_size)
    return input_data

data_number = 160
batch_number = 16
repeat_number = 200

ds_train = create_dataset(data_number, batch_size=batch_number, repeat_size=repeat_number)
dict_datasets = next(ds_train.create_dict_iterator())

class LinearNet(nn.Cell):
    def __init__(self):
        super(LinearNet, self).__init__()
        self.fc = nn.Dense(2, 1, 0.02, 0.02)

    def construct(self, x):
        x = self.fc(x)
        return x

net = LinearNet()
model_params = net.trainable_params()
print ('Param Shape is: {}'.format(len(model_params)))
for net_param in net.trainable_params():
    print(net_param, net_param.asnumpy())
net_loss = nn.loss.MSELoss()

optim = nn.Momentum(net.trainable_params(), learning_rate=0.01, momentum=0.6)
model = Model(net, net_loss, optim)

epoch = 1
model.train(epoch, ds_train, callbacks=[LossMonitor(10)], dataset_sink_mode=False)

for net_param in net.trainable_params():
    print(net_param, net_param.asnumpy())

此时我们需要改成用mpirun去运行这个代码：

dechin@ubuntu2004:~/projects/gitlab/dechin/src/mindspore$ mpirun -n 2 singularity exec --nv /home/dechin/tools/singularity/mindspore-gpu_1.2.0.sif python test_nonlinear.py

运行过程中我们就可以看到，有两个python的任务分别跑在两个不同的GPU卡上：

这个监控方式其实还是nvidia-smi的指令，只不过为了长期监视GPU状态，我是使用了watch -n 1 nvidia-smi的指令，每隔1s的时间就刷新一次nvidia-smi的状态。但是这里需要注意的也有两点：1. 这个代码直接用mindspore指令也可以跑在单GPU卡上，但是如果要用mpirun来运行，那么我们就不能使用刚才alias的mindspore指令，而需要手动写上完整的指令，除非把新的指令再alias一个；2. 上述的代码因为只是初始化了一下，所以虽然跑在两张卡上，但是实际上训练过程并没有互相通信，是两个独立的任务。那么如果要构造一个完整的自动化的分布式训练，就需要像如下代码一样，再加入context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL)这样的一个字样，以及checkpoint中的一些配置，新的完整代码如下所示：

# test_nonlinear.py

from mindspore import context
from mindspore.communication.management import init
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")
init()

import numpy as np
from mindspore import dataset as ds
from mindspore import nn, Tensor, Model
import time
from mindspore.train.callback import Callback, LossMonitor, ModelCheckpoint, CheckpointConfig
from mindspore.context import ParallelMode
import mindspore as ms
ms.common.set_seed(0)

start_time = time.time()
def get_data(num, a=2.0, b=3.0, c=5.0):
    for _ in range(num):
        x = np.random.uniform(-1.0, 1.0)
        y = np.random.uniform(-1.0, 1.0)
        noise = np.random.normal(0, 0.03)
        z = a * x ** 2 + b * y ** 3 + c + noise
        yield np.array([[x**2], [y**3]],dtype=np.float32).reshape(1,2), np.array([z]).astype(np.float32)

def create_dataset(num_data, batch_size=16, repeat_size=1):
    input_data = ds.GeneratorDataset(list(get_data(num_data)), column_names=['xy','z'])
    input_data = input_data.batch(batch_size)
    input_data = input_data.repeat(repeat_size)
    return input_data

data_number = 160
batch_number = 16
repeat_number = 20

context.set_auto_parallel_context(parallel_mode=ParallelMode.DATA_PARALLEL)
ds_train = create_dataset(data_number, batch_size=batch_number, repeat_size=repeat_number)
dict_datasets = next(ds_train.create_dict_iterator())

class LinearNet(nn.Cell):
    def __init__(self):
        super(LinearNet, self).__init__()
        self.fc = nn.Dense(2, 1, 0.02, 0.02)

    def construct(self, x):
        x = self.fc(x)
        return x

net = LinearNet()
model_params = net.trainable_params()
print ('Param Shape is: {}'.format(len(model_params)))
for net_param in net.trainable_params():
    print(net_param, net_param.asnumpy())
net_loss = nn.loss.MSELoss()

optim = nn.Momentum(net.trainable_params(), learning_rate=0.01, momentum=0.6)
ckpt_config = CheckpointConfig()
ckpt_callback = ModelCheckpoint(prefix='data_parallel', config=ckpt_config)

model = Model(net, net_loss, optim)

epoch = 10
model.train(epoch, ds_train, callbacks=[ckpt_callback], dataset_sink_mode=True)

for net_param in net.trainable_params():
    print(net_param, net_param.asnumpy())

print ('The total time cost is: {}s'.format(time.time() - start_time))

这个代码的运行结果如下所示：

dechin@ubuntu2004:~/projects/gitlab/dechin/src/mindspore$ mpirun -n 2 singularity exec --nv /home/dechin/tools/singularity/mindspore-gpu_1.2.0.sif python test_nonlinear.py 
Param Shape is: 2
Parameter (name=fc.weight, shape=(1, 2), dtype=Float32, requires_grad=True) [[0.02 0.02]]
Parameter (name=fc.bias, shape=(1,), dtype=Float32, requires_grad=True) [0.02]
Param Shape is: 2
Parameter (name=fc.weight, shape=(1, 2), dtype=Float32, requires_grad=True) [[0.02 0.02]]
Parameter (name=fc.bias, shape=(1,), dtype=Float32, requires_grad=True) [0.02]
[WARNING] ME(2528801:139843698521024,MainProcess):2021-06-10-09:45:37.603.010 [mindspore/train/callback/_checkpoint.py:428] OSError, failed to remove the older ckpt file /home/dechin/projects/gitlab/dechin/src/mindspore/data_parallel-1_200.ckpt.
[WARNING] ME(2528799:139709496722368,MainProcess):2021-06-10-09:45:37.713.232 [mindspore/train/callback/_checkpoint.py:428] OSError, failed to remove the older ckpt file /home/dechin/projects/gitlab/dechin/src/mindspore/data_parallel-2_200.ckpt.
[WARNING] ME(2528799:139709496722368,MainProcess):2021-06-10-09:45:37.824.271 [mindspore/train/callback/_checkpoint.py:428] OSError, failed to remove the older ckpt file /home/dechin/projects/gitlab/dechin/src/mindspore/data_parallel-3_200.ckpt.
[WARNING] ME(2528799:139709496722368,MainProcess):2021-06-10-09:45:37.943.749 [mindspore/train/callback/_checkpoint.py:428] OSError, failed to remove the older ckpt file /home/dechin/projects/gitlab/dechin/src/mindspore/data_parallel-4_200.ckpt.
[WARNING] ME(2528801:139843698521024,MainProcess):2021-06-10-09:45:38.433.85 [mindspore/train/callback/_checkpoint.py:428] OSError, failed to remove the older ckpt file /home/dechin/projects/gitlab/dechin/src/mindspore/data_parallel-5_200.ckpt.
Parameter (name=fc.weight, shape=(1, 2), dtype=Float32, requires_grad=True) [[0.12186428 0.21167319]]
Parameter (name=fc.bias, shape=(1,), dtype=Float32, requires_grad=True) [5.561276]
The total time cost is: 8.412446737289429s
Parameter (name=fc.weight, shape=(1, 2), dtype=Float32, requires_grad=True) [[0.12186428 0.21167319]]
Parameter (name=fc.bias, shape=(1,), dtype=Float32, requires_grad=True) [5.561276]
The total time cost is: 8.439369916915894s

虽然运行成功了，但是有一点需要注意的是，分布式的这个案例耗时为8.44s，而单卡的训练耗时为8.15s，分布式的训练速度甚至比单卡的要慢，我们在总结里面对这个现象进行一个解释。

总结概要

这篇文章我们主要探讨如何去部署一个基于MindSpore框架的分布式训练环境，在MindSpore环境已经配置好的情况下，我们只需要安装好openmpi和nccl这两个工具就可以实现分布式的训练，在文中我们已经给出了相应的示例。虽然分布式与并行技术的目的是为了提升性能，但不是说对所有的场景都能够起到加速的作用，比如文章中的案例就没有加速的效果。这其实是因为我们的用例场景太简单了，纵观整个训练过程，GPU的使用率不到10%，在这种情况下还要考虑上通信的开销，自然是没有放在同一个卡上去训练来得快。这也给我们一个启发，考虑使用分布式和并行计算的技术时，一定也要先评估好问题本身是否适用于并行化的处理，否则是达不到预期的加速的目的的。