一种基于强化学习的神经网络自适应分布式并行训练方法技术

本发明专利技术公开了一种基于强化学习的神经网络自适应分布式并行训练方法，为大规模复杂神经网络的切分和调度提供最优化解决方案。本发明专利技术首先分析神经网络模型结构和计算属性对执行性能的影响，在此基础上提取包括计算代价、通信代价和内存使用率等性能因子，构建能够全面体现分布式训练性能的多维性能评价模型，提高并行策略综合性能；其次利用前馈网络根据算子属性特征实现算子自适应分组，确定并行度，减小搜索空间的同时实现了端到端的策略搜索；最后基于重要性采样，采用近端策略梯度迭代优化强化学习模型，搜索最优切分和调度策略，扩展策略网络离线学习能力，提高算法平稳性、收敛速率以及策略搜索性能。敛速率以及策略搜索性能。敛速率以及策略搜索性能。

【技术实现步骤摘要】
一种基于强化学习的神经网络自适应分布式并行训练方法


[0001]本专利技术涉及一种基于强化学习的神经网络自适应分布式并行训练方法，为大规模复杂神经网络提供最优的模型并行训练方案。

技术介绍

[0002]近年来，受益于AI算法、硬件计算能力和数据集发展，深度神经网络技术在自然语言处理、计算机视觉和搜索推荐等领域应用广泛。由于这些领域不断迭代研发规模更大，结构更复杂的神经网络，因此“摩尔定律”已难以匹配计算需求，单设备已无法支持大规模深度网络训练。因此，研究切分神经网络计算图，并将切分后的网络调度至含有多个CPU和GPU集群，实现模型并行训练已成为解决大规模神经网络训练的通用方法。
[0003]业界通常将神经网络计算图横向按层划分、纵向跨层划分或随机划分并调度至不同设备执行，但这些方法严重依赖专家经验，切分方式不合理，集群利用率低而且通常具有较大通信开销。
[0004]基于“图搜索”的启发式方法是指导模型并行训练方法之一，F.Pellegrini等人提出Scotch静态图切分方法，通过计算各节点成本和流过每条边的数据量，指导静态图切分，平衡工作负载以最小化通信成本，但该方法无法直接适用于动态环境。Jia等人提出FlexFlow框架，使用有向图描述DNN所有计算和状态，从算子、样本、参数以及特征四个维度的并行空间中自动搜索有效策略；Zhang等人针对数据同步并行问题提出AutoSync，随机搜索同步并行空间，获得最佳同步并行策略；Peng等人提出Optimus框架针对PS框架建立绩效模型，准确预估机器学习训练速度，通过构建在线资源性能模型来最小化作业训练时间。这些方法虽然在特定场景具有较好效果，但受限于网络类型(如FlexFlow只适用于DNN)、场景(如AutoSync只适用于数据同步并行，且只针对参数服务器架构)限制，难以获得普适推广。
[0005]另一类方法是通过机器学习指导模型切分和资源调度，Kim等人研发的Parallex通过学习线性模型调整变量大小，以实现自适应张量划分；Peter等人使用贝叶斯优化确定资源调度的信用度大小；Uber团队研发Horovod利用贝叶斯对可调节部分进行优化。但是上述方法并行维度较为单一，主体调优模型简单，而且这些方法聚焦于分布式运行策略运行时间的优化，并不能保证策略除运行时间外的其他性能。
[0006]因此，本专利技术研究上述缺陷，设计并实现一种基于强化学习的神经网络自适应分布式并行训练方法。

技术实现思路

