ChatGLM2-6B的通透解析:从FlashAttention、Multi-Query Attention到GLM2的微调、源码解读

目录

前言

第一部分 相比第一代的改进点:FlashAttention与Multi-Query Attention

第二部分 FlashAttention:减少内存访问提升计算速度——更长上下文的关键

2.1 FlashAttention相关的背景知识

2.1.1  Transformer计算复杂度:​编辑——Self-Attention层与MLP层

2.1.1.1 Self-Attention层的计算复杂度:

2.1.1.2 MLP层的计算复杂度:

2.1.2 Transformer的空间复杂度:​编辑——Self-Attention层与MLP层

2.1.2.1 Self-Attention块的中间激活:

2.1.2.2 MLP块的中间激活:

2.1.2.3 两个layer norm需要保存的中间激活:

2.1.3 分析GPU的内存分析图:计算的瓶颈是显存访问

2.1.4 safe softmax

2.2 前向传递:Standard Attention/Memory-efficient Attention/Flash Attention

2.2.1 Standard Attention

2.2.2 Memory-efficient Attention:把显存复杂度从平方降低到线性,但HBM访问次数仍是平方

2.2.3 Flash Attention:避免频繁地从HBM中读写数据

第三部分 多查询注意力(Muti Query Attention):各自Query矩阵,但共享Key 和 Value 矩阵

3.1 Multi-Head Attention、Grouped-Query Attention、Muti Query Attention的区别

​3.2 MHA 和 MQA在代码实现上的差异

第四部分 模型的使用/部署、微调

4.1 模型的使用/部署

