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

Abstract

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

Description

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

技术领域

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

背景技术

近年来，受益于AI算法、硬件计算能力和数据集发展，深度神经网络技术在自然语言处理、计算机视觉和搜索推荐等领域应用广泛。由于这些领域不断迭代研发规模更大，结构更复杂的神经网络，因此“摩尔定律”已难以匹配计算需求，单设备已无法支持大规模深度网络训练。因此，研究切分神经网络计算图，并将切分后的网络调度至含有多个CPU和GPU集群，实现模型并行训练已成为解决大规模神经网络训练的通用方法。

业界通常将神经网络计算图横向按层划分、纵向跨层划分或随机划分并调度至不同设备执行，但这些方法严重依赖专家经验，切分方式不合理，集群利用率低而且通常具有较大通信开销。

基于“图搜索”的启发式方法是指导模型并行训练方法之一，F.Pellegrini等人提出Scotch静态图切分方法，通过计算各节点成本和流过每条边的数据量，指导静态图切分，平衡工作负载以最小化通信成本，但该方法无法直接适用于动态环境。Jia等人提出FlexFlow框架，使用有向图描述DNN所有计算和状态，从算子、样本、参数以及特征四个维度的并行空间中自动搜索有效策略；Zhang等人针对数据同步并行问题提出AutoSync，随机搜索同步并行空间，获得最佳同步并行策略；Peng等人提出Optimus框架针对PS框架建立绩效模型，准确预估机器学习训练速度，通过构建在线资源性能模型来最小化作业训练时间。这些方法虽然在特定场景具有较好效果，但受限于网络类型(如FlexFlow只适用于DNN)、场景(如AutoSync只适用于数据同步并行，且只针对参数服务器架构)限制，难以获得普适推广。

另一类方法是通过机器学习指导模型切分和资源调度，Kim等人研发的Parallex通过学习线性模型调整变量大小，以实现自适应张量划分；Peter等人使用贝叶斯优化确定资源调度的信用度大小；Uber团队研发Horovod利用贝叶斯对可调节部分进行优化。但是上述方法并行维度较为单一，主体调优模型简单，而且这些方法聚焦于分布式运行策略运行时间的优化，并不能保证策略除运行时间外的其他性能。

因此，本发明研究上述缺陷，设计并实现一种基于强化学习的神经网络自适应分布式并行训练方法。

发明内容

本发明设计并实现一种基于强化学习的神经网络自适应分布式并行训练方法。通过构建能够全面体现分布式训练性能的多维性能评价模型，提高并行策略综合性能；根据算子属性特征实现算子自适应分组，确定并行度，从而达到减小搜索空间，实现自适应训练的目的；采用近端策略梯度方法迭代优化强化学习模型，扩展策略网络离线学习能力，提高算法平稳性、收敛速率以及策略搜索性能。

一种基于强化学习的神经网络自适应分布式并行训练方法，步骤如下：

步骤1：构建多维性能评价模型R(π_g,π_s)，衡量分布式并行策略综合性能。首先分析影响神经网络执行性能的因素，包括神经网络模型结构、计算属性和集群拓扑等；其次提取计算代价E_i、通信代价C_i和内存使用率M_i等性能因子，最后通过以下两个步骤构建能够全面体现分布式训练性能的多维性能评价模型：

1)通过计算代价E_i和通信代价C_i的复杂线性关系f预测获得分布式并行策略单次迭代运行时间r_t＝f(E_i,C_i)；

2)建立多维性能评价模型R＝-log(αr_t+βr_c+γr_m)；其中r_t，r_c和r_m分别表示单次迭代运行时间、峰值通信代价和峰值内存占用率，α，β和γ分别表示权重超参数，权重超参数是基于专家经验的方式，通过手动调优获得最佳权重组合。

步骤2：构建分布式并行策略搜索空间，确定分组和并行度，并生成完整模型并行策略。

首先，根据神经网络模型结构特性和集群拓扑，抽象计算图

和集群设备资源组

构建分布式并行策略搜索空间。计算图

中，顶点O表示神经网络算子，E表示连接顶点O的有向边。

其次，提取模型及算子属性，生成算子特征向量。基于计算图

提取所有算子的相关属性，包括算子类型、输入输出以及相邻节点索引，基于上述属性对算子进行特征编码，生成算子特征向量。

再次，构建切分策略网络

