深度学习模型训练方法及显存优化装置

公开了一种深度学习模型训练方法及显存优化装置。所述方法包括：在第一训练阶段，获取该阶段中各个张量的显存生命周期信息和显存大小信息；根据显存生命周期信息进行张量分类，并根据张量分类结果和显存大小信息生成显存分配方案；在第二训练阶段，基于显存分配方案进行显存分配。本公开利用模型训练中可预测的分配模式先验知识来提前规划显存分配，在不影响训练性能的基础上有效地降低显存碎片，进而提升训练过程中的资源使用效率。本公开还可通过对临时张量按大小分组的空间聚合以及聚合前的张量合并来提升显存分配方案的生成效率，并通过引入静态结合动态的显存分配器使得本公开的显存优化方案在稠密和MoE模型上均实现高效显存分配。