4.2 基于 P-Tuning v2 的微调(官方


前言

本文最初和第一代ChatGLM-6B的内容汇总在一块,但为了阐述清楚FlashAttention、Multi-Query Attention等相关的原理,以及GLM2的微调、源码解读等内容,导致之前那篇文章越写越长,故特把ChatGLM2相关的内容独立抽取出来成本文

且本文会和本博客内其他大模型相关的文章一样,极其注重可读性,比如为了不断提高可读性,本文近期会不断反复修改,细抠标题的层级、措辞,甚至排版、标点符号,如果不通俗易懂,宁愿不写

第一部分 相比第一代的改进点:FlashAttention与Multi-Query Attention

ChatGLM2-6B(GitHub项目地址、HuggingFace地址)是开源中英双语对话模型 ChatGLM-6B 的第二代版本,相比第一代,第二点引入了如下新特性:

  • 数据集上
    经过了 1.4T 中英标识符的预训练与人类偏好对齐训练
  • 更长的上下文
    基于 FlashAttention 技术,将基座模型的上下文长度(Context Length)由 ChatGLM-6B 的 2K 扩展到了 32K,并在对话阶段使用 8K 的上下文长度训练,允许更多轮次的对话
    (当前版本的 ChatGLM2-6B 对单轮超长文档的理解能力有限,会在后续迭代升级中着重进行优化)
  • 更高效的推理
    基于 Multi-Query Attention 技术,ChatGLM2-6B 有更高效的推理速度和更低的显存占用:在官方的模型实现下,推理速度相比初代提升了 42%,INT4 量化下,6G 显存支持的对话长度由 1K 提升到了 8K
  • 模型架构上变成了decoder only的架构
    chatglm还是encoder架构,但是到了chatglm2 变成了decoder only的架构(这点很少有资料会提及到),何以见得呢?
    如七月黄老师所说,chatglm2仓库的modeling用了新版pytorch的这个函数:context_layerChatGLM2-6B的通透解析:从FlashAttention、Multi-Query Attention到GLM2的微调、源码解读_第1张图片

    context_layer 这个函数实现了attention机制的计算,入参 is_causal=True 表示遮后看前的mask(这种类型的注意力通常用在transformer的decoder部分,以确保当前位置只能关注到之前的位置,俗称“看不见未来”,从而使模型可以进行自回归预测 )

  • 允许商业使用
  • 准确性不足
    尽管模型在训练的各个阶段都尽力确保数据的合规性和准确性,但由于 ChatGLM2-6B 模型规模较小,且模型受概率随机性因素影响,无法保证输出内容的准确性,且模型易被误导

第二部分 FlashAttention:减少内存访问提升计算速度——更长上下文的关键

FlashAttention是斯坦福联合纽约州立大学在22年6月份提出的一种具有 IO 感知,且兼具快速、内存高效的新型注意力算法「对应论文为:FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,这是其GitHub地址,这是其解读之一,该解读也是本第二部分的重要参考之一」

它要解决一个什么样的问题呢?

首先,GPT3、LLaMA、ChatGLM、BLOOM等大语言模型输入输出的最大序列长度只有2048或4096,扩展到更长序列的难度在哪里呢?本质原因是,transformer模型的计算复杂度和空间复杂度都是 O(N^2)​的,其中N​为序列长度

2.1 FlashAttention相关的背景知识

2.1.1  Transformer计算复杂度:l *\left(24 b N h^{2}+4 b N^{2} h\right)——Self-Attention层与MLP层

当输入批次大小为 b​ ,序列长度为 N​ 时,
l​ 层transformer模型的计算量为 l *\left(24 b N h^{2}+4 b N^{2} h\right)​,h​是隐藏层维度通常等于词向量维度,可能不少同学都会疑问这个计算量是怎么一步一步计算得来的,下面详细拆解下这个计算过程

首先,我们知道,transformer模型由 l​ 个相同的层组成,每个层分为两部分:self-attention块和MLP块

2.1.1.1 Self-Attention层的计算复杂度:8bNh^2 + 4bN^2h

self-attention层的模型参数有两部分,一部分是Q​、K​、V​的权重矩阵W_Q​、W_K​、W_V​和偏置,另一部分是输出权重矩阵W_O​和偏置,最终为8bNh^2 + 4bN^2h

具体怎么计算得来的呢?

  1. 第一步是计算Q​、K​、V
    Q=x W_{Q}, K=x W_{K}, V=x W_{V}
    该矩阵乘法的输入和输出形状为 [b, N, h] \times[h, h] \rightarrow[b, N, h]
    计算量为:3 * 2 b N h^{2}=6 b N h^{2}
  2. 计算Q K^T
    该部分的输入和输出形状为
    \left[b, h e a d \_n u m, l, p e r \_h e a d \_h i d d e n \_s i z e\right]​ \times​ \left[b, h e a d \_n u m, p e r \_h e a d \_h i d d e n \_s i z e\right. , N] \rightarrow\left[b, h e a d \_n u m, N, l\right]
    计算量为:2bN^2h
  3. 计算在V​上的加权 score \cdot V
    该部分矩阵乘法的输入和输出形状为
    \left[b, h e a d \_n u m, l, l\right] \times\left[b, h e a d \_n u m, l, p e r \_h e a d \_h i d d e n \_s i z e\right]​ \rightarrow\left[b, h e a d \_n u m, N, p e r \_h e a d \_h i d d e n \_s i z e\right]
    计算量为:2bN^2h
  4.  attention后的线性映射,矩阵乘法的输入和输出形状为[b, N, h] \times[h, h] \rightarrow[b, N, h]
    计算量为2bNh^2

    最终自注意力层的输出结果为
    x_{o u t}=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{h}}\right) \cdot V \cdot W_{o}+x
2.1.1.2 MLP层的计算复杂度:16bNh^2

MLP块由2个线性层组成,最终是16bNh^2

怎么计算得来的呢?

  1. 一般地,第一个线性层是,第二个线性层再将维度从 4h​映射到h
    x=f_{\text {gelu }}\left(x_{\text {out }} W_{1}\right) W_{2}+x_{\text {out }}

    第一个线性层的权重矩阵 W_1​ 的形状为 [h,4h]​,相当于先将维度从 h​ 映射到4h​,矩阵乘法的输入和输出形状为[b, N, h] \times[h, 4 h] \rightarrow[b, N, 4 h]​,计算量为 8bNh^2
    第二个线性层的权重矩阵 W_2​ 的形状为 [4h,h]​,相当于再将维度从 4h​映射到 h​,矩阵乘法的输入和输出形状为[b, N, 4 h] \times[4 h, h] \rightarrow[b, N, h]​,计算量为 8bNh^2
  2. 将上述所有表粗所示的计算量相加,得到每个transformer层的计算量大约为
    24 b N h^{2}+4 b N^{2} h
  3. 此外,另一个计算量的大头是logits的计算(毕竟词嵌入矩阵的参数量也较多),将隐藏向量映射为词表大小,说白了,词向量维度通常等于隐藏层维度h​ ,词嵌入矩阵的参数量为Vh​,最后的输出层的权重矩阵通常与词嵌入矩阵是参数共享的「解释一下,如七月杜老师所说,这个是transformer中一个重要的点,参数共享可以减小参数量,词嵌入矩阵是[vocab_size,hidden_size],输出层矩阵是 [hidden_size,vocab_size],是可以共享的」
    其矩阵乘法的输入和输出形状为[b, N, h] \times[h, V] \rightarrow[b, N, V]​,计算量为 2bNhV
  4. 因此,对于一个 l​ 层的transformer模型,输入数据形状为 [b,N]​的情况下,一次训练迭代的计算量为
    l *\left(24 b N h^{2}+4 b N^{2} h\right)+2 b N h V

