Qwen模型架构深度解析：从Config到Attention

Qwen模型，作为国内自主研发的大语言模型，其架构与Meta的Llama2颇为相似，这并非简单的模仿，而是在借鉴成熟架构的基础上，进行创新与优化。本文将深入剖析Qwen模型的内部构造，重点解读其关键组件及其工作原理，帮助读者理解Qwen模型的设计思想和技术特点。

Qwen2Config：模型的配置蓝图

Qwen2Config类是Qwen模型的核心配置类，它定义了模型的各项参数，包括词汇表大小、隐藏层维度、注意力头数等。这些参数共同决定了模型的规模和性能。

1.1 Model：模型的骨架

1.1.1 初始化：搭建模型的基石

在Qwen2Model类的初始化阶段，首先会设置模型的两个关键属性：padding_idx和vocab_size。padding_idx用于指定填充标记的索引，在处理变长序列时，需要对序列进行填充，使其长度一致，padding_idx就是用来标识这些填充位置。vocab_size则定义了模型词汇表的大小，即模型能够处理的不同词汇的数量。

接下来，初始化模型的三个核心组件：嵌入层、解码器层和归一化层。

• 嵌入层（nn.Embedding）：嵌入层是模型的第一层，负责将输入的标记（tokens）映射成密集的向量表示。每个词汇都对应一个唯一的嵌入向量，这些向量捕捉了词汇的语义信息。嵌入层的权重是模型需要学习的参数之一，通过训练，模型可以学习到更有效的词汇表示。
• 解码器层（nn.ModuleList()）：Qwen模型包含多个解码器层，这些层是模型的核心计算单元。每个解码器层都由Qwen2DecoderLayer定义，负责对输入的隐藏状态进行处理，提取特征并生成新的隐藏状态。多个解码器层堆叠在一起，可以使模型学习到更复杂的语言模式。
• 归一化层（Qwen2RMSNorm）：归一化层用于对隐藏状态进行归一化处理，可以提高模型的训练稳定性和收敛速度。Qwen模型使用的归一化层是Root Mean Square Layer Normalization（RMSNorm），它是一种简单而有效的归一化方法。

此外，模型还设置了gradient_checkpoint属性，用于控制是否使用梯度检查点技术。梯度检查点是一种节省显存的技术，通过在反向传播时重新计算部分层的激活值，可以减少显存的占用，从而训练更大的模型。

最后，调用post_init()方法完成一些初始化和准备检查的代码。post_init()主要用于对参数进行初始化，以及初始化梯度检查点。