[0007]本专利技术设计并实现一种基于强化学习的神经网络自适应分布式并行训练方法。通过构建能够全面体现分布式训练性能的多维性能评价模型，提高并行策略综合性能；根据算子属性特征实现算子自适应分组，确定并行度，从而达到减小搜索空间，实现自适应训练的目的；采用近端策略梯度方法迭代优化强化学习模型，扩展策略网络离线学习能力，提高算法平稳性、收敛速率以及策略搜索性能。
[0008]一种基于强化学习的神经网络自适应分布式并行训练方法，步骤如下：
[0009]步骤1：构建多维性能评价模型R(π
g
,π
s
)，衡量分布式并行策略综合性能。首先分析影响神经网络执行性能的因素，包括神经网络模型结构、计算属性和集群拓扑等；其次提取计算代价E
i
、通信代价C
i
和内存使用率M
i
等性能因子，最后通过以下两个步骤构建能够全面体现分布式训练性能的多维性能评价模型：
[0010]1)通过计算代价E
i
和通信代价C
i
的复杂线性关系f预测获得分布式并行策略单次迭代运行时间r
t
＝f(E
i
,C
i
)；
[0011]2)建立多维性能评价模型R＝
‑
log(αr
t
+βr
c
+γr
m
)；其中r
t
，r
c
和r
m
分别表示单次迭代运行时间、峰值通信代价和峰值内存占用率，α，β和γ分别表示权重超参数，权重超参数是基于专家经验的方式，通过手动调优获得最佳权重组合。
[0012]步骤2：构建分布式并行策略搜索空间，确定分组和并行度，并生成完整模型并行策略。
[0013]首先，根据神经网络模型结构特性和集群拓扑，抽象计算图和集群设备资源组构建分布式并行策略搜索空间。计算图中，顶点O表示神经网络算子，E表示连接顶点O的有向边。
[0014]其次，提取模型及算子属性，生成算子特征向量。基于计算图提取所有算子的相关属性，包括算子类型、输入输出以及相邻节点索引，基于上述属性对算子进行特征编码，生成算子特征向量。
[0015]再次，构建切分策略网络对算子分组，确定并行度，生成切分策略π
g
，继而再根据算子分组情况生成组特征编码。切分策略网络采用含有SoftMax的前馈神经网络，将算子特征向量和集群设备资源组作为切分策略网络的输入，实现算子的粗粒度分组，确定并行度。其中，同一组内算子不适合并行执行，组间算子适合并行执行。完成分组后选取组内算子特征向量生成组特征编码序列。
[0016]最后，构建调度策略网络确定不同分组在设备上的调度策略π
s
。调度策略网络是含有注意力机制和长短时记忆的Seq2Seq双循环神经网络，采用编码器和解码器分别处理输入和输出不定长序列。调度策略网络将切分策略网络所得到的组特征编码作为输入序列，用于生成设备序列，生成的设备序列和输入的组特征编码序列呈一一对应的顺序关系，即第一分组中的所有算子将被调度至输出设备序列的第一个设备上执行，以此类推。这样的对应关系即为调度策略π
s
。切分策略π
g
和调度策略π
s
共同构成完整的模型并行策略(π
g
,π
s
)。如图1左侧所示，强化学习智能体Agent包含上述切分策略网络和调度策略网络双层策略网络。
[0017]步骤3：基于重要性采样，采用近端策略梯度迭代优化强化学习模型，搜索最优切分和调度策略。
[0018]首先，基于重要性采样的核心思想是分离采样和优化模型对象，完全拷贝一份切分策略网络和调度策略网络副本副本负责不断采样切分调度策略，负责模型参数的不断迭代更新，同时利用KL散度捕捉两模型间差异，当两模
型间差异过大时，同步一次副本参数
[0019]其次，将步骤2中得到的完整模型并行策略输入执行模拟器(图1右侧执行模拟器所示)，多维度采集策略执行性能，并通过计算代价E
i
和通信代价C
i
的复杂线性关系f预测获得分布式并行策略单次运行时间r
t
＝f(E
i
,C
i
)，通信开销，内存占用和计算负载在内的多种性能指标。在此基础上，以多维性能评价模型R(π
...

【技术保护点】

【技术特征摘要】
1.一种基于强化学习的神经网络自适应分布式并行训练方法，其特征在于，所述方法包括以下步骤：(1)分析神经网络模型结构、计算属性等对神经网络执行性能的影响，提取包括计算代价、通信代价和内存使用率的性能因子，构建能够全面体现分布式训练性能的多维性能评价模型；所述计算代价指参与运算的张量精度与设备计算密度比值，通信代价指通信量大小与通信带宽比值，内存使用率指当前设备模型参数物理存储大小与设备总内存比值；且所述性能因子还可扩展显存和吞吐量；(2)根据神经网络模型结构特性和集群拓扑，构建分布式并行策略搜索空间；提取模型算子属性，包括算子类型、输入输出以及相邻节点索引，并根据算子属性采用前馈网络对模型所有算子进行自适应分组来确定深度网络模型并行度，其中，同一算子组中的算子不适合并行执行，算子组之间适合并行执行；(3)基于重要性采样，采集强化学习训练样本，通过近端策略梯度方法迭代优化强化学习模型，寻找最优分布式并行策略。2.根据权利要求1所述的一种基于强化学习的神经网络自适应分布式并行训练方法，其特征在于：步骤(1)中所述多维性能评价模型通过以下两个步骤建立：(1)通过计算代价E
i
和通信代价C
i
的复杂线性关系f预测获得分布式并行策略单步运行时间r
t
＝f(E
i
,C
i
)；(2)建立多维性能评价模型R＝
‑
log...

【专利技术属性】
技术研发人员：吴吉央，曾艳，张纪林，袁俊峰，任永坚，周丽，
申请(专利权)人：杭州电子科技大学，
类型：发明
国别省市：