手搭一个大模型-part1-Qwen模型的介绍

1o0.0o1

已于 2024-05-18 21:17:35 修改

阅读量4.5k

点赞数 22

分类专栏：大模型文章标签：神经网络 python 自然语言处理 transformer

于 2024-05-16 19:27:07 首次发布

本文链接：https://blog.csdn.net/qq_60489376/article/details/138817599

版权

大模型专栏收录该内容

4 篇文章

订阅专栏

前置知识

Transformer的魅力

在2017年，一篇划时代的论文《Attention Is All You Need》发布，极大地推动了自然语言处理（NLP）领域的进展。自此以后，Transformer模型不仅在文本生成中得到了广泛应用，还在扩散模型等多个领域显示出了其强大的能力。接下来，我们将详细介绍Transformer的整体框架。
Transformer框架图
如上图:我们可以看到Transformer架构分为四个部分,分别是输入部分,编码器部分,解码器部分,输出部分。

输入部分：包括编码器输入和解码器输入。对于编码器输入，我们首先将文本通过词嵌入（Word Embedding）和位置嵌入（Position Embedding）转换为向量，然后送入编码器以提取特征。解码器输入则在编码器的基础上增加了一个掩码机制，该机制用于屏蔽未来的信息，防止信息提前泄露。
编码器部分：根据Transformer的原始论文，编码器由六个结构相同的编码层堆叠而成。这里我们重点分析单个编码层的结构。从结构图中可以看出，每个编码器层包括两个子模块：多头自注意力（Multi-Head Self-Attention）和前馈全连接层（Feed-Forward Neural Network）。每个子模块都是通过一个残差连接和随后的层归一化（Layer Normalization）来实现的。
解码器部分：解码器同样由六个结构相同的解码层堆叠而成，构成整个解码器。在解码器的每个层中，有三个主要的子模块：多头自注意力层、编码器-解码器注意力层（用于从编码器层抽取信息），以及前馈全连接层。与编码器类似，这些子模块也采用残差连接和层归一化。
输出部分：最后，输出通过一个全连接层进行特征提取，并通过Softmax函数生成最终的预测结果

我们以机器翻译为例子理解一下Transformer的训练全过程:

假设我们的任务是将我爱北京翻译成I love Beijing

预处理和词嵌入
输入处理：首先，中文句子“我爱北京”会被分词为单独的词或字符。假设我们使用字符级的分割，得到“我”、“爱”、“北”、“京”。这一部分是由tokenizer完成的
词嵌入：这些字符通过词嵌入层转换成向量。此外，由于Transformer不具备处理序列顺序的能力，我们还需为每个字符添加位置嵌入，以表示其在句子中的位置。对应结构中的Embedding
编码器操作
多头自注意力机制：在编码器中，多头自注意力层会评估句子中每个字符与其他字符的关系，这有助于捕获例如“我爱”（我和爱之间的直接关系）这样的局部依赖关系。(Multi-Head Self-Attention)
前馈全连接层机制(Feed-Forward Neural Network): 经过大量的实验表面,全连接层的特征提取能力是很强的,而且结构简单,为了防止多头注意力机制特征提取不够充分,所有加入了这一层,让模型进一步学习到词语词之间的依赖关系
层次结构处理：编码器的每一层都将之前层的输出作为输入，逐层提取更抽象的特征。每个层的输出都是一个加强了输入句子每个部分上下文信息的表示。
解码器操作
屏蔽未来信息：解码器在生成翻译时使用屏蔽技巧来避免“看到”未来的输出。例如，在预测单词“love”时，模型只能访问到“ I”，而不能访问到“Beijing”。
注意力机制：解码器的编码器-解码器注意力层使得每一步的生成都可以关注到输入句子的不同部分。例如，当生成“Beijing”时，模型可能会特别关注“北京”。
生成预测和训练
输出：每次解码步骤，模型都会输出一个词的概率分布，选择概率最高的词作为这一位置的翻译。例如，首先生成“I”，然后是“love”，最后是“Beijing”。
训练过程：在训练阶段，我们使用实际的目标句子“ I love Beijing ”作为训练目标。模型通过比较预测输出与实际输出，并通过反向传播优化其参数。

