Qwen2源码阅读——核心代码跟读

本文链接：https://blog.csdn.net/m0_59235245/article/details/143806903

我们继续使用上一节使用的样例文本：

["你好啊", "简单的机器学习是为了让机器学习变得更简单而存在的"]

这个样例产生的tokens结果为：

{'input_ids': tensor([[108386, 103924, 151643, 151643, 151643, 151643, 151643, 151643, 151643,            151643],           [105172, 102182, 100134, 104802,  99258, 102182, 100134, 112606, 100405,             68536]]), 'attention_mask': tensor([[1, 1, 0, 0, 0, 0, 0, 0, 0, 0],           [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}

即得到的shape为：[2, 10]

由上一节print(model)内容：

Qwen2ForCausalLM(     (model): Qwen2Model(       (embed_tokens): Embedding(151936, 896)       (layers): ModuleList(         (0-3): 4 x Qwen2DecoderLayer(           (self_attn): Qwen2SdpaAttention(             (q_proj): Linear(in_features=896, out_features=896, bias=True)             (k_proj): Linear(in_features=896, out_features=128, bias=True)             (v_proj): Linear(in_features=896, out_features=128, bias=True)             (o_proj): Linear(in_features=896, out_features=896, bias=False)             (rotary_emb): Qwen2RotaryEmbedding()           )           (mlp): Qwen2MLP(             (gate_proj): Linear(in_features=896, out_features=4864, bias=False)             (up_proj): Linear(in_features=896, out_features=4864, bias=False)             (down_proj): Linear(in_features=4864, out_features=896, bias=False)             (act_fn): SiLU()           )           (input_layernorm): Qwen2RMSNorm((896,), eps=1e-06)           (post_attention_layernorm): Qwen2RMSNorm((896,), eps=1e-06)         )       )       (norm): Qwen2RMSNorm((896,), eps=1e-06)       (rotary_emb): Qwen2RotaryEmbedding()     )     (lm_head): Linear(in_features=896, out_features=151936, bias=False)   )

我们看到Qwen2Model主体是由embed_tokens + 4*(self_attn + mlp + input_layernorm + post_attention_layernorm) + norm + rotary_emb组成的。

详情

embed_tokens层：

embed_tokens就是我们熟悉的nn.Embedding初始化得到的层。即：

self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)

其中：config.vocab_size=151936, config.hidden_size=896, self.padding_idx=config.pad_token_id=None。当self.padding_idx为None时候，默认取值就为0。

对于shape为[2, 10]的输入，经过embed_tokens层，可获得shape为[2, 10, 896]，记为inputs_embeds。

cache_position和position_ids：

因为cache_position和position_ids都是None(注：正对本样例而言)，所以cache_position直接是通过传进来的序列长度计算得到的，即为：tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])，position_ids为：tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]])

causal_mask：

causal_mask是由方法self._update_causal_mask产生的，它将产生四维的矩阵数据，shape为[2, 1, 10, 10]，即[bs, 1, seq_len, seq_len]，我们这里展示一下最后两维的数据，分别是causal_mask[0,0][:5, :5]和causal_mask[1,0][:5, :5]，如下：

# causal_mask[0,0][:5, :5]   tensor([[ 0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],           [ 0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38],           [ 0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38],           [ 0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38],           [ 0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38]])   # causal_mask[1,0][:5, :5]   tensor([[ 0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38, -3.4028e+38],           [ 0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38, -3.4028e+38],           [ 0.0000e+00,  0.0000e+00,  0.0000e+00, -3.4028e+38, -3.4028e+38],           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00, -3.4028e+38],           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00]])

可以看出下三角全为0，causal_mask[0,0][:5, :5]不全为0的原因是由attention_mask引起的，即pad部分是不用去计算的。

rotary_emb：

rotary_emb层只计算一次，然后运用到后面的各层，这一层是没有参数的，不参与训练。使用旋转位置编码最直接的好处有：

可以使用绝对位置编码来表示相对位置编码；
计算量是线性的；
通过配置，可以实现一定的长度往外延拓能力；

rotary_emb主要用于计算cos和sin的值，即计算公式：