2.1.2 Transformer的空间复杂度:l *\left(34 b N h+5 b N^{2} a\right)——Self-Attention层与MLP层

中间激活的显存大小为l *\left(34 b N h+5 b N^{2} a\right)​  ,其中 a​ 为注意力头数

大模型在训练过程中通常采用混合精度训练,中间激活值一般是float16或者bfloat16数据类型的。在分析中间激活的显存占用时,假设中间激活值是以float16或bfloat16数据格式来保存的,每个元素占了2个bytes。唯一例外的是,dropout操作的mask矩阵,每个元素只占1个bytes。在下面的分析中,单位是bytes,而不是元素个数。

每个transformer层包含了一个self-attention块和MLP块,并分别对应了一个layer normalization连接。

2.1.2.1 Self-Attention块的中间激活:11 b N h+5 b N^{2} a

self-attention块的计算公式如下:

Q=x W_{Q}, K=x W_{K}, V=x W_{V}
x_{o u t}=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{h}}\right) \cdot V \cdot W_{o}+x

最终,self-attention块的中间激活占用显存大小为:11 b N h+5 b N^{2} a

具体怎么计算得来的呢?

  1. 对于Q,K,V ,需要保存它们共同的输入 x ,这就是中间激活。输入 x的形状为[b, N, h],元素个数为 bNh占用显存大小为2 * b N h=2 b N h
  2. 对于 Q K^{T}矩阵乘法,需要保存中间激活 Q,K ,两个张量的形状都是[b,s,h]占用显存大小合计为2 * 2 * b N h=4 b N h
  3. 对于 \text { softmax () } 函数,需要保存函数的输入Q K^{T} ,占用显存大小为2 b N^{2} a,这里的a 表示注意力头数
    \text { score }=\operatorname{softmax}\left(\frac{Q K^{T}}{\sqrt{d_{k}}}\right)

    其中
    Q的形状为:\left[b, h e a d \_n u m, N, p e r \_h e a d \_h i d d e n \_s i z e\right]
    K^{T}的形状为:\left[b, h e a d \_n u m, p e r \_h e a d \_h i d d e n \_s i z e, N\right]
    Q K^{T}的形状为:\left[b, h e a d \_n u m, N, N\right],元素个数为b N^{2} a,占用显存大小为2 b N^{2} a
  4. 计算完 \text { softmax () }函数后,会进行dropout操作。需要保存一个mask矩阵,mask矩阵的形状与Q K^{T} 相同,占用显存大小为b N^{2} a
  5. 计算在 V上的attention,即\text { score } \cdot V,需要保存 \text { score } ,大小为 2 b N^{2} a ;以及 V,大小为 2 b N h二者占用显存大小合计为2 b N^{2} a+2 b N h
  6. 计算输出映射以及一个dropout操作。输入映射需要保存其输入,大小为2 b s h;dropout需要保存mask矩阵,大小为\text { bsh }二者占用显存大小合计为3 b N h

因此,将上述中间激活相加得到,self-attention块的中间激活占用显存大小为11 b N h+5 b N^{2} a

2.1.2.2 MLP块的中间激活:19 b N h

MLP块的计算公式如下:x=f_{\text {gelu }}\left(x_{\text {out }} W_{1}\right) W_{2}+x_{\text {out }},最终对于MLP块,需要保存的中间激活值为 19 b N h

具体怎么计算得来的呢?

  1. 第一个线性层需要保存其输入,占用显存大小为 2 b N h
  2. 激活函数需要保存其输入,占用显存大小为8 b N h
  3. 第二个线性层需要保存其输入,占用显存大小为 8 b N h
  4. 最后有一个dropout操作,需要保存mask矩阵,占用显存大小为\text { bNh }