了解GPT模型

为什么要介绍GPT模型呢?

现在很多大模型都是基于transform中的decoder模块发展的,而最典型的模型就是GPT。GPT模型本质上是利用了Transformer的解码器结构。与完整的Transformer相比，GPT仅使用了解码器部分，这是因为它主要被设计为一个生成模型。

下图是GPT模型的结构图
在这里插入图片描述
如上图所示, 经典的Transformer Decoder Block包含3个子层, 分别是Masked Multi-Head Attention层, encoder-decoder attention层, 以及Feed Forward层. 但是在GPT中取消了第二个encoder-decoder attention子层, 只保留Masked Multi-Head Attention层, 和Feed Forward层.

上图最后的输出部分还展示了GPT最经典的两个预训练任务,分别为:无监督训练和有监督训练。

无监督训练: 在这个阶段，模型主要通过大量的文本数据进行预训练，目的是使模型能够预测下一个词。这一过程称为自回归语言建模。通过这种方式，GPT不仅学习了词汇的共现（co-occurrence）和语法结构，而且还能把握到更复杂的语义信息和上下文关系。
有监督训练:GPT模型通常会通过有监督的微调（SFT）来适应具体的任务。这包括但不限于文本分类、情感分析、问答系统等。通过针对性的训练优化特定任务的表现。这种策略能够有效提高了模型的灵活性和效能，使其能够在多种NLP任务中表现出色

千问大模型框架

介绍

介绍完Transformer和GPT模型后,接下来就是正题,手搭一个千问大模型。首先我们先看看千问大模型的结构图
Qwen大模型结构图
我们现在看主干部分,是不是很熟悉。和我们GPT模型的结构是不是很相似? 现在基本上大模型都是基于GPT的结构来实现的,唯一不同的是可能对GPT每一个模块都有了一定程度的魔改,从而达到GPT模型所做不到的程度。

接着我们开始介绍每一部分的实现

Tokenizer部分

Tokenizer就是一个训练好的分词器,可以理解为,我现在有一张词汇表,上面每个单词都有一个序号,现在我将一个句子划分成一个个词,然后根据这张词汇表将这些词用一个序号表示。这就是Tokenizer的作用

Embedding部分

Qwen的Embedding部分和GPT相比在PositionEmbedding部分做出了一些改变,GPT中的PositionEmbedding通常使用固定的三角函数位置编码。这种位置编码是通过正余弦函数的变化来为模型的每个输入位置生成唯一的编码，从而帮助模型理解单词的顺序关系。
代码如下

