一种模型训练方法及装置、电子设备、存储介质和计算机程序产品

本公开涉及一种模型训练方法及装置、电子设备、存储介质和计算机程序产品，所述方法包括：确定目标模型中的检查点模块；在前向传播阶段，对检查点模块执行前向计算，以及在显存中仅保存检查点模块执行前向计算过程中的目标参数；在反向传播阶段，若检查点模块中包括的最后一层为Flash Attention层，跳过Flash Attention层的前向重计算，以及根据预设Flash Attention反向算法，确定Flash Attention层的梯度，并对检查点模块中其它各网络层执行前向重计算确定对应的梯度。根据本公开实施例能够节省Flash Attention层前向重计算的冗余计算开销，有效提高训练效率。