基于张量版本控制的内存优化方法及装置

专利2026-09-03  2

本发明涉及内存优化,尤其涉及一种基于张量版本控制的内存优化方法及装置。
背景技术
::1、传统的深度学习模型在训练过程中通常要将所有的模型参数和中间计算结果保存在gpu内存中,这可能会导致内存不足的问题。gpu内存容量不足已成为限制深度学习发展的主要原因之一。现有技术中gpu内存优化主要采用重计算、内存交换、模型剪枝和蒸馏方法,其中重计算方法是通过在前向传播的过程中丢弃掉部分中间计算结果,只保留部分检查点,然后在反向传播的过程中根据保存的检查点重新计算出所需要的结果,是一种以时间换取空间的办法,但是该类方法在训练规模较大的模型时,会引入较多的额外的时间开销。模型剪枝和蒸馏则是通过减少模型的参数量来达到节省内存占用的目的,但会降低模型训练的精度和性能。2、内存交换技术通过充分利用cpu内存资源,根据深度学习训练过程中的访存特征,在gpu内存和cpu内存之间有序地交换中间张量,在特定时刻将gpu中保存的数据从gpu内存中移到cpu内存,并在需要访问时从cpu内存重新移回gpu内存,从而大大降低深度学习训练过程中的瞬时显存需求,可以解决深度学习训练中的gpu内存不足问题。3、虽然内存交换技术能够通过计算与通信的重叠来减少引入的额外时间开销,但现有技术中内存交换技术大多数是基于神经网络的层级别进行交换实现的,该类优化的粒度较粗,容易引入较大的通信开销。且现有技术中内存交换方法通常是基于启发式算法或者张量的访问特征选择所需执行交换操作的张量以及交换操作的时机,即选择什么时候对哪些张量进行swap-in(换入)和swap-out(换出)操作,该类方式通常都是仅考虑swap-in所带来的时间开销,因而通常是使用预取的方法,以尽可能减少训练过程中由于通信而造成的训练阻塞时间,但是swap-out操作同样会带来一定的时间开销,导致swap-out阶段产生大量的通信开销。技术实现思路1、本发明要解决的技术问题就在于:针对现有技术存在的技术问题,本发明提供一种实现方法简单、成本低以及优化效率高的基于张量版本控制的内存优化方法及装置,能够在gpu内存优化过程中减少swap-out次数,大大降低由于swap-out阶段的同步操作会引入的同步阻塞开销。2、为解决上述技术问题,本发明提出的技术方案为:3、一种基于张量版本控制的内存优化方法,步骤包括:4、s01.抽取深度学习模型的计算图g={v,e},其中v表示节点,e表示边,并对抽取的计算图g中节点进行拓扑排序,得到算子的执行顺序;5、s02.对深度学习模型的训练过程,先训练一个iteration,记录张量的访问模式以及统计每个算子的计算时间c、内存开销u,并根据每个算子的内存开销u以及通信带宽β计算出每个算子的通信开销t;6、s03.在深度学习模型的训练过程中,使用预先定义的张量类型stensor对各张量的状态和版本进行管理,所述张量类型stensor中设置有用于标识张量是否可以被驱逐的flag字段、用于记录张量的存储位置的location字段、用于记录该张量的引用计数值的refcount字段以及用于记录最后一次修改当前张量的时间戳的timestamp字段戳;每当创建或操作张量时,更新张量对应的所述refcount字段,当对张量进行修改时,在张量对应的所述timestamp字段中更新修改时的时间戳,并记录张量的访问顺序;7、s04.当需要在gpu上分配内存时,根据gpu空闲内存的充足状态控制使用张量并对张量类型stensor的各字段进行更新,其中如果内存不足时,选择最近最少使用的张量,将选择的张量换出到cpu内存上,并修改对应的location字段以及flag字段。8、进一步的,所述flag字段中使用non-evictable表示当前张量不能被驱逐,使用evictable表示可以被驱逐;所述location字段中使用cpu-only表示当前张量只存储在cpu内存中,使用gpu-only表示只存储在gpu内存中,使用cpu-gpu表示cpu内存和gpu内存中分别存储了一个副本,所述refcount字段中当引用计数为0的时候进行垃圾回收。9、进一步的,步骤s04中,将选择的张量换出到cpu内存上后,修改对应的location字段为gpu-only以表示只存储在gpu内存中,且交换期间张量的flag字段被标记为non-evictable以表示当前张量不能被驱逐,直到cuda流执行完毕后再将flag字段修改为evictable以表示当前张量能被驱逐。10、进一步的,步骤s04中,若训练过程需要使用目标张量但副本只保存在cpu内存上,则提前将目标张量复制回gpu内存上,并修改对应的location字段为cpu-gpu以表示只存储在gpu内存中,且复制期间张量的flag字段被标记为non-evictable以表示当前张量不能被驱逐,直到cuda流执行完毕后再将flag字段修改为evictable以表示当前张量能被驱逐。11、进一步的,步骤s04的步骤包括:12、步骤s401.判断gpu空闲内存是否充足;13、步骤s402.如果判断到gpu空闲内存充足,则直接分配内存;14、步骤s403.如果判断到如果gpu空闲内存不足,则先检查location字段为cpu-gpu的张量,若对应的flag字段为evictable,且gpu副本的时间戳与cpu副本的时间戳保持一致,cpu-gpu表示只存储在gpu内存中,evictable表示当前张量能被驱逐,则判定该张量当前未被修改过,直接将对应的gpu副本从内存中删除后返回执行步骤s401,如果不满足转入步骤s404。15、步骤s404.检查是否存在引用计数值为0的张量,若存在则运行垃圾回收机制,将该引用计数值为0的张量从gpu内存中删除后返回步骤s401,如果不满足转入步骤s405;16、步骤s405.判断当前是否完成所有cuda流,如果有未完成的cuda流则等待cuda流完成,如果全部完成则使用lru机制执行换出操作以换出一部分最近最少使用张量,返回执行步骤s401。17、进一步的,步骤s03中,通过维护一个列表list存储张量的访问顺序,每当使用张量时,将对应张量移到列表底部,以使得列表list中张量的使用顺序按照从上到下进行排序。18、进一步的,步骤s02中,按照式计算出每个算子的通信开销t。19、一种基于张量版本控制的内存优化装置,包括:20、模型抽取模块,用于抽取深度学习模型的计算图g={v,e},其中v表示节点,e表示边,并对抽取的计算图g中节点进行拓扑排序,得到算子的执行顺序;21、profile模块,用于对深度学习模型的训练过程,先训练一个iteration,记录张量的访问模式以及统计每个算子的计算时间c、内存开销u,并根据每个算子的内存开销u以及通信带宽β计算出每个算子的通信开销t;22、张量管理模块,用于在深度学习模型的训练过程中,使用预先定义的张量类型stensor对各张量的状态和版本进行管理,所述张量类型stensor中设置有用于标识张量是否可以被驱逐的flag字段、用于记录张量的存储位置的location字段、用于记录该张量的引用计数值的refcount字段以及用于记录最后一次修改当前张量的时间戳的timestamp字段戳;每当创建或操作张量时,更新张量对应的所述refcount字段,当对张量进行修改时,在张量对应的所述timestamp字段中更新修改时的时间戳,并记录张量的访问顺序;23、内存优化模块,用于当需要在gpu上分配内存时,根据gpu空闲内存的充足状态控制使用张量并对张量类型stensor的各字段进行更新,其中如果内存不足时,选择最近最少使用的张量,将选择的张量换出到cpu内存上,并修改对应的location字段以及flag字段。24、一种计算机装置,包括处理器以及存储器,所述存储器用于存储计算机程序,所述处理器用于执行所述计算机程序以执行如上述方法。25、一种存储有计算机程序的计算机可读存储介质,所述计算机程序被处理器执行时实现如上述方法。26、与现有技术相比,本发明的优点在于:27、1、本发明利用cpu内存作为扩展提升可训练模型规模的基础上,充分考虑对传统优化方式中swap-out阶段的优化,通过采用一种算子级的内存交换策略,在模型训练过程中,定义一个张量类型stensor,使用四个字段分别对应表示张量是否可以被驱逐、记录张量的存储位置、引用计数值以及修改张量的时间戳,利用该张量类型对训练过程中张量的状态和版本进行管理,使得在gpu内存优化过程中能够减少冗余的张量swap-out操作次数,可以大大降低由于swap-out操作引入的同步开销,从而降低模型训练gpu内存开销,提高内存利用率和计算效率。28、2、本发明通过基于张量类型stensor记录和管理张量的状态和版本,通过张量的标记字段和时间戳的比较,判断gpu内存中的张量版本是否和cpu内存中的一致,可以避免不必要的数据交换,从而在降低模型训练的gpu内存开销的同时,进一步减少内存交换引入的时间开销,提高模型的训练效率。当前第1页12当前第1页12
技术特征:

1.一种基于张量版本控制的内存优化方法,其特征在于,步骤包括:

2.根据权利要求1所述的基于张量版本控制的内存优化方法,其特征在于,所述flag字段中使用non-evictable表示当前张量不能被驱逐,使用evictable表示可以被驱逐;所述location字段中使用cpu-only表示当前张量只存储在cpu内存中,使用gpu-only表示只存储在gpu内存中,使用cpu-gpu表示cpu内存和gpu内存中分别存储了一个副本,所述refcount字段中当引用计数为0的时候进行垃圾回收。

3.根据权利要求2所述的基于张量版本控制的内存优化方法,其特征在于,步骤s04中,将选择的张量换出到cpu内存上后,修改对应的location字段为gpu-only以表示只存储在gpu内存中,且交换期间张量的flag字段被标记为non-evictable以表示当前张量不能被驱逐,直到cuda流执行完毕后再将flag字段修改为evictable以表示当前张量能被驱逐。

4.根据权利要求2所述的基于张量版本控制的内存优化方法,其特征在于,步骤s04中,若训练过程需要使用目标张量但副本只保存在cpu内存上,则提前将目标张量复制回gpu内存上,并修改对应的location字段为cpu-gpu以表示只存储在gpu内存中,且复制期间张量的flag字段被标记为non-evictable以表示当前张量不能被驱逐,直到cuda流执行完毕后再将flag字段修改为evictable以表示当前张量能被驱逐。

