混合专家模型的专家网络训练方法、装置、介质以及设备与流程

本申请涉及人工智能、混合专家模型和专家网络训练等，尤其涉及一种混合专家模型的专家网络并行训练方法、装置、可读存储介质和电子设备。

背景技术：

1、混合专家模型是一种大规模神经网络架构，其将传统网络层中的单一变换替换为多个并行的专家网络，并引入路由网络，由路由网络为每个输入token动态选择其中若干个专家网络进行处理。通过这种方式，混合专家模型能够在总参数量大幅增加的情况下保持计算量基本不变。

2、在混合专家模型的分布式训练中，通常采用并行训练专家网络的训练策略：将同一混合专家层中的全部专家网络分散部署至多个计算节点，每个计算节点负责存储并计算其中一部分专家网络。每个训练微批次进入混合专家层时，各计算节点上的路由网络根据本地输入token的特征为各token选择目标专家网络，各计算节点据此将本地token发送至其目标专家网络所在的计算节点。这使得不同专家网络被分配的token数量存在显著差异。

3、在并行训练专家网络的训练策略下，各计算节点持有不同的专家网络集合，token分配的不均衡表现为部分计算节点接收到的token数量远大于其他节点。由于各计算节点在同步训练模式下需相互等待，接收token数量最多的计算节点成为性能瓶颈，其计算耗时决定了当前混合专家层的整体处理时长，拖累训练进度。

技术实现思路

1、本申请提供了一种混合专家模型的专家网络训练方法、装置、可读存储介质和电子设备，通过使各计算节点的当前节点负载均恰好达到预设节点负载，消除各计算节点间的负载不均衡，避免因部分计算节点负载过高而拖慢整体训练进度，有利于降低混合专家模型的训练耗时并提升训练稳定性。

2、根据本申请的第一方面，提供了混合专家模型的专家网络训练方法，所述方法包括：

3、响应于检测到训练微批次进入混合专家层，通过第一节点基于各计算节点上各本地专家网络的当前专家负载确定各计算节点的当前节点负载，并根据预设节点负载和各计算节点的当前节点负载在所述计算节点中确定超载节点和欠载节点；所述计算节点的共享存储中设置有第一数量的冗余专家槽位；所述第一数量等于混合专家层的总专家网络数与总计算节点数之比；所述预设节点负载为均衡状态下各计算节点应承担的负载数量；所述第一节点在所述计算节点中产生；

4、基于所述预设节点负载、所述欠载节点的当前节点负载、所述超载节点的当前节点负载以及所述超载节点上各本地专家网络的当前专家负载，生成负载调度信息；其中，所述负载调度信息指定至少一个待迁移节点对；一个待迁移节点对包括由一个作为迁出节点的超载节点和一个作为迁入节点的欠载节点；负载调度信息指定将所述迁出节点上至少一个本地专家网络及其对应的任务负载迁移到所述迁入节点以使该迁入节点的当前节点负载恰好等于所述预设节点负载；

5、将所述负载调度信息发送给各计算节点，以供各计算节点基于所述负载调度信息确定自身角色并执行负载迁移。

6、根据本申请的第二方面，提供了混合专家模型的专家网络训练装置，所述装置包括：

7、节点负载确定模块，用于响应于检测到训练微批次进入混合专家层，通过第一节点基于各计算节点上各本地专家网络的当前专家负载确定各计算节点的当前节点负载，并根据预设节点负载和各计算节点的当前节点负载在所述计算节点中确定超载节点和欠载节点；所述计算节点的共享存储中设置有第一数量的冗余专家槽位；所述第一数量等于混合专家层的总专家网络数与总计算节点数之比；所述预设节点负载为均衡状态下各计算节点应承担的负载数量；所述第一节点在所述计算节点中产生；

8、调度信息生成模块，用于基于所述预设节点负载、所述欠载节点的当前节点负载、所述超载节点的当前节点负载以及所述超载节点上各本地专家网络的当前专家负载，生成负载调度信息；其中，所述负载调度信息指定至少一个待迁移节点对；一个待迁移节点对包括由一个作为迁出节点的超载节点和一个作为迁入节点的欠载节点；负载调度信息指定将所述迁出节点上至少一个本地专家网络及其对应的任务负载迁移到所述迁入节点以使该迁入节点的当前节点负载恰好等于所述预设节点负载；

9、调度信息发送模块，用于将所述负载调度信息发送给各计算节点，以供各计算节点基于所述负载调度信息确定自身角色并执行负载迁移。

10、根据本发明的第三方面，本申请实施例提供了一种计算机可读存储介质，其上存储有计算机程序，该程序被处理器执行时实现如本申请实施例所述的混合专家模型的专家网络训练方法。

