一种面向人工智能大模型高效训练的并行策略搜索方法

摘要

本发明公开了一种面向人工智能大模型高效训练的并行策略搜索方法，步骤如下：输入大模型；获取系统执行性能信息和大模型的执行性能信息；利用代价模型估计大模型在当前系统上的执行时间开销和存储开销；以流水线并行训练的每次迭代所花费的时间为优化目标，建立用于自动搜索大模型的并行训练策略的混合整数二次规划数学模型并求解；输出最优的大模型的并行训练策略。本发明支持流水线并行、数据并行、张量并行和优化器并行，策略空间大，搜索时间短，可应用于大模型的单机多卡并行训练和多机集群的分布式训练，提升大模型的训练效率。

说明书

技术领域

本发明涉及人工智能技术，特别是一种面向人工智能大模型(以下简称大模型)高效训练的并行策略搜索方法。

背景技术

近年来，人工智能大模型的参数量快速上涨，并在计算机视觉、自然语言处理等各个领域的大部分任务上取得了最佳的效果。与之相对的，计算硬件发展的速度非常缓慢。因此不得不利用单机多卡或集群进行并行训练。

当前，并行训练大模型主要有流水线并行、数据并行、张量并行和优化器并行等并行模式。流水线并行主要对大模型的不同层之间进行切分，然后将若干相邻的层放在同一个流水线阶段上。在计算时，流水线并行把一个小批量划分为多个微批量，然后把这些微批量组织为流水线，从而提升整体的训练效率。数据并行、张量并行和优化器并行则在大模型的各个层的内部进行并行计算，提升了层内计算的效率。其中数据并行需要在不同的节点上拷贝完整的大模型和优化器状态，拥有较小的通信量和较大的显存开销；张量并行和优化器并行分别切分大模型的张量和优化器状态，虽然显存开销较小，但通信量相对数据并行而言更大。需要注意的是，这些并行计算策略在不同的模型和不同的计算环境下各有优劣。因此，针对特定的模型和给定的计算环境，如何组合这些并行策略以最大化训练效率成为了大模型的并行训练中的一个重要研究方向。

但是由于并行策略众多，人工组合和评测的代价较高。同时，对于不同的模型和计算环境，选用合适的并行策略组合依赖专家级别的知识，难以应用于现实场景。因此，有研究人员提出了自动搜索大模型的并行训练策略方法。但是截至目前，已有方法要么只关注上述四种并行策略中的一种或其中几种，要么把上述四种并行策略分层考虑，逐层求解。这些方法的策略空间受限，容易得到局部最优的并行训练策略，而非全局最优的并行训练策略。又因为训练一个参数规模较大的大模型往往需要花费数周乃至数月的时间，消耗大量能源，所以如何自动搜索到一个全局最优的并行训练策略来最大化地加速大模型训练成为了一个重要问题。

发明内容

发明目的：本发明的目的是提供一种面向人工智能大模型高效训练的并行策略搜索方法，从而提供一种基于混合整数二次规划的面向大模型高效训练的并行策略搜索方法，以解决现有技术所搜索的并行策略的局部最优性问题，从而进一步提高大模型的训练效率，降低能源消耗。

技术方案：本发明所述的一种面向人工智能大模型高效训练的并行策略搜索方法，包括以下步骤：

步骤S1、输入大模型；

步骤S2、获取系统执行性能信息和大模型的执行性能信息；

设定All-Reduce的通信量a

设定P2P的通信量a

设定All-Reduce的通信量a

统计大模型G(V,E)每一层u(u∈V)的前向执行时间fp

统计除模型和优化器占用的显存以外，大模型G(V,E)每一层的激活显存

步骤S3中所述的执行时间开销包括大模型每一层的所有候选策略的计算及通信时间之和，相邻层的所有候选策略对的同一流水线阶段的通信时间和跨流水线阶段的P2P通信时间的一种或多种。

步骤S3、利用代价模型估计大模型在当前系统上的执行时间开销和存储开销；

步骤S3中所述的执行时间开销具体包括以下步骤：

利用代价模型估计大模型G(V,E)的层u(u∈V)的所有候选数据并行、张量并行和优化器并行策略的计算及通信时间之和A

利用代价模型估计大模型G(V,E)的边＜u,v>(∈E)所对应的层u和层v的所有候选数据并行、张量并行和优化器并行策略对之间的同一流水线阶段的通信时间R

利用代价模型估计大模型G(V,E)的边(＜u,v＞∈E)所对应的层u和层v的所有候选数据并行、张量并行和优化器并行策略对之间的跨流水线阶段的P2P通信时间R′

步骤S3中所述的存储开销包括大模型每一层的所有候选策略的显存开销。

步骤S3中所述的存储开销具体包括以下步骤：利用代价模型估计大模型G(V,E)的层u(u∈V)的所有候选数据并行、张量并行和优化器并行策略的显存开销M

