一种模型分布式训练显存优化方法、系统及可读存储介质技术方案

本发明专利技术公开了一种模型分布式训练显存优化方法、系统及可读存储介质，包括：获取模型第一轮训练过程中每个算子的参数量和访问顺序，对算子进行唯一性编号后生成执行顺序图并存储；对每个算子进行均匀切分，并记录分片映射关系；在训练过程中，根据算子的执行顺序，通过计算设备获取最初运行算子的完整参数；进行当前算子的计算，并根据算子的执行顺序图，通过异步通信预取即将运行的下一个算子的参数；在当前算子执行完毕后，仅保留原始的部分算子参数，并释放其余部分的算子参数；同步算子参数通信，验证分片参数的完整性后执行训练迭代，循环上述过程直至训练完成。利用本发明专利技术，可以在降低单个计算设备显存占用的同时，保证模型的高效训练。

【技术实现步骤摘要】

本专利技术涉及深度学习领域，尤其是涉及一种模型分布式训练显存优化方法、系统及可读存储介质。

技术介绍

1、深度学习和 ai 的快速发展在计算机视觉、自然语言处理、推荐系统等领域取得了显著成果。然而，随着任务复杂度的提升，模型的规模也急剧膨胀，参数数量从百万级跃升至万亿级，显存需求随之激增。例如，一个 100 亿参数的模型在 fp16 精度下需要至少20gb 显存，而训练过程中激活函数、梯度和优化器状态的存储需求可能使显存占用达到80gb 以上。尽管 gpu 性能不断提升，但其内存容量仍然有限，成为模型规模和训练速度的主要瓶颈。

2、针对显存优化，有以下方法：量化（压缩、混合精度）、重计算、换入换出。量化是将模型中的高精度数值转换为低精度数值。重计算是一种策略，在前向传播过程中不存储所有中间激活值，而是在反向传播时重新计算这些值。换入换出是一种通过在显存和主内存之间动态转移数据的方式来管理显存使用。如公开号为cn112329834a的中过专利文献公开了循环网络模型训练时显存空间的分配方法和装置，有效的压缩网络计算中使用到的显存，从而提高训练速度。本文档来自技高网...

【技术保护点】

1.一种模型分布式训练显存优化方法，其特征在于，包括：

2.根据权利要求1所述的模型分布式训练显存优化方法，其特征在于，步骤（2）中，对每个算子进行均匀切分后，每个计算设备保留切割后的算子参数的部分值，对于参数的其它部分进行释放，并记录为分片映射关系。

3.根据权利要求1所述的模型分布式训练显存优化方法，其特征在于，步骤（2）中，还包括为每个算子增加属性信息，算子属性信息包括：算子编号、算子分片参数、参数完整性状态以及下一个算子的编号；

4.根据权利要求3所述的模型分布式训练显存优化方法，其特征在于，步骤（4）中，当参数完整性状态指示当前算子的参数完整...

【技术特征摘要】

1.一种模型分布式训练显存优化方法，其特征在于，包括：

2.根据权利要求1所述的模型分布式训练显存优化方法，其特征在于，步骤（2）中，对每个算子进行均匀切分后，每个计算设备保留切割后的算子参数的部分值，对于参数的其它部分进行释放，并记录为分片映射关系。

3.根据权利要求1所述的模型分布式训练显存优化方法，其特征在于，步骤（2）中，还包括为每个算子增加属性信息，算子属性信息包括：算子编号、算子分片参数、参数完整性状态以及下一个算子的编号；

4.根据权利要求3所述的模型分布式训练显存优化方法，其特征在于，步骤（4）中，当参数完整性状态指示当前算子的参数完整后，启动计算设备执行正向计算与反向传播操作。

5.根据权利要求3所述...

【专利技术属性】
技术研发人员：解臣林，李石坚，张犁，潘纲，
申请(专利权)人：浙江大学，
类型：发明
国别省市：