【AI设计模式】06-分布式并行模式（Distribution Strategy）

作者：王磊
更多精彩分享，欢迎访问和关注：https://www.zhihu.com/people/wldandan

在前一篇文章分享了《检查点模式（CheckPoints）》，本篇将带大家了解分布式并行模式。互联网应用通常面临的最大挑战是流量，而传统的单体架构无法应对流量的变化，因此需要采用分布式的微服务架构进行改造。同理，在AI领域，当前最大的挑战也是数据量和模型参数规模不断的增长，单卡的资源要么无法支撑训练，要么效率较低。因此，也需要采用分布式的方式，将负载分摊到不同的硬件设备上来应对大数据/参数规模化，从而提升模型训练效率。这种通过分布式方式解决模型训练效率的模式被称为分布式并行模式。

模式定义

分布式并行模式就是将训练的过程扩展到多个worker，并借助缓存、硬件加速、并行等技术手段，加速模型训练效率，通常有数据并行和模型并行以及混合并行的方式。

问题

AI领域的模型参数和训练的数据规模正呈现越来越大的趋势。从2018年Google发布Bert模型开始，参数规模从几亿到百亿、千亿（如华为盘古大模型）甚至万亿规模（Google Switch Transformer）。同时，训练数据的规模也越来越大，以GPT系列为例：

1. GPT-1是上亿规模的参数量，数据集使用了1万本书的BookCorpus，25亿单词量；
2. GPT-2参数量达到了15亿规模，其中数据来自于互联网，使用了800万在Reddit被链接过的网页数据，清洗后越40GB（WebText)；
3. GPT-3参数规模首次突破百亿，数据集上将语料规模扩大到570GB的CC数据集(4千亿词)+WebText2(190亿词)+BookCorpus(670亿词)+维基百科(30亿词)。

引自： 2022，大模型还能走多远

对于如此大规模的模型及训练数据，使用单卡的方式完全无法完成训练。以GPT-3模型训练为例 ，使用 8 张 V100 显卡，训练时长预计要 36 年， 512 张 V100显卡 ，训练时间接近 7 个月，而1024 张A100的训练时长可以减少到 1 个月。时间越长，意味着成本越高，大模型的训练可能是普通人难以负担的。因此，需要分布式并行的方式来增强算力、加速数据处理和模型训练。

解决方案

在模式定义中提到分布式并行可以分为数据并行、模型并行、混合并行三种方式。其中数据并行是将训练的样本数据分成不同的批量，在不同的硬件设备上运行相同的模型进行训练，而后聚合梯度，更新参数，进行下一轮的训练。数据并行根据通信的方式又可以分为同步和异步（参数服务器）方式。混合并行指数据并行和模型并行结合的方式。

数据并行

如下图，假设我们使用8张GPU卡或者昇腾的NPU卡来训练图片分类的模型，训练的批量为160，那么每张卡上面分到的批量数据（min-batch）为20，每张卡基于样本数据完成训练。因为各张卡上处理的数据样本不同，所以获得的梯度会有些差别。因此，需要通过对梯度进行聚合（求和、均值）等计算来保持和单卡训练相同的结果，最后再更新参数。

数据并行依赖于集合通信的操作，上面的例子中使用到了Broadcast和AllReduce通信原语，其中Broadcast将数据分发到不同的卡，而AllReduce操作则完成不同卡上的梯度聚合操作。

除了BroadCast和AllReduce算子外，还有AllGather（将每张卡的输入Tensor在第0维度上进行拼接，最终每张卡输出是相同的数值）、ReduceScatter（将每张卡的输入先进行求和，然后在第0维度按卡数切分，将数据分发到对应的卡上）、AlltoAll（将输入数据在特定的维度切分成特定的块数，并按顺序发送给其它rank，同时从其它rank接收输入，按顺序在特定的维度拼接数据）等算子，方便多种数据并行的需求。

上面提到的都是同步的方式完成数据并行的训练，即不同分片上的梯度计算完成后会汇总一起更新模型参数，然后再启动下一轮的训练，而异步的方式是模型的权重和参数异步更新，就是单卡上的worker不需要等待其它worker，梯度聚合更新模型参数后即可开始下一轮训练。异步方式的典型应用是参数服务器架构，即有一个参数服务器管理模型权重，并根据worker训练的梯度，持续更新模型，并推送新模型，触发worker进行下一轮训练。

