[发明专利]一种显存优化方法、装置、设备及存储介质

说明书

本申请实施例提供一种显存优化方法、装置、设备及存储介质，其中，所述方法包括：基于预设网络模型，生成第一计算图；确定所述第一计算图的显存峰值与运行数据之间的关联关系；基于所述关联关系，对所述第一计算图进行调整，生成至少一个第二计算图；基于所述至少一个第二计算图的显存峰值和运行时长，在所述至少一个第二计算图中确定目标计算图；基于所述目标计算图，确定所述预设网络模型所需的显存空间。

技术领域

本申请实施例涉及数据处理技术领域，涉及但不限于一种显存优化方法、装置、设备及存储介质。

背景技术

随着深度学习领域的快速发展，训练拥有超多参数的大模型甚至超大模型逐渐步入人们的视野。训练模型逐渐变大变深的同时必然会使显存的开销加剧增长，当模型进一步加大，或者加大批大小(Batch size)后，模型训练的显存占用也会随之增长，最后高于显卡的显存容量，使得模型无法训练。

发明内容

本申请实施例提供一种显存优化技术方案。

本申请实施例的技术方案是这样实现的：

本申请实施例提供一种显存优化方法，所述方法包括：

基于预设网络模型，生成第一计算图；

确定所述第一计算图的显存峰值与运行数据之间的关联关系；

基于所述关联关系，对所述第一计算图进行调整，生成至少一个第二计算图；

基于所述至少一个第二计算图的显存峰值和运行时长，在所述至少一个第二计算图中确定目标计算图；

基于所述目标计算图，确定所述预设网络模型所需的显存空间。

在一些实施例中，所述基于预设网络模型，生成第一计算图，包括：基于所述预设网络模型，生成数据交换格式的计算图信息；基于所述计算图信息中的算子队列，生成所述计算图信息匹配的第一计算图。如此，通过采用深度学习框架将需要训练的模型生成JSON格式的计算图，能够得到第一计算图，进而通过多种优化方案对第一计算图进行调整生成多个第二计算图，便于从中筛选显存花销最优的计算图。

在一些实施例中，所述确定所述第一计算图的显存峰值与运行数据之间的关联关系，包括：确定所述第一计算图中显存峰值的出现时刻；确定所述第一计算图中算子的运行数据；确定所述运行数据的生成时刻和所述运行数据在所述第一计算图中的应用时刻；确定所述生成时刻和所述应用时刻，与所述显存峰值的出现时刻之间的时序关系，为所述关联关系。如此，通过分析算子生成的运行数据的时刻和应用该运行数据的时刻，与达到显存峰值的时刻之间的时序关系，能够进一步确定是否需要移动该算子以降低峰值。

在一些实施例中，基于所述关联关系，对所述第一计算图进行调整，生成至少一个第二计算图，包括：在所述第一计算图中，确定所述关联关系满足预设条件的目标运行数据；基于所述目标运行数据，对所述第一计算图进行调整，生成所述至少一个第二计算图。如此，通过在第一计算图中，移动目标运行数据对应的算子，得到第二计算图，从而能够降低第二计算图的峰值。

在一些实施例中，所述在所述第一计算图中，确定所述关联关系满足预设条件的目标运行数据，包括：在所述第一计算图中的算子的运行数据中，确定生成时刻在所述显存峰值的出现时间之前且应用时刻在所述显存峰值的出现时刻之后的运行数据，为满足所述预设条件的目标运行数据。

在一些实施例中，所述基于所述目标运行数据，对所述第一计算图进行调整，生成所述至少一个第二计算图，包括：在所述第一计算图中，确定所述目标运行数据对应的目标算子；基于所述第一计算图中所述显存峰值的出现时刻，调整所述第一计算图中的所述目标算子，生成所述至少一个第二计算图。如此，针对第一计算图，通过按照该第一计算图中显存峰值的出现时刻，移动目标算子，从而能够生成多个第二计算图。