class Qwen2Model(Qwen2PreTrainedModel):
    def __init__(self, config: Qwen2Config): #传入一个配置对象，它包含了模型的所有配置参数
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size #设置词汇表的大小

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) #创建一个嵌入层，用于将词汇表中的每个单词映射到一个隐藏向量
        self.layers = nn.ModuleList( #创建一个模块列表，包含多个 `Qwen2DecoderLayer`，每个层对应模型的一个解码器层
            [Qwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        self.gradient_checkpointing = False
        # Initialize weights and apply final processing
        self.post_init()

post_init方法在构造函数执行完毕后调用，用于执行一些依赖于构造函数中已执行步骤的初始化工作。这通常包括权重初始化和梯度检查点设置。

def post_init(self):
    """
    A method executed at the end of each Transformer model initialization, to execute code that needs the model's
    modules properly initialized (such as weight initialization).
    """
    self.init_weights()#初始化模型的权重
    self._backward_compatibility_gradient_checkpointing()#向后兼容而设置的私有方法(梯度检查点)

1.1.2 Forward：数据流动的引擎

forward方法定义了模型的前向传播过程，即数据在模型中的流动路径。这是自然语言处理模型中至关重要的部分，特别是对于Transformer模型而言。

inputs_embeds = self.embed_tokens(input_ids)#这行代码使用 `embed_tokens` 方法将输入的 `input_ids`转换为嵌入向量
hidden_states = inputs_embeds#初始化 `hidden_states` 为输入嵌入，它将在后续的解码器层中被更新

for idx, decoder_layer in enumerate(self.layers):
    # 将所有的hidden_states保存成tuple
    if output_hidden_states: #判断是否需要输出所有层的隐藏状态
        all_hidden_states += (hidden_states,)
    # 将hs送入每一层decoder_layer
    # 调用当前解码器层的前向传播方法，传入当前的 `hidden_states` 和其他必要的参数
    layer_outputs = decoder_layer(
        hidden_states,
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_value=past_key_value,
        output_attentions=output_attentions,
        use_cache=use_cache,
    )
    # 取出上一层decoder_输出的hs,再传入下一个layer
    # 只要第一个,第二个是cache的一个类，然后进入下一个layer
    hidden_states = layer_outputs[0]

hidden_states = self.norm(hidden_states)

if output_hidden_states:
    all_hidden_states += (hidden_states,)

首先，input_ids（输入的标记序列）通过嵌入层转换为嵌入向量inputs_embeds。然后，hidden_states被初始化为inputs_embeds，并在后续的解码器层中不断更新。

模型遍历每一层解码器层，将hidden_states输入到当前层，并获取输出。每一层解码器层都会返回新的hidden_states，以及一些其他信息，如注意力权重等。hidden_states会被传递到下一层解码器层，直到最后一层。

如果output_hidden_states为True，则模型会将每一层解码器层的hidden_states保存到all_hidden_states列表中。这在某些情况下是有用的，例如，当需要分析模型每一层的表示时。

最后，模型对最后一层解码器层输出的hidden_states进行归一化处理，并将其作为模型的最终输出。

1.2 Qwen2DecoderLayer：解码器的核心

Qwen2DecoderLayer是Qwen模型解码器层的核心组件，它包含了自注意力机制（self-attention）和多层感知机（MLP），以及两种归一化层。这些组件协同工作，共同完成对输入序列的解码。

1.2.1 初始化：构建解码器的积木

解码器层的初始化主要包括三个部分：自注意力机制、多层感知机和归一化层。

QWEN2_ATTENTION_CLASSES = {
    "eager": Qwen2Attention,  # 默认的注意力实现，一般情况下是这个
    "flash_attention_2": Qwen2FlashAttention2, # 一个优化的注意力实现
    "sdpa": Qwen2SdpaAttention, # 一种特殊的注意力实现
}

class Qwen2DecoderLayer(nn.Module):
    def __init__(self, config: Qwen2Config):
        super().__init__()
        self.hidden_size = config.hidden_size # 设置隐藏层的大小
        self.self_attn = QWEN2_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)# 根据配置中的 `_attn_implementation` 键来选择使用哪种自注意力实现，并初始化它

        self.mlp = Qwen2MLP(config)# 初始化一个多层感知机（MLP），用于在自注意力之后处理隐藏状态
        self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)# 初始化两个归一化层，分别用于自注意力之前和之后。这两个层都是 RMS 归一化层，使用配置中的 `hidden_size` 和 `rms_norm_eps` 参数
        self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

• 自注意力机制（Qwen2Attention）：自注意力机制是Transformer模型的核心，它允许模型在处理序列时，关注到序列中不同位置之间的关系。Qwen模型支持多种自注意力实现，包括Qwen2Attention（默认实现）、Qwen2FlashAttention2（优化实现）和Qwen2SdpaAttention（特殊实现）。
• 多层感知机（Qwen2MLP）：多层感知机用于在自注意力机制之后处理隐藏状态，它可以学习到更复杂的特征表示。Qwen模型使用的MLP包含两个线性层和一个激活函数。
• 归一化层（Qwen2RMSNorm）：Qwen模型使用RMSNorm对隐藏状态进行归一化处理，以提高训练稳定性和收敛速度。解码器层包含两个RMSNorm层，分别位于自注意力机制之前和之后。

1.2.2 Forward：解码器的运作

residual = hidden_states# 保存原始的 `hidden_states` 到 `residual` 变量中，用于后面的残差连接
hidden_states = self.input_layernorm(hidden_states)  # 将 `hidden_states` 通过输入归一化层（`input_layernorm`），RMSNorm标准化
hidden_states, self_attn_weights, present_key_value = self.self_attn(
    hidden_states=hidden_states,
    attention_mask=attention_mask,
    position_ids=position_ids,
    past_key_value=past_key_value,
    output_attentions=output_attentions,
    use_cache=use_cache,
    **kwargs,
)

hidden_states = residual + hidden_states

residual = hidden_states
hidden_states = self.post_attention_layernorm(hidden_states)
hidden_states = self.mlp(hidden_states)
hidden_states = residual + hidden_states

outputs = (hidden_states,)

return outputs

解码器层的前向传播过程如下：

1. 残差连接：将原始的hidden_states保存到residual变量中，用于后面的残差连接。
2. 归一化：将hidden_states通过输入归一化层（input_layernorm）进行归一化处理。
3. 自注意力：将归一化后的hidden_states输入到自注意力机制（self_attn）中，计算自注意力权重，并生成新的hidden_states。
4. 残差连接：将自注意力机制输出的hidden_states与residual相加，实现残差连接。
5. 归一化：将残差连接后的hidden_states通过后注意力归一化层（post_attention_layernorm）进行归一化处理。
6. 多层感知机：将归一化后的hidden_states输入到多层感知机（mlp）中，学习更复杂的特征表示。
7. 残差连接：将多层感知机输出的hidden_states与residual相加，实现残差连接。
8. 输出：将最终的hidden_states作为解码器层的输出。