步骤S4、以流水线并行训练的每次迭代所花费的时间为优化目标，建立用于自动搜索大模型的并行训练策略的混合整数二次规划数学模型并求解；

步骤S4.1、枚举所有合法的流水线阶段数量deg和合法的微批量数量c，选取最小的每次迭代所花费时间TPI，及其对应的流水线并行策略、数据并行策略、张量并行策略和优化器并行策略的一种或多种；

步骤S4.2、基于一个非流水线并行策略(deg＝1)，建立并求解以数据并行、张量并行和优化器并行每次迭代所花费时间TPI为优化目标的混合整数二次规划数学模型，搜索的内容为数据并行策略、张量并行策略和优化器并行策略的一种或多种；

步骤S4.3、基于一个给定的流水线并行策略(deg>1)，建立并求解以流水线并行、数据并行、张量并行和优化器并行训练的每次迭代所花费时间TPI为优化目标的混合整数二次规划数学模型，搜索的内容为流水线并行策略、数据并行策略、张量并行策略和优化器并行策略的一种或多种。

步骤S5、输出最优的大模型的并行训练策略。

基于混合整数二次规划数学模型的各求解结果，输出大模型在当前系统下的最优流水线并行策略、数据并行策略、张量并行策略和优化器并行策略的一种或多种组合。

一种计算机存储介质，其上存储有计算机程序，该计算机程序被处理器执行时实现上述的一种面向人工智能大模型高效训练的并行策略搜索方法。

一种计算机设备，包括储存器、处理器及存储在存储器上并可在处理器上运行的计算机程序，所述处理器执行所述计算机程序时实现上述的一种面向人工智能大模型高效训练的并行策略搜索方法。

有益效果：与现有技术相比，本发明具有如下优点：

1、本发明支持流水线并行、数据并行、张量并行和优化器并行等四类并行训练策略，其策略空间大；

2、本发明通过利用混合整数二次规划，其搜索时间短；

3、本发明可应用于大模型的单机多卡并行训练和多机集群的分布式并行训练，有助于提升大模型的训练效率，从而减少训练时间，降低能源消耗。

附图说明

图1为本发明所述方法的步骤流程图；

图2为本发明实施例的系统环境示意图；

图3为本发明实施例的步骤S4.1的流程图。

具体实施方式

下面结合附图对本发明的技术方案作进一步说明。

本发明提供的面向人工智能大模型高效训练的并行策略搜索方法，支持各类大模型的流水线并行、数据并行、张量并行和优化器并行训练，可应用于大模型的单机多卡并行训练和多机集群的分布式并行训练。

本实施例使用两台服务器Server0和Server1，每台服务器有四张GPU卡，依次编号为GPU0，GPU1，GPU2和GPU3。两台服务器互为同构，即包括但不限于：硬件环境完全相同，软件环境完全相同和网络环境及协议完全相同等。于本实施例中，定义GPU的全局编号globalID＝serverID*gpuPerServer+localID，其中serverID为服务器编号，gpuPerServer为每台服务器的GPU数量，localID为当前GPU的本地编号。例如Server1的GPU1的全局编号为1*4+1＝5。

为帮助理解本实施例的两台服务器之间的关系，图2展示了本实施例的两台服务器的拓扑结构。两台服务器的GPU0和GPU1通过PCIe总线与CPU0相连，GPU2和GPU3通过PCIe总线与CPU1相连。两台服务器的CPU0和CPU1之间均通过QPI总线相连。两台服务器之间通过以太网相连。本实施例将应用面向人工智能大模型高效训练的并行策略搜索方法，搜索得到可以在Server0和Server1上用Adam优化器分布式训练BERT-Huge模型的并行策略。于本实施例中，本发明的方法具体工作流程如下所述：

面向人工智能大模型高效训练的并行策略搜索方法的工作流程如图1所示。

步骤S1、输入大模型。

本实施例中，所述的大模型为BERT-Huge模型。

步骤S2、获取系统执行性能信息和大模型的执行性能信息。

本实施例中，所述的获取系统执行性能信息和大模型的执行性能信息，包括：

设定All-Reduce的通信量a

设定P2P的通信量a

设定All-Reduce的通信量a

在Server0上，利用GPU0，设定批量大小为1，统计BERT-Huge模型G(V,E)每一层u(u∈V)的前向执行时间fp

在Server0上，利用GPU0，设定批量大小为1，统计除BERT-Huge模型和Adam优化器占用的显存以外，BERT-Huge模型G(V,E)每一层的激活显存

步骤S3、利用代价模型估计大模型在当前系统上的执行时间开销和存储开销。

本实施例中，所述的大模型在当前系统上的执行时间开销，包括：