5.根据权利要求1所述的基于张量版本控制的内存优化方法,其特征在于,步骤s04的步骤包括:

6.根据权利要求1所述的基于张量版本控制的内存优化方法,其特征在于,步骤s03中,通过维护一个列表list存储张量的访问顺序,每当使用张量时,将对应张量移到列表底部,以使得列表list中张量的使用顺序按照从上到下进行排序。

7.根据权利要求1~6中任意一项所述的基于张量版本控制的内存优化方法,其特征在于,步骤s02中,按照式计算出每个算子的通信开销t。

8.一种基于张量版本控制的内存优化装置,其特征在于,包括:

9.一种计算机装置,包括处理器以及存储器,所述存储器用于存储计算机程序,其特征在于,所述处理器用于执行所述计算机程序以执行如权利要求1~7中任意一项所述方法。

10.一种存储有计算机程序的计算机可读存储介质,其特征在于,所述计算机程序被处理器执行时实现如权利要求1~7中任意一项所述的方法。


技术总结
本发明公开一种基于张量版本控制的内存优化方法及装置,该方法步骤包括:S01.抽取模型的计算图进行拓扑排序,得到算子的执行顺序;S02.对模型训练过程进行profile,先训练一个iteration,记录张量的访问模式以及统计每个算子的计算时间、内存开销,并根据每个算子的内存开销以及通信带宽计算出通信开销;S03.在深度学习模型的训练过程中,使用预先定义的张量类型对各张量的状态和版本进行管理;S04.当需要在GPU上分配内存时,根据GPU空闲内存的充足状态控制使用张量并对各字段进行更新,如果内存不足选择最近最少使用的张量换出到CPU内存。本发明能够在GPU内存优化过程中减少swap‑out次数,降低同步开销。

技术研发人员:乔林波,杨智琳,阚志刚,高翊夫,李东升
受保护的技术使用者:中国人民解放军国防科技大学
技术研发日:
技术公布日:2024/6/26
转载请注明原文地址:https://doc.8miu.com/read-1833850.html

最新回复(0)