从零手搓大模型之路（一、学习Qwen模型架构）

最新推荐文章于 2025-03-06 13:48:25 发布

晓光尘梦

最新推荐文章于 2025-03-06 13:48:25 发布

阅读量1.3w

点赞数 73

分类专栏：大模型学习文章标签：语言模型

本文链接：https://blog.csdn.net/qq_37021523/article/details/138901191

版权

大模型学习专栏收录该内容

10 篇文章

订阅专栏

写在前面

五月，依旧是给自己挖坑，参加了DataWhale的五月自学课堂（从零手搓大模型实战）。

说是从零手搓，但深知自己远远没有大佬水平，达不到研究透透的程度，所以也就是看源码了解下内部逻辑，简简单单写写心得，期望有朝一日也能成为大佬。

So，今天挖的坑让明天的自己哭着来填~~（哭？哭也是要算时间的！）~~。

饭得一口一口吃，路得一点一点走。

下面是课程的链接：

从零手搓大模型

本篇博文的所研究学习的源码是Qwen2，链接如下：

https://github.com/huggingface/transformers/tree/v4.39.3/src/transformers/models/qwen2

同时也链接一位大佬的知乎，对Qwen2源码逐行分析讲解，本博文也有一部分借鉴了该作：

Qwen2逐行代码分析

本博文里面图来源于课程链接。

开卷开卷！~~(怎么感觉东西越学越多)~~

1、Qwen2整体介绍

这张图第一眼看上去确实很唬人，如果从左往右一点点耐心看，还是能看出些端倪。

其核心流程实际上只是最左边Qwen2主干的虚线框区域，文本Text经过词向量化处理后，得到一个向量（这么做的原因是每个文本长度不一样，为了让机器更好学习，便统一编码将其转化为某个固定长度的向量来互相区分，即文本转序列），至此便完成了Tokenizer和Embedding。

在Layers阶段，图中是进行了三次，每一次的Decoder中，输入都会先copy一份用于残差连接（将输入直接加入到输出，为的是防止神经网路的梯度爆炸和消失，也有一定原因是加强输入数据的影响）。

输入经过RMSNorm（通常矩阵计算要进行归一化防止不同特征的取值过大，常用的是layer norm，也就是每一项减去样本的均值，再除以样本的方差；而RMS则是去除了减去均值的操作，以便提升效率）处理后便来到Attn注意力机制。

（关于注意力机制部分之前看过的沐神讲论文，视频链接如下，后续也会补一文写心得）

Transformer论文逐段精读【论文精读】_哔哩哔哩_bilibili

之后再经过MLP数据处理，得到最后的结果。

大致框架如上，打开源码的路径，下面有四个文件：

configuration_qwen2	qwen2里面参数的配置文件
modeling_qwen2	qwen2的主体模型源码
tokenization_qwen2	qwen2的词向量化代码
tokenization_qwen2_fast	qwen2的词向量化代码

注：run_demo是我按照视频教程敲的运行文件，并不是源码里带的。

源码架构如下:

- Qwen2RMSNorm: RMS归一化层

- Qwen2RotaryEmbedding: 旋转位置编码

- Attention

- Qwen2Attention: 注意力层

- Qwen2FlashAttention2: 使用Flash Attention 2.0版本加速的注意力层

- Qwen2SdpaAttention: 使用Sdpa(pytorch自带的加速, Scaled Dot-Product Attention)加速的注意力层

- Qwen2DecoderLayer: 编码层,核心结构,之后就是堆叠

- Qwen2PreTrainedModel: 预训练类

- Qwen2Model: 不带head的Qwen2模型

- Qwen2ForCausalLM: 带Causal LM head的Qwen2模型

- Qwen2ForSequenceClassification: 带序列分类头的Qwen2模型

其它函数简介:

- _get_unpad_data: 在flash attention的数据预处理中会用到.主要是对attention mask进行一些操作.

- rotate_half: 在旋转位置编码中用到