# 定义位置编码器
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        """位置编码器类的初始化函数, 共有三个参数, 分别是d_model: 词嵌入维度, 
           dropout: 置0比率, max_len: 每个句子的最大长度"""
        super(PositionalEncoding, self).__init__()
        # 实例化nn中预定义的Dropout层, 并将dropout传入其中, 获得对象self.dropout
        self.dropout = nn.Dropout(p=dropout)
        # 初始化一个位置编码矩阵, 它是一个0阵，矩阵的大小是max_len x d_model.
        pe = torch.zeros(max_len, d_model)
        # 初始化一个绝对位置矩阵, 在我们这里，词汇的绝对位置就是用它的索引去表示. 
        # 所以我们首先使用arange方法获得一个连续自然数向量，然后再使用unsqueeze方法拓展向量维度使其成为矩阵， 
        # 又因为参数传的是1，代表矩阵拓展的位置，会使向量变成一个max_len x 1 的矩阵， 
        position = torch.arange(0, max_len).unsqueeze(1)
        # 绝对位置矩阵初始化之后，接下来就是考虑如何将这些位置信息加入到位置编码矩阵中，
        # 最简单思路就是先将max_len x 1的绝对位置矩阵， 变换成max_len x d_model形状，然后覆盖原来的初始位置编码矩阵即可， 
        # 要做这种矩阵变换，就需要一个1xd_model形状的变换矩阵div_term，我们对这个变换矩阵的要求除了形状外，
        # 还希望它能够将自然数的绝对位置编码缩放成足够小的数字，有助于在之后的梯度下降过程中更快的收敛.  这样我们就可以开始初始化这个变换矩阵了.
        # 首先使用arange获得一个自然数矩阵， 但是细心的同学们会发现， 我们这里并没有按照预计的一样初始化一个1xd_model的矩阵， 
        # 而是有了一个跳跃，只初始化了一半即1xd_model/2 的矩阵。 为什么是一半呢，其实这里并不是真正意义上的初始化了一半的矩阵，
        # 我们可以把它看作是初始化了两次，而每次初始化的变换矩阵会做不同的处理，第一次初始化的变换矩阵分布在正弦波上， 第二次初始化的变换矩阵分布在余弦波上， 
        # 并把这两个矩阵分别填充在位置编码矩阵的偶数和奇数位置上，组成最终的位置编码矩阵.
        # 跳跃式初始化
        div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        # 将pe注册乘buffer
        # 最后把pe位置编码矩阵注册成模型的buffer，什么是buffer呢，
        # 我们把它认为是对模型效果有帮助的，但是却不是模型结构中超参数或者参数，不需要随着优化步骤进行更新的增益对象. 
        # 注册之后我们就可以在模型保存后重加载时和模型结构与参数一同被加载.
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        """forward函数的参数是x, 表示文本序列的词嵌入表示"""
        # 在相加之前我们对pe做一些适配工作， 将这个三维张量的第二维也就是句子最大长度的那一维将切片到与输入的x的第二维相同即x.size(1)，
        # 因为我们默认max_len为5000一般来讲实在太大了，很难有一条句子包含5000个词汇，所以要进行与输入张量的适配. 
        # 最后使用Variable进行封装，使其与x的样式相同，但是它是不需要进行梯度求解的，因此把requires_grad设置成false.
        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)
        return x

我们可以通过代码查看一下,单词在不同位置下,对应的波长变化

import matplotlib.pyplot as plt
import numpy as np

# 创建一张15 x 5大小的画布
plt.figure(figsize=(15, 5))

# 实例化PositionalEncoding类得到pe对象, 输入参数是20和0
pe = PositionalEncoding(20, 0)

# 然后向pe传入被Variable封装的tensor, 这样pe会直接执行forward函数, 
# 且这个tensor里的数值都是0, 被处理后相当于位置编码张量
# 100个词,每个词表达成20维的词向量
y = pe(Variable(torch.zeros(1, 100, 20)))

# 然后定义画布的横纵坐标, 横坐标到100的长度, 纵坐标是某一个词汇中的某维特征在不同长度下对应的值
# 因为总共有20维之多, 我们这里只查看4，5，6，7维的值.
plt.plot(np.arange(100), y[0, :, 4:8].data.numpy())

# 在画布上填写维度提示信息
plt.legend(["dim %d" % p for p in [4, 5, 6, 7]])

在这里插入图片描述
通过这种方式能够对词向量引入位置关系。

而Qwen的Position Embedding也是用的这种编码。

hidden_stage

hidden_stage就是我们词汇经过tokenizer变成索引张量后,然后经过Embedding层变成一个词向量,这个时候我们就称其为一个hidden_stage

Qwen attention

Qwen attention的结构图如下:
在这里插入图片描述
接下来我将从注意力机制开始,介绍每一个框架的作用

注意力机制

我们先介绍一下什么是注意力机制,首先我们先看一下下面的图片:
在这里插入图片描述
大家看到这图后第一时间的关注点在哪里呢(手动狗头)?

我们观察事物时，之所以能够快速判断一种事物(当然允许判断是错误的), 是因为我们大脑能够很快把注意力放在事物最具有辨识度的部分从而作出判断，而并非是从头到尾的观察一遍事物后，才能有判断结果. 正是基于这样的理论，就产生了注意力机制.

