一种面向扩散模型并行训练的通信优化方法

本发明公开了一种面向扩散模型并行训练的通信优化方法，主要解决现有并行训练方法通信占比高、易出现落后者的问题。方案包括：1)获取训练样本集，选用任意一个现有的扩散模型；2)获取扩散模型的损失，进行反向传播计算梯度，更新一阶动量和二阶动量，并对其进行修正；3)判断当前训练步长是否达到设定的全局或节点内通信周期，确定通信组；4)根据通信组和类型计数更新通信类型为FP32或BF16；5)使用AllReduce操作分组同步模型参数，对其进行平均；5)重复上述过程直到模型收敛，得到训练好的模型。本发明能够有效降低通信开销，缓解落后者问题，显著提升并行训练效率，可用于大规模扩散模型的深度学习任务。