Qwen 个人笔记
Qwen的整体架构与Llama2类似,如下图所示:
1 Qwen2Config
1.1 Model
1.1.1 初始化
- 设置了模型的两个属性:
padding_idx
(用于指定填充标记的索引),vocab_size
(词汇表的大小) - 初始化了模型的嵌入层、解码器层、归一化层
- 嵌入层(
nn.Embedding
):模型使用嵌入层将输入的标记映射成密集的向量表示。 - 解码器层(
nn.ModuleList()
):模型包含多个解码器层,这些层都是由Qwen2DecoderLayer
定义 - 归一化层
Qwen2RMSNorm
:归一化层使用的是 Root Mean Square Layer Normalization - 设置了是否使用
gradient_checkpoint
主要是用来节省显存 - 调用
post_init()
完成一些初始化和准备检查的代码 (主要是对参数进行初始化,以及初始化梯度检查点作用)jie
整个 Qwen2Model
类的目的是为了构建一个完整的Qwen2模型,它可以通过传入不同的配置参数来适应不同的任务和数据集
python
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
这样的方法通常用于设置一些在构造函数中无法完成的初始化工作,因为这些工作依赖于构造函数中已经执行的步骤
python
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
自然语言处理模型中的前向传播过程的一部分(重要部分),Transformer 模型中的典型操作
python
inputs_embeds = self.embed_tokens(input_ids)#这行代码使用 `embed_tokens` 方法将输入的 `input_ids`转换为嵌入向量
# embed positions
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]
# 将最后layers输出后的hidden_states进行标准化
hidden_states = self.norm(hidden_states)
# 如果需要输出所有层的隐藏状态,则将最后一层的 `hidden_states` 添加到 `all_hidden_states` 中
if output_hidden_states:
all_hidden_states += (hidden_states,)
1.2 Qwen2DecoderLayer
1.2.1 初始化
layer
三件套:attn
+MLP
+norm
这段代码定义了一个名为 Qwen2DecoderLayer
的类,它继承自 PyTorch 的 nn.Module
。这个类代表了一个 Transformer 模型中的解码器层,它包含了自注意力机制(self-attention)和多层感知机(MLP),以及两种归一化层
python
# 字典 映射了不同的注意力实现类
QWEN2_ATTENTION_CLASSES = {
"eager": Qwen2Attention, # 默认的注意力实现,一般情况下是这个
"flash_attention_2": Qwen2FlashAttention2, # 一个优化的注意力实现
"sdpa": Qwen2SdpaAttention, # 一种特殊的注意力实现
}
# 解码器层的类,它包含自注意力、MLP、输入归一化和后注意力归一化
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)
1.2.2 Forward
这段代码体现了 Transformer 模型中的两个关键特性:残差连接和归一化。[[残差连接]]有助于避免深层网络中的梯度消失问题,而归一化则有助于稳定训练过程,加快收敛速度。代码中的 **kwargs
表示函数还可能接收其他未明确列出的关键字参数
python
residual = hidden_states# 保存原始的 `hidden_states` 到 `residual` 变量中,用于后面的残差连接
# 标准化后送入attn
hidden_states = self.input_layernorm(hidden_states) # 将 `hidden_states` 通过输入归一化层(`input_layernorm`),RMSNorm标准化
# 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,
**kwargs,
)
# 再次应用残差连接,将 MLP 之前的 `hidden_states`(保存在 `residual` 中)与 MLP 的输出相加
hidden_states = residual + hidden_states
# Fully Connected
residual = hidden_states
# 同样的RMSNorm标准化
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.3 Qwen2Attention
1.3.1 初始化
这段代码定义了一个名为 Qwen2Attention
的类,它继承自 PyTorch 的 nn.Module
。这个类实现了多头自注意力机制,这是 Transformer 架构中的一个核心组件
这个自注意力类的实现包括了对多头注意力的支持,以及对旋转嵌入的使用。这种自注意力机制可以用于处理序列数据,如自然语言或时间序列,并且是许多现代 NLP 模型的基础
python
# 定义了一个自注意力类的蓝图
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,
)
-
max_position_embeddings
确定了模型能够编码的位置索引的最大值,通常对应于模型能够处理的最长序列长度 -
rope_theta
决定了旋转矩阵的周期性,它影响着位置编码的周期性变化
python
max_position_embeddings (`int`, *optional*, defaults to 32768):
The maximum sequence length that this model might ever be used with.
rope_theta (`float`, *optional*, defaults to 10000.0):
The base period of the RoPE embeddings.
1.3.2 Forward
这段代码是一个 Transformer 模型中自注意力(Self-Attention)机制的实现,具体来说是多头自注意力(Multi-Head Self-Attention)的计算过程
这个过程是 Transformer 模型中自注意力机制的核心,它允许模型在处理序列数据时考虑到不同位置之间的关系。旋转位置嵌入(RoPE)是一种特殊的位置编码方式,它通过旋转变换将位置信息融入到查询和键中,从而提高模型对序列顺序的感知能力
python
# 获取形状信息,hidden_states输入的为(bs,T,hd)
# 获取 `hidden_states` 的形状信息,其中 `bsz` 是批次大小,`q_len` 是序列长度,`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),为了将输入的隐藏状态转换为适合多头自注意力计算的形式,每个头可以独立地处理序列的一部分,从而实现并行处理和更细粒度的表示学习
# `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和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
# 计算查询和键的点积除以缩放因子(`head_dim` 的平方根),得到注意力权重
attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
# 然后 attn_weights 加上 attention_mask,实现读取顺序
# 对注意力权重进行 softmax 操作以进行归一化,然后应用 dropout 以防止过拟合
attn_weights = attn_weights + attention_mask
# 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)
# 转置,修改形状等reshape操作
# 对注意力输出进行转置和重塑,以恢复到原始的 `bsz x q_len x hd` 形状
attn_output = attn_output.transpose(1, 2).contiguous()
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
# 最后在进行一次o_proj,通过输出投影层 `o_proj` 对注意力输出进行最后的线性变换
attn_output = self.o_proj(attn_output)
# 返回结果,注意力权重和可能使用的 `past_key_value`(在缓存注意力权重时使用)
return attn_output, attn_weights, past_key_value
1.3.3 细节Debug
1.3.3.1 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)定义初始张量
python
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
相同的维度。方便后续计算
python
# 定义输入x, n_rep是需要重复的次数,在这里一般是组数
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
相乘了
python
#(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
的类,它继承自 PyTorch 的 nn.Module
。这个类实现了一个多层感知机(MLP)结构,通常用于 Transformer 模型中的前馈网络(Feed-Forward Network,FFN)
这种 MLP 结构在 Transformer 模型中很常见,它通过引入非线性激活函数和中间层来增加模型的表达能力。在 Qwen2MLP
中,使用了一种特殊的结构,其中 gate_proj
的输出与 up_proj
的输出相乘,这种结构有助于模型学习输入数据的复杂特征
python
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
1.5 Qwen2RMSNorm
计算公式: 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。
这个类实现了 RMS(Root Mean Square)归一化
RMS 归一化与传统的 LayerNorm 不同,它只使用方差(不包括均值)来进行归一化,这使得它在某些情况下可能更简单或更有效。在 Qwen2RMSNorm
中,通过对方差进行归一化,然后将结果乘以一个可学习的权重,这有助于模型调整归一化后的特征表示
python
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)