对算子分组，确定并行度，生成切分策略π_g，继而再根据算子分组情况生成组特征编码。切分策略网络采用含有SoftMax的前馈神经网络，将算子特征向量和集群设备资源组

作为切分策略网络

的输入，实现算子的粗粒度分组，确定并行度。其中，同一组内算子不适合并行执行，组间算子适合并行执行。完成分组后选取组内算子特征向量生成组特征编码序列。

最后，构建调度策略网络

确定不同分组在设备上的调度策略π_s。调度策略网络是含有注意力机制和长短时记忆的Seq2Seq双循环神经网络，采用编码器和解码器分别处理输入和输出不定长序列。调度策略网络

将切分策略网络

所得到的组特征编码作为输入序列，用于生成设备序列，生成的设备序列和输入的组特征编码序列呈一一对应的顺序关系，即第一分组中的所有算子将被调度至输出设备序列的第一个设备上执行，以此类推。这样的对应关系即为调度策略π_s。切分策略π_g和调度策略π_s共同构成完整的模型并行策略(π_g,π_s)。如图1左侧所示，强化学习智能体Agent包含上述切分策略网络

和调度策略网络

双层策略网络。

步骤3：基于重要性采样，采用近端策略梯度迭代优化强化学习模型，搜索最优切分和调度策略。

首先，基于重要性采样的核心思想是分离采样和优化模型对象，完全拷贝一份切分策略网络和调度策略网络

副本

副本

负责不断采样切分调度策略，

负责模型参数的不断迭代更新，同时利用KL散度捕捉两模型间差异，当两模型间差异过大时，同步一次副本参数

其次，将步骤2中得到的完整模型并行策略输入执行模拟器(图1右侧执行模拟器所示)，多维度采集策略执行性能，并通过计算代价E_i和通信代价C_i的复杂线性关系f预测获得分布式并行策略单次运行时间r_t＝f(E_i,C_i)，通信开销，内存占用和计算负载在内的多种性能指标。在此基础上，以多维性能评价模型R(π_g,π_s)为指导计算强化学习奖励值R。

最后，根据奖励指导切分策略网络和调度策略网络

迭代优化。根据奖励值R，采用近端策略梯度方法，不断迭代协同优化切分策略网络和调度策略网络

参数，激励

策略网络在下一轮训练生成综合性能更优的模型并行策略。最后生成最优的并行策略，在真实分布式环境中执行，如图1右侧所示。

本发明具有的有益效果是：多维性能评价模型的建立，提高搜索到并行策略的综合性能；切分和调度双层策略网络的建立，实现模型并行策略完全端到端的搜索，无需人为介入；使用执行模拟器接管真实分布式执行环境，加快策略执行和性能评测速度，提高采样效率；使用基于重要性采样的分布式近端策略优化算法改进原有策略梯度算法，扩展策略网络离线学习能力，提高算法平稳性和收敛速率。于此同时，本发明中的多维性能评价模型和执行模拟器均可推广至其它自适应搜索场景。

附图说明

图1是基于强化学习搜索最优并行策略示意图；

图2是切分网络和策略网络联合架构图。

具体实施方式

下面将结合附图和具体实施步骤对本发明做出进一步说明：

如图1所示，一种基于强化学习的神经网络自适应分布式并行训练方法，包括以下步骤：

步骤1：构建多维性能评价模型R(π_g,π_s)，衡量策略综合性能。首先分析影响神经网络执行性能的因素，包括神经网络模型结构、计算属性和集群拓扑等；其次提取计算代价E_i、通信代价C_i和内存使用率M_i等性能因子，其中计算代价E_i、通信代价C_i和内存使用率M_i定义如下：

计算代价，采用参与运算的张量精度除以设备计算密度联合表示，可衡量设备训练过程中计算负载情况，记T_j为该张量第j维深度，

为张量精度计算，c_i为第i个设备计算密度，计算代价定义如下式：

通信代价，采用通信量大小除以通信带宽表示，可衡量设备训练过程的通信负载情况，记

表示张量大小，b_i表示设备间通信带宽，通信代价表示如下式：

内存使用率，内存使用率M_d采用当前设备i内存中模型参数物理存储大小和设备总内存m_i的比值，内存使用率表示如下式：

通过上述定义，本发明将从运行时间、峰值内存和峰值通信代价三个维度对调度策略性能评价进行建模。建立起基于强化学习的多维性能评价模型，并将R作为强化学习奖励值，用于不断迭代提升

