Tiny-universe学习笔记1:Qwen-blog
本文是参与Datawhale Tiny-universe组队学习的第一篇学习笔记,参考链接:https://github.com/datawhalechina/tiny-universe
Tiny-universe学习笔记1:Qwen-blog
Qwen整体架构与Llama2类似,具体如下图所示:
其中:
tokenizer
将文本转为词表里面得数值(文本转为对应得index)- 数值经过
embedding
得到一一对应得向量 attention_mask
是用来看见左边、右边,双向等待来设定- 各类下游任务,
Casual
、seqcals
等,基本都是基础模型model
后面接对应得Linear
层,还有损失函数不一样。
1. Qwen2Config
Qwen2Config
继承自PretrainedConfig
,PretrainedConfig
中有一些是可以被子类重写的类属性:
model_type(str)
:模型的类型attribute_map(Dict[str, str])
:存储模型特定的属性名称和transformers属性标准名称之间的映射关系vocab_size(int)
:词表中token的数量hidden-size(int)
:模型隐藏层的维度num_attention_deads
:每个 Encoder 中 attention 层的 head 数
Qwen2Config
中除了从PretrainedConfig
中继承的参数外,需要注意的参数有
num_hidden_layers(int)
:隐藏层的数量,由于Qwen是Only-Decoder架构的,所以这个参数也可以被描述为transformer decoder中的隐藏层的数量rms_norm_eps(float)
:用于防止RMS中发生除0的情况
1.1 Qwen2Model
1.1.1 初始化
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, layer_idx) for layer_idx in range(config.num_hidden_layers)]
)
self._attn_implementation = config._attn_implementation
self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.gradient_checkpointing = False
# Initialize weights and apply final processing
self.post_init()
-
__init__
方法根据传入的Qwen2Config
初始化父类的属性和模型属性:- 首先设置了模型的两个属性:
padding_idx
:词表中用于指定填充标记的索引id
vocab_size
:词表的大小
- 初始化了模型的嵌入层、解码器层、归一化层
- 嵌入层(
nn.Embedding
):在嵌入层中,模型将输入的标记映射成密集的向量表示。 - 解码器层(
nn.ModuleList
):模型包含多个解码器层,这些层都由Qwen2DecoderLayer
定义的 - 归一化层(
Qwen2RMSNorm
):归一化层使用的是 Root Mean Square Layer Normalization
- 嵌入层(
- 设置了是否使用
gradient_checkpoint
主要是用来节省显存 - 调用
post_init()
完成一些初始化和准备检查的代码
- 首先设置了模型的两个属性:
-
post_init()
函数主要是对参数进行初始化,以及对初始化梯度检查点作用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
以下代码为formard核心计算部分:
# input_ids经过嵌入层后得到inputs_embeds
inputs_embeds = self.embed_tokens(input_ids)
# embed positions
hidden_states = inputs_embeds
for idx, decoder_layer in enumerate(self.layers):
# 将所有的hidden_states保存成tuple
if output_hidden_states:
all_hidden_states += (hidden_states,)
# 将hs送入每一层decoder_layer
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
if output_hidden_states:
all_hidden_states += (hidden_states,)
整体流程为将第一个hidden_states
,也就是inputs_embeds
,输入到每一层decoder-layer
中,再对最后一个decoder-layer
的输出,利用归一化层进行处理得到最后的输出。
如果保存output_hidden_states的话,就是第一个为input_ids进行emb,然后保存到n-1层的decoder_layer的输出hidden_states
,再加上最后一层layer的输出hs进行过norm后的hidden_states
.
1.2 Qwen2DecoderLayer
1.2.1 Qwen2DecoderLayer初始化
Qwen2DecoderLayer
分为三个模块:
- attention:一般为
Qwen2Attention
- MLP:Qwen2MLP
- norm:Qwen2RMSNorm
QWEN2_ATTENTION_CLASSES = {
"eager": Qwen2Attention, # 一般情况下是这个
"flash_attention_2": Qwen2FlashAttention2,
"sdpa": Qwen2SdpaAttention,
}
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
内容是一样的,都用到了Qwen2RMSNorm
,只是应用的顺序不一样。
1.2.2 Forward
Qwen2DecoderLayer
流程图如下所示:
核心代码如下:
residual = hidden_states
# 标准化后送入attn
hidden_states = self.input_layernorm(hidden_states) # 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,
)
# 残差与新的hidden_states相加
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 初始化
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
# num_key_value_heads表示键值对的头数
self.num_key_value_heads = config.num_key_value_heads
# num_key_value_groups表示键值对的组数
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_proj、k_proj、v_proj、o_proj为四个Linear操作,Lora微调基本都对它们进行更改
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,
)
1.3.2 Forward
# 获取形状信息,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)
# 将旋转位置嵌入应用于查询和键张量。使用了旋转位置嵌入的余弦和正弦部分,将它们与查询和键张量相乘,并将结果相加,从而实现旋转位置嵌入的效果
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
attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
# 然后 attn_weights 加上 attention_mask,实现读取顺序
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操作
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)
# 返回结果
return attn_output, attn_weights, past_key_value
1.4 Qwen2 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
1.5 Qwen2RMSNorm
计算公式:
R
M
S
N
o
r
m
(
x
)
=
x
1
n
∑
i
=
1
n
w
i
2
+
ϵ
RMSNorm(x)=\frac{x}{\sqrt{\frac{1}{n}\sum_{i=1}^n w_i^2 + ϵ}}
RMSNorm(x)=n1∑i=1nwi2+ϵx
- x是层的输入的
hidden_state
- wi 表示的是
hidden_state
的最后一个维度的值 - n 表示上面输入的最后一个维度的数量。
- ϵ 表示是很小的数,防止除0。
代码实现:
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)