一种基于显存状态感知的AI模型可靠性训练方法及装置

本发明公开了一种基于显存状态感知的AI模型可靠性训练方法及装置，该方法首先在每轮训练算子执行前，通过当前训练步数和保存频率判断当前训练轮次是否保存。其次若触发保存操作，分析当前训练的AI模型的网络结构，计算可与训练前反向进行并行传输的检查点大小，并生成切分策略。然后基于当前训练数据生成的切分策略对检查点进行切分操作，完成切分后将部分检查点暂存到空余显存中，在暂存时使用显存复用技术暂存检查点。最后暂存执行完成后，通过多流并行执行异步拷贝操作，并更新参数，完成训练。本发明减少大模型训练时的暂停时间和检查点保存引入的开销，减少端到端训练时间，并实现了保存性能的显著提升。