注意力机制有三种计算方法:
将𝑄𝐾进行纵轴拼接，做一次线性变化，再使用 softmax 处理获得结果最后与做张量乘法：
$\text{Attention}(Q, K, V) = \text{Softmax}(\text{Linear}([Q, K])) \cdot V$
将𝑄𝐾进行纵轴拼接，做一次线性变化后再使用 tanh 函数激活，然后再进行内部求和，最后使用 softmax 处理获得结果最后与 $V$ 做张量乘法
$\text{Attention}(Q, K, V) = \text{Softmax}(\sum (\text{tanh}(\text{Linear}([Q, K])))) \cdot V$
将 𝑄𝐾的转置做点积运算，然后除以一个缩放系数，再使用 softmax 处理获得结果最后与 𝑉做张量乘法：
$\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{Q \cdot K^T}{\sqrt{d_k}}\right) \cdot V$

我们可以将注意力机制的计算公式分成两个部分,一个是Q和K相乘这一部分,另一个是V。我们将Q和K相乘这一部分(包括缩放,或者线性变化之类的操作)最后得出的结果称为得分函数,然后对得分函数做一个softmax,将分数映射到0-1之间。这个就是我们的权重矩阵。最后，将权重矩阵与值（Value, V）相乘，得到最终的注意力输出。如下图,权重越大,线的颜色越深。
在这里插入图片描述
自注意力机制即K=Q=V

多头注意力机制

多头自注意力机制的结构图如下:

多头注意力机制即将词向量的维度分为多个头去管理,每个头只分析属于自己部分的维度。这种结构设计能让每个注意力机制去优化每个词汇的不同特征部分，从而均衡同一种注意力机制可能产生的偏差，让词义拥有来自更多元的表达，实验表明可以提升模型效果.

GQA机制

Qwen的多头自注意力层使用了GQA（Group Query Attention）机制，在原来的多头注意力机制结构上做了一些改进。具体来说，GQA机制的关键点在于：在GQA中，每个查询（Query）不再与单个键（Key）对应，而是与一组键进行交互。这种机制可以更有效地利用计算资源，提高模型的并行计算能力。

在每个组内，计算查询与键的注意力得分，然后进行归一化（如Softmax）并计算加权和，类似于标准的多头注意力机制。这样，每个组的注意力计算是独立的，可以并行进行。

最后,将各组的注意力结果进行汇总，形成最终的输出。这种机制有助于模型在保持计算效率的同时，提高注意力计算的灵活性和表达能力
在这里插入图片描述

引入了GQA的代码:

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
        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=config.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        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,
        )

我们还注意到了在Qwen Attention的图中,出现了一个repeat-kv模块,下面我将介绍一下这个模块

旋转编码RoPE机制

旋转位置编码_rotary_pos_emb,通过这种编码方式,不仅可以为Token添加绝对位置的信息,还能添加到相对位置的信息

具体公式如下:
在这里插入图片描述

RoPE机制通过在计算注意力权重时引入旋转编码，使模型在处理长序列时更具优势。具体来说，RoPE将位置信息编码到Query和Key向量中，通过在这些向量上应用旋转变换，使得注意力分数中包含了位置信息。这种方法使得模型能够更好地捕捉序列中的相对位置关系。

具体代码实现:

class Qwen2RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()
        # 定义初始值
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        # 定义旋转角
        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)

        # Build here to make `torch.jit.trace` work.
        self._set_cos_sin_cache(
            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
        )
    # 为seq里面的每个token形成独一无二的旋转角嵌入(外积)
    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)
        # 生成角度信息(利用注册机制生成self.cos_cached与sin_cached
        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)

    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

        return (
            self.cos_cached[:seq_len].to(dtype=x.dtype),
            self.sin_cached[:seq_len].to(dtype=x.dtype),
        )

# 后半部分和前半部分进行了交换，并且将后半部分的符号取反。
def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
    """Applies Rotary Position Embedding to the query and key tensors.

    query and key tensors rotated using the Rotary Position Embedding.
    """
    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

repeat-kv

在标准的多头自注意力机制中，输入的特征会被投影成 Query (Q), Key (K) 和 Value (V)，然后分成多个头独立计算注意力。每个头都有自己的 Q, K 和 V，并在每个头内部进行自注意力计算。这种方法虽然有效，但计算量较大，特别是对于大规模模型和长序列数据

