权重流水并行的神经网络分布式训练方法和装置

本发明涉及人工智能，尤其涉及一种权重流水并行的神经网络分布式训练方法和装置。

背景技术：

1、以神经网络为中心的人工智能技术发展离不开算力的支持。神经网络，尤其是拥有巨大潜力的llm，往往需要长时间的训练过程才能达到较好的性能水平。为了应对这样的挑战，人们往往采用多芯片/服务器集群的方式，通过分布式并行训练来加速整个训练过程，流水线并行方案可以有效减少训练时间，但是会消耗大量诸如gpu的计算芯片。常见的神经网络分布式并行技术包括数据并行、张量并行、流水并行和混合并行等，这些并行技术还会在计算、通信、存储等各种资源做出权衡。

2、流水线并行的一个明显缺点是设备容易出现bubble空闲状态（因为后一个阶段需要等待前一个阶段执行完毕），导致计算资源的浪费。为了解决这个问题，一些研究提出了基于micro-batch的流水方案，该方案中，将一个batch的多个样本进一步拆分到多个micro-batch。每个micro-batch分别执行，在这种策略下bubble的占比会进一步减小。

3、然而一个batch的训练过程中，设备之间的通信量和batch内包含的样本数量成正比。当一个batch包含很多样本的时候，通信量就会很大。

4、综上，现有的流水线并行方案存在通信量大的问题。

技术实现思路

1、本发明提供一种权重流水并行的神经网络分布式训练方法和装置，用以解决现有技术中通信量大的缺陷，实现通信量小的流水并行的神经网络分布式训练。

2、本发明提供一种权重流水并行的神经网络分布式训练方法，包括如下步骤：

3、将神经网络拆分为多个计算阶段，将所述计算阶段按照前/反向操作表示为前反向结点并按照顺序进行环形排列，得到目标轮盘；

4、对所述目标轮盘的所述前反向结点所在的位置进行设备框标记，以将所述前反向结点所在的位置平均分配给计算设备；

5、根据所述目标轮盘的旋转状态和所述设备框得到每个所述计算设备的当前前反向结点，对于每个所述计算设备，利用训练数据根据所述当前前反向结点对应的前/反向操作进行计算并传输相应权重，进而得到训练结果。

6、根据本发明提供的一种权重流水并行的神经网络分布式训练方法，所述目标轮盘的前反向结点还包括对应前/反向操作所需的权重标记。

7、根据本发明提供的一种权重流水并行的神经网络分布式训练方法，所述对于每个所述计算设备，利用训练数据根据所述当前前反向结点对应的前/反向操作进行计算并传输相应权重，具体包括：

8、根据所述计算设备的数量将训练数据拆分为若干个批次子数据，将所述批次子数据发送给所述计算设备；

9、对于每个所述计算设备，利用所述批次子数据根据所述当前前反向结点对应的前/反向操作进行计算并向下一设备进行权重传输。

10、根据本发明提供的一种权重流水并行的神经网络分布式训练方法，所述计算设备之间的权重传输的顺序是根据所述目标轮盘的旋转方向以及所述设备框的位置确定的。

11、根据本发明提供的一种权重流水并行的神经网络分布式训练方法，所述根据所述目标轮盘的旋转状态和所述设备框得到每个所述计算设备的当前前反向结点，对于每个所述计算设备，利用训练数据根据所述当前前反向结点对应的前/反向操作进行计算并传输相应权重，进而得到训练结果，具体包括：

12、s1：在初始时刻初始化所述目标轮盘，根据每个所述计算设备的设备框内包括的前反向结点得到所述计算设备的当前前反向结点；

13、s2：在所述计算设备能够进行计算的情况下，利用所述训练数据进行所述当前前反向结点对应的前/反向操作的计算，并传输相应权重；在所述计算设备不能进行计算的情况下，将所述计算设备置于初始空闲状态，并传输相应权重；

14、s3：旋转所述目标轮盘，更新所述计算设备的当前前反向结点，跳转步骤s2，直至训练结束，输出训练结果。

15、根据本发明提供的一种权重流水并行的神经网络分布式训练方法，所述步骤s2还包括：

16、在所述计算设备完成一个计算阶段的所有反向操作的情况下，将所述计算阶段更新的权重传输给下一个设备，以使之后的设备在更新的权重的基础上进行更新。

17、本发明还提供一种权重流水并行的神经网络分布式训练装置，包括如下模块：

18、轮盘单元，用于将神经网络拆分为多个计算阶段，将所述计算阶段按照前/反向操作表示为前反向结点并按照顺序进行环形排列，得到目标轮盘；

19、设备框单元，用于对所述目标轮盘的所述前反向结点所在的位置进行设备框标记，以将所述前反向结点所在的位置平均分配给计算设备；