2.1.2.3 两个layer norm需要保存的中间激活:4bNh

另外,self-attention块和MLP块分别对应了一个layer normalization。每个layer norm需要保存其输入,大小为2bNh2个layer norm需要保存的中间激活为 4bNh

综上,每个transformer层需要保存的中间激活占用显存大小为34 b N h+5 b N^{2} a

对于 l 层transformer模型,还有embedding层、最后的输出层。embedding层不需要中间激活。总的而言,当隐藏维度 h 比较大,层数 l 较深时,这部分的中间激活是很少的,可以忽略

因此,对于 l 层transformer模型,中间激活占用的显存大小可以近似为\left(34 b N h+5 b N^{2} a\right) * l  「更多分析见此文《分析transformer模型的参数量、计算量、中间激活、KV cache》

通过上面两小节的内容,可以看到,transformer模型的计算量和储存复杂度随着序列长度 N 呈二次方增长。这限制了大语言模型的最大序列长度 N​ 的大小

其次,GPT4将最大序列长度 N​ 扩大到了32K,Claude更是将最大序列长度 N​ 扩大到了100K,这些工作一定采用了一些优化方法来降低原生transformer的复杂度,那具体怎么优化呢?
我们知道,每个transformer层分为两部分:self-attention块和MLP块,但上面计算量中的 4bN^2h​项和中间激活中的5bN^2a​ 项都是self-attention块产生的,与MLP块无关

如此,FlashAttention提出了一种加速计算、节省显存和IO感知的精确注意力,可以有效地缓解上述问题

Meta推出的开源大模型LLaMA,阿联酋推出的开源大模型Falcon都使用了Flash Attention来加速计算和节省显存。目前,Flash Attention已经集成到了pytorch2.0中,另外triton、xformer等开源框架也进行了整合实现

2.1.3 分析GPU的内存分析图:计算的瓶颈是显存访问

通过上文可知,transformer的核心组件self-attention块的计算复杂度和空间复杂度是序列长度 N的二次方

  1. 对于self-attention块,除了大矩阵乘法是计算受限的,其他操作(计算softmax、dropout、mask)都是内存受限的。
    尽管已经有许多近似注意力的方法尝试减少attention的计算和内存要求。例如,稀疏近似和低秩近似的方法,将计算复杂度降低到了序列长度的线性或亚线性
  2. 但这些近似注意力方法方法并没有得到广泛应用。因为这些方法过于关注FLOPs(浮点数计算次数)的减少,而忽略了IO读写的内存访问开销,导致这并没有效减少运行时间(wall-clock time)
    总之,在现代GPU中,计算速度已经远超过了显存访问速度,transformer中的大部分计算操作的瓶颈是显存访问。对于显存受限的操作,IO感知是非常重要的,因为显存读写占用了大部分的运行时间
  3. 而Flash Attention则是IO感知的,通过减少内存访问,来计算精确注意力,从而减少运行时间,实现计算加速

GPU的内存由多个不同大小和不同读写速度的内存组成。内存越小,读写速度越快。对于A100-40GB来说,内存分级图如下所示

ChatGLM2-6B的通透解析:从FlashAttention、Multi-Query Attention到GLM2的微调、源码解读_第2张图片

  • SRAM内存分布在108个流式多处理器上,每个处理器的大小为192K。合计为 192 * 108 K B=20,736 K M=20 M B​ 
  • 高带宽内存HBM(High Bandwidth Memory),也就是我们常说的显存,大小为40GB。SRAM的读写速度为19TB/s,而HBM的读写速度只有1.5TB/s,不到SRAM的1/10

所以,上面讲到计算注意力的主要瓶颈是显存访问,因此减少对HBM的读写次数,有效利用更高速的SRAM来进行计算是非常重要的,而GPU有大量的线程来执行某个操作,称为kernel。GPU执行操作的典型方式分为三步:

  1. 每个kernel将输入数据从低速的HBM中加载到高速的SRAM中
  2. 在SRAM中,进行计算
  3. 计算完毕后,将计算结果从SRAM中写入到HBM中