残差连接和归一化是Transformer模型中的两个关键技术，残差连接可以避免深层网络中的梯度消失问题，归一化可以稳定训练过程，加快收敛速度。

1.3 Qwen2Attention：自注意力的引擎

Qwen2Attention类实现了多头自注意力机制，这是Transformer架构中的一个核心组件。自注意力机制允许模型在处理序列数据时考虑到不同位置之间的关系。

1.3.1 初始化：构建自注意力的蓝图

class Qwen2Attention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: Qwen2Config):
        super().__init__()
        self.config = config# 保存传入的配置对象，它包含了自注意力层所需的所有配置参数。
        self.layer_idx = layer_idx# 保存索引
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.hidden_size // self.num_heads# 计算每个注意力头的维度大小，它是隐藏层大小除以头的数量
        self.num_key_value_heads = config.num_key_value_heads# 设置键值对头的数量
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.max_position_embeddings = config.max_position_embeddings# 设置最大位置嵌入的大小，这通常用于位置编码
        self.rope_theta = config.rope_theta# 设置旋转嵌入（Rotary Positional Embedding）的参数
        self.is_causal = True# 指示是否使用因果自注意力（即在生成下一个 token 时只能使用之前的 token）
        self.attention_dropout = config.attention_dropout# 设置注意力权重的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=config.attention_bias)# Query
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)# Key
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)# Value
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)# 计算最终输出的投影

        # 初始化一个旋转嵌入层，这是一种特殊的位置编码，可以增强模型对序列顺序的感知能力
        self.rotary_emb = Qwen2RotaryEmbedding(
            self.head_dim,
            max_position_embeddings=self.max_position_embeddings,
            base=self.rope_theta,
        )

自注意力类的初始化包括以下步骤：

1. 配置参数：保存传入的配置对象，包括隐藏层大小、注意力头数、键值对头数、最大位置嵌入大小、旋转嵌入参数等。
2. 线性投影：初始化四个线性层，分别用于计算查询（Query）、键（Key）、值（Value）和输出。
3. 旋转嵌入：初始化旋转嵌入层，这是一种特殊的位置编码，可以增强模型对序列顺序的感知能力。

max_position_embeddings确定了模型能够编码的位置索引的最大值，通常对应于模型能够处理的最长序列长度。rope_theta决定了旋转矩阵的周期性，它影响着位置编码的周期性变化。

1.3.2 Forward：自注意力的计算

bsz, q_len, _ = hidden_states.size()

query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)

 # reshape多头处理--分块--(bs,T,heads,hd_d),为了将输入的隐藏状态转换为适合多头自注意力计算的形式,每个头可以独立地处理序列的一部分，从而实现并行处理和更细粒度的表示学习
     # `ranspose` 函数用于交换张量的两个维度
query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

key_states = repeat_kv(key_states, self.num_key_value_groups)
value_states = repeat_kv(value_states, self.num_key_value_groups)

attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

attn_weights = attn_weights + attention_mask

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)
attn_output = torch.matmul(attn_weights, value_states)

attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)

attn_output = self.o_proj(attn_output)

return attn_output, attn_weights, past_key_value

自注意力机制的计算过程如下：

1. 线性投影：将输入的hidden_states通过线性层投影到查询（Query）、键（Key）和值（Value）空间。
2. 多头处理：将查询、键和值张量reshape成多个头，每个头独立地处理序列的一部分，从而实现并行处理和更细粒度的表示学习。
3. 旋转嵌入：将旋转位置嵌入应用于查询和键张量，以增强模型对序列顺序的感知能力。
4. 注意力权重计算：计算查询和键的点积，并除以缩放因子（head_dim的平方根），得到注意力权重。
5. 注意力权重归一化：对注意力权重进行softmax操作以进行归一化，然后应用dropout以防止过拟合。
6. 注意力输出计算：使用归一化后的注意力权重和值张量进行点积，得到注意力输出。
7. 输出投影：通过输出投影层o_proj对注意力输出进行最后的线性变换。

1.3.3 细节Debug：深入理解GQA和位置编码

1.3.3.1 GQA：减少显存占用

GQA（Grouped-Query Attention）是一种多头注意力机制的变体，旨在减少推理过程中的显存占用。在传统的自注意力机制中，每个注意力头都需要存储一个完整的键值对缓存（KV Cache），这在处理长序列时会占用大量的显存。GQA通过将多个注意力头共享同一个键值对缓存，从而减少了显存的占用。