相比同步的数据并行模式，异步并行的吞吐会高一些，因为不用等待训练比较慢的worker，但如果过程中有worker失败，部分数据可能没有训练，导致有时候无法准确的知道处理了多少个epoch数据。配合检查点模式，通过定期保存或者利用断点续训的功能，避免这样的问题。参数服务器的具体实现也可以使用同步的方式。

模型并行

对于GPT-3上千亿、万亿参数的模型，数据并行无法解决模型过大，单卡无法承载的问题。因此，需要通过层间并行（模型以层为单位切分到多个硬件设备）或者层内并行（每层的模型参数切分到多个硬件设备）的模型并行方式来解决。

如上图，层间模型并行，每卡执行的网络模型存在差异，如P0卡上执行embedding层，p1，p2，p3分别计算后面的全连接层，而层内模型并行不会改变网络结构，而是将每层的模型参数切分到不同的设备上实现并行，如右图中把每层的矩阵计算切分到p0-p3卡中。

案例

AI框架普遍都提供了分布式并行的能力，以MindSpore框架为例，它支持数据并行、模型并行和混合并行三种分布式并行模式，同时在模型并行中，根据自动化程度又支持自动，半自动和手动的方式。

在MindSpore中可以通过context模块来设置不同的并行策略：

from mindspore.communication import init, get_rank, get_group_size
from mindspore import context
init()
device_num = get_group_size()
rank = get_rank()
print("rank_id is {}, device_num is {}".format(rank, device_num))
context.reset_auto_parallel_context()
# 下述的并行配置用户只需要配置其中一种模式
# 数据并行模式
context.set_auto_parallel_context(parallel_mode=context.ParallelMode.DATA_PARALLEL)
# 半自动并行模式
# context.set_auto_parallel_context(parallel_mode=context.ParallelMode.SEMI_AUTO_PARALLEL)
# 全并行模式
# context.set_auto_parallel_context(parallel_mode=context.ParallelMode.AUTO_PARALLEL)
# 手动并行模式
# context.set_auto_parallel_context(parallel_mode=context.ParallelMode.HYBRID_PARALLEL)

完成策略配置后，就可以根据不同的场景来使用MindSpore提供的各种分布式并行模式。

数据并行

如下为一个简单网络的数据并行代码样例。不难看出，开发者如需在数据并行时对网络有特殊的修改，只需要通过context设置ParallelMode.DATA_PARALLEL 为数据并行模式即可。

import numpy as np
from mindspore import Tensor, context, Model, Parameter
from mindspore.communication import init
from mindspore import ops, nn

class DataParallelNet(nn.Cell):
    def __init__(self):
        super(DataParallelNet, self).__init__()
        # 初始化权重
        weight_init = np.random.rand(512, 128).astype(np.float32)
        self.weight = Parameter(Tensor(weight_init))
        self.fc = ops.MatMul()
        self.reduce = ops.ReduceSum()

    def construct(self, x):
        x = self.fc(x, self.weight)
        x = self.reduce(x, -1)
        return x

init()
# 设置并行模式为数据并行，其他方式一致
context.set_auto_parallel_context(parallel_mode=context.ParallelMode.DATA_PARALLEL)
net = DataParallelNet()
model = Model(net)
model.train(*args, **kwargs)

在MindSpore 1.7版本中，引入了数据并行自动加速的功能，它可以在训练过程中，主动监测训练瓶颈是否在数据侧，如果在数据侧，则根据资源（内存，CPU）状态，自动调节数据处理算子的并行度以及算子队列长度（影响内存占用）来加速数据并行处理过程。启用自动调优只需要引入dataset模块，使用1行代码启用该功能：

import mindspore.dataset as ds
ds.config.set_enable_autotune(True)

如下为Resnet-50引入自动调优后的效果，单epoch耗时从72s降低到了17s：

[WARNING] [auto_tune.cc:297 Analyse] Op (MapOp(ID:3)) is slow, input connector utilization=0.975806, output connector utilization=0.298387, diff= 0.677419 > 0.35 threshold.
[WARNING] [auto_tune.cc:253 RequestNumWorkerChange] Added request to change "num_parallel_workers" of Operator: MapOp(ID:3)From old value: [2] to new value: [4].
[WARNING] [auto_tune.cc:309 Analyse] Op (BatchOp(ID:2)) getting low average worker cpu utilization 1.64516% < 35% threshold.
[WARNING] [auto_tune.cc:263 RequestConnectorCapacityChange] Added request to change "prefetch_size" of Operator: BatchOp(ID:2)From old value: [1] to new value: [5].
epoch: 1 step: 1875, loss is 1.1544309
epoch time: 72110.166 ms, per step time: 38.459 ms
epoch: 2 step: 1875, loss is 0.64530635
epoch time: 24519.360 ms, per step time: 13.077 ms
epoch: 3 step: 1875, loss is 0.9806979
epoch time: 17116.234 ms, per step time: 9.129 ms