利用代价模型估计BERT-Huge模型G(V,E)的层u(u∈V)的所有候选数据并行、张量并行和优化器并行策略在批量大小为bs情况下的计算及通信时间之和A

利用代价模型估计BERT-Huge模型G(V,E)的边(∈E)所对应的层u和层

利用代价模型估计BERT-Huge模型G(V,E)的边(∈E)所对应的层u和层v的所有候选数据并行、张量并行和优化器并行策略对之间在批量大小为bs情况下的跨流水线阶段的P2P通信时间R′

本实施例中，所述的大模型在当前系统上的存储开销，包括：

利用代价模型估计BERT-Huge模型G(V,E)的层u(u∈V)的所有候选数据并行、张量并行和优化器并行策略在批量大小为bs情况下的显存开销M

步骤S4、以流水线并行训练的每次迭代所花费的时间为优化目标，建立用于自动搜索大模型的并行训练策略的混合整数二次规划数学模型并求解：

本实施例中，所述的以流水线并行训练的每次迭代所花费的时间为优化目标，建立用于自动搜索大模型的并行训练策略的混合整数二次规划数学模型并求解，包括：

步骤S4.1、枚举所有合法的流水线阶段数量deg和合法的微批量数量c，选取最小的每次迭代所花费时间TPI，及其对应的流水线并行策略、数据并行策略、张量并行策略和优化器并行策略的一种或多种；

步骤S4.2、基于一个非流水线并行策略(deg＝1)，建立并求解以数据并行、张量并行和优化器并行每次迭代所花费时间TPI为优化目标的混合整数二次规划数学模型，搜索的内容为数据并行策略、张量并行策略和优化器并行策略的一种或多种；

步骤S4.3、基于一个给定的流水线并行策略(deg>1)，建立并求解以流水线并行、数据并行、张量并行和优化器并行训练的每次迭代所花费时间TPI为优化目标的混合整数二次规划数学模型，搜索的内容为流水线并行策略、数据并行策略、张量并行策略和优化器并行策略的一种或多种。

为便于理解，图3展示了本实施例中的步骤S4.1的执行流程。于本实施例中，所述的流水线为GPipe模式的同步流水线。

步骤S4.1.1、输入大模型的每一层的计算及通信时间之和A和显存开销M，相邻层的统一流水线阶段的通信时间R和不同流水线阶段的P2P通信时间R′，以及GPU数量n和小批量数量B。

步骤S4.1.2、令流水线阶段数deg＝1，最小每次迭代执行时间minTPI＝∞，及其对应的最小流水线阶段数minDeg＝-1、微批量数量minC＝-1、模型每一层的放置策略

步骤S4.1.3、判断流水线阶段数量deg是否不大于GPU数量n。如果S4.1.3判断成立，则进入外层循环(入口为步骤S4.1.4)。

步骤S4.1.4、判断流水线度数deg是否为1。如果是，则进入步骤S4.2，建立并求解以数据并行、张量并行和优化器并行每次迭代所花费时间TPI为优化目标的混合整数二次规划数学模型。求解完成后，进入步骤S4.1.5、判断优化目标TPI是否小于minTPI，如果步骤S4.1.5判断成立，程序将执行步骤S4.1.6、令minTPI＝TPI，minDeg＝deg，minC＝c，minP＝P，minS＝S。最后，程序将执行步骤S4.1.7，令deg＝deg*2，再对步骤S4.1.3中外层循环的条件进行判断。

如果步骤S4.1.4的判断结果为否，执行步骤S4.1.8、令微批量数量c为2，然后执行步骤S4.1.9、判断微批量数量c是否不超过小批量数量B。如果判断不成立，那么程序执行步骤S4.1.7、deg＝deg*2；如果判断成立，那么程序进入内层循环。在内层循环中，程序首先执行步骤S4.1.10、判断B是否不被c整除。如果判断成立，那么执行步骤S4.1.14、c＝c+1，继续判断步骤S4.1.9中内层循环的条件；如果步骤S4.1.10的判断不成立，那么执行步骤S4.1.11、微批量大小b＝B/c，然后执行步骤S4.3、建立并求解以流水线并行、数据并行、张量并行和优化器并行训练的每次迭代所花费时间TPI为优化目标的混合整数二次规划数学模型。在得到求解结果后，执行步骤S4.1.12、判断优化目标TPI是否小于minTPI，如果判断成立，那么执行步骤S4.1.13、令minTPI＝TPI，minDeg＝deg，minC＝c，minP＝P，minS＝S，然后程序将执行步骤S4.1.14、c＝c+1，继续步骤S4.1.9中内层循环条件的判断。

如果步骤S4.1.3判断不成立，进入步骤S4.1.15、输出最小的每次迭代执行时间minTPI及其对应的最小流水线阶段数minDeg、微批量数量minC、模型每一层的放置策略minP和模型每一层的并行策略minS，结束。