GQA（图形问答）和MQA（多步问答）不需要在推理的过程存储那么多的kv cache（键值对缓存）, 那么kv cache占用的显存就变小，那么我们LLM serving可以处理的请求数量就更多

• GQA（Graphical Question Answering）：这是一种机器理解视觉和语言的任务，其中模型需要理解图像内容并回答有关图像的问题。
• MQA（Multi-turn Question Answering）：这是一种交互式问答任务，模型需要在多轮对话中回答用户的问题，通常需要根据之前的对话历史来理解上下文。
• KV Cache（Key-Value Cache）：在自回归语言模型中，为了加速生成过程，会缓存之前计算过的键（Key）和值（Value），这样在生成下一个词时可以重用这些信息，而不是重新计算。这可以显著提高推理速度，但同时也占用了显存。

（1）定义初始张量

import torch

## shape:(batch, seq_len, head, head_dim)
## 批次大小 序列长度 头的数量 每个头的维度
query = torch.randn(10, 128, 8, 128)
key = torch.randn(10, 128, 2, 128)
value = torch.randn(10, 128, 2, 128)

## 在此设置组数为4
groups = query.shape[-2] // key.shape[-2]

（2）之后进行扩展key，value的操作

在GQA中，key和value都要比query小group倍，但是为在后续做矩阵乘法时方便，我们需要先把key和value的head利用expand扩展张量到和query相同的维度。方便后续计算

def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:

    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    # dont need repeat here means multi head attention
    if n_rep == 1:
        return hidden_states
    # first we expand x to (bs, seq_len, head, group, head_dim)
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    # reshape make head -> head * group
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

（3） 矩阵乘法得到score与output 后面就是征程的kqv相乘了

#(bs, head, seq_len, head_dim)
query = query.transpose(1, 2)
key = repeat_kv(key.transpose(1, 2), 4)
value = repeat_kv(value.transpose(1, 2), 4)
scores = torch.matmul(query, key.transpose(2, 3)) / math.sqrt(head_dim)
scores = torch.nn.functional.softmax(scores, dim=-1)

out = torch.matmul(scores, value)
#上一步转置了，还得转回去
out = out.transpose(1, 2)

1.3.3.2 apply_rotary_pos_emb：旋转位置编码

1.3.3.3 读取顺序attention_mask

这两部分直接看原文（见参考资料）吧
ヽ(゜▽゜　)－C<(/;◇;)/~

1.4 Qwen2MLP：多层感知机的实现

Qwen2MLP类实现了一个多层感知机（MLP）结构，通常用于Transformer模型中的前馈网络（Feed-Forward Network，FFN）。

class Qwen2MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        # 这俩不必多说
        self.config = config
        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)# 定义一个线性层，用于将输入投影到中间层的大小。这个层的权重在训练过程中是不变的（`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]# 根据配置中的激活函数类型，选择相应的激活函数。`ACT2FN` 是一个将激活函数名称映射到 PyTorch 激活函数的字典

    # 定义了前向传播函数，它是模型的输入数据流经网络的路径
    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj

多层感知机包含三个线性层和一个激活函数。gate_proj和up_proj将输入投影到中间层的大小，down_proj将中间层的输出投影回隐藏层的大小。act_fn是激活函数，用于引入非线性。

1.5 Qwen2RMSNorm：RMS归一化的实现

RMS（Root Mean Square）归一化是一种简单而有效的归一化方法，与传统的LayerNorm不同，它只使用方差（不包括均值）来进行归一化。

class Qwen2RMSNorm(nn.Module):  # 标准化层
    def __init__(self, hidden_size, eps=1e-6):
        """
        Qwen2RMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        # 将归一化后的隐藏状态乘以学习到的权重，并将张量转换回原始的数据类型，然后返回
        return self.weight * hidden_states.to(input_dtype)

RMS归一化的计算公式如下：

R M S N o r m ( x ) = x 1 n ∑ i = 1 n ω i 2 + ϵ RMSNorm(x)=\frac{x}{ \sqrt{\frac{1}{n}\sum[i = 1]^{n}\omega^2[i+\epsilon }} RMSNorm(x) =n1∑i=1nωi2+ϵ x

其中：

• x是层的输入的hidden_state
• wi 表示的是hidden_state的最后一个维度的值
• n 表示上面输入的最后一个维度的数量。
• ϵ 表示是很小的数，防止除0。

总结

Qwen模型在架构上借鉴了Llama2，但在实现细节上进行了创新和优化。通过深入理解Qwen模型的内部构造，我们可以更好地掌握其设计思想和技术特点，并为后续的研究和应用奠定基础。