- apply_rotary_pos_emb: 对数据主要是注意力运算中的q,k做旋转位置编码

- repeat_kv: 主要是在MQA(Multi-Query Attention)和GQA(Group-Query Attention)中用到,因为q head数量是k,v head的数量的整数倍

后面的内容基本是围绕modeling_qwen2。

2、源码学习笔记

整体的代码很长，不过有一部分是关于几种注意力机制的不同实现。

按照前面的框架调用流程，逐一看一下源码。

研究学习下面的源码前，应该对照上面调用整体流程的demo示例，也就是run_demo。

1）Qwen2Model

这是Qwen2的主体类，继承了源码里面的Qwen2PreTrainedModel，父类中主要是一些参数设定以及模型中各个层级的权重初始化。

说白了，在这里初始化过程中，就要完成一系列参数设定，便于在forward中使用。

class Qwen2Model(Qwen2PreTrainedModel):
    """
    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2DecoderLayer`]

    Args:
        config: Qwen2Config，来源于前面提到的configuration_qwen2.py以及自己实际的传参
    """

    def __init__(self, config: Qwen2Config):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        # 构建主体的embedding
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        
        # 构建主体的layers层
        self.layers = nn.ModuleList(
            [Qwen2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self._attn_implementation = config._attn_implementation
        
        # 创建RMS归一化对象
        self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        # 梯度检查点，一种训练策略
        self.gradient_checkpointing = False

        # Initialize weights and apply final processing
        # 主要是对参数进行初始化
        self.post_init()

其中，gradient_checkpointing（梯度检查点）是一种训练优化策略，会降低显卡的内存占用，但相对应的，会延长训练时间。其工作原理是用时间换空间。检查点不保存整个计算所有中间结果帮助进行反向传播的计算，而是在反向传播的过程中重新计算中间结果。

在该类的forward函数中，前面一大堆都是对环境变量的设定与配置，如输入数据的shape，是否梯度检查，是否添加位置信息（注意力机制中的数据位置编码），注意力机制的实现方式等等。配合示意图查看很容易理解。

2）Qwen2DecoderLayer

# 注意力机制的实现方式字典
QWEN2_ATTENTION_CLASSES = {
    "eager": Qwen2Attention,
    "flash_attention_2": Qwen2FlashAttention2,
    "sdpa": Qwen2SdpaAttention,
}


class Qwen2DecoderLayer(nn.Module):
    def __init__(self, config: Qwen2Config, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size

        # 确定注意力机制的实现方式
        if config.use_sliding_window and config._attn_implementation != "flash_attention_2":
            logger.warning_once(
                f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "
                "unexpected results may be encountered."
            )
        self.self_attn = QWEN2_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)

        # 多层感知机
        self.mlp = Qwen2MLP(config)
        # 创建两个对象，RMS标准化
        self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

其中，在设定注意力机制的实现方式时，Qwen2Attention是本博文后面会注释学习的；Qwen2FlashAttention2是使用flash attention进行加速，目前该加速已经是attention计算的必备，只要有GPU加速；Qwen2SdpaAttention为Scaled Dot Product Attention，也是论文《Attention Is All You Need》中提到的。

对于forword函数，配合这部分的示意图更好理解：

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: Optional[bool] = False,
        use_cache: Optional[bool] = False,
        **kwargs,
    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
        if "padding_mask" in kwargs:
            warnings.warn(
                "Passing `padding_mask` is deprecated and will be removed in v4.37. "
                "Please make sure use `attention_mask` instead.`"
            )

        residual = hidden_states

        # 进入注意力前的RMS标准化
        hidden_states = self.input_layernorm(hidden_states)

        # Self Attention
        
        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,
        )
        # 残差连接
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        # MLP前面的RMS标准化
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (self_attn_weights,)

        if use_cache:
            outputs += (present_key_value,)

        return outputs

3）Qwen2Attention

配合框架图看更容易理解，这个类是注意力机制的实现形式，这部分先直接copy了

class Qwen2Attention(nn.Module):

    def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        if layer_idx is None:
            logger.warning_once(
                f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
                "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
                "when creating this class."
            )

        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})."
            )

        # 构建注意力机制的q，k，v，以及框架中新填的o四个线性层
        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,
        )

