大语言模型的显存管理方法、装置、电子设备及存储介质

专利检索2025-06-13  2


本公开涉及机器学习领域,尤其涉及一种大语言模型的显存管理方法、装置、电子设备及存储介质。


背景技术:

1、大语言模型(large language model,llm)在规模增大到一定程度后的“涌现”现象已经在各种领域展示了惊人的应用潜力,例如文本生成、机器翻译、编程辅助等。然而大模型对内存资源的需求远远高于传统的卷积神经网络(convolutional neural networks,cnns)。常用的cnn参数量一般为百万(million)量级,例如alexnet(60 million)、vgg(138million)。但常用的llm参数量一般为十亿(billion)量级,例如llama2的多个版本(7billion、13 billion、70 billion)。甚至gpt4的参数量可能有几千billion。超大规模的参数给gpu(图形处理单元)的内存带来了极大的压力。相关技术对大语言模型过程中产生的数据进行管理时存在内存浪费以及数据不连续等问题。


技术实现思路

1、有鉴于此,本公开提出了一种大语言模型的显存管理方法、装置、电子设备及存储介质,旨在合理管理内存避免内存空间浪费。

2、根据本公开的第一方面,提供了一种大语言模型的显存管理方法,所述方法包括:

3、确定当前需要进行大语言模型推理的输入文本所对应的文本单元数量,所述大语言模型推理过程包括预填充阶段和解码阶段;

4、根据所述文本单元数量确定预填充阶段产生的激活向量对应的存储位置;

5、根据所述存储位置利用显存中的kv缓存区域的尾部作为激活向量缓存区域;

6、响应于进入解码阶段,从所述激活向量缓存区域中取出激活向量进行解码,得到增加的kv数据和解码阶段产生的激活向量;

7、将增加的kv数据存储至所述激活向量缓存区域,将所述解码阶段产生的激活向量存储在所述增加的kv数据之后。

8、在一种可能的实现方式中,根据所述存储位置利用显存中的kv缓存区域的尾部作为激活向量缓存区域,包括:

9、根据所述解码阶段产生的激活向量占用的最大内存空间确定所述显存中在kv缓存区域外部的外部存储空间;

10、根据所述外部存储空间修正所述存储位置;

11、根据修正后的所述存储位置利用显存中的kv缓存区域的尾部作为激活向量缓存区域。

12、在一种可能的实现方式中,所述激活向量缓存区域的形状为[bs, seqlen, 2×dim+hdim] ,其中,bs为批处理大小,seqlen为文本单元数量,dim为每个文本单元对应的特征维度,hdim为扩展后每个文本单元对应的特征维度。

13、在一种可能的实现方式中,所述激活向量缓存区域包括第一区域、第二区域和第三区域;

14、所述第一区域和所述第二区域的形状为[bs, seqlen, dim],所述第三区域的形状为[bs, seqlen, hdim]。

15、在一种可能的实现方式中,所述第一区域存储的激活向量包括在注意力层处理阶段产生的第一输入参数、第一残差参数和第一输出参数,以及在前馈神经网络处理阶段产生的第二输入参数、第二残差参数和第二输出参数。

16、在一种可能的实现方式中,所述第二区域存储的激活向量包括在注意力层处理阶段产生的第一输入参数、第三输出参数、第四输出参数和第三输入参数,以及在前馈神经网络处理阶段产生的第五输出参数、第四输入参数和第六输出参数。

17、在一种可能的实现方式中,所述第三区域存储的激活向量包括在注意力层处理阶段产生的第七输出参数、第五输入参数和第八输出参数,以及在前馈神经网络处理阶段产生的第九输出参数和第六输入参数。

18、根据本公开的第二方面,提供了一种大语言模型的显存管理装置,所述装置包括:

19、信息确定模块,用于确定当前需要进行大语言模型推理的输入文本所对应的文本单元数量,所述大语言模型推理过程包括预填充阶段和解码阶段;

20、内存计算模块,用于根据所述文本单元数量确定预填充阶段产生的激活向量对应的存储位置;

21、缓存分割模块,用于根据所述存储位置利用显存中的kv缓存区域的尾部作为激活向量缓存区域;

22、数据解码模块,用于响应于进入解码阶段,从所述激活向量缓存区域中取出激活向量进行解码,得到增加的kv数据和解码阶段产生的激活向量;

23、数据存储模块,用于将增加的kv数据存储至所述激活向量缓存区域,将所述解码阶段产生的激活向量存储在所述增加的kv数据之后。

24、在一种可能的实现方式中,所述缓存分割模块,进一步用于:

25、根据所述解码阶段产生的激活向量占用的最大内存空间确定所述显存中在kv缓存区域外部的外部存储空间;

26、根据所述外部存储空间修正所述存储位置;

27、根据修正后的所述存储位置利用显存中的kv缓存区域的尾部作为激活向量缓存区域。

28、在一种可能的实现方式中,所述激活向量缓存区域的形状为[bs, seqlen, 2×dim+hdim] ,其中,bs为批处理大小,seqlen为文本单元数量,dim为每个文本单元对应的特征维度,hdim为扩展后每个文本单元对应的特征维度。