本实施例提供的步骤S4.2中，为帮助理解所述以数据并行、张量并行和优化器并行每次迭代所花费时间TPI为优化目标的混合整数二次规划数学模型，本实施例将例举一个混合整数二次规划数学模型：

假设目标函数最小化每次迭代所花费时间TPI可以表示为公式(1)：

min TPI＝p

在公式(1)中，TPI是每次迭代所花费时间，因此该式等价于最大化每单位时间所迭代的次数，即等价于最大化吞吐量。在公式(1)中，p

为了更好地理解，下面结合公式介绍各约束条件：

约束条件1：流水线的第1个阶段的计算时间包括BERT-Huge所有层u的计算及通信时间之和，以及BERT-Huge的所有边(∈E)所对应的层u和层v之间的同一流水线阶段的通信代价之和。约束条件1可以表达为公式(2)：

其中S

约束条件2：流水线的第1个阶段的存储开销包括BERT-Huge所有层u(u∈V)的所有候选数据并行、张量并行和优化器并行策略的显存开销之和。约束条件2可以表达为公式(3)：

本实施例中，m代表GPU0的显存容量减去加载BERT-Huge模型G(V,E)所需要的额外显存m

约束条件3：BERT-Huge所有层u(u∈V)能且仅能选择一个候选数据并行、张量并行和优化器并行策略。约束条件3可以表达为公式(4)(5)：

其中g

本实施例提供的步骤S4.3中，为帮助理解所述以流水线并行、数据并行、张量并行和优化器并行训练的每次迭代所花费时间TPI为优化目标的混合整数二次规划数学模型，本实施例将例举一个混合整数二次规划数学模型：

假设目标函数最小化每次迭代所花费时间TPI可以表示为公式(6)：

在公式(6)中，p

为了更好地理解，下面结合公式介绍各个约束条件：

约束条件1：流水线的第i个阶段的计算和通信时间，包括：(1)放置在第i个阶段的BERT-Huge层的计算及通信时间之和；(2)BERT-Huge的所有放置在第i个阶段的层u和层v之间的边(∈E)所对应的同一流水线阶段的通信代价之和。约束条件1可以表达为公式(7)：

其中P

约束条件2：流水线并行的第j个跨流水线阶段的通信时间。约束条件2可以表达为公式(8)：

约束条件3：流水线的第i个阶段的存储开销包括放置在第i个阶段的BERT-Huge层u(u∈V)的所有候选数据并行、张量并行和优化器并行策略的显存开销之和。约束条件3可以表达为公式(9)：

本实施例中，m代表GPU0的显存容量减去加载BERT-Huge模型G(V,E)所需要的额外显存m

约束条件4：每个流水线阶段必须包含BERT-Huge模型的连续部分。约束条件4可以表达为公式(10)(11)(12)：

其中Z

约束条件5：BERT-Huge所有层u(u∈V)能且只能被放置在其中一个流水线阶段上，并且所有流水线阶段上至少有一个层。约束条件5可以表达为公式(13)(14)(15)：

约束条件6：BERT-Huge所有层u(u∈V)能且仅能选择一个候选数据并行、张量并行和优化器并行策略。约束条件6可以表达为公式(4)(5)。

最后，基于以上混合整数二次规划模型的求解结果，输出最优的大模型的并行训练策略(步骤S5):

本实施例中，所述输出最优的大模型的并行训练策略，具体包括：

基于步骤S4的输出步骤S4.1.15，输出BERT-Huge模型在当前系统下的最优流水线并行策略、数据并行策略、张量并行策略和优化器并行策略的一种或多种组合。

更具体地，minTPI代表最小的每次迭代执行时间，minDeg代表当前minTPI对应的最小流水线阶段数，minC代表当前minTPI对应的微批量数量，minP代表当前minTPI对应的模型每一层的放置策略，minS代表当前minTPI对应的模型每一层的并行策略。所以最终的流水线并行策略为：设定流水线的阶段数为minDeg，流水线计算时拆分小批量为minC个微批量。模型的各层u应当放置在minP中对应的每一层放置的流水线阶段上。同时，模型每一层u的数据并行、张量并行和优化器并行策略为minS中对应每一层所采用的并行策略。此外，对于并行计算策略和系统的映射关系，系统中全局编号为l的GPU将负责计算流水线的第

本发明的方法在多个模型和系统上进行了实验。实验过程中，在编号为0的服务器上统计面向人工智能大模型高效训练的并行策略搜索方法的搜索时间。在所述方法搜索得到的并行训练策略后，在编号为0的服务器上统计在该并行策略在当前系统下的吞吐量。在BERT-Huge、T5-Large、ViT-Huge、Swin-Huge四个模型上的实验结果表明，相比已有的最佳方法，本发明的方法最快可以提升约16倍的策略搜索速度，最大可以提升约1.7倍的训练吞吐量。