而对于性能受限于内存带宽的操作,进行加速的常用方式就是kernel融合。kernel融合的基本思想是:避免反复执行“从HBM中读取输入数据,SRAM执行计算,最后将计算结果写入到HBM中”,将多个操作融合成一个操作,减少读写HBM的次数(需要注意的是,模型训练通常会影响到算子融合的效果,因为为了后向传递计算梯度,通常需要将某些中间结果写入到HBM中)

2.1.4 safe softmax

继续行文之前,先补充两个背景知识,一个是safe softmax,一个是Standard Attention

对于第一个背景知识:safe softmax而言

  1. 考虑到向量 \left[x_{1}, x_{2}, \cdots, x_{d}\right]​,原生softmax的计算过程如下:
    \operatorname{softmax}\left(x_{i}\right)=\frac{e^{x_{i}}}{\sum_{j=1}^{d} e^{x_{j}}}
  2. 在实际硬件中,浮点数表示的范围是有限的
    对于float32和bfloat16来说,当 x \geq 89​ 时,e^x​就会变成inf,发生数据上溢的问题
    故为了避免发生数值溢出的问题,保证数值稳定性,计算时通常会“减去最大值”,称为“safe softmax”

    即现在所有的深度学习框架中都采用了“safe softmax”这种计算方式
    m=\max _{i}\left(x_{i}\right) ; \quad \operatorname{softmax}\left(x_{i}\right)=\frac{e^{x_{i}-m}}{\sum_{j=1}^{d} e^{x_{j}-m}}
  3. 在训练语言模型时,通常会采用交叉熵损失函数。交叉熵损失函数等价于先执行log_softmax函数,再计算负对数似然函数
    且在计算log_softmax时,同样会执行“减去最大值”,这不仅可以避免数值溢出,提高数值稳定性,还可以加快计算速度
    \log \left(\operatorname{softmax}\left(x_{i}\right)\right)=\log \left(\frac{e^{x_{i}-m}}{\sum_{j=1}^{d} e^{x_{j}-m}}\right)=x_{i}-m-\log \left(\sum_{j=1}^{d} e^{x_{j}-m}\right)

2.2 前向传递:Standard Attention/Memory-efficient Attention/Flash Attention

2.2.1 Standard Attention

  1. 首先,transformer中注意力机制的计算过程为:
    \operatorname{Attention}(Q, K, V)=\operatorname{softmax}\left(\frac{Q K^{\top}}{\sqrt{d}}\right) V
    其中, Q, K, V \in R^{N \times d}​,其中 N​ 是序列长度, d​ 是每个注意力头的维度,输出可以记为 O \in R^{N \times d}​ 
  2. 上面的式子可以拆解为:​
    S=Q K^{\top} \in R^{N \times N}, P=\operatorname{softmax}(S) \in R^{N \times N}, O=P V \in R^{N \times d}
    在标准注意力实现中, S, P \in R^{N \times N}​ 都要写回到HBM中,占用了 O\left(N^{2}\right)​的内存,通常 N \gg d
    例如,对于GPT2, N = 1024​,d = 64​ ;对于GPT3,N = 1028​,d = 128​ 

    总之,注意力矩阵P, S​ 需要的内存 O\left(N^{2}\right)​远大于Q, K, V, O​ 所需要的内存O(N d)
    相当于,self-attention中,大部分操作都是内存受限的逐点运算,例如,对 S​ 的mask操作、 S​ 的softmax操作、对 P​的dropout操作,这些逐点操作的性能是受限于内存带宽的,会减慢运行时间
  3. 下图展示了标准注意力的实现过程
    ChatGLM2-6B的通透解析:从FlashAttention、Multi-Query Attention到GLM2的微调、源码解读_第3张图片
    标准注意力实现存在两个问题:
    1. 显存占用多,过程中由于实例化了完整的注意力矩阵P, S \in R^{N \times N}​ ,导致了 O\left(N^{2}\right)​ 的内存要求
    2. HBM读写次数多,减慢了运行时间(wall- clock time)

    接下来的Memory-efficient Attention、Flash Attention,便是要分别解决上述这两个问题

2.2.2 Memory-efficient Attention:把显存复杂度从平方降低到线性,但HBM访问次数仍是平方