至于forword，也是直接copy，不过这部分注意力机制新添了旋转位置嵌入操作，在本博文后面进行理解学习。

关于注意力部分实现方式理解呢，更多是结合之前看的沐神讲论文，详细内容还是打算放在新的博文章节里了，之后会重新编辑此博文放出链接。~~（才不是要交作业了，内容太多来不及编辑呢）~~

源码里面的步骤是：

首先将hidden_states送入Linear中得到query、key与value。
使用旋转位置嵌入操作rotary_emb，使用了旋转位置嵌入的余弦和正弦部分，将他们与query和key相乘，并将结果相加，从而实现旋转位置嵌入的效果。
将key_states和value_states重复group次，再执行dot attn操作。
在dot attn操作后得到attn_weights,加上attention_mask从而实现读取掩盖操作，在经过softmax与value_states相乘。得到attn_output。
再将上述的attn_output进行reshape操作，送入o_proj，得到最终的输出。

    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,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        if "padding_mask" in kwargs:
            warnings.warn(
                "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
            )
        # 获取形状信息,hidden_states输入的为(bs,T,hd)
        bsz, q_len, _ = hidden_states.size()
        
        # 对hidden_states进行Linear生成query、key、value
        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)
        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)

        kv_seq_len = key_states.shape[-2]
        if past_key_value is not None:
            if self.layer_idx is None:
                raise ValueError(
                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
                    "with a layer index."
                )
            kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)


        # 将旋转位置嵌入应用于查询和键张量。使用了旋转位置嵌入的余弦和正弦部分，将它们与查询和键张量相乘，并将结果相加，从而实现旋转位置嵌入的效果
        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)

        if past_key_value is not None:
            cache_kwargs = {"sin": sin, "cos": cos}  # 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
        # 先将key_states和value_states重复了num_key_value_groups次
        key_states = repeat_kv(key_states, self.num_key_value_groups)
        value_states = repeat_kv(value_states, self.num_key_value_groups)
        
        # 使用dot attn实现q*kT/hd_d^0.5
        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
            raise ValueError(
                f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
                f" {attn_weights.size()}"
            )

        if attention_mask is not None:
            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
                )
            # 然后 attn_weights 加上 attention_mask，实现读取顺序
            attn_weights = attn_weights + attention_mask

        # upcast attention to fp32
        # softmax + dropout + values_states相乘
        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)

        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()}"
            )
        
        # 转置，修改形状等reshape操作
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
        
        # 最后在进行一次o_proj
        attn_output = self.o_proj(attn_output)

        if not output_attentions:
            attn_weights = None

        return attn_output, attn_weights, past_key_value

4）Qwen2RMSNorm

配合公式，就是计算均方根标准化。这里也就把教程里面的copy过来了。

其中，x为输入的矩阵数据，wi为最后一个维度的值，n表示最后一个维度的数量。

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))

        # 设置这部分主要是为了防止除数为0，所以加一个小数
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype

        # 将数据转为torch.float32，保证足够的精度下，也有较快的运行速度
        # 尤其时数据量比较大的情况下
        hidden_states = hidden_states.to(torch.float32)

        # 计算方差，这里keepdim是均值时保留维度，确保下面与hidden_states进行矩阵运算
        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)

5）Qwen2MLP

配合结构示意图，这个类平平无奇，很传统的多层感知机，copy一下再额外加些注释：

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)

        # ACT2FN是自定义的一系列激活函数字典，config的定义来获取这里用到的激活函数类型
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        # 两个尺寸相同的网络，其中一个经过激活后与另一个结果相乘，得到的结果过第三层
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))