切分和调度策略网络综合性能，自适应寻找最优分布式并行切分和调度策略(π_g,π_s)。多维性能评价模型建模如下：

R(π_g,π_s)＝-log(αr_t+βr_c+γr_m) (4)

其中，奖励值R越大表示搜索到的策略越好。值得注意的是，上式中的自然对数运算也支持采用开方运算替换，其中α，β和γ分别表示权重超级参数，通过手动调优获得最佳权重组合。而r_t，r_c和r_m分别表示运行时间、峰值通信负载和峰值内存使用率三个性能因子，各性能因子定义方式如下：

其中r_t表示模拟或真实环境下策略单次迭代运行时间，利用模拟器间接通过计算代价E_i和通信代价C_i间的非线性关系f(E_i,C_i)预测获得，R_c表示多个设备的峰值通信代价，通过max{C₁,C₂,...,C_k}计算获得，其中C_i对应硬件设备i的通信代价，

表示用户可允许的最大通信量，如果超过最大通信量，将做出相应的处罚。同理R_m表示多个设备的峰值内存使用率，通过max{M₁,M₂,...,M_k}计算获得，其中M_i为内存使用率，

表示用户可容忍的最大内存占用率。

步骤2：如图2所示，构建分布式并行策略搜索空间，确定分组和并行度，并生成模型并行策略。

首先，根据神经网络模型结构特性和集群拓扑，抽象计算图

和集群设备资源组

构建分布式并行策略搜索空间。抽象计算图

和集群设备资源组

其中顶点O表示神经网络算子，记为O＝{o₁,o₂,…,o_N}。O为连接顶点P的有向边，表示计算节点间通信数据依赖，记为E＝{e₁₁,e₁₂,…,e_ij}。针对集群设备资源组

将集群中CPU和GPU设备抽象为一系列资源组

现假设集群具有M个可用资源，那么可用资源组

可抽象为

其次，提取模型及算子属性，生成算子特征向量。提取计算图

中所有算子属性，包括算子类型(type)、输入输出(output shape)以及相邻节点索引(adj)，对算子属性进行特征编码，生成算子特征向量。编码方式如下：提取计算图中算子三部分属性作为算子属性特征向量，(1)算子类型(type)：例如卷积(Conv2D)、最大池化(MaxPool)或矩阵乘法(MalMul)等算子类型定义。本发明采用自然语言处理方法，收集TensorFlow API中常用的200个算子词汇，建立词汇表，将词汇作为自然语言处理任务，得到大小尺寸为20的词向量作为算子类型部分的嵌入特征；(2)算子输出张量大小(output shape)：假设当前算子为o_i，获取计算图顶点连边{e_i1,e_i2,…,e_ij}作为该算子输出张量列表。为了统一输出张量数，本发明统一输出张量数为6，并且指定输出张量的最大维度为4，如有不足则通过-1填充向量。例如，现有卷积算子输出一个形状为(2,2,1,64)的四维张量，那么它的大小为：256＝2×2×1×64，其余输出张量位置均填充-1；(3)邻接算子索引(adj)：按照拓扑排序遍历所有算子得到算子索引，将每个算子的入度和出度大小限制为6，并将该算子的父节点算子和子节点算子填充索引集合，依然采用-1填充入度和出度小于6的情况。

再次，构建切分策略网络

(如图2第一层架构所示)对算子分组，确定并行度，生成切分策略π_g，再根据算子分组情况生成组特征编码。切分策略网络

是含有大小分别为64和128两层隐藏层的前馈神经网，并且在隐藏层之间引入30％丢弃层以防止过拟合。切分策略网络

采用带有SoftMax的前馈神经网络对计算图算子进行粗粒度分组。将前文所述算子特征向量和集群设备资源组

作为切分策略网络输入，产生粗粒度分组，同时确定并行度，其中，同一组内算子不适合并行执行，组间算子适合并行执行。分组完成后，切分策略网络

将组内算子特征编码合成组特征编码输出，包括三个部分：(1)将该组内所有算子类型编码求平均值后作为该组编码第一部分；(2)将该组内所有算子输出张量大小计算平均值后作为该组编码第二部分，注意，该部分编码不存在-1值；(3)分组的独热编码，表示组与组之间的相连关系，该部分向量长度为分组数量(例如将算子切分为256组，则向量长度即为256)。如果该组中的某个算子与第i组中的某个算子相连，则该向量第i个位置设置为1，否则为0。

最后，构建调度策略网络

