Datawhale------Tiny-universe学习笔记——Qwen(2)

最新推荐文章于 2024-09-25 09:27:18 发布

人生百态，人生如梦

最新推荐文章于 2024-09-25 09:27:18 发布

阅读量71

点赞数

分类专栏： Datawhale打卡文章标签：学习笔记人工智能

原文链接：https://github.com/datawhalechina/tiny-universe/tree/main/content/Qwen-blog

版权

Datawhale打卡专栏收录该内容

4 篇文章 0 订阅

订阅专栏

这一节基本转载于（稍有改变）:

tiny-universe/content/Qwen-blog at main · datawhalechina/tiny-universe · GitHub

2. Qwen2Attention

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

config里面的参数可直接看Qwen2Config里面的介绍

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.

其中有几个核心参数：

num_key_value_heads:表示在注意力机制中使用的“头”（Heads）的数量。每个头可以学习序列的不同部分或不同特征。多头注意力（Multi-Head Attention）通过将查询（Query）、键（Key）、值（Value）通过不同的线性变换分割成多个头，然后并行处理，最后再将结果合并起来，以增强模型的表达能力。

num_key_value_groups:表示键值对的组数，这通常与num_key_value_heads有关。在一些模型中，可能会将头分组，每组头共享相同的参数。这样可以在保持多头注意力的同时减少模型的参数量。计算为num_heads // num_key_value_headsGQA的实现！！

q_proj,k_proj,v_proj:这些是三个线性变换（通常使用nn.Linear实现），分别用于将输入序列映射到查询（Query）、键（Key）和值（Value）的表示空间。在自注意力机制中，模型会计算查询与所有键的相似度，然后使用这个相似度来加权值（Value）。

o_proj:这是另一个线性变换，用于将多头注意力的输出（即多个头的输出合并后的结果）映射回原始的表示空间。这个步骤是必要的，因为它允许模型在注意力层之后继续处理信息。

后续LoRa也基本都对q得四个操作动的刀子.

2.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

首先将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，得到最终的输出.

想具体了解注意力机制得话的话可以看我另一篇：

AI大模型开发——transformer模型超全完整版（0基础可懂)_大数据 transformer 与大模型技术学习 ppt-CSDN博客

3 细用Debug

3.1 GQA

主旨:GQA和MQA不需要在推理的过程存储那么多的kv cache, 那么kv cache占用的显存就变小，那么我们LLM serving可以处理的请求数量就更多.

定义初始张量：

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]

之后进行扩展key，value的操作：

在GQA中，key和value都要比query小group倍，但是为在后续做矩阵乘法时方便，我们需要先把key和value的head利用expand扩展张量到和query相同的维度。方便后续计算。

# 定义输入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)

矩阵乘法得到score与output 后面就是征程的kqv相乘了

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

为什么要用expand之后再reshape而不能直接用tensor自带的repeat?

expand 方法用于对张量进行扩展，但不实际分配新的内存。它返回的张量与原始张量共享相同的数据

repeat 方法通过实际复制数据来扩展张量。它返回的新张量不与原始张量共享数据，扩展后的张量占用了更多的内存。

3.2 apply_rotary_pos_emb

位置编码的含义是对每一个token的每一个dim赋予不同的位置信息。公式定义:

概念：通过旋转编码，使得每个token既有相对位置信息，又有绝对位置信息。

既能以自注意力矩阵偏置的形式作用于 At,s,直接反映两个token的相对位置信息，又能拆解到向量 qt 和 ks 上，通过直接编码token的绝对位置实现。

RoPE本质是实现对特征向量的旋转操作，如果以二维特征向量举例，对于相邻两个token来说，其对应同一个 θ,其定义为:

可得，其本质就是: qt, ks 旋转后的结果，就是 qt, ks乘上cos再加上 qt, ks翻转维度并取反一维后乘上sin。

对于高纬向量，由于奇、复数维度两两交错实现较为复杂，则现在可简化为将特征维度一切二，如下图所示，在实现过程中对前后各半进行的操作即为rotate_half操作：

代码实现：
先定义旋转角度

class Qwen2RotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()
        # 定义初始值
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        # 定义旋转角
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # Build here to make `torch.jit.trace` work.
        self._set_cos_sin_cache(
            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
        )
    # 为seq里面的每个token形成独一无二的旋转角嵌入(外积)
    def _set_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len
        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)

        freqs = torch.outer(t, self.inv_freq)
        # 生成角度信息(利用注册机制生成self.cos_cached与sin_cached
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

        return (
            self.cos_cached[:seq_len].to(dtype=x.dtype),
            self.sin_cached[:seq_len].to(dtype=x.dtype),
        )

首先要先生成角度: $\theta = \left(\frac{1}{10000^{2n/d}}\right)$

其中，n表示维度数，其取值范围为[0, 1, ..., d/2-1]

然后将上述生成角度与每一个位置乘积，区分一个seq中各个词：其实等价于: θ=(i100002n/d)
其中: i为行数。

` emb将二者cat起来，得到dim维度，每dim/2一循环：

然后，在取出位置编码信息cos与sin的时候，就是将seq的部分切出来，原先设置的1024是最大pos编码，每次用的时候只取当下seq_len的即可.之前求得外积，是为了保证seq里面得每一个词都能有不同的1024个位置编码。

进行旋转嵌入:

# 后半部分和前半部分进行了交换，并且将后半部分的符号取反。
def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
    """Applies Rotary Position Embedding to the query and key tensors.

    query and key tensors rotated using the Rotary Position Embedding.
    """
    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

对应公式:

其中，下标t则表示位于同一行，也就是同一seq_len，对于相邻的两个元素，

3.3 读取顺序attention_mask

第一步的时候只能看到自己，第二步只能看到0、1...其余的都是负无穷。

经过softmax，对应负无穷的位置权重为0，从而实现只能从左往右。

attn形状为(bs,heads,T,T),values的形状为(bs,heads,T,hd),最终落实到[30,30]×[30,128]上，30表示一句话的步长，也就是总词数。

value里面每一个词有128个维度来描述，对于第一个词，由于attn为下三角，所以每一个维度都只有第一个非零元素1进行相乘，其他的都是×0。

对于第二行，则是前两个有不同的权值，让value的128个维度分别依据这两个的权重，在128个维度上根据前两行，计算得出output的第二个词(第二步或者第二行)的128个维度.... 这种加权，体现出考虑前词关系。

第n步则对应有n个权重，用来描述从1到n个步之间的各个关系，进而计算出各个维度。

每一个矩阵乘法的结果相当于是下一个词的dim，那么score则是根据mask来判断，能通过前几个词对应dim的值从而进行加权，进而得到下一个词的该dim上的值

对于推理的过程，问询不一样长没关系，因为所有的权重都是dim-dim，得到的attention_score是一个seq,seq的，权重跟seq的大小没关系。其都是后面的dim维度的参数。 - 推理过程的attention_mask可有可无，是一个一个吐，循环cat到下一个，每一次都取最后一个，代表着预测的是下一个token.