11、根据本发明的第四方面，本申请实施例提供了一种电子设备，包括存储器，处理器及存储在存储器上并可在处理器运行的计算机程序，所述处理器执行所述计算机程序时实现如本申请实施例所述的混合专家模型的专家网络训练方法。

12、根据本申请的第五方面，本申请实施例提供了一种计算机程序产品，包括计算机程序，所述计算机程序在被处理器执行时实现如本申请实施例所述的混合专家模型的专家网络训练方法。

13、本申请技术方案，通过第一节点根据各计算节点的当前节点负载与预设节点负载的差异确定超载节点和欠载节点，并生成指定由超载节点向欠载节点迁移本地专家网络及其对应任务负载的负载调度信息，使欠载节点的当前节点负载恰好达到预设节点负载，从而消除各计算节点间的负载不均衡，避免因部分计算节点负载过高而成为性能瓶颈拖慢整体训练进度。由于各计算节点的当前节点负载均被调整为预设节点负载，每个计算节点在混合专家层处理过程中承担的计算量保持恒定，各计算节点接收的任务负载固定为预设值，使得专家网络处理的输入数据量在编译期即可确定，无需在运行时动态分配显存，从而避免显存碎片；同时，由于每个混合专家层的输入数据量已知，各专家网络输入数据的起始地址和长度无需在运行时计算，无需在每层处理时将各计算节点的任务计数从计算节点拷贝至主机，消除了主机与节点之间的同步开销，降低了训练耗时。另外，第一数量被设置为等于混合专家层的总专家网络数与总计算节点数之比，使得各计算节点预留的冗余专家槽位数量足以容纳任意路由分布下所需迁移的专家网络，保证负载调度信息的生成在任意路由分布下均有可行解，避免因槽位不足导致训练中断，提升了训练稳定性。

14、应当理解，本部分所描述的内容并非旨在标识本申请的实施例的关键或重要特征，也不用于限制本申请的范围。本申请的其它特征将通过以下的说明书而变得容易理解。

1.一种混合专家模型的专家网络训练方法，其特征在于，所述方法包括：

2.根据权利要求1所述的方法，其特征在于，所述基于所述预设节点负载、所述欠载节点的当前节点负载、所述超载节点的当前节点负载以及所述超载节点上各本地专家网络的当前专家负载，生成负载调度信息，包括：

3.根据权利要求2所述的方法，其特征在于，所述基于所述预设节点负载、所述欠载节点的当前节点负载和所述超载节点的当前节点负载，构建所述待迁移节点对并为所述待迁移节点对确定待迁移的负载数量，包括：

4.根据权利要求2所述的方法，其特征在于，所述基于所述待迁移的负载数量和所述待迁移节点对中迁出节点上各本地专家网络的当前专家负载，在所述迁出节点中确定待迁移的目标专家网络以及各目标专家网络需迁移的负载数量，包括：

5.根据权利要求4所述的方法，其特征在于，所述基于已遍历的总负载量、所述待迁移的负载数量和所述目标专家网络的当前专家负载，确定各目标专家网络需迁移的负载数量，包括：

6.根据权利要求1所述的方法，其特征在于，所述各计算节点基于所述负载调度信息确定自身角色并执行负载迁移，包括：

7.根据权利要求6所述的方法，其特征在于，所述基于所述负载调度信息所指定的迁移数量在分配给所述目标专家网络的任务负载中确定待迁移的目标负载，并将所述目标负载发送给所述迁入节点，包括：

8.根据权利要求6所述的方法，其特征在于，所述迁入节点从所述迁出节点中拉取所述目标专家网络的权重参数，与所述迁出节点将所述目标负载发送给所述迁入节点，在时间上存在部分重叠。

9.根据权利要求1所述的方法，其特征在于，所述响应于检测到训练微批次进入混合专家层，通过第一节点基于各计算节点上各本地专家网络的当前专家负载确定各计算节点的当前节点负载，并根据预设节点负载和各计算节点的当前节点负载在所述计算节点中确定超载节点和欠载节点，包括：

10.根据权利要求1所述的方法，其特征在于，所述将所述负载调度信息发送给各计算节点，包括：

11.根据权利要求1所述的方法，其特征在于，所述方法还包括：

12.一种混合专家模型的专家网络训练装置，其特征在于，所述装置包括：

13.一种计算机可读存储介质，其上存储有计算机程序，其特征在于，该程序被处理器执行时实现如权利要求1-11中任一项所述的混合专家模型的专家网络训练方法。

14.一种电子设备，包括存储器，处理器及存储在存储器上并可在处理器运行的计算机程序，其特征在于，所述处理器执行所述计算机程序时实现如权利要求1-11中任一项所述的混合专家模型的专家网络训练方法。