解密 Qwen: 探索类 Llama2 架构的技术内幕

会飞的Anthony

于 2024-09-17 09:00:00 发布

阅读量594

点赞数 15

分类专栏：信息系统自然语言处理人工智能文章标签：自然语言处理人工智能

本文链接：https://blog.csdn.net/ljd939952281/article/details/142307857

版权

人工智能同时被 3 个专栏收录

133 篇文章 2 订阅

订阅专栏

信息系统

100 篇文章 0 订阅

订阅专栏

自然语言处理

24 篇文章 0 订阅

订阅专栏

在快速发展的人工智能领域，Qwen 模型是新一代的大规模语言模型，它基于 Transformer 架构，专为处理自然语言任务而设计。作为一款与 Llama2 架构相似的模型，Qwen 不仅具有强大的生成能力，还能支持各种下游任务，如语言生成、文本分类等。Qwen 通过优化注意力机制（如 Flash Attention 和 GQA）以及高效的内存管理，使其在推理速度、计算效率和扩展性方面都有出色的表现。

本文将深入探讨 Qwen 模型的内部架构，从核心模块到注意力机制的详细实现，帮助大家理解这款先进语言模型背后的技术原理。如果你对大型语言模型感兴趣，那么这篇文章将是一个非常好的起点。

Qwen 的整体架构

Qwen的模型架构与Llama2非常相似，采用了Transformer架构作为基础。整个模型的工作流程如下：

Tokenizer：将输入文本转换为词表中的数值。
Embedding：将词表中的数值映射为向量表示。
Attention Mask：用于设定哪些位置可以互相“看到”，例如双向、单向等。
模型输出：通过不同的线性层和损失函数，可以应用于各种下游任务，如Casual、SeqCls等。

下面我们详细讲解Qwen模型的各个组成部分，并提供相关代码。

1. Qwen2Config: 模型配置

Qwen的Qwen2Config是所有模型组件的配置文件，它定义了模型的核心超参数，例如词汇大小、隐藏层维度和注意力头数。通过修改这些参数，我们可以调整模型的大小和复杂度。

代码示例：

class Qwen2Config:
    def __init__(self, vocab_size=30522, hidden_size=768, num_hidden_layers=12, num_attention_heads=12, **kwargs):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.pad_token_id = kwargs.get('pad_token_id', 0)

# 实例化配置
config = Qwen2Config(vocab_size=50257, hidden_size=1024, num_hidden_layers=24, num_attention_heads=16)

2. Qwen2Model: 模型初始化与前向传播

Qwen2Model是Qwen的核心模型类。它负责初始化嵌入层、多层解码器和归一化层，并在前向传播时，将输入文本逐层处理并生成输出。

2.1 初始化模型

在初始化时，Qwen会创建一个词嵌入层来将输入的词ID转换为词向量。随后，它会构建多个解码器层，每一层包含注意力机制、MLP 和标准化层。

代码示例：

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(config) for _ in range(config.num_hidden_layers)])
        
        # RMSNorm 标准化
        self.norm = Qwen2RMSNorm(config.hidden_size, eps=1e-6)

    # 初始化权重
    self.post_init()

2.2 前向传播

在前向传播过程中，输入的词ID首先通过嵌入层转化为词向量，接着通过每一层的解码器进行处理，最后输出经过标准化得到最终结果。

代码示例：

def forward(self, input_ids, attention_mask=None):
    # 将输入词ID转为词向量
    inputs_embeds = self.embed_tokens(input_ids)
    hidden_states = inputs_embeds

    for decoder_layer in self.layers:
        # 逐层传递隐藏状态
        hidden_states = decoder_layer(hidden_states, attention_mask=attention_mask)[0]

    # 最后一层标准化
    hidden_states = self.norm(hidden_states)
    
    return hidden_states

3. Qwen2DecoderLayer: 解码器层

Qwen的每个解码器层由三个核心模块组成：自注意力机制、MLP模块和归一化层。每一层负责处理输入并产生新的隐藏状态。

3.1 初始化

每个解码器层初始化时会创建一个自注意力模块（Qwen2Attention）和一个MLP模块，并将输入状态标准化。

代码示例：

class Qwen2DecoderLayer(nn.Module):
    def __init__(self, config: Qwen2Config):
        super().__init__()
        self.self_attn = Qwen2Attention(config)
        self.mlp = Qwen2MLP(config)
        self.input_layernorm = Qwen2RMSNorm(config.hidden_size)
        self.post_attention_layernorm = Qwen2RMSNorm(config.hidden_size)

3.2 前向传播

在前向传播时，输入状态首先经过标准化，然后进入自注意力模块，接着再经过MLP层处理，最后输出。

代码示例：

def forward(self, hidden_states, attention_mask=None):
    # 先进行归一化
    residual = hidden_states
    hidden_states = self.input_layernorm(hidden_states)
    
    # 自注意力机制
    hidden_states, _ = self.self_attn(hidden_states, attention_mask=attention_mask)
    hidden_states = residual + hidden_states

    # 再次归一化
    residual = hidden_states
    hidden_states = self.post_attention_layernorm(hidden_states)
    
    # MLP 层
    hidden_states = self.mlp(hidden_states)
    hidden_states = residual + hidden_states

    return (hidden_states,)

4. Qwen2Attention: 自注意力机制

Qwen中的自注意力模块支持多头注意力机制，并能通过不同的实现来优化性能（如GQA或Flash Attention）。

4.1 初始化

初始化自注意力模块时，主要创建了线性变换用于生成查询（Query）、键（Key）和值（Value），并设定多头注意力的维度。

代码示例：

class Qwen2Attention(nn.Module):
    def __init__(self, config: Qwen2Config):
        super().__init__()
        self.num_heads = config.num_attention_heads
        self.head_dim = config.hidden_size // config.num_attention_heads
        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim)
        self.k_proj = nn.Linear(config.hidden_size, self.head_dim)
        self.v_proj = nn.Linear(config.hidden_size, self.head_dim)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size)

4.2 前向传播

在前向传播时，自注意力机制将查询、键和值通过点积计算得到注意力分数，再通过softmax将这些分数转化为权重，并加权生成最终输出。

代码示例：

def forward(self, hidden_states, attention_mask=None):
    query = self.q_proj(hidden_states)
    key = self.k_proj(hidden_states)
    value = self.v_proj(hidden_states)

    # 计算注意力分数
    attn_scores = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(self.head_dim)

    if attention_mask is not None:
        attn_scores += attention_mask

    # softmax 得到注意力权重
    attn_weights = nn.functional.softmax(attn_scores, dim=-1)
    attn_output = torch.matmul(attn_weights, value)
    attn_output = self.o_proj(attn_output)
    
    return attn_output, attn_weights

5. Qwen2MLP: 多层感知机

Qwen2MLP模块由三层线性变换组成，用于处理非线性转换。

代码示例：

class Qwen2MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size)
        self.act_fn = nn.GELU()

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

6. Qwen2RMSNorm: 标准化层

RMSNorm是Qwen中的标准化层，用于对每一层的输出进行规范化处理，确保稳定的训练和推理效果。

代码示例：

class Qwen2RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, hidden_states):
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        normed_hidden_states = hidden_states / torch.sqrt(variance + self.eps)
        return self.weight * normed_hidden_states