20、计算单元，用于根据所述目标轮盘的旋转状态和所述设备框得到每个所述计算设备的当前前反向结点，对于每个所述计算设备，利用训练数据根据所述当前前反向结点对应的前/反向操作进行计算并传输相应权重，进而得到训练结果。

21、本发明还提供一种电子设备，包括存储器、处理器及存储在存储器上并可在处理器上运行的计算机程序，所述处理器执行所述程序时实现如上述任一种所述权重流水并行的神经网络分布式训练方法。

22、本发明还提供一种非暂态计算机可读存储介质，其上存储有计算机程序，该计算机程序被处理器执行时实现如上述任一种所述权重流水并行的神经网络分布式训练方法。

23、本发明还提供一种计算机程序产品，包括计算机程序，所述计算机程序被处理器执行时实现如上述任一种所述权重流水并行的神经网络分布式训练方法。

24、本发明提供的权重流水并行的神经网络分布式训练方法和装置，通过将神经网络拆分为多个计算阶段，将所述计算阶段按照前/反向操作表示为前反向结点并按照顺序进行环形排列，得到目标轮盘；对所述目标轮盘的所述前反向结点所在的位置进行设备框标记，以将所述前反向结点所在的位置平均分配给计算设备；根据所述目标轮盘的旋转状态和所述设备框得到每个所述计算设备的当前前反向结点，对于每个所述计算设备，利用训练数据根据所述当前前反向结点对应的前/反向操作进行计算并传输相应权重，进而得到训练结果。本发明将神经网络根据计算阶段以及每个阶段的前/反向操作进行拆分，根据拆分得到的前反向结点组成目标轮盘，使每个计算设备按照目标轮盘的旋转状态进行前向推理或反向传播，由于本发明的各个设备之间主要进行权重的传递，从而降低了神经网络训练过程中设备之间的数据传输量，提高整体训练效率。

技术特征：

1.一种权重流水并行的神经网络分布式训练方法，其特征在于，包括：

2.根据权利要求1所述的权重流水并行的神经网络分布式训练方法 ，其特征在于，所述目标轮盘的前反向结点还包括对应前/反向操作所需的权重标记。

3.根据权利要求1所述的权重流水并行的神经网络分布式训练方法 ，其特征在于，所述对于每个所述计算设备，利用训练数据根据所述当前前反向结点对应的前/反向操作进行计算并传输相应权重，具体包括：

4.根据权利要求1所述的权重流水并行的神经网络分布式训练方法 ，其特征在于，所述计算设备之间的权重传输的顺序是根据所述目标轮盘的旋转方向以及所述设备框的位置确定的。

5.根据权利要求1或3所述的权重流水并行的神经网络分布式训练方法 ，其特征在于，所述根据所述目标轮盘的旋转状态和所述设备框得到每个所述计算设备的当前前反向结点，对于每个所述计算设备，利用训练数据根据所述当前前反向结点对应的前/反向操作进行计算并传输相应权重，进而得到训练结果，具体包括：

6.根据权利要求5所述的权重流水并行的神经网络分布式训练方法 ，其特征在于，所述步骤s2还包括：

7.一种权重流水并行的神经网络分布式训练装置，其特征在于，包括：

8.一种电子设备，包括存储器、处理器及存储在所述存储器上并可在所述处理器上运行的计算机程序，其特征在于，所述处理器执行所述程序时实现如权利要求1至6任一项所述权重流水并行的神经网络分布式训练方法。

9.一种非暂态计算机可读存储介质，其上存储有计算机程序，其特征在于，所述计算机程序被处理器执行时实现如权利要求1至6任一项所述权重流水并行的神经网络分布式训练方法。

10.一种计算机程序产品，包括计算机程序，其特征在于，所述计算机程序被处理器执行时实现如权利要求1至6任一项所述权重流水并行的神经网络分布式训练方法。

技术总结
本发明提供一种权重流水并行的神经网络分布式训练方法和装置，本包括：将神经网络拆分为多个计算阶段，将所述计算阶段按照前/反向操作表示为前反向结点并按照顺序进行环形排列，得到目标轮盘；对所述目标轮盘的所述前反向结点所在的位置进行设备框标记，以将所述前反向结点所在的位置平均分配给计算设备；根据所述目标轮盘的旋转状态和所述设备框得到每个所述计算设备的当前前反向结点，对于每个所述计算设备，利用训练数据根据所述当前前反向结点对应的前/反向操作进行计算并传输相应权重，进而得到训练结果。本发明使每个计算设备按照目标轮盘的旋转状态进行前向推理或反向传播，降低了神经网络训练过程中设备之间的数据传输量，提高整体训练效率。

技术研发人员：施路平,林俊峰,赵蓉
受保护的技术使用者：清华大学
技术研发日：
技术公布日：2024/12/5