repeat-kv 技术主要是通过重复使用键 (K) 和值 (V) 来优化计算。在这种技术中，K 和 V 在每个分组之间是共享的，而查询 (Q) 仍然是独立计算的。这可以与卷积神经网络和全连接神经网络进行对比。

在 CNN 中，卷积层代替了全连接层，卷积层中的卷积核实现了权重共享。每次特征提取时，它们都使用相同的卷积核参数进行卷积运算，并通过反向传播来更新这些权重。同时，每个卷积层有多个卷积核，每个卷积核维护着自己的权重矩阵。这里的每个卷积核可以类比为不同分组中的 K 和 V 矩阵。

通过这种方式，相较于全连接神经网络，CNN 大大减少了参数量。repeat-kv 技术在多头自注意力机制中起到了类似卷积核的作用，通过共享 K 和 V 矩阵，减少了重复计算的次数，从而优化了计算效率和参数量。

Attention mask机制

我们可以想象一下，当我们向一个大模型提问时，大模型是如何生成答案的。大模型在生成答案时是一个个token逐步生成的，然后将当前生成的token拼接到输入中，再继续生成下一个token。我们称这种机制为自回归。

自回归机制要求模型只能利用当前和之前的 token 信息，而不能看到未来的 token 信息。可以想象一下，如果模型能够看到未来的 token 信息，那直接告诉你答案就可以了，还需要一次次地重复输入吗？然而，在训练模型时，我们实际上是已经知道整个句子的所有信息的。如果在训练时将整个句子都输入进去，就无法达到我们希望的预测效果。为了模拟生成过程中的逐步预测，我们引入了 Attention Mask 机制。

Attention Mask 通过遮蔽不应该被注意到的未来信息，确保每个 token 只能看到当前和之前的 token。例如，在自回归生成任务中，我们通常使用一个下三角矩阵作为 Attention Mask。这个矩阵确保每个 token 只能关注到自己及之前的 token，而无法看到之后的 token

假设我们有一个句子“I love beijing”，在生成过程中，Attention mask会确保：
生成第一个token “I” 时，只能看到起始符(通常是SOS)。
生成第二个token “love” 时，只能看到 “I ”。
生成第三个token “beijing” 时，只能看到 “I love”

这个过程我们可以用下面这个下三角矩阵实现
$\begin{bmatrix} 1 & 0 & 0 \\ 1 & 1 & 0 \\ 1 & 1 & 1 \\ \end{bmatrix}$
预测I的时候, love和beijing被mask遮掩,值变成无限小(或等于0),然后在使用softmax计算权重矩阵时,就会把love和beijing这两个词的值映射成0,从而做到让模型对这两个词的注意力为0,做到了遮挡未来信息的作用。

这就是Attention mask的作用

Qwen Attention层的代码实现

class Qwen2Attention(nn.Module)
    def __init__(self, config: Qwen2Config):
        super().__init__()
        self.config = config
        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  # 旋转位置嵌入的基数
        self.is_causal = True  # 是否为因果注意力机制
        self.attention_dropout = config.attention_dropout  # 注意力的dropout概率

        # 检查 hidden_size 是否能被 num_heads 整除
        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})."
            )

        # 定义线性层，用于生成 Query, Key 和 Value
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        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,
        )

    def forward(self, hidden_states, attention_mask, position_ids, layer_idx, past_key_value):
        # 获取输入的形状信息
        bsz, q_len, _ = hidden_states.size()

        # 通过线性变换生成 Query, Key 和 Value
        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        # 将生成的 Query, Key 和 Value 进行 reshape，适应多头处理
        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)

        # 应用旋转位置嵌入到 Query 和 Key 上
        cos, sin = self.rotary_emb(value_states, seq_len=key_states.size(2))
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

        # 重复 Key 和 Value，使其在每个分组中共享
        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

        # 对注意力权重进行 softmax 和 dropout，然后与 Value 相乘
        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

# 配置类，用于定义模型配置
class Qwen2Config:
    hidden_size = 512
    num_attention_heads = 8
    num_key_value_heads = 2
    max_position_embeddings = 1024
    rope_theta = 10000
    attention_dropout = 0.1
    attention_bias = False

