张量切分方法、电子设备、存储介质和计算机程序产品与流程

本发明涉及芯片，尤其涉及一种张量切分方法、电子设备、存储介质和计算机程序产品。

背景技术：

1、在现代高性能计算和深度学习模型训练中，数据并行和模型并行是常见的策略，用于加速计算任务。现有的并行计算方案主要依赖于张量切分(tensor splitting)，这对性能优化具有重要影响。

2、目前的张量切分方案可以分为两大类：单轴切分和多轴切分。单轴切分是指根据张量的维度和形状尺寸自动选择需要切分的轴(此处的轴可以理解为需要切分的维度)，沿着张量的一个轴切分，但由于只能沿一个轴切分，当形状尺寸不均衡或较大时，无法充分利用所有计算核，导致部分计算核闲置。多轴切分指沿着多个轴同时进行切分，从而更灵活地分配计算任务，但多轴切分需要人工配置具体的切分策略，需额外高昂的人工成本且泛化性较差。

技术实现思路

1、本发明提供一种张量切分方法、电子设备、存储介质和计算机程序产品，用以解决相关技术中存在的缺陷。

2、本发明提供一种张量切分方法，包括如下步骤。

3、基于目标张量的属性以及计算核数，确定所述目标张量的多个候选切分模式；

4、基于各候选切分模式的评分，从各候选切分模式中确定目标切分模式；

5、基于所述目标切分模式，对所述目标张量进行切分。

6、根据本发明提供的一种张量切分方法，所述基于目标张量的属性以及计算核数，确定所述目标张量的多个候选切分模式，包括：

7、枚举步骤：在当前计算核数下，基于所述目标张量的属性，确定所述目标张量是否至少有一个切分模式；

8、若是，则将所述至少一个切分模式作为所述候选切分模式；若否，则减小所述当前计算核数，并返回执行所述枚举步骤。

9、根据本发明提供的一种张量切分方法，首个计算核数为最大可用计算核数。

10、根据本发明提供的一种张量切分方法，所述基于各候选切分模式的评分，从各候选切分模式中确定目标切分模式，包括：

11、若有且只有一个最高评分，则将所述最高评分对应的候选切分模式作为所述目标切分模式；若存在多个最高评分，则将所述最高评分中批处理轴的切分数量最大的候选切分模式作为所述目标切分模式。

12、根据本发明提供的一种张量切分方法，所述从各候选切分模式中确定目标切分模式，之后还包括：

13、若所述目标切分模式对应的计算核数为最大可用计算核数，则确定所述目标张量是否为权重张量，若是，则开启权重张量广播后，将所述目标切分模式绑定至所述目标张量；

14、若所述目标切分模式对应的计算核数不为最大可用计算核数，则不开启权重张量广播，并将所述目标切分模式绑定至所述目标张量。

15、根据本发明提供的一种张量切分方法，所述基于目标张量的属性以及计算核数，确定所述目标张量的多个候选切分模式，包括：

16、基于所述目标张量的属性、所述目标张量中各轴的最小切分粒度以及所述计算核数，确定所述目标张量的多个候选切分模式。

17、根据本发明提供的一种张量切分方法，各候选切分模式的评分基于各候选切分张量的形状大小确定，各候选切分张量指基于各候选切分模式对所述目标张量进行切分后的张量。

18、本发明还提供一种张量切分装置，包括如下模块。

19、确定单元，用于基于目标张量的属性以及计算核数，确定所述目标张量的多个候选切分模式；

20、评分单元，用于基于各候选切分模式的评分，从各候选切分模式中确定目标切分模式；

21、切分单元，用于基于所述目标切分模式，对所述目标张量进行切分。

22、本发明还提供一种电子设备，包括存储器、处理器及存储在存储器上并可在处理器上运行的计算机程序，所述处理器执行所述程序时实现如上述任一种所述张量切分方法。

23、本发明还提供一种非暂态计算机可读存储介质，其上存储有计算机程序，该计算机程序被处理器执行时实现如上述任一种所述张量切分方法。

24、本发明还提供一种计算机程序产品，包括计算机程序，所述计算机程序被处理器执行时实现如上述任一种所述张量切分方法。

25、本发明提供的张量切分方法、电子设备、存储介质和计算机程序产品，在确定候选切分模式时，考虑了计算核数，进而实现在充分利用所有计算核的基础上确定候选切分模式，并从各候选切分模式中确定最优的目标切分模式，最大化利用计算核，避免计算核闲置。此外，本发明基于不同目标张量的属性和计算核数自动确定候选切分模式，而不需要人工配置切分策略，降低人工成本，且本发明中不同任务对应的目标张量均可以采用该方法进行张量切分，泛化性较强。

技术特征：

1.一种张量切分方法，其特征在于，包括：

2.根据权利要求1所述的张量切分方法，其特征在于，所述基于目标张量的属性以及计算核数，确定所述目标张量的多个候选切分模式，包括：

3.根据权利要求2所述的张量切分方法，其特征在于，首个计算核数为最大可用计算核数。

4.根据权利要求1至3任一项所述的张量切分方法，其特征在于，所述基于各候选切分模式的评分，从各候选切分模式中确定目标切分模式，包括：

5.根据权利要求1至3任一项所述的张量切分方法，其特征在于，所述从各候选切分模式中确定目标切分模式，之后还包括：

6.根据权利要求1至3任一项所述的张量切分方法，其特征在于，所述基于目标张量的属性以及计算核数，确定所述目标张量的多个候选切分模式，包括：

7.根据权利要求1至3任一项所述的张量切分方法，其特征在于，各候选切分模式的评分基于各候选切分张量的形状大小确定，各候选切分张量指基于各候选切分模式对所述目标张量进行切分后的张量。

8.一种电子设备，包括存储器、处理器及存储在所述存储器上并可在所述处理器上运行的计算机程序，其特征在于，所述处理器执行所述程序时实现如权利要求1至7任一项所述张量切分方法。

9.一种非暂态计算机可读存储介质，其上存储有计算机程序，其特征在于，所述计算机程序被处理器执行时实现如权利要求1至7任一项所述张量切分方法。

10.一种计算机程序产品，包括计算机程序，其特征在于，所述计算机程序被处理器执行时实现如权利要求1至7任一项所述张量切分方法。

技术总结本发明提供一种张量切分方法、电子设备、存储介质和计算机程序产品，所述方法包括：基于目标张量的属性以及计算核数，确定目标张量的多个候选切分模式；基于各候选切分模式的评分，从各候选切分模式中确定目标切分模式；基于目标切分模式，对目标张量进行切分。本发明提供的张量切分方法、电子设备、存储介质和计算机程序产品，能够在在充分利用所有计算核的基础上进行张量切分，避免计算核闲置，且不同任务对应的目标张量均可以采用本发明进行张量切分，泛化性较强。技术研发人员：请求不公布姓名,请求不公布姓名,请求不公布姓名,请求不公布姓名受保护的技术使用者：上海壁仞科技股份有限公司技术研发日：技术公布日：2024/9/9