数据并行的另一种实践是参数服务器，MindSpore基于自身通信原语，实现了同步的参数服务器。参数服务器通常包含三个组件，Server、Worker和Scheduler，作用分别是：

• Server：保存模型的权重和反向计算的梯度值，并使用优化器通过Worker上传的梯度值对模型进行更新。
• Worker：执行网络的正反向计算，反向计算的梯度值通过Push接口上传至Server中，通过Pull接口把Server更新好的模型下载到Worker本地。
• Scheduler：用于建立Server和Worker的通信关系。

下面介绍如何使用MindSpore在昇腾910上通过参数服务器的方式完成LeNet的分布式训练，其基本流程如下：

1. 准备训练脚本：直接使用MindSpore模型代码仓的LeNet训练代码，使用MNIST数据集.

2. 设置参数服务器训练模式的参数：通过如下代码，先启动参数服务器模式，然后初始化网络，最后控制训练权重通过参数服务器更新。

context.set_ps_context(enable_ps=True)
network = LeNet5(cfg.num_classes)
network.set_param_ps()

3. 配置三个组件的拉起的脚本。

Scheduler组件脚本（scheduler.sh）：

#!/bin/bash
export MS_SERVER_NUM=1
export MS_WORKER_NUM=1
export MS_SCHED_HOST=XXX.XXX.XXX.XXX
export MS_SCHED_PORT=XXXX
export MS_ROLE=MS_PSERVER
python train.py --device_target=Ascend --data_path=path/to/dataset

Server组件脚本（server.sh）：

Server组件脚本（server.sh）：
#!/bin/bash
export MS_SERVER_NUM=1
export MS_WORKER_NUM=1
export MS_SCHED_HOST=XXX.XXX.XXX.XXX
export MS_SCHED_PORT=XXXX
export MS_ROLE=MS_PSERVER
python train.py --device_target=Ascend --data_path=path/to/dataset

Worker组件脚本（worker.sh）：

#!/bin/bash
export MS_SERVER_NUM=1
export MS_WORKER_NUM=1
export MS_SCHED_HOST=XXX.XXX.XXX.XXX
export MS_SCHED_PORT=XXXX
export MS_ROLE=MS_WORKER
python train.py --device_target=Ascend --data_path=path/to/dataset

4. 启动训练并查看结果：分别执行sh Scheduler.sh > scheduler.log 2>&1 &、sh Server.sh > server.log 2>&1 &以及sh Worker.sh > worker.log 2>&1 &启动三个组件，然后在scheduler.log可以看到Server和Worker的通信日志：

The server node id:b5d8a47c-46d7-49a5-aecf-d29d7f8b6124,node ip: 10.90.53.118,node port:46737 assign rank id:0
The worker node id:55e86d4b-d717-4930-b414-ebd80082f541 assign rank id:1
Start the scheduler node is successful！

在worker.log可以看到训练结果，说明参数服务器方式运行LeNet训练成功：

epoch: 1 step: 1, loss is 2.302287
epoch: 1 step: 2, loss is 2.304071
epoch: 1 step: 3, loss is 2.308778
epoch: 1 step: 4, loss is 2.301943
...

模型并行

在【AI工程】05-基于MindSpore的Resnet-50模型分布式训练实践 一文中，我们已经详细介绍了基于MindSpore，使用模型自动并行策略完成Resnet-50模型的分布式训练的案例，这里不再赘述。

总结

按照当前模型参数和训练数据发展的趋势，分布式并行将成为开发者在加速模型训练时唯一的选择。分布式训练时，可以从硬件和实践的角度进一步优化，如使用专用硬件，NPU（如昇腾），TPU（Google）。分布式并行模式的选择也会影响到训练的速度，此外，训练的数据批量的大小也需要合理选择，如果min-batch批量数据大，可以减少训练迭代次数，但会影响梯度下降收敛的速度以及最终的进度。在实际的模型训练过程中，需要选择合适的硬件、并行模式和超参来加速模型训练的过程。

参考资料

[1] 什么是大模型？超大模型？Foundation Model?

[2] MindSpore分布式自动并行训练

[3] 机器学习设计模式

[4] 分布式通信原语

[5] 分布式并行训练案例

说明：严禁转载本文内容，否则视为侵权。