29、在一种可能的实现方式中,所述激活向量缓存区域包括第一区域、第二区域和第三区域;

30、所述第一区域和所述第二区域的形状为[bs, seqlen, dim],所述第三区域的形状为[bs, seqlen, hdim]。

31、在一种可能的实现方式中,所述第一区域存储的激活向量包括在注意力层处理阶段产生的第一输入参数、第一残差参数和第一输出参数,以及在前馈神经网络处理阶段产生的第二输入参数、第二残差参数和第二输出参数。

32、在一种可能的实现方式中,所述第二区域存储的激活向量包括在注意力层处理阶段产生的第一输入参数、第三输出参数、第四输出参数和第三输入参数,以及在前馈神经网络处理阶段产生的第五输出参数、第四输入参数和第六输出参数。

33、在一种可能的实现方式中,所述第三区域存储的激活向量包括在注意力层处理阶段产生的第七输出参数、第五输入参数和第八输出参数,以及在前馈神经网络处理阶段产生的第九输出参数和第六输入参数。

34、根据本公开的第三方面,提供了一种电子设备,包括:处理器;用于存储处理器可执行指令的存储器;其中,所述处理器被配置为在执行所述存储器存储的指令时,实现上述方法。

35、根据本公开的第四方面,提供了一种非易失性计算机可读存储介质,其上存储有计算机程序指令,其中,所述计算机程序指令被处理器执行时实现上述方法。

36、根据本公开的第五方面,提供了一种计算机程序产品,包括计算机可读代码,或者承载有计算机可读代码的非易失性计算机可读存储介质,当所述计算机可读代码在电子设备的处理器中运行时,所述电子设备中的处理器执行上述方法。

37、在本公开实施例中,通过确定当前需要进行大语言模型推理的文本单元数量进一步确定预填充阶段产生的第一激活向量内存,然后根据第一激活向量内存利用显存中的kv缓存区域的尾部作为激活向量缓存区域。在进入解码阶段后,从激活向量缓存区域中取出激活向量进行解码,得到增加的kv数据和解码阶段产生的激活向量。再将增加后的kv数据存储至激活向量缓存区域,将解码阶段产生的激活向量存储在kv数据之后。本公开可以通过利用kv缓存区域尾部存储激活向量的方式实现kv缓存区域的复用,解决了kv数据不连续的问题,且避免内存资源浪费。


技术特征:

1.一种大语言模型的显存管理方法,其特征在于,所述方法包括:

2.根据权利要求1所述的方法,其特征在于,根据所述存储位置利用显存中的kv缓存区域的尾部作为激活向量缓存区域,包括:

3.根据权利要求1或2所述的方法,其特征在于,所述激活向量缓存区域的形状为[bs,seqlen, 2×dim+hdim] ,其中,bs为批处理大小,seqlen为文本单元数量,dim为每个文本单元对应的特征维度,hdim为扩展后每个文本单元对应的特征维度。

4.根据权利要求3所述的方法,其特征在于,所述激活向量缓存区域包括第一区域、第二区域和第三区域;

5.根据权利要求4所述的方法,其特征在于,所述第一区域存储的激活向量包括在注意力层处理阶段产生的第一输入参数、第一残差参数和第一输出参数,以及在前馈神经网络处理阶段产生的第二输入参数、第二残差参数和第二输出参数。

6.根据权利要求4所述的方法,其特征在于,所述第二区域存储的激活向量包括在注意力层处理阶段产生的第一输入参数、第三输出参数、第四输出参数和第三输入参数,以及在前馈神经网络处理阶段产生的第五输出参数、第四输入参数和第六输出参数。

7.根据权利要求4所述的方法,其特征在于,所述第三区域存储的激活向量包括在注意力层处理阶段产生的第七输出参数、第五输入参数和第八输出参数,以及在前馈神经网络处理阶段产生的第九输出参数和第六输入参数。

8.一种大语言模型的显存管理装置,其特征在于,所述装置包括:

9.一种电子设备,其特征在于,包括:

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


技术总结
本公开涉及机器学习领域,尤其涉及一种大语言模型的显存管理方法、装置、电子设备及存储介质,通过确定当前需要进行大语言模型推理的文本单元数量进一步确定预填充阶段产生的第一激活向量内存,然后根据第一激活向量内存在显存中的KV缓存区域的尾部分割激活向量缓存区域。在进入解码阶段后,从激活向量缓存区域中取出激活向量进行解码,得到增加的KV数据和解码阶段产生的激活向量。再将增加后的KV数据存储至激活向量缓存区域,将解码阶段产生的激活向量存储在KV数据之后。本公开可以通过分割KV缓存区域给激活向量的方式实现KV缓存区域的复用,解决了KV数据不连续的问题,且避免内存资源浪费。

技术研发人员:汪玉,毛秋力,洪可
受保护的技术使用者:清华大学
技术研发日:
技术公布日:2024/5/29
转载请注明原文地址:https://win.8miu.com/read-1155112.html

最新回复(0)