DW大模型学习笔记

Qwen 模型学习笔记

1. Qwen 整体架构

Qwen的整体架构与Llama2类似，如下所示：

• Tokenizer: 将文本转为词表里面的数值。
• Embedding: 数值经过embedding得到一一对应的向量。
• Attention Mask: 用来看见左边、右边，双向等等来设定。
• 下游任务: 各类下游任务 (Casual, seqcls等) 基本都是基础模型model后面接对应的Linear层，损失函数不同。

2. Qwen2Config

Qwen2Config包含一些自定义的超参数，例如 vocab_size, hidden_size, num_hidden_layers, num_attention_heads等。类似于dict可以调用里面的超参数，比如 config.pad_token_id。

3. Qwen2Model

3.1 初始化

设置了模型的两个属性：

• padding_idx: 用于指定填充标记的索引
• vocab_size: 词汇表的大小

初始化了模型的嵌入层、解码器层、归一化层：

• 嵌入层（nn.Embedding）：将输入的标记映射成密集的向量表示。
• 解码器层（nn.ModuleList）：包含多个解码器层，这些层由 Qwen2DecoderLayer 定义。
• 归一化层（Qwen2RMSNorm）：使用 Root Mean Square Layer Normalization。

设置了是否使用 gradient_checkpoint 以节省显存。调用 post_init() 完成一些初始化和准备检查的代码。

3.2 Forward

核心主干讲解：

inputs_embeds = self.embed_tokens(input_ids)
hidden_states = inputs_embeds

for idx, decoder_layer in enumerate(self.layers):
    if output_hidden_states:
        all_hidden_states += (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,
    )
    hidden_states = layer_outputs[0]
    
hidden_states = self.norm(hidden_states)
    
if output_hidden_states:
    all_hidden_states += (hidden_states,)

最终以 BaseModelOutputWithPast 的形式输出。

4. Qwen2DecoderLayer

4.1 初始化

DecoderLayer的三件套：attn + MLP + norm

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)
        self.mlp = Qwen2MLP(config)
        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)

input_layernorm 和 post_attention_layernorm 的内容相同，但应用的顺序不同。

4.2 Forward

residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
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

5. Qwen2Attention

5.1 初始化

核心参数解析：

• num_key_value_heads: 表示键值对的头数
• num_key_value_groups: 表示键值对的组数，计算为 num_heads // num_key_value_heads (GQA的实现)
• q_proj, k_proj, v_proj, o_proj 四个 Linear 操作。

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

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

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

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

6. Qwen2MLP

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

7. Qwen2RMSNorm

class Qwen2RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        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)

8. 关键细节解析

8.1 GQA

• repeat_kv 方法用于扩展 key 和 value 的头维度，使其与 query 的头维度匹配。
• expand 方法用于对张量进行扩展，不实际分配新的内存。
• repeat 方法通过实际复制数据来扩展张量，占用更多内存。

8.2 Rotary Position Embedding

• 通过旋转位置嵌入，使每个 token 既有相对位置信息，又有绝对位置信息

。

• apply_rotary_pos_emb 函数将旋转位置嵌入应用于查询和键张量。

总结

Qwen的整体架构和实现细节为我们提供了一个强大的自然语言处理模型。通过对其结构和实现细节的深入了解，我们可以更好地理解和应用Qwen模型来解决各种自然语言处理任务。