`class Qwen2RotaryEmbedding(nn.Module):       def __init__(           self,           dim=None,           max_position_embeddings=2048,           base=10000,           device=None,           scaling_factor=1.0,           rope_type="default",           config: Optional[Qwen2Config] = None,       ):           super().__init__()           # TODO (joao): remove the `if` below, only used for BC           self.rope_kwargs = {}           if config is None:               logger.warning_once(                   "`Qwen2RotaryEmbedding` can now be fully parameterized by passing the model config through the "                   "`config` argument. All other arguments will be removed in v4.46"               )               self.rope_kwargs = {                   "rope_type": rope_type,                   "factor": scaling_factor,                   "dim": dim,                   "base": base,                   "max_position_embeddings": max_position_embeddings,               }               self.rope_type = rope_type               self.max_seq_len_cached = max_position_embeddings               self.original_max_seq_len = max_position_embeddings           else:               # BC: "rope_type" was originally "type"               if config.rope_scaling is not None:                   self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))               else:                   self.rope_type = "default"               self.max_seq_len_cached = config.max_position_embeddings               self.original_max_seq_len = config.max_position_embeddings              self.config = config           self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]       # 这里会获取到一个shape为[config.hidden_size // config.num_attention_heads//2]的inv_freq           # 因为是多头，所以实际上每个头的维度是config.hidden_size // config.num_attention_heads           # 再除以2是由公式确定的，具体看下面的公式矩阵.           inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, **self.rope_kwargs)           self.register_buffer("inv_freq", inv_freq, persistent=False)           self.original_inv_freq = self.inv_freq          def _dynamic_frequency_update(self, position_ids, device):           """           dynamic RoPE layers should recompute `inv_freq` in the following situations:           1 - growing beyond the cached sequence length (allow scaling)           2 - the current sequence length is in the original scale (avoid losing precision with small sequences)           """           seq_len = torch.max(position_ids) + 1           if seq_len > self.max_seq_len_cached:  # growth               inv_freq, self.attention_scaling = self.rope_init_fn(                   self.config, device, seq_len=seq_len, **self.rope_kwargs               )               self.register_buffer("inv_freq", inv_freq, persistent=False)  # TODO joao: may break with compilation               self.max_seq_len_cached = seq_len              if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len:  # reset               self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)               self.max_seq_len_cached = self.original_max_seq_len          @torch.no_grad()       def forward(self, x, position_ids):          """          x: shape为[bs, seq_len, hidden_size]          position_ids: shape为[1, seq_len]          """           if "dynamic" in self.rope_type:               self._dynamic_frequency_update(position_ids, device=x.device)              # Core RoPE block           # self.inv_freq本身的shape为[32], 经过下面的操作可获得[1, 32, 1]           inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)           position_ids_expanded = position_ids[:, None, :].float()  # shape为[1, 1, seq_len]           # Force float32 (see https://github.com/huggingface/transformers/pull/29285)           device_type = x.device.type           device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"           with torch.autocast(device_type=device_type, enabled=False):              # 经过[1, 32, 1]和[1, 1, seq_len]矩阵乘法之后可以得到[1, 32, seq_len]               # 再经过变换可以得到[1, seq_len, 32]               freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)               emb = torch.cat((freqs, freqs), dim=-1)  # [1, seq_len, 64]               cos = emb.cos()               sin = emb.sin()              # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention           cos = cos * self.attention_scaling           sin = sin * self.attention_scaling              return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)   `

原位置编码公式：

代码里面得到的是：

注：更多文档请参考Transformer升级之路：2、博采众长的旋转式位置编码[1]

经过rotary_emb可以得到position_embeddings，它是一个元组，分别是(cos, sin)对应的值，它们的shape都是[1, seq_len, 64]。

**self_attn: **

这里使用Qwen2SdpaAttention来计算self_attention，下面我们仔细介绍一下这个模块。

首先是Qwen2SdpaAttention继承自Qwen2Attention，然后修改了其forward方法。而Qwen2Attention初始化方案主要初始化了4个可训练参数权重，分别是self.q_proj、self.k_proj、self.v_proj、self.o_proj，如下代码：

self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=True)   self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)   self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=True)   self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)

self.hidden_size=config.hidden_size=896；
self.num_heads=config.num_attention_heads=14；
self.head_dim=self.hidden_size // self.num_heads=64；
self.num_key_value_heads=config.num_key_value_heads=2；
注意这里的q, k, v偏置全部设为了True，即bias=True；

接着我们看一下Qwen2Attention中的forward部分：