(如图2第二层架构所示)，确定不同分组的设备的调度策略π_s。调度网络是含有注意力机制(Attention)和长短时记忆的Seq2Seq双循环神经网络，采用编码器(Encode)和解码器(Decode)分别处理输入和输出不定长序列。(1)将步骤2中获得的组特征编码作为调度策略网络

的输入，并采用编码器处理输入分组的不定长序列，每次读取一个分组g_i的组特征编码向量并产生k个隐藏状态(Hidden)，其中k是等于分组数量的超参数。(2)采用解码器处理输出设备的不定长序列，每次预测得到一个调度设备d_j，得到的设备序列和输入分组序列具有一一对应的顺序关系，即第一组中的所有算子将被调度至解码器输出的第一个设备上，以此类推。值得注意的是，每个设备都具有其自己的可训练特征编码，而且前一个设备的特征编码将输入到下一步解码器中。调度策略网络π_s采用注意力机制注意编码器状态，在训练过程中解码器将在第t步从SoftMax层采样获得设备d_t。为了使得调度策略网络激活函数u_t更加平坦，本发明引入SoftMax温度和分对数裁剪，使用温度T和tanh常数C来表达调度策略网络激活函数u_t，即通过如下方式采样：

d_t～softmax(Ctanh(u_t/T)) (6)

解码器输出的设备序列即为对应输入分组的调度策略π_s，切分策略π_g和调度策略π_s共同形成完整模型的并行策略(π_g,π_s)。

步骤3：基于重要性采样，采用近端策略梯度迭代优化强化学习模型，搜索最优切分和调度策略。

首先，基于重要性采样的核心思想是分离采样和优化模型对象，完全拷贝一份切分策略网络和调度策略网络

副本

副本

负责不断采样切分调度策略，生成切分调度策略样本，用于后续强化学习算法迭代优化；而

只负责模型参数的不断迭代更新。近端策略梯度方法利用KL散度捕捉两模型间差异，如果模型差异过大，将会同步一次副本参数

反之则继续使用陈旧模型副本完成采样。

其次，将步骤2中得到的完整模型并行策略输入执行模拟器，采集策略各方面执行性能，并通过计算代价E_i和通信代价C_i的复杂线性关系f预测获得分布式并行策略单步运行时间r_t＝f(E_i,C_i)。在此基础上，以多维性能评价模型R(π_g,π_s)为指导计算强化学习奖励值R。

其中执行模拟器设计方法如下：策略的首次执行将在真实分布式环境中执行，用于采集给定模型在所有设备上的运行性能。之后由执行模拟器接管真实分布式环境，将计算代价建模为E_i，将通信代价建模为C_i，本发明假设通信代价与设备间通信数据流大小呈线性关系。每个设备采用双线程维护以下两个先进先出度列。

表示设备d上算子运行队列，记录即将运行算子的先后顺序；

表示将从设备d上传输至其他设备的张量队列，记录即将通信至其他设备上的张量集合。

执行模拟器在设计过程中遵循如下三个原则：(1)只有当操作o所有依赖的父节点均执行完毕，且依赖的张量已全部获取到后才可以在设备d上开始运行；(2)计算与通信过程可相互重叠，即具备流水线执行机制；(3)同一设备上算子以串行方式执行。本发明中模拟器采用以下三个触发机制生成时间流水线：

触发机制一：触发执行。此处以当前算子o_i正在设备d_k上执行为例，完整触发执行过程将完成以下流程：1)完成当前算子o_i在设备d_k上执行过程，获取依赖该节点的所有子节点o_j。2)获取o_j节点所在设备d_l，如果d_j≠d_i，那么将操作o_i输出张量t_i入队列

如果d_j＝d_i，检查o_j是否符合运行原则，如果符合，将其入队

3)如果

空闲，则将触发设备d_k空闲状态；如果不空闲，则选择下一个队列

中算子执行触发。

触发机制二：触发通信。此处以张量t_i从设备d_k传输到d_j为例，完整触发通信过程将执行以下流程：1)完成当前张量t_i从设备d_k传输到d_j的完整过程。2)检查依赖于t_i张量的算子o_k是否符合运行原则，如果符合，将其入队

3)判断

是否为空，如果为空将设备d_k的通信触发状态为空闲。如果不是，选择

下一传输张量触发通信流程。

触发机制三：触发状态。首先判断

和

是否为空，如果为空则触发执行空闲或通信空闲状态，如果不为空，则出队实现下一个算子触发执行或张量触发传输。