在注意力计算过程中,节省显存的主要挑战是softmax与K,V的列是耦合的。其方法是单独计算softmax的归一化因子,来实现解耦

  1. 为了简化分析,忽略计算softmax时“减去最大值”的步骤
    记 Q 的第 i 列为 q_{i} \in R^{d} , K 的第 j 列为 K_{j} \in R^{d},有S_{i j}=q_{i}^{\top} k_{j} \in R
    定义softmax的归一化因子为:
    L_{i}=\sum_{j} e^{q_{i}^{\top} k_{j}} \in R
  2. 记 v_{j} \in R^{d} 为 V的第 j 个列向量,则输出 O 的第 i个列向量 o_i 为:
    o_{i}=P_{i:} V=\sum_{j} P_{i j} v_{j}=\sum_{j} \frac{e^{q_{i}^{\top} k_{j}}}{L_{i}} v_{j}
  3. 在计算得到归一化因子L_i 后,就可以通过反复累加 \frac{e^{q_{i}^{\top} k_{j}}}{L_{i}} v_{j}来得到 o_i

如此,节省内存(memory-efficient)的注意力机制,改变了计算顺序,相比于Standard Attention,节省显存的注意力机制将显存复杂度从 O(N^2) 降低到了O(N) 

这种方法在《Online normalizer calculation for softmax》和《Self-attention Does Not Need O\left(n^{2}\right) Memory》中已经使用过,称其为“lazy softmax”,这种方法避免了实例化完整的注意力矩阵 S,P,从而达到了节省显存的目的。然而HBM访问次数仍然是 O(N^2)的,因此运行时间并没有减少

2.2.3 Flash Attention:避免频繁地从HBM中读写数据

// 待更..


第三部分 多查询注意力(Muti Query Attention):各自Query矩阵,但共享Key 和 Value 矩阵

多查询注意力(Muti Query Attention)是 19 年Google一研究者提出的一种新的 Attention 机制(对应论文为:Fast Transformer Decoding: One Write-Head is All You Need、这是其解读之一),其能够在保证模型效果的同时加快 decoder 生成 token 的速度

3.1 Multi-Head Attention、Grouped-Query Attention、Muti Query Attention的区别

那其与17年 Google提出的transformer中多头注意力机制(简称MHA)有啥本质区别呢?有意思的是,区别在于:

ChatGLM2-6B的通透解析:从FlashAttention、Multi-Query Attention到GLM2的微调、源码解读_第4张图片

  • 我们知道MHA的每个头都各自有一份不同的Key、Query、Value矩阵
  • 而MQA 让所有的头之间 共享 同一份 Key 和 Value 矩阵,每个头只单独保留了一份 Query 参数,从而大大减少 Key 和 Value 矩阵的参数量
    总之,MQA 实际上是将 head 中的 key 和 value 矩阵抽出来单独存为一份共享参数,而 query 则是依旧保留在原来的 head 中,每个 head 有一份自己独有的 query 参数

下图对比了多头注意力(Multi-Head Attention)、LLaMA2中分组查询注意力(Grouped-Query Attention)、多查询注意力(Muti Query Attention)的差别

ChatGLM2-6B的通透解析:从FlashAttention、Multi-Query Attention到GLM2的微调、源码解读_第5张图片

​3.2 MHA 和 MQA在代码实现上的差异

总之,MHA 和 MQA 之间的区别只在于建立 Wqkv Layer 上

# Multi Head Attention
self.Wqkv = nn.Linear(                        # 【关键】Multi-Head Attention 的创建方法
    self.d_model, 
    3 * self.d_model,                         # 有 query, key, value 3 个矩阵, 所以是 3 * d_model
    device=device
)

query, key, value = qkv.chunk(                # 【关键】每个 tensor 都是 (1, 512, 768)
    3, 
    dim=2
)


# Multi Query Attention
self.Wqkv = nn.Linear(                                # 【关键】Multi-Query Attention 的创建方法
    d_model,
    d_model + 2 * self.head_dim,                      # 只创建 query 的 head 向量,所以只有 1 个 d_model
    device=device,                                    # 而 key 和 value 不再具备单独的头向量
)

query, key, value = qkv.split(                        # query -> (1, 512, 768)
    [self.d_model, self.head_dim, self.head_dim],     # key   -> (1, 512, 96)
    dim=2                                             # value -> (1, 512, 96)
)

