本发明涉及大模型训练,尤其涉及一种大模型并行训练通信优化方法。
背景技术:
1、传统的数据并行性在所有数据并行进程中静态地维护整个训练过程中所需的所有模型状态,导致内存冗余。zero数据并行通过对模型状态进行分区而不是复制来消除数据并行进程中的内存冗余,这样每个计算设备上的内存消耗与数据并行维度成反比,zero数据并行也因此在学术界以及工业界流行起来。zero有主要的三个优化阶段,分别将优化器状态、梯度和模型参数分割存储在多个计算设备上,减少了传统数据并行中每个计算设备上存储这些状态所需的内存。但是,这样带来了更多的训练期间的通信开销。
技术实现思路
1、为了解决现有技术中存在的问题,本申请实施例提供了一种大模型并行训练通信优化的方法、装置、计算设备、计算机存储介质及包含计算机程序的产品,能够降低通信开销,提高并行训练的效率。
2、第一方面,本申请实施例提供一种大模型并行训练通信优化方法,应用于多个计算设备组成的系统,每个计算设备配置有gpu,gpu包含至少一个进程,进程用于大模型并行训练中的一个任务分区,方法用于对zero数据的zero-1阶段的优化,方法包括:执行模型的向前计算以及反向计算,确定梯度信息;基于梯度信息,更新本地优化器状态中的动量;通过1-bit allreduce通信对更新后的动量进行处理,并更新压缩过程中产生的误差;利用聚合后的动量以及方差,更新每个计算设备上存储的模型参数分区。
3、在一些可能的实现方式中,1-bit allreduce通信包括:使用all-to-all通信原语将压缩的动量数据块发送给其他计算设备,并接收来自其他计算设备的压缩动量数据块;对所有接收到的动量数据块进行规约操作;使用allgather通信原语收集所有计算设备的动量数据块,完成动量的全局聚合。
4、在一些可能的实现方式中,更新压缩过程中产生的误差,包括:记录每次压缩产生的误差,并在下一次迭代中用于补偿压缩后的动量。
5、在一些可能的实现方式中,方法应用于对zero数据的zero-2阶段的优化,确定梯度信息之后,方法还包括:使用1-bit reduce-scatter进行通信,确定聚合后的动量;基于聚合后的动量以及固定方差,完成模型参数的更新。
6、在一些可能的实现方式中,方法还包括:通过allgather对模型参数进行通信。
7、在一些可能的实现方式中,方法应用于对zero数据的zero-3阶段的优化,方法包括:每个进程执行模型的前向计算,使用allgather通信操作收集所有计算设备上的完整模型信息;完成前向计算后,每个计算设备将不属于当前参数分区的部分进行丢弃操作;执行模型的反向计算,确定梯度信息,使用allgather通信操作收集所有计算设备上的梯度信息;进行1-bit reduce-scatter通信以及更新压缩误差;计算设备再次丢弃非当前分区的参数信息,只保留当前分区必要的数据;利用聚合后的梯度信息和方差,计算每个进程独立更新其分配的模型参数分区。
8、第二方面,本申请实施例提供一种大模型并行训练通信优化装置,部署于多个计算设备组成的系统,每个计算设备配置有gpu,gpu包含至少一个进程,进程用于大模型并行训练中的一个任务分区,装置用于对zero数据的zero-1阶段的优化,装置包括:梯度信息获取模块,用于执行模型的向前计算以及反向计算,确定梯度信息;动量更新模块,用于基于梯度信息,更新本地优化器状态中的动量;处理模块,用于通过1-bit allreduce通信对更新后的动量进行处理,并更新压缩过程中产生的误差;处理模块,还用于利用聚合后的动量以及方差,更新每个计算设备上存储的模型参数分区。
9、第三方面,本申请实施例提供一种计算机可读存储介质,包括计算机可读指令,当计算机读取并执行所述计算机可读指令时,使得计算机执行如第一方面任一项所述的方法。
10、第四方面,本申请实施例提供一种计算设备,包括处理器和存储器,其中,所述存储器中存储有计算机程序指令,所述计算机程序指令被所述处理器运行时,执行如第一方面任一项所述的方法。
11、第五方面,本申请实施例提供一种包含计算机程序的产品,当计算机程序产品在处理器上运行时,使得所述处理器执行如第一方面任一项所述的方法。
1.一种大模型并行训练通信优化方法,其特征在于,应用于多个计算设备组成的系统,每个所述计算设备配置有gpu,所述gpu包含至少一个进程,所述进程用于大模型并行训练中的一个任务分区,所述方法用于对zero数据的zero-1阶段的优化,所述方法包括:
2.根据权利要求1所述的方法,其特征在于,所述1-bit allreduce通信包括:
3.根据权利要求1所述的方法,其特征在于,所述更新压缩过程中产生的误差,包括:记录每次压缩产生的误差,并在下一次迭代中用于补偿压缩后的动量。
4.根据权利要求1所述的方法,其特征在于,所述方法应用于对zero数据的zero-2阶段的优化,所述确定梯度信息之后,所述方法还包括:
5.根据权利要求4所述的方法,其特征在于,所述方法还包括:通过allgather对模型参数进行通信。
6.根据权利要求1所述的方法,其特征在于,所述方法应用于对zero数据的zero-3阶段的优化,所述方法包括:
7.一种大模型并行训练通信优化装置,其特征在于,部署于多个计算设备组成的系统,每个所述计算设备配置有gpu,所述gpu包含至少一个进程,所述进程用于大模型并行训练中的一个任务分区,所述装置用于对zero数据的zero-1阶段的优化,所述装置包括:
