张量卸载方法技术

本申请涉及一种张量卸载方法

【技术实现步骤摘要】
张量卸载方法、装置、计算机设备及存储介质


[0001]本申请涉及人工智能
，特别是涉及一种张量卸载方法
、
装置
、
计算机设备及存储介质
。

技术介绍

[0002]人工智能技术（
Artificial Intelligence
，
AI
）作为智能学科重要的组成部分，在图像识别
、
语音处理
、
机器翻译
、
游戏
、
医疗保健等诸多领域得到了广泛应用，
AI
模型也随之不断发展
。
但是随着
AI
模型规模的不断增加，模型参数
、
中间变量等张量的显存需求也越来越高
。
另一方面，随着
AI
芯片计算能力的快速增加，需要不断通过增加批量（
batch
）大小来充分发挥
AI
芯片的算力，但是批量大小的增加又受限于显存大小
。
因此，如何提高显存利用率对于
AI
模型的训练有着重要意义
。
[0003]在传统技术中，通常使用剪枝
、
量化等模型压缩方法来减少显存的使用量，从而减少显存需求
。
但是，模型压缩会造成模型训练和推理精度的降低，并且会引入额外的计算开销
。
[0004]基于此，传统技术中亟需一种在对模型训练精度影响较小的情况下，增加显存有效容量的方法
。

技术实现思路

[0005]基于此，有必要针对上述技术问题，提供一种能够在对
AI
模型训练精度影响较小的情况下，增加显存有效容量的张量卸载方法
、
装置
、
计算机设备及存储介质
。
[0006]第一方面，本申请提供了一种张量卸载方法
。
所述方法包括：获取张量特征信息，所述张量特征信息包括待训练模型每层的显存容量需求
、
计算间隔以及计算延迟信息；基于预设卸载策略和所述显存容量需求确定显存约束条件，基于所述预设卸载策略和计算间隔确定卸载时间约束条件，所述预设卸载策略包括主存卸载和重计算卸载；基于所述显存约束条件和卸载时间约束条件筛选所述预设卸载策略，确定多个候选卸载策略；基于所述计算延迟信息确定每个候选卸载策略的额外计算延迟，基于所述额外计算延迟确定目标卸载策略，并基于所述目标卸载策略卸载张量
。
[0007]在其中一个实施例中，所述计算间隔基于以下方式确定：确定目标层的前向计算延迟以及反向计算延迟之和，作为当前层的计算延迟信息，其中，所述目标层包括沿前向计算方向当前层之后的所有层，所述前向计算方向包括所述待训练模型的输入端至输出端的方向
。
[0008]在其中一个实施例中，所述显存容量需求基于所述张量的数据信息确定，所述数据信息包括数据类型信息
、
数据形状信息和数据维度信息
。
[0009]在其中一个实施例中，所述基于预设卸载策略和所述显存容量需求确定显存约束
条件包括：基于所述预设卸载策略确定所述待训练模型的张量卸载总量；基于所述待训练模型每层的显存容量需求确定所述待训练模型显存容量的总需求；所述显存约束条件为，所述张量卸载总量不小于所述显存容量总需求
。
[0010]在其中一个实施例中，所述基于预设卸载策略和所述显存容量需求确定显存约束条件还包括：基于所述预设卸载策略确定显存剩余容量；确定多个显存容量需求中的最大单层显存容量需求；所述显存约束条件为，所述显存剩余容量不小于所述最大单层显存容量需求
。
[0011]在其中一个实施例中，所述基于所述预设卸载策略和计算间隔确定卸载时间约束条件包括：基于所述预设卸载策略确定所述待训练模型每层的张量主存卸载量；基于所述张量主存卸载量确定第一时间阈值；所述卸载时间约束条件为，所述计算间隔不小于所述第一时间阈值
。
[0012]在其中一个实施例中，所述基于所述预设卸载策略和计算间隔确定卸载时间约束条件还包括：基于所述预设卸载策略确定所述待训练模型的张量主存卸载总量；基于所述张量主存卸载总量确定第二时间阈值；基于所述预设卸载策略确定主存卸载层，并基于所述计算间隔确定主存卸载层计算间隔，所述主存卸载层包括所述待训练模型中至少有一个张量卸载至主存的层；所述卸载时间约束条件为，所述主存卸载层计算间隔不小于所述第二时间阈值
。
[0013]在其中一个实施例中，所述张量包括常驻张量以及临时张量，所述常驻张量包括训练完成之前计算所需的持续张量，所述临时张量包括单次前后向计算所需张量，所述预设卸载策略基于以下方式确定：确定所述常驻张量的预设卸载策略包括主存卸载；确定所述临时张量的预设卸载策略包括主存卸载和
/
或重计算卸载
。
[0014]在其中一个实施例中，所述计算延迟信息包括重计算延迟，所述预设卸载策略还基于以下方式确定：若所述计算间隔不小于卸载传输延迟，则确定所述张量的预设卸载策略包括主存卸载，所述卸载传输延迟包括显存至主存的传输延迟；若所述计算间隔不小于所述重计算延迟，且小于所述卸载传输延迟，则确定所述张量的预设卸载策略包括重计算卸载
。
[0015]在其中一个实施例中，在所述基于所述目标卸载策略卸载张量之后，还包括：确定当前可用显存容量与每轮次延迟时间，所述每轮次延迟时间包括所述待训练模型单次前后向计算的实际延迟时间；基于所述每层的计算延迟信息确定每轮次延迟需求，所述每轮次延迟需求包括所述待训练模型单次前后向计算的延迟需求；若所述每轮次延迟时间大于所述每轮次延迟需求，且需求显存容量不大于所述可
用显存容量，则基于所述需求显存容量释放显存，其中，所述需求显存容量基于所述每轮次延迟时间和每轮次延迟需求确定
。
[0016]在其中一个实施例中，所述确定当前可用显存容量包括：基于所述目标卸载策略确定当前显存剩余容量；若所述每轮次延迟时间小于所述每轮次延迟需求，则基于所述每轮次延迟时间和每轮次延迟需求确定可释放显存容量；基于所述当前显存剩余容量和可释放显存容量确定所述可用显存容量
。
[0017]第二方面，本申请还提供了一种张量卸载装置
。
所述装置包括：张量特征信息获取模块，用于获取张量特征信息，所述张量特征信息包括待训练模型每层的显存容量需求
、
计算间隔以及计算延迟信息；约束条件确定模块，用于基于预设卸载策略和所述显存容量需求确定显存约束条件，基于所述预设卸载策略和计算间隔确定卸载时间约束条件，所述预设卸载策略包括主存卸载和重计算卸载；候选卸载策略确定模块，用于基于所述显存约束条件和卸载时间约束条件筛选所述预设卸载策略，确定多个候选卸载策略；目标卸载策略确定模块，用于基于所述计算延迟信息确定每个候选卸载策略的额外计算延迟，本文档来自技高网...

【技术保护点】

【技术特征摘要】
1.
一种张量卸载方法，其特征在于，所述方法包括：获取张量特征信息，所述张量特征信息包括待训练模型每层的显存容量需求
、
计算间隔以及计算延迟信息；基于预设卸载策略和所述显存容量需求确定显存约束条件，基于所述预设卸载策略和计算间隔确定卸载时间约束条件，所述预设卸载策略包括主存卸载和重计算卸载；基于所述显存约束条件和卸载时间约束条件筛选所述预设卸载策略，确定多个候选卸载策略；基于所述计算延迟信息确定每个候选卸载策略的额外计算延迟，基于所述额外计算延迟确定目标卸载策略，并基于所述目标卸载策略卸载张量
。2.
根据权利要求1所述的方法，其特征在于，所述计算间隔基于以下方式确定：确定目标层的前向计算延迟以及反向计算延迟之和，作为当前层的计算延迟信息，其中，所述目标层包括沿前向计算方向当前层之后的所有层，所述前向计算方向包括所述待训练模型的输入端至输出端的方向
。3.
根据权利要求1所述的方法，其特征在于，所述显存容量需求基于所述张量的数据信息确定，所述数据信息包括数据类型信息
、
数据形状信息和数据维度信息
。4.
根据权利要求1所述的方法，其特征在于，所述基于预设卸载策略和所述显存容量需求确定显存约束条件包括：基于所述预设卸载策略确定所述待训练模型的张量卸载总量；基于所述待训练模型每层的显存容量需求确定所述待训练模型显存容量的总需求；所述显存约束条件为，所述张量卸载总量不小于所述显存容量总需求
。5.
根据权利要求1所述的方法，其特征在于，所述基于预设卸载策略和所述显存容量需求确定显存约束条件还包括：基于所述预设卸载策略确定显存剩余容量；确定多个显存容量需求中的最大单层显存容量需求；所述显存约束条件为，所述显存剩余容量不小于所述最大单层显存容量需求
。6.
根据权利要求1所述的方法，其特征在于，所述基于所述预设卸载策略和计算间隔确定卸载时间约束条件包括：基于所述预设卸载策略确定所述待训练模型每层的张量主存卸载量；基于所述张量主存卸载量确定第一时间阈值；所述卸载时间约束条件为，所述计算间隔不小于所述第一时间阈值
。7.
根据权利要求1所述的方法，其特征在于，所述基于所述预设卸载策略和计算间隔确定卸载时间约束条件还包括：基于所述预设卸载策略确定所述待训练模型的张量主存卸载总量；基于所述张量主存卸载总量确定第二时间阈值；基于所述预设卸载策略确定主存卸载层，并基于所述计算间隔确定主存卸载层计算间隔，所述主存卸载层包括所述待训练模型中至少有一个张量卸载至主存的层；所述卸载时间约束条件为，所述主存卸载层计算间隔不小于所述第二时间阈值
。8.
根据权利要求1所述的方法，其特征在于，所述张量包括常驻张量以及临时张量，所述常驻张量包括训练完成之前计算所需的持续张量，所述临时张量包...

【专利技术属性】
技术研发人员：李勇，程稳，刘懿，吴运翔，朱健，吕波，常璟飞，黄章敏，陈光，曾令仿，
申请(专利权)人：之江实验室，
类型：发明
国别省市：