对比上面的代码,你可以发现

  • 在 MHA 中,query, key, value 每个向量均有 768 维度
  • 而在 MQA 中,只有 query 是 768 维,而 key 和 value 均只剩下 96 维了,恰好是 1 个 head_dim 的维度

因此,可以确认:在 MQA 中,除了 query 向量还保存着 8 个头,key 和 value 向量都只剩 1 个「公共头」了,这也正好印证了论文中所说的「所有 head 之间共享一份 key 和 value 的参数」

剩下的问题就是如何将这 1 份参数同时让 8 个头都使用,代码里使用矩阵乘法 matmul 来广播,使得每个头都乘以这同一个 tensor,以此来实现参数共享:

def scaled_multihead_dot_product_attention(
        query,
        key,
        value,
        n_heads,
        multiquery=False,
    ):
    q = rearrange(query, 'b s (h d) -> b h s d', h=n_heads)         # (1, 512, 768) -> (1, 8, 512, 96)
    kv_n_heads = 1 if multiquery else n_heads
    k = rearrange(key, 'b s (h d) -> b h d s', h=kv_n_heads)        # (1, 512, 768) -> (1, 8, 96, 512) if not multiquery 
                                                                    # (1, 512, 96) -> (1, 1, 96, 512)  if multiquery
    v = rearrange(value, 'b s (h d) -> b h s d', h=kv_n_heads)      # (1, 512, 768) -> (1, 8, 512, 96) if not multiquery 
                                                                    # (1, 512, 96) -> (1, 1, 512, 96)  if multiquery
    
    attn_weight = q.matmul(k) * softmax_scale                       # (1, 8, 512, 512)
    attn_weight = torch.softmax(attn_weight, dim=-1)                # (1, 8, 512, 512)

    out = attn_weight.matmul(v)                                     # (1, 8, 512, 512) * (1, 1, 512, 96) = (1, 8, 512, 96)
    out = rearrange(out, 'b h s d -> b s (h d)')                    # (1, 512, 768)

    return out, attn_weight, past_key_value

第四部分 模型的使用/部署、微调

4.1 模型的使用/部署

  1. 首先需要下载本仓库:
    git clone https://github.com/THUDM/ChatGLM2-6B
    cd ChatGLM2-6B
  2. 然后使用 pip 安装依赖:
    pip install -r requirements.txt
    
    其中 transformers 库版本推荐为 4.30.2torch 推荐使用 2.0 及以上的版本,以获得最佳的推理性能
  3. 代码调用
    可以通过如下代码调用 ChatGLM2-6B 模型来生成对话:
    >>> from transformers import AutoTokenizer, AutoModel
    >>> tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
    >>> model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True, device='cuda')
    >>> model = model.eval()
    >>> response, history = model.chat(tokenizer, "你好", history=[])
    >>> print(response)
    你好!我是人工智能助手 ChatGLM2-6B,很高兴见到你,欢迎问我任何问题。
    >>> response, history = model.chat(tokenizer, "晚上睡不着应该怎么办", history=history)
    >>> print(response)
    晚上睡不着可能会让你感到焦虑或不舒服,但以下是一些可以帮助你入睡的方法:
    1. 制定规律的睡眠时间表:保持规律的睡眠时间表可以帮助你建立健康的睡眠习惯,使你更容易入睡。尽量在每天的相同时间上床,并在同一时间起床。
    2. 创造一个舒适的睡眠环境:确保睡眠环境舒适,安静,黑暗且温度适宜。可以使用舒适的床上用品,并保持房间通风。
    3. 放松身心:在睡前做些放松的活动,例如泡个热水澡,听些轻柔的音乐,阅读一些有趣的书籍等,有助于缓解紧张和焦虑,使你更容易入睡。
    4. 避免饮用含有咖啡因的饮料:咖啡因是一种刺激性物质,会影响你的睡眠质量。尽量避免在睡前饮用含有咖啡因的饮料,例如咖啡,茶和可乐。
    5. 避免在床上做与睡眠无关的事情:在床上做些与睡眠无关的事情,例如看电影,玩游戏或工作等,可能会干扰你的睡眠。
    6. 尝试呼吸技巧:深呼吸是一种放松技巧,可以帮助你缓解紧张和焦虑,使你更容易入睡。试着慢慢吸气,保持几秒钟,然后缓慢呼气。

    如果这些方法无法帮助你入睡,你可以考虑咨询医生或睡眠专家,寻求进一步的建议

