Minimind项目源码详细解析（2）Attention机制

Attention机制代码详细解析既然大家开始看LLM相关了内容了那么大家一定对attention机制有了一定的了解在此我就不对attention机制进行过于细致的讲解了在此主要讲解一些具体实现和一些扩展attention机制简要讲解在大语言模型里我们本质上是让 LLM 学习 token 与 token 之间的依赖关系而提取这种关系的核心正是 Attention 注意力机制。Attention 的计算主要依靠三个关键矩阵 —— 也就是大家熟知的 Q、K、V 矩阵QQuery查询向量代表当前 token 想要 “查找” 什么信息KKey键向量代表每个 token 自身携带的 “特征标识”VValue值向量代表每个 token 真正要传递出去的信息内容具体过程可以简单理解为用当前 token 的 Q去和其他 token 的 K 做点积运算算出它们之间的相似度 / 关联强度。在 decoder 里为了保证生成顺序合法我们只会让当前 token 关注它之前出现的 token而看不到未来的 token。再用这个关联权重去加权求和所有对应的 V最终得到的结果就是当前 token 从所有相关 token 那里收集到的有效信息。这就是注意力机制最核心的原理。Attention(Q,K,V)softmax(QKTdk)V\text{Attention}(Q, K, V) \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)VAttention(Q,K,V)softmax(dk​​QKT​)V这里的softmax是为了将所有值变成一个概率分布没学过可以在B站上搜一下讲解非常多。话不多说直接上代码我会在代码中进行比较详细的讲解。attention的具体实现代码不过这里用到的其实是多头注意力Multi-Head Attention。肯定有人会问什么是 “头”其实非常直观我们把 token 映射成 Q、K、V 时用到的一组线性变换矩阵 W就可以看作一个注意力头。而多头注意力就是同时使用多组这样的独立矩阵也就是多个 “头”。这么多头有什么用举个简单例子假设一个 token 向量是 16 维我们用 4 个注意力头就可以把这 16 维平均分成 4 份每份 4 维让每个注意力头只负责其中一小段维度。这样一来不同的头可以关注不同的语义信息有的头关注语法结构有的头关注语义关联有的头捕捉局部依赖最后再把所有头的结果拼接起来模型就能更细、更全面、更丰富地提取 token 之间的语义关系。此外这里为了节省参数我们将KV头进行复用GQA分组查询注意力怎么理解呢还是假设一个token有16个维度我们要4个头也就是4个Q头4个K头4个V头。也就是将这16个维度分成了4个小组。但是这样太耗费参数了怎么办呢我们可以把这4个小组再组合一下分成2个大组不同小组之间使用不同的Q不同大组之间使用不同的KV这样我们就能减少参数的使用了具体可以看一下代码的注释classAttention(nn.Module):def__init__(self,args:MiniMindConfig):super().__init__()self.num_key_value_headsargs.num_attention_headsifargs.num_key_value_headsisNoneelseargs.num_key_value_headsassertargs.num_attention_heads%self.num_key_value_heads0self.n_local_headsargs.num_attention_heads#我们需要的Q注意力头数self.n_local_kv_headsself.num_key_value_heads#我们需要的K头和V头的数量self.n_repself.n_local_heads//self.n_local_kv_heads#每个大组中的小组数量self.head_dimargs.hidden_size//args.num_attention_heads#每个头需要关注的维度数量self.q_projnn.Linear(args.hidden_size,args.num_attention_heads*self.head_dim,biasFalse)#Wq这里我们将多个头拼接在一起下面同理。self.k_projnn.Linear(args.hidden_size,self.num_key_value_heads*self.head_dim,biasFalse)self.v_projnn.Linear(args.hidden_size,self.num_key_value_heads*self.head_dim,biasFalse)self.o_projnn.Linear(args.num_attention_heads*self.head_dim,args.hidden_size,biasFalse)#输出矩阵用来将计算出来的注意力分数进行整合输出。self.attn_dropoutnn.Dropout(args.dropout)self.resid_dropoutnn.Dropout(args.dropout)0self.dropoutargs.dropout self.flashhasattr(torch.nn.functional,scaled_dot_product_attention)andargs.flash_attn#flash Attention 建议自己学习一下这里不细讲# print(WARNING: using slow attention. Flash Attention requires PyTorch 2.0)defforward(self,x:torch.Tensor,position_embeddings:Tuple[torch.Tensor,torch.Tensor],# 修改为接收cos和sinpast_key_value:Optional[Tuple[torch.Tensor,torch.Tensor]]None,use_cacheFalse,attention_mask:Optional[torch.Tensor]None):bsz,seq_len,_x.shape xq,xk,xvself.q_proj(x),self.k_proj(x),self.v_proj(x)xqxq.view(bsz,seq_len,self.n_local_heads,self.head_dim)xkxk.view(bsz,seq_len,self.n_local_kv_heads,self.head_dim)xvxv.view(bsz,seq_len,self.n_local_kv_heads,self.head_dim)#重塑一下尺寸cos,sinposition_embeddings xq,xkapply_rotary_pos_emb(xq,xk,cos,sin)#应用RoPE旋转编码具体可以看我讲RoPE的文章# kv_cache实现ifpast_key_valueisnotNone:xktorch.cat([past_key_value[0],xk],dim1)xvtorch.cat([past_key_value[1],xv],dim1)past_kv(xk,xv)ifuse_cacheelseNone#为什么要有KV cache呢考虑一下LLM推理时的原理LLM的推理是自回归的也就是每次会生成一个token然后再根据该token返回查询与以前所有token的关系从而生成下一个token那么如果我们不保存之前token的xk和xv我们在每次查询的时候都要再算一遍那如果把之前的token的k和v全部保存下来我们是不是就不用再算了也就大大加快了推理速度。xq,xk,xv(xq.transpose(1,2),repeat_kv(xk,self.n_rep).transpose(1,2),repeat_kv(xv,self.n_rep).transpose(1,2))为什么要这样呢自己用笔写一下如果不换维度那么最后乘出来后两维是head_dim * head_dim的因为我们要的其实是token和token之间的关系所以我们要的实际上是seq_len * seq_len我们通过维度的变换就能够做到这样的结果了。ifself.flashand(seq_len1)and(past_key_valueisNone)and(attention_maskisNoneortorch.all(attention_mask1)):outputF.scaled_dot_product_attention(xq,xk,xv,dropout_pself.dropoutifself.trainingelse0.0,is_causalTrue)#flash_attn 自己学一下我也不太会else:scores(xq xk.transpose(-2,-1))/math.sqrt(self.head_dim)公式除以根号是为了防止数值过大如果不这样的话假设有一个很大的数字1e9那么他经过softmax后的权重就是1,会导致梯度消失scores[:,:,:,-seq_len:]torch.triu(torch.full((seq_len,seq_len),float(-inf),devicescores.device),diagonal1)因果注意力掩码 会生成一个下三角矩阵从而让每个token只能关注到自己前面的tokenifattention_maskisnotNone:extended_attention_maskattention_mask.unsqueeze(1).unsqueeze(2)extended_attention_mask(1.0-extended_attention_mask)*-1e9scoresscoresextended_attention_maskattention_mask可以忽略掉某些无用的token比如为了将一句话填充为seq_len长度的padding(填充)他只是为了充长度的没有实际作用我们就给他去除scoresF.softmax(scores.float(),dim-1).type_as(xq)scoresself.attn_dropout(scores)outputscores xv outputoutput.transpose(1,2).reshape(bsz,seq_len,-1)outputself.resid_dropout(self.o_proj(output))#复原并输出returnoutput,past_kv