机器学习模型混合并行策略自动搜索的方法及系统

本发明涉及机器学习训练并行，具体地，涉及一种机器学习模型混合并行策略自动搜索的方法及系统。

背景技术：

1、机器学习通常在多个领域都有着重要应用，在使用机器学习过程中，通常需要对机器学习模型进行训练，而当前的机器学习模型的规模越来越大，例如meta提出的llama3模型便含有7b到70b左右的参数量，这通常需要上百g的gpu内存。为实现快速的机器学习训练过程，需将机器学习的模型训练过程并行化，对于不同的模型，通常有不同的适合其并行训练的方法，为找到不同模型的最优并行策略，就需要去自动搜索机器学习训练并行的策略。

2、现有的机器学习模型通常包含数据并行，张量并行以及流水线并行。数据并行是指对训练的数据样本进行划分，分别传输到不同的节点上进行训练的过程。张量并行是指对算子涉及到的张量进行多维度划分，分配给不同计算设备，由多个计算设备并行计算得到该算子的计算结果。而流水线并行是指将机器学习模型的神经网络中不同的算子划分给流水线不同的阶段，将机器学习的训练过程变成流水线形式，提高训练过程的吞吐量。

3、现有的自动搜索机器学习训练并行策略的工作主要有两种：未考虑流水线并行的并行策略自动搜索方法以及考虑了流水线并行的并行策略自动搜索方法。未考虑流水线并行的并行策略自动搜索方法首先对计算图中的每一个点，即算子，统计不同并行策略的计算开销以及算子内通信开销，并对计算图中的每一条边，即张量的流转过程，统计算子间通信开销，然后调用一种基于全局动态规划的方法对计算图进行每一个算子的并行策略选择。具体来说，尝试将计算图中的每一个点或者每一条边删去，将该点或者该边的计算以及通信开销存进邻接点中；将该计算图不断变小，在变小后的计算图上遍历每一个节点，选择该节点的并行策略，使得最后的计算图的整体开销最小，并在选择完每一个结点的并行策略后，将原先删去的点和边按照删去的顺序逆序地添加回来，并在添加的过程中为每一个添加回来的节点选择其适合的并行策略，在恢复成原先图后便完成了整个算法的搜索过程。这种不考虑流水线并行的策略自动搜索方式拥有较快的搜索速度，并且能找到全局最优的并行策略，但是由于其不考虑流水线并行，缺少了流水线并行的维度，难以达到更好的优化效果，并且由于其不考虑不同并行策略所需要占用的gpu内存大小，其依据开销找到的并行策略很有可能无法在相应的设备上执行。

4、现有的考虑了流水线并行的并行策略自动搜索方法将机器学习模型划分成不同流水线阶段，并在每一个阶段内进行数据并行和张量并行的策略搜索。该方法相关的工作例如alpa，将流水线阶段的开销模型为一个批次经过整个流水线所需要的耗时，即在流水线并行策略的自动搜索过程中要尽可能地减少一个批次经过整个流水线所需要的耗时。该方法将数据并行和张量并行的策略自动搜索过程抽象成一个整数线性规划的问题，并调用整数线性规划的求解器进行求解。这种考虑流水线并行的策略自动搜索方式结合了三种不同的并行策略，拓宽了并行训练的优化空间，能够找到更多的优化机会。但是这种方法消耗了更多的搜索时间，并且不能找到全局最优的并行训练策略。

5、现如今的并行训练策略搜索方法虽然能够找到一种混合了三种并行策略的并行训练策略，但是所需要花的时间比较久，不能够高效地找到这种策略。

6、综上所述，针对上述现有技术的问题，研究一种机器学习模型混合并行策略自动搜索的方法及系统，能够快速地找出质量较优的并行策略，成为当下亟待解决的关键任务。

技术实现思路

1、针对现有技术中的缺陷，本发明的目的是提供一种机器学习模型混合并行策略自动搜索的方法及系统。

2、根据本发明提供的一种机器学习模型混合并行策略自动搜索的方法，包括如下步骤：

3、算子间搜索步骤：对输入的计算图和设备拓扑情况，进行流水线阶段的划分，得到划分后的流水线阶段；