`def forward(       self,       hidden_states: torch.Tensor,       attention_mask: Optional[torch.Tensor] = None,       position_ids: Optional[torch.LongTensor] = None,       past_key_value: Optional[Cache] = None,       output_attentions: bool = False,       use_cache: bool = False,       cache_position: Optional[torch.LongTensor] = None,       position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46   ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:       """       参数shape说明：        hidden_states: [bs, seq_len, hidden_size]        attention_mask: [bs, 1, seq_len, seq_len]        position_ids: [1, seq_len]        cache_position: [seq_len]        position_embeddings: 元组数据，即(cos, sin)，shape都是[1, seq_len, self.head_dim]       """       bsz, q_len, _ = hidden_states.size()            # [bs, seq_len, hidden_size]=[bs, seq_len, 896]       query_states = self.q_proj(hidden_states)       # [bs, seq_len, self.num_key_value_heads * self.head_dim]=[bs, seq_len, 128]        key_states = self.k_proj(hidden_states)       # [bs, seq_len, self.num_key_value_heads * self.head_dim]=[bs, seq_len, 128]       value_states = self.v_proj(hidden_states)     # [bs, self.num_heads, seq_len, self.head_dim]=[bs, 14, sql_len, 64]       query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)       # [bs, self.num_key_value_heads, seq_len, self.head_dim] = [bs, 2, sql_len, 64]       key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)       # [bs, self.num_key_value_heads, seq_len, self.head_dim] = [bs, 2, sql_len, 64]       value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)          if position_embeddings is None:           logger.warning_once(               "The attention layers in this model are transitioning from computing the RoPE embeddings internally "               "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "               "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "               "removed and `position_embeddings` will be mandatory."           )           cos, sin = self.rotary_emb(value_states, position_ids)       else:          # cos: [1, seq_len, self.head_dim]=[1, seq_len, 64]           # sin: [1, seq_len, self.head_dim]=[1, seq_len, 64]           cos, sin = position_embeddings       # 针对query_states和key_states运用旋转位置编码，即使用下面的公式。       # 得到的shape为[bs, self.num_heads, seq_len, self.head_dim]=[bs, 14, seq_len, 64]       # 和 [bs, self.num_key_value_heads, seq_len, self.head_dim] = [bs, 2, seq_len, 64]       query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)          if past_key_value is not None:           cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}  # Specific to RoPE models           key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)          # repeat k/v heads if n_kv_heads < n_heads       # 这里的self.num_key_value_groups=self.num_heads // self.num_key_value_heads=7       # num_key_value_groups作用请看下面注释。              # 得到的shape为[bs, self.num_heads, seq_len, self.head_dim]=[bs, 14, seq_len, 64]       key_states = repeat_kv(key_states, self.num_key_value_groups)       value_states = repeat_kv(value_states, self.num_key_value_groups)            # 计算attn_weights，其shape是[bs, self.num_head, seq_len, seq_len]       attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)       if attention_mask is not None:  # no matter the length, we just slice it           causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]           attn_weights = attn_weights + causal_mask          # upcast attention to fp32       attn_weights = nn.functional.softmax(attn_weights, dim=-1,dtype=torch.float32).to(query_states.dtype)       attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)       # [bs, self.num_head, seq_len, seq_len]矩阵乘以[bs, self.num_head, seq_len, head_dim]       # 得到[bs, self.num_head, seq_len, head_dim]       attn_output = torch.matmul(attn_weights, value_states)          if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):           raise ValueError(               f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"               f" {attn_output.size()}"           )            # [bs, seq_len, self.num_head, head_dim]       attn_output = attn_output.transpose(1, 2).contiguous()       # [bs, seq_len, self.hidden_size]       attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)     # [bs, seq_len, self.hidden_size]       attn_output = self.o_proj(attn_output)          if not output_attentions:           attn_weights = None          return attn_output, attn_weights, past_key_value   `

注：

在Transformer模型中，num_key_value_groups 是分组查询注意力（Grouped-query attention, GQA）的一个概念。分组查询注意力是多头注意力的一种改进形式，它在保持一定数量的query头的同时，减少key和value头的数量，以此来提高计算效率。

具体来说，num_key_value_groups 表示的是将key和value头分组的数量。在标准的多头注意力中，每个query头都会与一个对应的key和value头配对。但在分组查询注意力中，多个query头会共享一组key和value头。这样做可以减少模型的参数数量和计算量，从而提高效率。

例如，如果我们有8个query头，但在分组查询注意力中，我们可能只有4个key-value组，那么num_key_value_groups 就是4。这意味着每两个query头会共享一个key和value头。在实际计算中，这组key和value会被复制（或者说广播）到与query头相同的数量，以便进行注意力权重的计算。

这种方法在保持多头注意力的优势的同时，减少了参数数量和计算复杂度，有助于提升模型的推理速度，尤其是在解码阶段。但是，它也需要仔细设计，以避免对模型性能产生负面影响。

在实际的代码实现中，num_key_value_groups 通常是通过将总的query头数除以key-value头数来计算的。例如，如果num_heads（query头的数量）是8，而num_key_value_heads（key-value头的数量）是4，那么num_key_value_groups 就是2，意味着每两个query头共享一个key-value头。

mlp：

对于这一层，其实直接看代码就可以理解了，没有特别难的内容在里面。所以这里就不进行介绍了。

总结

本篇文章主要集中在介绍数据在流转的过程中，各个矩阵的shape，通过shape的变化，来理解整个过程。其实如果对Bert本身有理解的情况下，整篇内容只需要理解旋转位置编码的实现以及分组查询注意力的理解就好了，其它内容和Bert相比，并没有本质的变化（除了attention_mask部分）。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述