最后，完成执行模拟，得到奖励值R后，根据奖励值R，指导切分策略网络和调度策略网络

迭代优化，并输出最优并行策略在真实分布式环境中执行。在优化算法的选择上，本发明采用近端策略优化方法协同优化切分策略网络和调度策略网络，本质上我们希望最大化优化如下代价函数

令θ_g和θ_s分别为切分策略网络和调度策略网络参数，

为模型副本参数，将θ_g和θ_s联合表示为θ，

和

联合表示为θ^k。p(g,s；θ)表示切分和调度策略网络

策略的联合概率分布，q(g,s；θ^k)表示模型副本

采样获得的联合概率分布，R为奖励值，b为平均移动基线。

对上式的切分和策略网络参数进行拆分也可以写作如下形式：

其中

其中p(g；θ_g)表示切分策略网络

生成策略的概率分布，p(s|g；η_s)表示调度策略网络

生成调度策略的条件概率分布；

则表示切分策略网络副本

采样切分策略的概率分布，

表示调度策略网络副本

采样调度策略的条件概率分布；η为KL散度权重超参数。b为指数移动平均奖励基线，基线的引入可有效提高采样效率。如果令N为代表周期的超参数，那么指数移动平均奖励基线EMA_N(b_n)的递推计算式如下：

代价函数

的求导过程如下：

本发明采用Adam完成梯度下降过程，近端策略优化方法引入KL散度以避免θ和θ^k差距过大导致采样效率低下的问题。

Claims (5)

1.一种基于强化学习的神经网络自适应分布式并行训练方法，其特征在于，所述方法包括以下步骤：

(1)分析神经网络模型结构、计算属性等对神经网络执行性能的影响，提取包括计算代价、通信代价和内存使用率的性能因子，构建能够全面体现分布式训练性能的多维性能评价模型；

所述计算代价指参与运算的张量精度与设备计算密度比值，通信代价指通信量大小与通信带宽比值，内存使用率指当前设备模型参数物理存储大小与设备总内存比值；且所述性能因子还可扩展显存和吞吐量；

(2)根据神经网络模型结构特性和集群拓扑，构建分布式并行策略搜索空间；提取模型算子属性，包括算子类型、输入输出以及相邻节点索引，并根据算子属性采用前馈网络对模型所有算子进行自适应分组来确定深度网络模型并行度，其中，同一算子组中的算子不适合并行执行，算子组之间适合并行执行；

(3)基于重要性采样，采集强化学习训练样本，通过近端策略梯度方法迭代优化强化学习模型，寻找最优分布式并行策略。

2.根据权利要求1所述的一种基于强化学习的神经网络自适应分布式并行训练方法，其特征在于：步骤(1)中所述多维性能评价模型通过以下两个步骤建立：

(1)通过计算代价E_i和通信代价C_i的复杂线性关系f预测获得分布式并行策略单步运行时间r_t＝f(E_i,C_i)；

(2)建立多维性能评价模型R＝-log(αr_t+βr_c+γr_m)；其中；r_t，r_c和r_m分别表示单步运行时间、峰值通信代价和峰值内存占用率，α，β和γ分别表示权重超参数，基于经验调参，获得最佳权重组合。

3.根据权利要求1所述的一种基于强化学习的神经网络自适应分布式并行训练方法，其特征在于：步骤(2)中所述确定深度网络模型并行度是指：对算子属性进行特征编码，基于多维性能评价模型计算所得奖励值，利用强化学习指导含有SoftMax的前馈神经网络完成算子分组，将互相依赖且通信代价较高的算子分为一组，从而确定深度网络模型并行度。

4.根据权利要求1所述的一种基于强化学习的神经网络自适应分布式并行训练方法，其特征在于：步骤(3)中所述近端策略梯度方法是指：分离采样和优化模型对象，利用KL散度捕捉两模型间差异，实现强化学习算法的离线迭代更新，引入指数移动平均奖励基线b降低收敛方差。

5.根据权利要求2所述的一种基于强化学习的神经网络自适应分布式并行训练方法，其特征在于：所述复杂线性关系f通过如下方式构造：每个设备采用双线程维护执行

和通信

两个先进先出队列构造策略执行过程，遵循如下原则：(1)只有当算子o所有依赖父节点均执行完毕，且获取到所有依赖张量后才可以在当前设备执行；(2)计算与通信过程可相互重叠；(3)同一设备算子以串行方式执行。

Images