从本地加载模型

以上代码会由 transformers 自动下载模型实现和参数

完整的模型实现在 Hugging Face Hub。如果你的网络环境较差,下载模型参数可能会花费较长时间甚至失败。此时可以先将模型下载到本地,然后从本地加载。

从 Hugging Face Hub 下载模型需要先安装Git LFS,然后运行

git clone https://huggingface.co/THUDM/chatglm2-6b

如果你从 Hugging Face Hub 上下载 checkpoint 的速度较慢,可以只下载模型实现

GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/THUDM/chatglm2-6b

然后从这里手动下载模型参数文件,并将下载的文件替换到本地的 chatglm2-6b 目录下

将模型下载到本地之后,将以上代码中的 THUDM/chatglm2-6b 替换为你本地的 chatglm2-6b 文件夹的路径,即可从本地加载模型。

模型的实现仍然处在变动中。如果希望固定使用的模型实现以保证兼容性,可以在 from_pretrained 的调用中增加 revision="v1.0" 参数。v1.0 是当前最新的版本号,完整的版本列表参见 Change Log

最后,可以通过以下命令启动基于 Gradio 的网页版 demo:

python web_demo.py

在这里插入图片描述

4.2 基于 P-Tuning v2 的微调(官方)

P-Tuning v2 将需要微调的参数量减少到原来的 0.1%,再通过模型量化、Gradient Checkpoint 等方法,最低只需要 7GB 显存即可运行(当然,我司杜老师也会在七月类ChatGPT微调实战课上录一个ChatGLM2-6B的微调视频)

  1. 环境配置
    在原chatglm-6b的环境中安装以下依赖
    pip install rouge_chinese nltk jieba datasets
  2. 微调数据准备
    ADGEN 数据集任务为根据输入(content)生成一段广告词(summary)
    { “content”: “类型#上衣版型#宽松版型#显瘦图案#线条衣样式#衬衫衣袖型#泡泡袖衣款式#抽绳”, “summary”:
    “这件衬衫的款式非常的宽松,利落的线条可以很好的隐藏身材上的小缺点,穿在身上有着很好的显瘦效果。领口装饰了一个可爱的抽绳,漂亮的绳结展现出了十足的个性,配合时尚的泡泡袖型,尽显女性甜美可爱的气息。”
    }
    从 Google Drive 或者 Tsinghua Cloud 下载处理好的 ADGEN 数据集,将解压后的 AdvertiseGen 目录放到本 ptuning 目录下即可
  3. 微调
    修改train.sh文件
    去掉最后的 --quantization_bit 4( 去掉后为FP16 精度加载)
    修改模型路径,THUDM/chatglm-6b修改为/data/sim_chatgpt/chatglm2-6b
    目前专业级GPU Tesla P100也不支持INT4或8量化

    执行train.sh文件
    bash train.sh
    如遇报错:
    wandb.errors.UsageError: api_key not configured (no-tty). call wandb.login(k…
    解决方法:
    在main.py文件中加入下面两行,禁用wandb即可
  4. import os
    os.environ["WANDB_DISABLED"] = "true"
    其中,train.sh 中的 PRE_SEQ_LEN 和 LR 分别是 soft prompt 长度和训练的学习率,可以进行调节以取得最佳的效果。

微调过程显存使用情况如下:

在这里插入图片描述

 微调完成后,在./output/adgen-chatglm2-6b-pt-128-2e-2 下回生成微调好的模型文件。

我们可以对比下微调前后的效果
以命令行 Demo为例,只需修改ptuning路径下web_demo.sh中的模型路径为/data/sim_chatgpt/chatglm2-6b,运行 web_demo.py即可:

bash web_demo.sh

Input:
类型#上衣材质#牛仔布颜色#白色风格#简约图案#刺绣衣样式#外套衣款式#破洞
Label:
简约而不简单的牛仔外套,白色的衣身十分百搭。衣身多处有做旧破洞设计,打破单调乏味,增加一丝造型看点。衣身后背处有趣味刺绣装饰,丰富层次感,彰显别样时尚。

Output[微调前]:

在这里插入图片描述

 Output[微调后]:

乱码,待解决......

 // 待更

你可能感兴趣的:(论文,代码,实战,ChatGLM2-6B,FlashAttention,Multi-Query注意力)