# 位置嵌入类，用于定义旋转位置嵌入
class Qwen2RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, max_position_embeddings, base):
        super().__init__()
        self.head_dim = head_dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base

    def forward(self, value_states, seq_len):
        # 计算 cos 和 sin 嵌入
        return cos, sin

# 用于应用旋转位置嵌入
def apply_rotary_pos_emb(query, key, cos, sin, position_ids):
    # 计算旋转位置嵌入
    return query, key

# 用于重复键和值
def repeat_kv(states, num_groups):
    # 重复键和值
    return states.repeat(1, num_groups, 1, 1)

Qwen MLP层

经过了Attention层后,我们来到了下一个子层,即MLP层。首先我们先看一下它的结构图
在这里插入图片描述
在我们transformer中的这一层被称为前馈全连接层,使用的是两层全连接神经网络。但是全连接神经网络有一个缺点,那就是只能线性拟合,对于非线性特征拟合效果不算太好。而MLP则是借助了LSTM中的门值的思想,通过添加一个激活函数,引入非线性变化。从而提高了模型对非线性数据特征抽取能力。

这里的HS经过了Linear层后,经过了一个Act层,这里的Act代指激活函数,常见的有ReLU, Sigmoid等,得到一个类似门值的矩阵(可以理解为权重矩阵), 在和另一边经过绿色部分的linear拟合的数据做乘积,从而引入了非线性能力。后面再经过一个线性层对特征进行抽取,最后输出。这就是Qwen的MLP层。

实现代码

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)
        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, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj

Qwen中的RMSNorm

RMSNorm的公式为:

$\text{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{n} \sum_{i=1}^{n} w_i^2 + \epsilon}}$
其中:

$x$ 表示我们输入的向量,也就是Qwen 结构图中的HS
$w_i$ 表示我们输入的向量x中,第i个位置的数值

RMSNorm通过计算输入向量每个元素的平方和的均值，再取平方根，进行归一化。这种方式不改变输入的均值，而是基于输入向量的范数来进行标准化。

在HS输入到子层中时,经过RMSNorm能够加快模型收敛的速度,且在处理具有较大数值波动的输入时，能够较好地保持输入的相对大小。

代码实现

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)

最后的输出层

在词向量经过了多个Decoder层堆叠后,来到了最后(用 $H S$ 来代指这个经过了多层处理后的词向量),这个时候,我们将HS通过一个线性层拟合,然后进入到Loss部分,这一块我个人理解应该是类似GPT模型那样的做法:将这个HS和词表做乘积,得出一个概率分布。然后用贪心算法,拿到概率最大的做预测。最后计算损失,这就是这一步命名为Loss的原因(可能)

对于上面的这一步骤:"将这个 $H S$ 和词表做乘积,得出一个概率分布。"可以这样理解:

$H S$ 是隐藏状态矩阵，其尺寸为 $1,1,hidden\_dim)$
$V$ 是词表,尺寸为: $vocab\_size, hidden\_dim)$

隐藏状态矩阵 $H S$ 降维：

从 $hidden_dim ) (1, 1, \text{hidden\_dim})$ 降维到 $hidden_dim ) (1,\text{hidden\_dim})$ ，记作 $\mathbf{h}$ ：

$\mathbf{h} = \mathbf{HS}[0, 0, :]$
词表矩阵 $\mathbf{V}$ 的转置：

词表矩阵 $\mathbf{V}$ 的转置记作 $\mathbf{V}^\top$ ，其尺寸为 $hidden_dim , vocab_size ) (\text{hidden\_dim}, \text{vocab\_size})$ ：
矩阵相乘：

计算 $\mathbf{h}$ 和 $\mathbf{V}^\top$ 的矩阵乘积，得到概率分布 $\mathbf{p}$ ，其尺寸为 $vocab_size ) (1,\text{vocab\_size})$ ：

$\mathbf{p} = \mathbf{h} \cdot \mathbf{V}^\top$
最大值的索引作为预测的值：

找出 $\mathbf{p}$ 中最大的值对应的索引，记作 $\hat{y}$ ：

$\hat{y} = \arg\max(\mathbf{p})$