Qwen2 源码解析

本文将深入分析千问模型推理代码，详细解析其内部实现机制, 简单阅读下官方推理代码:

from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "qwen/Qwen2___5-7B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
prompt = "Give me a short introduction to large language model."
messages = [
    {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
    {"role": "user", "content": prompt},
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
generated_ids = model.generate(
    **model_inputs,
     max_new_tokens=512,
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

批量推理

from transformers import AutoModelForCausalLM, AutoTokenizer
model_name = "qwen/Qwen2___5-7B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name,padding_side='left')  # 批量推理，要设置 左 padding
batch_prompt = [
    "Give me a short introduction to large language model.",
    "请介绍下您自己？"
         ]
batch_massages = []
for prompt in batch_prompt:
    messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": prompt}]  
    batch_massages.append(messages)

text = tokenizer.apply_chat_template(
    batch_massages,
    tokenize=False,
    add_generation_prompt=True,
)
model_inputs = tokenizer(text, padding=True, return_tensors="pt").to(model.device)
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=512,
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

• model  实例化粗探

  AutoModelForCausalLM.from_pretrained 方法执行过程如下:

  config, kwargs = AutoConfig.from_pretrained(
                  pretrained_model_name_or_path,
                  return_unused_kwargs=True,
                  trust_remote_code=trust_remote_code,
                  code_revision=code_revision,
                  _commit_hash=commit_hash,
                  **hub_kwargs,
                  **kwargs,
              )
  

  即根据config.json文件里字段 "model_type": "qwen2"   加载 实例类型:

  """
  type(config):
        <class 'transformers.models.qwen2.configuration_qwen2.Qwen2Config'>
  """        
  

  根据 config  类型匹配模型类:

  model_class = _get_model_class(config, cls._model_mapping)
  """
  type(model_class)
       <class 'transformers.models.qwen2.modeling_qwen2.Qwen2ForCausalLM'>
  """
  

  所以: AutoModelForCausalLM.from_pretrained 本质调用是 Qwen2ForCausalLM.from_pretrained方法.外层封装绕一大圈, 换个写法也会你会一目了然内在逻辑. 简化版本, 即:

  # 1.加载配置文件
  with open('config.json', "r", encoding="utf-8") as reader:
           text = reader.read()
  config_dict = json.loads(text)
  # 2.实例化配置类
  config = Qwen2Config(**config_dict)
  # 3.实例化模型
  model = Qwen2ForCausalLM.from_pretrained(pretrained_model_name_or_path, config=config, **kwarg)
  # 即使 不写初始化 config 实例代码, PreTrainedModel.from_pretrained也会自动完成. 
  model = Qwen2ForCausalLM.from_pretrained(pretrained_model_name_or_path, **kwarg)
  

1. 首先会在本地缓存中加载所需的模型文件。若文件不存在，则会自动从远程仓库下载。**(cached_file)**.

2. 该方法会读取模型的配置信息json文件，并据此创建一个配置对象，确保模型结构.

  - vocab_size ：Qwen2模型的词汇表大小。定义调用[`Qwen2Model`]时`inputs_ids`可以表示的不同token的数量。
  - hidden_size ：隐藏表示的维度。
  - intermediate_size ：MLP表示的维度。
  - num_hidden_layers ：Transformer编码器中的隐藏层数量。
  - num_attention_heads ：Transformer编码器中每个注意力层的注意力头数量。
  - num_key_value_heads ：应使用的key_value头的数量以实现分组查询注意力（Grouped Query Attention）。如果`num_key_value_heads=num_attention_heads`，模型将使用多头注意力（MHA），如果`num_key_value_heads=1`，模型将使用多查询注意力（MQA），否则使用GQA。在将多头检查点转换为GQA检查点时，每个分组的关键和值头应由该组内的所有原始头的平均池化构建。
  - hidden_act (`str` 或 `function`, *可选*, 默认为`"silu"`)：解码器中的非线性激活函数（函数或字符串）。
  - max_position_embeddings ：模型可能使用的最大序列长度。
  - initializer_range ：用于初始化所有权重矩阵的截断正态初始化器的标准差。
  - rms_norm_eps：rms规范化层使用的epsilon。
  - use_cache：模型是否应返回最后的key/values注意力（并非所有模型都使用）。仅当`config.is_decoder=True`时相关。
  - tie_word_embeddings ：模型输入和输出词嵌入是否应绑定。
  - rope_theta ：RoPE嵌入的基本周期。
  - use_sliding_window：是否使用滑动窗口注意力。
  - sliding_window ：滑动窗口注意力（SWA）窗口大小。如果未指定，将默认为`4096`。
  - max_window_layers ：使用SWA（滑动窗口注意力）的层数。底部图层使用SWA，而顶部图层使用全部注意力。
  - attention_dropout ：注意力概率的dropout比率。

• Qwen2RotaryEmbedding

  RoPE的基本思想是将输入的token embeddings视为复数向量，然后根据其位置信息进行旋转变换。具体来说，对于每个token，都会根据其位置生成一个旋转矩阵，然后将这个旋转矩阵应用于token embedding。这样，不同位置的token就会具有不同的旋转角度，从而实现了位置编码。在RoPE中，对于输入序列中的每个位置 (i)，其嵌入向量 (x_i) 会被一个旋转矩阵 (R_i) 进行变换：其中，旋转矩阵 (R_i) 是根据位置 (i) 计算得到的。具体来说，旋转矩阵 (R_i) 的形式通常是：

  import torch
  from  torch import nn
  class Qwen2RotaryEmbedding(nn.Module):
        def __init__(self,dim = 128 , base = 10000,max_max_position_embeddings = 2048, device = None):
            super().__init__()
            self.dim = dim
            self.base = base
            self.max_position_embeddings = max_position_embeddings
            inv_freq = 1.0 / (self.base ** (torch.arange(0,self.dim, 2,dtype = torch.int64).float().to(device))  / self.dim )   
            self.register_buffer("inv_freq", inv_freq, persistent=False)
            self._set_cos_sin_cache(
                   seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
                   )   
        def _set_cos_sin_cache(self, seq_len, device, dtype):
              self.max_seq_len_cached = seq_len
              t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
              freqs = torch.outer(t, self.inv_freq)
              # Different from paper, but it uses a different permutation in order to obtain the same calculation
              emb = torch.cat((freqs, freqs), dim=-1)
              self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
              self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)   
  

  上述旋转矩阵应用于k,v 向量，很巧妙，先看论文实现方法。

  向量元素(element) 两两分组旋转，即原本 (q_0,q_1)为旋转一组，代码实现中(q_0,   q_d/2) 设置一组，构思很巧：

  def rotata_half(x):
      x1 = x[..., :x.shape[-1] //2]
      x2 = x[..., x.shape[-1] // 2 :]
      retuwn torch.cat((-x2,x1),dim = -1)
  def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
      cos = cos[position_ids].unsqueeze(unsqueeze_dim)
      sin = sin[position_ids].unsqueeze(unsqueeze_dim)
      q_embed = (q * cos) + (rotate_half(q) * sin)
      k_embed = (k * cos) + (rotate_half(k) * sin)
      return q_embed, k_embed   
  

• Qwen2Attention

`Qwen2Model` 正在使用 `Qwen2SdpaAttention`，但 `torch.nn.functional.scaled_dot_product_attention` 不支持 `output_attentions=True`。因此回退到手动实现注意力机制，即外层封装`Qwen2SdpaAttention` 类，继承`Qwen2Attention`.

  class Qwen2Attention(nn.Module):
      def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
          super().__init__()
          self.config = config                              # 配置类
          self.layer_idx = layer_idx                        # 模型层id
          self.hidden_size = config.hidden_size                    # 隐藏层维度 28 * 128
              self.num_heads = config.num_attention_heads          # q 注意力头: 28
          self.head_dim = self.hidden_size // self.num_heads       # 注意力头向量：128 
          self.num_key_value_heads = config.num_key_value_heads    # k,v 注意力头: 4 ，广播 28 
          self.num_key_value_groups = self.num_heads // self.num_key_value_heads     # GQA 分 7 组, 每组 k,v 共享。
          self.max_position_embeddings = config.max_position_embeddings 
          self.rope_theta = config.rope_theta
          self.is_causal = True
          self.attention_dropout = config.attention_dropout
          if (self.head_dim * self.num_heads) != self.hidden_size:
              raise ValueError(
                  f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
                  f" and `num_heads`: {self.num_heads})."
              )
          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.rotary_emb = Qwen2RotaryEmbedding(
              self.head_dim,
              max_position_embeddings=self.max_position_embeddings,
              base=self.rope_theta,
          )
  ```

  在  `forward` 里 执行 `attention`  ，经`q_proj、k_proj、v_proj` 矩阵投影后，获取 `Q,K,V` 状态信息，分别是:`[batch_size, seq_len, hidden]`, 按`config`配置信息将 `hidden` 进行切头处理，`Q` ：`[ num_heads * head_dim]`;   `k,v` ：`[ num_key_value_heads* head_dim]`;  每个头认为是`channel`，转化后维度信息：

  ```latex
  Q,K,V:  [bs, channel, seq_len,head_dimhead_dim]
  ```
  
  `K,V` 作用于旋转矩阵后保存，加速解码过程。代码实现缓存相对简单，key_cache、value_cache 列表实现。
  
  ```python
  if layer_idx == 0:
       self._seen_tokens += key_states.shape[-2]
  # Update the cache
  if len(self.key_cache) <= layer_idx:
      self.key_cache.append(key_states)
      self.value_cache.append(value_states)
  else:
      self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
      self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)
  ```
  
  `K,V` 在 `channel` 维度上广播，对齐 `Q ` 维度，执行 `torch.nn.functional.scaled_dot_product_attention` 操作。
  
  ```python
   attn_output = torch.nn.functional.scaled_dot_product_attention(
              query_states,
              key_states,
              value_states,
              attn_mask=causal_mask,
              dropout_p=self.attention_dropout if self.training else 0.0,
              is_causal=is_causal,
          )
  ```
  
  该方法 `Torch` 实现：
  
  ```python
  def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None) -> torch.Tensor:
      L, S = query.size(-2), key.size(-2)
      scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
      attn_bias = torch.zeros(L, S, dtype=query.dtype)
      if is_causal:
          assert attn_mask is None
          temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
          attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
          attn_bias.to(query.dtype)
      if attn_mask is not None:
          if attn_mask.dtype == torch.bool:
              attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
          else:
              attn_bias += attn_mask
      attn_weight = query @ key.transpose(-2, -1) * scale_factor
      attn_weight += attn_bias
      attn_weight = torch.softmax(attn_weight, dim=-1)
      attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
      return attn_weight @ value

推理预填充阶段, is_causal 设置 True；在解码阶段，is_causal 设置 False；整个过程涉及`transpose`，`view` 操作可能会导致`tensor` 在内存里不连续，注意`contiguous`方法使用。

• Qwen2MLP

    class Qwen2MLP(nn.Module):
        def __init__(self, config):
            super().__init__()
            self.hidden_size = config.hidden_size
            self.intermediate_size = config.intermediate_size
            self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
            self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
            self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
            self.act_fn = ACT2FN[config.hidden_act]
    
        def forward(self, hidden_state):
            return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))

    Qwen2ForCausalLM(
      (model): Qwen2Model(
        (embed_tokens): Embedding(152064, 3584)
        (layers): ModuleList(
          (0-27): 28 x Qwen2DecoderLayer(
            (self_attn): Qwen2SdpaAttention(
              (q_proj): Linear(in_features=3584, out_features=3584, bias=True)
              (k_proj): Linear(in_features=3584, out_features=512, bias=True)
              (v_proj): Linear(in_features=3584, out_features=512, bias=True)
              (o_proj): Linear(in_features=3584, out_features=3584, bias=False)
              (rotary_emb): Qwen2RotaryEmbedding()
            )
            (mlp): Qwen2MLP(
              (gate_proj): Linear(in_features=3584, out_features=18944, bias=False)
              (up_proj): Linear(in_features=3584, out_features=18944, bias=False)
              (down_proj): Linear(in_features=18944, out_features=3584, bias=False)
              (act_fn): SiLU()
            )
            (input_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
            (post_attention_layernorm): Qwen2RMSNorm((3584,), eps=1e-06)
          )
        )
        (norm): Qwen2RMSNorm((3584,), eps=1e-06)
      )
      (lm_head): Linear(in_features=3584, out_features=152064, bias=False)
    )
    ```

    查看模型参数命名：

    ```python
     for name, parameters in model.named_parameters():
         print(name)
    """
    model.embed_tokens.weight
    model.layers.0.self_attn.q_proj.weight
    model.layers.0.self_attn.q_proj.bias
    model.layers.0.self_attn.k_proj.weight
    model.layers.0.self_attn.k_proj.bias
    model.layers.0.self_attn.v_proj.weight
    model.layers.0.self_attn.v_proj.bias
    model.layers.0.self_attn.o_proj.weight
    model.layers.0.mlp.gate_proj.weight
    model.layers.0.mlp.up_proj.weight
    model.layers.0.mlp.down_proj.weight
    model.layers.0.input_layernorm.weight
    model.layers.0.post_attention_layernorm.weight
    ...
    model.norm.weight
    lm_head.weight
    """    
    ```

    模型参数名 和 `model.safetensors.index.json` 文件里 `weight_map`字典下key 值一致，根据参数名去加载相关的权重矩阵。

+ `Load_checkpoints` **加载权重:**

  调用方法：`Qwen2ForCausalLM._load_pretrained_model`， 入参信息：

  ```python
  def _load_pretrained_model(
          cls,
          model,                # 实例化模型
          state_dict,
          loaded_keys,           # 模型参数的命名
          resolved_archive_file, # 模型权重文件路径,eg suffix.safetensors
          **kwargs）

"""摘自官方部分，仅供参考"""
def set_module_tensor(module, tensor_name, value):
    if "." in tensor_name:
        splits = tensor_name.split(".")
        for split in splits[:-1]:
            new_module = getattr(module, split)
            module = new_module
        tensor_name = splits[-1]
    '''
    举例：origin：
       module      : Qwen2ForCausalLM 
       tensor_name : 'model.embed_tokens.weight'
    执行上述递归后更新为：
       module      : Embedding(152064, 3584) 
       tensor_name : 'weight'
    '''  
    param = module._parameters[tensor_name] # 初始化参数
    param_cls = type(param)  # param类型：orch.nn.parameter.Parameter'>
    # value 是个 <class 'torch.Tensor'> 需要类型转化为<class 'torch.nn.parameter.Parameter'>
    new_value = param_cls(value, requires_grad = param.requires_grad)
    # 赋值替换,更新参数
    module._parameters[tensor_name] = new_value

for shard_file in resolved_archive_file:
    state_dict = load_state_dict(shard_file, is_quantized=is_quantized)
    for param_name, param in state_dict.items():
        set_module_tensor(model, param_name, param)
        

<transformers.generation.logits_process.RepetitionPenaltyLogitsProcessor object at 0x7f55df1d6390>
<transformers.generation.logits_process.TemperatureLogitsWarper object at 0x7f571c31bed0>
<transformers.generation.logits_process.TopKLogitsWarper object at 0x7f571c319cd0>
<transformers.generation.logits_process.TopPLogitsWarper object at 0x7f55db6c3910>

<transformers.generation.stopping_criteria.MaxLengthCriteria object at 0x7f5744a25c50>
<transformers.generation.stopping_criteria.EosTokenCriteria object at 0x7f5744b2ead0>

eos_token_id  = 151643
# kv cache 解码阶段输入:
model_inputs['input_ids']
tensor([[ 77235],
        [151643]], device='cuda:0')
# 输出:
next_tokens
tensor([ 11, 43959], device='cuda:0')

# finished sentences should have their next token be a padding token
# has_eos_stopping_criteria = [1,0],  1表示未完成，0表示完成
if has_eos_stopping_criteria:
    next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)