4、算子间开销生成步骤：将所述划分后的流水线阶段中最耗时阶段的时间作为性能指标，为所述算子间搜索步骤提供算子间性能开销模型；

5、算子内搜索步骤：对输入的计算图子图进行数据并行和张量并行结合的并行策略搜索；使用基于全局动态规划方法的缩点消边算法进行策略搜索，得到每个算子在计算图子图中的最优并行策略及对应的已选择并行策略的节点信息，其中每个节点相应的策略对应的计算开销、通信开销由算子内开销生成步骤给出；

6、算子内开销生成步骤：对所述已选择并行策略的节点信息，计算得到算子内的计算开销、通信开销以及算子间的通信开销。

7、优选地，所述算子间搜索步骤中，对计算图的算子进行拓扑排序，采用动态规划算法将算子划分到不同的流水线阶段中。

8、优选地，所述动态规划算法依赖于算子间开销生成步骤，即调用算子间性能开销模型，将不同的算子划分到不同的流水线阶段中，采取以减少流水线中最耗时流水线阶段的方式作为动态规划算法的开销模型进行搜索，最终得到划分后的流水线阶段。

9、优选地，所述算子间开销生成步骤中，对所述划分后的流水线阶段进行性能开销的预估，通过调用算子内搜索步骤得到一个流水线阶段内最优的数据并行和张量并行结合的并行策略开销。

10、优选地，所述算子间开销生成步骤中，将流水线中最耗时的阶段时间作为主要考量因素，应用搜索算法时应尽量降低该阶段耗时。

11、优选地，所述算子内搜索步骤中，并行策略搜索时，对算子内进行数据并行和算子模型并行的策略搜索，采用动态规划的方式，将算子组成的计算图进行点和边的删去，同时将点和边的开销信息保存到邻接点和邻接边中，从而在更小的图上进行策略的搜索。

12、优选地，所述算子内搜索步骤中，在进行点和边的删去操作时，对任意一个计算图子图应用消点和消边策略，其中消点包含了三个点，四个点以及大于四个点的情况，消边包含了两个点间有多条边的情况。

13、优选地，所述算子内开销生成步骤中，所述计算开销由实际测试得出，所述通信开销由张量大小除以带宽得出。

14、本发明还提供了一种机器学习模型混合并行策略自动搜索的系统，包括：

15、算子间并行策略搜索模块：根据计算图以及集群设备情况进行流水线阶段的划分；

16、算子间性能开销模型：根据划分的流水线情况以及划分后流水线阶段使用的设备情况依据算子内并行策略搜索模块进行算子间策略开销的计算；

17、算子内并行策略搜索模块：根据同一个阶段内的计算图子图选择每个算子的最优策略；

18、算子内性能开销模型：根据同一个阶段内的计算图子图每个算子选择的策略计算相应的开销。

19、优选地，所述算子内并行策略搜索模块包括：

20、缩点消边规则注册器单元：注册缩点消边过程中可用的规则；

21、缩点消边规则选择器单元：依据目前的计算图拓扑结构选择相应的缩点消边规则；

22、缩点消边算法执行器单元：根据选择的缩点消边规则执行相应的动态规划算法。

23、与现有技术相比，本发明具有如下的有益效果：

24、1、本发明提出并实现了基于新的流水线规划开销模型的流水线阶段搜索算法。该算法通过对流水线规划开销模型的优化，精准地评估和规划流水线阶段，显著减少了流水线阶段搜索的时延。与传统算法相比，它能更高效地根据计算图和设备情况进行流水线阶段的划分。

25、2、本发明在已有的缩点消边算法基础上，提出了新的无关方向性的缩点消边规则，通过启发式方法提出了针对大图的缩点消变规则，扩展了缩点消边能力，提升了搜索效率。对于大规模的机器学习模型计算图，这些新规则能够更有效地简化图结构，减少搜索空间的复杂度。在实际应用中，当处理包含大量算子和复杂连接关系的计算图时，新规则能够快速找到更优的并行策略搜索路径，使得搜索过程更加高效和准确，进一步提升了整个机器学习模型混合并行策略自动搜索的性能和效果。