LLaMa-1/2/3 原理+源码——拆解 (KV-Cache, RoPE, RMSNorm, GQA, SwiGLU)

原理

vanilla transformer 与 llama 的区别

主流的大语言模型都采用了transformer架构，它是一个基于多层self-attention的神经网络模型。

原始的transformer由编码器（encoder）和解码器（decoder）两个部分构成，同时，这两个部分也可以独立使用。

llama模型与gpt-2类似，也是采用了基于decoder-only的架构。在原始vanilla transformer decoder的基础上，llama进行了如下改动：

• 为了增强训练稳定性，前置了层归一化（pre-normalization），并使用rmsnorm 作为层归一化方法。
• 为了提高模型性能，采用swiglu 作为激活函数。
• 为了更好地建模长序列数据，采用rope 作为位置编码。
• 为了平衡效率和性能，部分模型采用了gqa分组查询注意力机制（grouped-query attention, gqa）。
• 并且将self-attention改进为使用kv-cache的grouped query。

每个版本的llama由于其隐藏层的大小、层数的不同，均有不同的变体。接下来，我们将展开看下每个版本的不同变体。

llama-1/2/3 演进

llama-1系列的演进：llama-1 -> alpaca -> vicuna 等：

• llama-1：meta开源的pre-trained model，模型参数从 7b、13b、32b、65b 不等，llama-7b在大多数基准测试上超过了text-davinci-003（即gpt3-173b），相比于chatgpt或者gpt4来说，llama可能效果上还有差距，目前hugging face已集成了llama的代码实现和开源模型。学术界和工业界都可以在此基础上进行学习和研究。
  

• alpaca：斯坦福在llama-7b的基础上监督微调出来的模型，斯坦福是用openai的text-davinci-003（即gpt3-173b）的api配合self-instruct技术，使用175个提示语种子自动生成了52k条提示-回复的指示数据集，在llama-7b上微调得到的模型，在8张80g的a100上训练了3小时。

• vicuna：在llama-13b的基础上使用监督微调得到的模型，数据集来自于sharegpt 产生的用户对话数据，共70k条。使用pytorch fsdp在8张a100上训练了一天。相较于alpaca，vicuna在训练中将序列长度由512扩展到了2048，并且通过梯度检测和flash attention来解决内存问题；调整训练损失考虑多轮对话，并仅根据模型的输出进行微调。通过gpt4来打分评测，vicuna可以达到chatgpt 90%的效果。

llama2系列的演进：

• llama2：采用了llama 1的大部分预训练设置和模型架构，有 7b、13b、34b、70b四个参数量版本。llama2和llama1的最大差别是：llama-2将预训练的语料扩充到了 2t token语料，同时将模型的上下文长度从2,048翻倍到了4,096，并在训练34b、70b的模型中引入了分组查询注意力机制（grouped-query attention, gqa）。
  
• llama-2 chat：有了更强大的基座模型llama-2，meta通过进一步的有监督微调（supervised fine-tuning, sft）、基于人类反馈的强化学习（reinforcement learning with human feedback, rlhf）等技术对模型进行迭代优化（pertrain -> sft -> rlhf），并发布了面向对话应用的微调系列模型 llama-2 chat。
• code-llama：得益于llama-2的优异性能，meta在2023年8月发布了专注于代码生成的code-llama，共有7b、13b、34b和70b四个参数量版本。

llama3系列的演进：

• llama3：包括8b和70b两个参数量版本。除此之外，meta还透露，400b的llama-3还在训练中。相比llama-2，llama-3支持8k长文本，并采用了一个编码效率更高的tokenizer，词表大小为128k。在预训练数据方面，llama-3使用了超过15t token的语料，这比llama 2的7倍还多。
  • 小型模型具有8b参数，其性能略优于mistral 7b和gemma 7b；
  • 中型模型则拥有70b参数，其性能介于chatgpt 3.5和gpt 4之间；
  • 大型模型规模达到400b，目前仍在训练中，旨在成为一个多模态、多语言版本的模型，预期性能应与gpt 4或gpt 4v相当。

embedding

embedding的过程：word -> token_id -> embedding_vector，其中第一步转化使用tokenizer的词表进行，第二步转化使用 learnable 的 embedding layer。

rms norm

对比batch norm 和 layer norm：都是减去均值mean，除以方差var，最终将归一化为正态分布n(0,1)。只不过两者是在不同的维度(batch还是feature)求均值和方差，（其中，减均值：re-centering 将均值mean变换为0，除方差：re-scaling将方差varance变换为1）。

rms norm（root mean layer norm）：rms norm认为，layer norm成功的原因是re-scaling，因为方差var计算的过程中使用了均值mean，因此rms norm不再使用均值mean，而是构造了一个特殊的统计量rms代替方差var。为什么使用rms norm？（1）rms norm计算量更小。（2）rms的效果和layer norm一样好。

针对输入向量 a 的rms norm 函数计算公式如下：

此外，rmsnorm 还可以引入可学习的缩放因子gi 和偏移参数bi，从而得到

rmsnorm 在huggingface transformer 库中代码实现如下所示：

class llamarmsnorm(nn.module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        llamarmsnorm is equivalent to t5layernorm
        """
        super().__init__()
        self.weight = nn.parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps # eps 防止取倒数之后分母为0
    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=true)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        # weight 是末尾乘的可训练参数, 即g_i
        return (self.weight * hidden_states).to(input_dtype)

为了使得模型训练过程更加稳定，gpt-2 相较于gpt 就提出了将layer norm前置，将第一个层归一化移动到多头自注意力层之前，第二个层归一化也移动到了全连接层之前，同时残差连接的位置也调整到了多头自注意力层与全连接层之后。层归一化中也采用了rmsnorm 归一化函数。

rope(rotary positional encodding)

绝对positional encodding的使用过程：word -> token_id -> embedding_vector + position_encodding -> encoder_input，其中第一步转化使用tokenizer的词表进行，第二步转化使用 learnable 的 embedding layer。将得到的embedding_vector 和 position_encodding 进行element-wise的相加，然后才做为input送入llm的encoder。

对比absolute pe 和 relative pe：

• absolute pe 绝对位置编码：每次单独1个token的pe，每个token的pe之间没有关系，是一组固定的vector，反映每个token在sequence中的绝对位置。计算 query, key 和 value 向量之前加在输入序列x上，$q/k/v=w_{q/k/v}(x+p)$，经典的位置编码向量 p 的计算方式是使用 sinusoidal 函数。
• relative pe 相对位置编码：每次处理2个token的pe，只在计算attention时使用（在query@key时加在key上），反映2个token的相关度。

和相对位置编码(relative pe)相比，旋转位置编码(rope) 具有更好的外推性，目前是大模型相对位置编码中应用最广的方式之一。

备注：什么是大模型外推性？（length extrapolation）

外推性是指大模型**在训练时和预测时的输入长度不一致，导致模型的泛化能力下降的问题。**例如，如果一个模型在训练时只使用了512个 token 的文本，那么在预测时如果输入超过512个 token，模型可能无法正确处理。这就限制了大模型在处理长文本或多轮对话等任务时的效果。

• 旋转位置编码（rope）：rope 借助了复数的思想，出发点是通过绝对位置编码的方式，实现token间的相对位置编码。其目标是通过下述 f 运算，计算self-attention前，来给q，k 添加，其在sequence中的绝对位置信息m和n，得到qm 和kn，然后进行qm@kn：

实际上，我们借助了复数的思想，寻找了一个 g 运算来合并 f 运算（分别给q和k嵌入绝对位置信息）和q@k（query和key内积）这两个操作，这样只需要给g运算输入：query，key，以及两者的在各自seqence中的绝对位置m和n，即可：

接下来的目标就是找到一个等价的位置编码方式，从而使得上述关系成立。

假定现在词嵌入向量的维度是两维 dim=2 ，这样就可以利用上2维度平面上的向量的几何性质，然后论文中提出了一个满足上述关系的 f 和 g 的形式如下(re 表示复数的实部)：

看到这里会发现，$f_q(x_m,m)=$qm这不就是 q乘以了一个旋转矩阵吗？这就是为什么叫做旋转位置编码的原因。($f_k(x_n,n)=$kn也是同样)

为什么叫旋转位置编码？因为使用欧拉公式构造旋转矩阵，将分别将q和k旋转到空间中对应的位置，实现对计算结果添加位置信息，再计算qm@kn。

上面是2维的例子， xm和xn只有2个token，llama中是n维的，n个token也是一样操作，内积满足线性叠加性，因此任意偶数维的rope，我们都可以表示为二维情形的拼接，即：

由于上述旋转矩阵rn 具有稀疏性，有大量元素是0，直接用矩阵乘法来实现会很浪费算力，因此可以使用逐位相乘⊗ 操作进一步加快计算速度。从下面这个qm的实现也可以看到，rope 可以视为是乘性位置编码的变体。

rope 在huggingface transformer 库中代码实现如下所示：

class llamarotaryembedding(torch.nn.module):

    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=none):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
        self.register_buffer("inv_freq", inv_freq)
        # build here to make `torch.jit.trace` work.
        self.max_seq_len_cached = max_position_embeddings
        t = torch.arange(self.max_seq_len_cached, device=self.inv_freq.device,
        dtype=self.inv_freq.dtype)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        # different from paper, but it uses a different permutation
        # in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        dtype = torch.get_default_dtype()
        self.register_buffer("cos_cached", emb.cos()[none, none, :, :].to(dtype), persistent=false)
        self.register_buffer("sin_cached", emb.sin()[none, none, :, :].to(dtype), persistent=false)
        
    def forward(self, x, seq_len=none):
        # x: [bs, num_attention_heads, seq_len, head_size]
        # this `if` block is unlikely to be run after we build sin/cos in `__init__`.
        # keep the logic here just in case.
        if seq_len > self.max_seq_len_cached:
            self.max_seq_len_cached = seq_len
            t = torch.arange(self.max_seq_len_cached, device=x.device, dtype=self.inv_freq.dtype)
            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            # different from paper, but it uses a different permutation
            # in order to obtain the same calculation
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
            self.register_buffer("cos_cached", emb.cos()[none, none, :, :].to(x.dtype),
            persistent=false)
            self.register_buffer("sin_cached", emb.sin()[none, none, :, :].to(x.dtype),
            persistent=false)
    
        return (
        self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
        self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
        )
    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):
        # the first two dimensions of cos and sin are always 1, so we can `squeeze` them.
        cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]
        sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]
        cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
        sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]
        q_embed = (q * cos) + (rotate_half(q) * sin)
        k_embed = (k * cos) + (rotate_half(k) * sin)
        return q_embed, k_embed

swiglu function

swiglu 激活函数是shazeer 在文献中提出，并在palm等模中进行了广泛应用，并且取得了不错的效果，相较于relu 函数在大部分评测中都有不少提升。在llama 中全连接层使用带有swiglu 激活函数的ffn（position-wise feed-forward network）的计算公式如下：

其中，σ(x) 是sigmoid 函数。下图给出了swish 激活函数在参数β 不同取值下的形状。可以看到当β 趋近于0 时，swish 函数趋近于线性函数y = x，当β 趋近于无穷大时，swish 函数趋近于relu 函数，β 取值为1 时，swish 函数是光滑且非单调。

huggingface 的transformer 库中 $swis{h}_{\beta =1}$函数使用 silu 函数 代替。

kv-cache

首先来了解一下llama的训练（下词预测任务）：seq2seq的生成，但迭代t次，seq_len逐渐增加。

下句预测时的self-attention：

• timpstep=1时seq_len=1，给[sos]时，预测love；
  
• timpstep=2时seq_len=2，给[sos] 和 love时，预测that
  
• timpstep=4时seq_len=4，给[sos] 和 love 和 can 和 quickly时，预测seize…
  

每个timestep我们只关注生成的最后一个token，但因为llama是一个seq2seq的model，每次必须重新计算和生成前面的token，因此我们希望能将之前timestep计算生成过的token给缓存起来，下个timestep不用再次计算，这样的背景下，kv-cache就产生了。

再来分析一下，每次个timestep的self-attention中我们到底需要哪些：因为我们只关注最后一个token的attention_output，如下图timestep=4，我们只需要attention_output的第4个token。

因此我们只需要q的最后一个token和k的所有token相乘，得到最后一个token的attention_score，然后用v的所有token再与attention_score点积(相乘求和)，得到最后一个token的attention_output：

由上分析可知，每个timestep，我们的q只需要新增的那个token即可，而k和v要缓存之前timestep的token，保证token是全的。每次计算出来的attention_output就是那个新增的token的attention。 这样就可以节省大量计算开销。

grouped multi-query attention

回顾原始的多头注意力multi-head attention：时间开销的瓶颈在于矩阵的运算matrix computation。

当我们使用kv-cache后：时间开销的瓶颈在于内存的访问memory access。

multi query attention（mqa）

为了提升attention计算效率，多查询注意力（multi query attention，mqa） 是多头注意力的一种变体。其主要区别在于，在多查询注意力中多个不同的注意力head共享一个k和v的集合，每个head只单独保留了一份q参数。 具体操作上，去除 k和v 的head维度，只为q保留head维度。因此这就是被叫做multi query attention的原因。

因此k和v的矩阵的数量仅为1个（不分head），大幅度减少了显存占用，使其更高效。由于多查询注意力改变了注意力机制的结构，因此模型通常需要从训练开始就支持多查询注意力。

研究结果表明，可以通过对已经训练好的模型进行微调来添加多查询注意力支持，仅需要约 5% 的原始训练数据量就可以达到不错的效果。包括falcon、santacoder、starcoder等在内很多模型都采用了多查询注意力机制。

grouped multi-query attention(gmqa)

就是在 multi-query attention的基础上，对input进行分组，如下图2个head分为1组，每组都有自己的k，v，每个组包含2个q。 （与mqa的区别在于：mqa的kv只有1份；gqa的kv有group份(llama-70b中是kv_heads=8，即每个kv对应8个q)）

import torch
import torch.nn as nn
from torch.nn import functional as f
from typing import optional

x = torch.rand(1, 512, 768)
batch, seq_len, d_model = x.shape
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
query = xq.view(batch, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
key = xk.view(batch, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
value = xv.view(batch, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)

# repeat kv heads if n_kv_heads < n_heads
key = key.reshape(1, 1, self.n_heads//self.n_kv_heads, 1)
value = value.reshape(1, 1, self.n_heads//self.n_kv_heads, 1)

# excute scaled dot product attention
attn_score = torch.matmul(query, key.transpose(-2, -1))
attn_score = attn_score / self.scale
attn_score = torch.softmax(attn_score, dim=-1)
attn_score = self.dropout(attn_score)
attn_output = torch.matmul(attn_score, value)
attn_output = attn_output.transpose(1, 2).reshape(batch, seq_len, self.head_dim)

源码

[llms 实践] 01 llama、alpaca、vicuna 整体介绍及 llama 推理过程

rms norm

import numpy as np
import torch
from torch import nn

bs, seq_len, emb_dim = 20, 5, 10  # 20个样本， 每个样本5个token(word/patch)，每个token(word/patch)是长度为10的embedding
x = torch.randn(bs, seq_len, emb_dim)

• ln (layer norm) 先re-centering(减均值)，再re-scaling(除方差)：$x=\frac{x-mean}{std}$。ln 作用在emb_dim维度上，使得每个embedding的均值=0，标准差=1。

ln = nn.layernorm(emb_dim)  # ln 作用在emb_dim维度
x_ln = ln(x)
print(x_ln[1, 4, :].mean())
print(x_ln[1, 4, :].std())
# tensor(-1.6391e-08, grad_fn=<meanbackward0>)
# tensor(1.0541, grad_fn=<stdbackward0>)

• rms norm （省略了ln的re-centering）只进行re-scaling：对于向量x，$x_i=\frac{x_i}{rms(x)}\cdot g_i$，where $rms(x)=\sqrt{\frac{1}{n}{\sum }_{i=1}^n{a}_i^2}$。rms norm 也作用在emb_dim维度上，使得每个embeddin的标准差=1。

import torch
import torch.nn as nn
class rmsnorm(torch.nn.module):
    def __init__(self, dim, eps=1e-8):
        super().__init__()
        self.eps = eps
        self.weight = nn.parameter(torch.ones(dim))  # 缩放因子g

    def _norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=true) + self.eps)
    
    def forward(self, x):
        output = self._norm(x.float().type_as(x))
        return output * self.weight
    
rms_norm = rmsnorm(emb_dim)
x_rms = rms_norm(x)
print(x_rms[1, 4, :].mean())
print(x_rms[1, 4, :].std())
# tensor(-0.0725, grad_fn=<meanbackward0>)
# tensor(1.0513, grad_fn=<stdbackward0>)

torch.rsqrt(x)就是 $\frac{1}{\sqrt{x}}$

rope (rotary positional encodding)

huggingface和meta实现了两版不同的rope，就是把复数位置信息快速融入query和key中，破坏transformer结构中sequence的完全对称性，使得对token位置敏感。

sinusoidal绝对位置编码（正余弦）：分别计算奇数和偶数的位置编码，然后拼接在一起。
$$pe(t,i)=sin(\frac{t}{1000^{\frac{i}{dim}}}),iiseven；pe(t,i)=cos(\frac{t}{1000^{\frac{i}{dim}}}),iisodd$$

import torch.nn as nn
import math
class sinusoidalpositionencoding(nn.module):
    def __init__(self, dim, seq_len=5000):
        super(sinusoidalpositionencoding, self).__init__()       
        pe = torch.zeros(seq_len, dim)  # [seq_len, dim]
        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)  # [seq_len, 1]
        # div_term 先算偶数，再用偶数反推奇数
        div_term = torch.exp(torch.arange(0, dim, 2).float() * (-math.log(10000.0) / dim)) # [dim/2]
        pe[:, 0::2] = torch.sin(position * div_term)  # even_pe: [seq_len, dim/2]
        pe[:, 1::2] = torch.cos(position * div_term)  # odd_pe: [seq_len, dim/2]
        # torch.size([seq_len, dim])
        pe = pe.unsqueeze(0).transpose(0, 1)
        # torch.size([seq_len, 1, dim])
        self.register_buffer('pe', pe)

    def forward(self, x):
        return x + self.pe[:x.size(0), :]

sinusoidal_pe = sinusoidalpositionencoding(6, 10)
print(sinusoidal_pe.pe.shape) # torch.size([10, 1, 6])

rope旋转位置编码（改进版）：

• 预计算旋转矩阵precompute_freqs_cis：计算出了i=[1,2,...,seqlen]位置的正余弦${\theta }_i$，并以复数形式表示，实部代表偶数，虚部代表奇数。 其中cis(x)=cos(x)+i·sin(x)

import torch
import torch.nn as nn
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
    """
    args:
        dim (int): token的维度
        end (int): 通常是seqlen
        theta (float, optional): 根据公式，默认为10000

    returns:
        torch.tensor: 预先计算好的复数张量
    """
	# 计算词向量元素两两分组之后，每组元素对应的旋转角度\theta_i
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))  # freqs.shape = [seq_len, dim // 2] 
    # # 生成 token 序列索引 t = [0, 1,..., seq_len-1]
    t = torch.arange(end, device=freqs.device)  # t.shape = [end]
    # 这里.outer()是外积，freqs是类似正弦编码的形式，它只有dim/2维，t是包括了最大position，也就是end变量，例如在bert中一般是512
    # 因此两个向量进行外积得到一个矩阵，维度为[end, dim/2]，对应论文中的mθ，其中m代表绝对位置
    freqs = torch.outer(t, freqs).float()  # 计算 m * \theta 
    # 转为复数向量，polar()函数用的很少，这一步之后实际上就算出了每一个位置m处的cosθ和sinθ的值，方便后续用于快速计算，这里具体形式可以参考上面发的文章
    # 简单来说就是每一行对应一个位置m，总共有dim/2列，每一列是该位置处的cosθi + isinθi，是一个复数向量
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # # torch.size([10, 3]) 
    return freqs_cis
    
precompute_freqs_cis(dim=6, end=10)
# # torch.size([10, 3])

可以看出旋转位置编码使用复数表示，对比前面的sincos位置编码，实部代表偶数，虚部代表奇数。 cos和sin都在一个数中，因此旋转位置编码比绝对的sincos位置编码要少一半。

• 施加rope信息apply_rotary_emb：计算query_rope时直接使用xq_ * freqs_cis，也就是按位相乘(rope是乘性位置编码x*p，不同于sinusoidal那种加性位置编码x+p)。apply_rotary_emb得到含有位置信息的sequence xq和xk。

from typing import tuple

def apply_rotary_emb(
    xq: torch.tensor,  # [bs, seq_len, heads, head_dim]
    xk: torch.tensor,  # [bs, seq_len, heads, head_dim]
    freqs_cis: torch.tensor,  # precompute_freqs_cis函数的输出 [seq_len, head_dim/2]
) -> tuple[torch.tensor, torch.tensor]:
    # 因为位置编码freqs_cis是复数,因此需要将query和key复数化,具体就是将dim维度分成两半, 每一半是dim/2维, 分别用做实数部分和虚数部分
    # [bs, seq_len, heads, head_dim] => [bs, seq_len, heads, head_dim/2]
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    # [bs, seq_len, heads, head_dim] => [bs, seq_len, heads, head_dim/2]
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))

    # 需要按位计算,因此维度要对齐 
    # [seq_len, head_dim/2] => [1, seq_len, 1, head_dim/2]
    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)

    # 这里先是xq_和freqs_cis两个复数张量按位相乘,这里直接相乘之后仍然是一个复数,然后再展开成实数形式,也就从dim/2维转到dim维,保持输出维度不变
    # [bs, seq_len, heads, head_dim/2]*[1, seq_len, 1, head_dim/2] => [bs, seq_len, heads, head_dim/2, 2] => [bs, seq_len, heads, head_dim]
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    # [bs, seq_len, heads, head_dim/2], [1, seq_len, 1, head_dim/2] => [bs, seq_len, heads, head_dim/2, 2] => [bs, seq_len, heads, head_dim]
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)

    # 需要注意的是大模型领域经常涉及到fp32和fp16的交互
    return xq_out.type_as(xq), xk_out.type_as(xk)

在后面计算attention_score=q@k时，先进行线性映射xq=wq(x),xk=wk(x)，然后同时假如相对位置信息xq, xk = apply_rotary_emb(xq, xk, freqs_cis) 其旋转矩阵相乘k根据sincos的数学性质，本质上等价于位置信息m和n的相对差值：

swiglu function

swiglu并不是一种全新的算法或理论，而是对现有transformer架构中的ffn层的一种改进。在transformer中，ffn是实现前馈传播的关键部分，通过两层全连接层和relu激活函数，实现从输入到输出的映射。

然而，swiglu对这一结构进行了优化，将第一层全连接和relu激活函数替换为两个权重矩阵和输入的变换，再配合swish激活函数进行哈达马积操作。

$$swiglu=swi+glu$$

• swi指的是swish非线性激活函数，
  

• glu指的是 gated linear unit，输入的向量x分别经过两个linear层，其中一个需要经过非线性激活函数，然后将两者对应元素相乘。

组合起来的swiglu就是把glu中的 $\sigma (\cdot )=swish(\cdot )$

llm中的tansformer block中的swiglu ffn除了上述操作外，另外接了一个output的线性层，整体的向量维度变化如下：

(13696-->5120)((swish(5120-->13696) * (5120-->13696) ))

class swiglu(torch.nn.module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        hidden_act: str,
    ):
        super().__init__()
        self.w1 = torch.nn.linear(hidden_size, intermediate_size, bias=false) ###（5120*13696）
        self.w2 = torch.nn.linear(intermediate_size, hidden_size, bias=false) ###（13696*5120）
        self.w3 = torch.nn.linear(hidden_size, intermediate_size, bias=false)  ###（5120*13696）
        self.act_fn = act2fn[hidden_act] ###可以是swish, silu等非线性激活函数

    def forward(self, x):
        return self.w2(self.act_fn(self.w1(x)) * self.w3(x))
        # 或 return self.w2(f.silu(self.w1(x)) * self.w3(x))

kv-cache

长文本输入的情况下，占用的显存线性随句长增长，此时需要做cache优化，kv-cache就是用来缓存自回归模型中attention中的key和value的，只出现在transformer-decoder中（如gpt/llama中，t5的decoder中，而bert就没有），kv-cache是做llm模型层面的推理加速，而flashattention则是gpu硬件算法优化。

clm (causal language model) task：

• ntp (next token prediction) training:

input: [sos] love that can quickly seize the gentle heart
target: love that can quickly seize the gentle heart [eos]

• ntp inference

sos: start of sentence
eos: end of sentence 

	input	output
t=1	[sos]	love
t=2	[sos] love	love that
t=3	[sos] love that	love that can
t=4	[sos] love that can	love that can quickly
t=5	[sos] love that can quickly	love that can quickly seize
t=6	[sos] love that can quickly seize	love that can quickly seize the
t=7	[sos] love that can quickly seize the	love that can quickly seize the gentle
t=8	[sos] love that can quickly seize the gentle	love that can quickly seize the gentle heart
t=9	[sos] love that can quickly seize the gentle heart	love that can quickly seize the gentle heart [eos]

就像上面描述的一样，自回归模型上次的输出，作为本次的输入，历史信息会作为k和v与q进行attention运算，但是随着历史信息内容的增加，k和v可能是非常庞大的matrix，而q还是一个vector。但存在一个问题，每次transformer都会重复推理前面已经有的token，这样会重复执行之前发生过的矩阵运算。

推理时，要预测下个token只需要当前最尾的一个q。我们可以得出一个结论,q[-1], k[:], v[:]就可以计算next_token！

• without kv-cache， 每次需要计算全wq(x),wk(x), wv(x), 每次需要计算全量attn。
  

• with kv-cache，第一步计算完整attn，将kv保存成kv_cache。第二步，取第一步的next token $x_n$计算q=wq($x_n$),k=wk($x_n$),v=wv($x_n$)， 将[kv_cache, kv]拼接，计算出qkv。kv-cache每个循环累增，的memory 量=2*(n层*l长度*d维度)。
  

使用kv-cache就是缓存前面已经有的kv的tokens（缓存的是k=wk@x），减少了x先前已经计算过了token再与wk相乘。每次只将x中新的token与w计算得到

attention with kv-cache源码：不断在max_seq_len维度上append新的kv

class attention(nn.module):
    """multi-head attention module."""
    def __init__(self, args: modelargs):
    
        self.wq = columnparallellinear(args.dim, args.n_heads * self.head_dim)
        self.wk = columnparallellinear(args.dim, self.n_kv_heads * self.head_dim)
        self.wv = columnparallellinear(args.dim, self.n_kv_heads * self.head_dim)
        self.wo = rowparallellinear(args.n_heads * self.head_dim, args.dim)
        
        
        # [8, 1024, 32, 128]
        self.cache_k = torch.zeros(
            (
                args.max_batch_size,    # 8
                args.max_seq_len,       # 1024, 不断地在这个维度上append keys
                self.n_local_kv_heads,  # 32
                self.head_dim,          # 128
            )
        ).cuda()
        
        #  [8, 1024, 32, 128]
        self.cache_v = torch.zeros(
            (
                args.max_batch_size,    # 8
                args.max_seq_len,       # 1024, 不断地在这个维度上append values
                self.n_local_kv_heads,  # 32
                self.head_dim,          # 128
            )
        ).cuda()
    
    def forward(
        self,
        x: torch.tensor,
        start_pos: int,
        freqs_cis: torch.tensor,
        mask: optional[torch.tensor],
    ):
        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
        
        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)

        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)

        self.cache_k = self.cache_k.to(xq)
        self.cache_v = self.cache_v.to(xq)
        
        
        self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
        self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv
        
        # 这里在复用之前的计算, all_past
        keys = self.cache_k[:bsz, : start_pos + seqlen]
        values = self.cache_v[:bsz, : start_pos + seqlen]

grouped query attention(gqa)

以llm foundry 为例，分组查询注意力实现代码如下，与llm foundry 中实现的多头自注意力代码相对比，其区别仅在于建立wqkv 层上：

import torch
import torch.nn as nn
from torch.nn import functional as f
from typing import optional


class multiqueryattention(nn.module):
    """multi-query self attention.
    using torch or triton attention implemetation enables user to also use
    additive bias.
    """
    def __init__(
        self,
        d_model: int,
        n_heads: int,
        device: optional[str] = none,
    ):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.wqkv = nn.linear( # multi-query attention 创建
            d_model,
            d_model + 2 * self.head_dim, # 只创建查询的头向量，所以只有1 个d_model
            device=device, # 而键和值则共享各自的一个head_dim 的向量
        )
        self.out_proj = nn.linear(
            self.d_model,
            self.d_model,
            device=device
        )
        self.out_proj._is_residual = true # type: ignore
    def forward(
        self,
        x,
    ):
        qkv = self.wqkv(x) # (1, 512, 960)
        query, key, value = qkv.split( # query -> (1, 512, 768)
            [self.d_model, self.head_dim, self.head_dim], # key -> (1, 512, 96)
            dim=2 # value -> (1, 512, 96)
        )
        context, attn_weights, past_key_value = f.scaled_dot_product_attention(query, key, value,
            self.n_heads, multiquery=true)
        return self.out_proj(context), attn_weights, past_key_value

llama finetine——alpaca、vicuna

• llama：是pertrained model，类似text-davini-003（gpt3），训练数据是公开的1t tokens。支持4k context window的长度。
• self instruction （sft）：在llama权重基础上进行instruction tuning。（性能：vicuna>alpaca>llama）
  • alpaca：prompt和answer都来自chatgpt，52k samples。
  • vicuna：prompt来自互联网，answer都来自sharegpt，70k samples；更长的context window；vicuna weights是基于llama weights的二次权重。

llama lora sft

lora全称low-rank adaption of llm，在不破坏原始参数$w$的基础上(freezed llm weights)，实现对参数$w$的等价微调，本质是微调1个linear layer参数矩阵$w$的2个低秩矩阵$a$和$b$（超参数：低秩的rank=$r\ll min(d,k)$，svd分解减少参数量），实现模型的高效微调(peft)得到推理时新的参数矩阵$w^{\prime }$，通常微调的linear layer是attention的$w_q,w_k,w_v,w_o$的$\delta w$，其中$\alpha$用于控制lora权重的比例：
$$w^{\prime }=w+\alpha \delta w=w+\alpha ba,wherew\in r^{dxk},b\in r^{dxr},a\in r^{rxk}$$

如d=100，k=500，r=5，对比$\delta w$ 和 $w$矩阵的参数量：（5100+5500)/(100*500)=3k/5w=6%，这就叫做parameter efficiency参数高效的微调！超参数$r$越小，参数约节省，但微调出来的模型性能也可能越差。

linear layer的lora伪代码：(这部分可以用peft库的get_peft_model实现）

import math
import torch
from torch import nn

input_dim = 768  # in_feature dim
output_dim = 768  # out_feature dim
rank = 8

w = ...  # from pretrained model with shape (input_dim, output_dim)
w_a = nn.parameter(torch.empty(input_dim, rank))
w_b = nn.parameter(torch.empty(rank, output_dim))

nn.init.kaiming_uniform_(w_a, a=math.sqrt(5))
nn.init.zeros_(w_b, a=math.sqrt(5))

def lora_forward(self, x, w, w_a, w_b):
    h = x @ w  + alpha * x @ (w_a @ w_b) 
    return h

原始llama模型结构：

llamaforcausallm(
  (model): llamamodel(
    (embed_tokens): embedding(32000, 4096, padding_idx=31999)
    (layers): modulelist(
      (0-31): 32 x llamadecoderlayer(
        (self_attn): llamaattention(
          (q_proj): linear8bitlt(in_features=4096, out_features=4096, bias=false)
          (k_proj): linear8bitlt(in_features=4096, out_features=4096, bias=false)
          (v_proj): linear8bitlt(in_features=4096, out_features=4096, bias=false)
          (o_proj): linear8bitlt(in_features=4096, out_features=4096, bias=false)
          (rotary_emb): llamarotaryembedding()
        )
        (mlp): llamamlp(
          (gate_proj): linear8bitlt(in_features=4096, out_features=11008, bias=false)
          (down_proj): linear8bitlt(in_features=11008, out_features=4096, bias=false)
          (up_proj): linear8bitlt(in_features=4096, out_features=11008, bias=false)
          (act_fn): siluactivation()
        )
        (input_layernorm): llamarmsnorm()
        (post_attention_layernorm): llamarmsnorm()
      )
    )
    (norm): llamarmsnorm()
  )
  (lm_head): linear(in_features=4096, out_features=32000, bias=false)
)

huggingface的trl库有专用于模型指令微调的sfttrainer，封装度较高，上手难度小，整个微调流程大约分为三步: 1. 模型和tokenizer载入。2. sft数据准备。3. 模型训练和保存（上传huggingface）。其中device_map="auto"会进行模型并行。

import torch
from datasets import load_dataset
from transformers import (
    automodelforcausallm,
    autotokenizer,
    bitsandbytesconfig,
    trainingarguments,
    pipeline)
from peft import loraconfig
from trl import sfttrainer



# model and tokenizer names
base_model_name = "/data3/huggingface/llm/llama-2-7b-chat-hf"
new_model_name = "llama-2-7b-enhanced" #you can give your own name for fine tuned model
# tokenizer
llama_tokenizer = autotokenizer.from_pretrained(base_model_name, trust_remote_code=true)
llama_tokenizer.pad_token = llama_tokenizer.eos_token
llama_tokenizer.padding_side = "right"
# model
base_model = automodelforcausallm.from_pretrained(
    base_model_name,
    device_map="auto"
)
base_model.config.use_cache = false
base_model.config.pretraining_tp = 1



# load dataset
data_name = "mlabonne/guanaco-llama2-1k"
training_data = load_dataset(data_name, split="train")
# check the data
print(training_data.shape)
# #11 is a qa sample in english
print(training_data[11])



# training params
train_params = trainingarguments(
    output_dir="./results_modified",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=1,
    optim="paged_adamw_32bit",
    save_steps=50,
    logging_steps=50,
    learning_rate=4e-5,
    weight_decay=0.001,
    fp16=false,
    bf16=false,
    max_grad_norm=0.3,
    max_steps=-1,
    warmup_ratio=0.03,
    group_by_length=true,
    lr_scheduler_type="constant",
    report_to="tensorboard"
)

from peft import get_peft_model
# lora config
peft_parameters = loraconfig(
    lora_alpha=8,
    lora_dropout=0.1,
    r=8,
    bias="none",
    task_type="causal_lm"
)
model = get_peft_model(base_model, peft_parameters)
model.print_trainable_parameters()

# trainer with lora configuration
fine_tuning = sfttrainer(
    model=base_model,
    train_dataset=training_data,
    peft_config=peft_parameters,
    dataset_text_field="text",
    tokenizer=llama_tokenizer,
    args=train_params
)

# training
fine_tuning.train()
# save model
fine_tuning.model.save_pretrained(new_model_name)



# reload model in fp16 and merge it with lora weights
base_model = automodelforcausallm.from_pretrained(
    base_model_name,
    low_cpu_mem_usage=true,
    return_dict=true,
    torch_dtype=torch.float16,
    device_map="auto"
)
from peft import loraconfig, peftmodel
model = peftmodel.from_pretrained(base_model, new_model_name)
model = model.merge_and_unload()

# reload tokenizer to save it
tokenizer = autotokenizer.from_pretrained(base_model_name, trust_remote_code=true)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"
from huggingface_hub import login
# you need to use your hugging face access tokens
login("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx")
# push the model to hugging face. this takes minutes and time depends the model size and your
# network speed.
model.push_to_hub(new_model_name, use_temp_dir=false)
tokenizer.push_to_hub(new_model_name, use_temp_dir=false)

# generate text using base model
query = "what do you think is the most important part of building an ai chatbot?"
text_gen = pipeline(task="text-generation", model=base_model_name, tokenizer=llama_tokenizer, max_length=200)
output = text_gen(f"<s>[inst] {query} [/inst]")
print(output[0]['generated_text'])

# generate text using fine-tuned model
query = "what do you think is the most important part of building an ai chatbot?"
text_gen = pipeline(task="text-generation", model=new_model_name, tokenizer=llama_tokenizer, max_length=200)
output = text_gen(f"<s>[inst] {query} [/inst]")
print(output[0]['generated_text'])

alpaca inference

alpaca-lora-7b ：基于llama-7b 在 stanford alpaca dataset上进行lora微调得到。下面config中的lora target modules就是指q_proj、k_proj、v_proj、o_proj中会包含lora_a和lora_b。lora rank=16。

epochs: 10 (load from best epoch)
batch size: 128
cutoff length: 512
learning rate: 3e-4
lora r: 16
lora target modules: q_proj, k_proj, v_proj, o_proj

4bit或8bit推理：使用 4 比特量化的不同变体，例如 nf4 (normalfloat4 (默认) ) 或纯 fp4 量化。从理论分析和实证结果来看，我们建议使用 nf4 量化以获得更好的性能。其他选项包括 bnb_4bit_use_double_quant ，它在第一轮量化之后会进行第二轮量化，为每个参数额外节省 0.4 比特。最后是计算类型，虽然 4 比特 bitsandbytes 以 4 比特存储权重，但计算仍然以 16 或 32 比特进行，这里可以选择任意组合 (float16、bfloat16、float32 等)。如果使用 16 比特计算数据类型 (默认 torch.float32)，矩阵乘法和训练将会更快。用户应该利用 transformers 中最新的 bitsandbytesconfig 来更改这些参数。下面是使用 nf4 量化加载 4 比特模型的示例，例子中使用了双量化以及 bfloat16 计算数据类型以加速训练。

import torch
import transformers
from transformers import llamatokenizer, llamaforcausallm, generationconfig, bitsandbytesconfig

# bitsandbytesconfig设置推理精度
nf4_config = bitsandbytesconfig(
   load_in_4bit=true,
   bnb_4bit_quant_type="nf4",
   bnb_4bit_use_double_quant=true,
   bnb_4bit_compute_dtype=torch.bfloat16
)

# load base_model and tokenizer
model = llamaforcausallm.from_pretrained("/data3/huggingface/llm/llama-2-7b-chat-hf", quantization_config=nf4_config, device_map="auto")
tokenizer = llamatokenizer.from_pretrained("/data3/huggingface/llm/llama-2-7b-chat-hf")
# load lora using peft
from peft import peftmodel
model = peftmodel.from_pretrained(model, "/data3/huggingface/llm/alpaca-lora-7b")


from peft import mapping
from peft.utils import other
print('model_type', model.config.model_type)
print(model.peft_config['default'].target_modules)
#默认的 target module
print(other.transformers_models_to_lora_target_modules_mapping)


def generate_prompt(instruction, input=none):
    if input:
        return f"""below is an instruction that describes a task, paired with an input that provides further context. write a response that appropriately completes the request.

### instruction:
{instruction}

### input:
{input}

### response:"""
    else:
        return f"""below is an instruction that describes a task. write a response that appropriately completes the request.

### instruction:
{instruction}

### response:"""
    
generation_config = generationconfig(
    temperature=1.5,
    # nucleus sampling
    top_p=0.8,
    num_beams=4,
)

def inference(instruction, input=none):
    prompt = generate_prompt(instruction, input)
#     print(prompt)
    inputs = tokenizer(prompt, return_tensors="pt")
    input_ids = inputs["input_ids"].cuda()
    generation_output = model.generate(
        input_ids=input_ids,
        generation_config=generation_config,
        return_dict_in_generate=true,
        output_scores=true,
        max_new_tokens=256
    )
    for s in generation_output.sequences:
        output = tokenizer.decode(s)
        print("response:", output.split("### response:")[1].strip())


inference(input("instruction: "))

加载lora后的模型结构：可以看到q_proj、k_proj、v_proj、o_proj中会包含lora_a和lora_b。

peftmodelforcausallm(
  (base_model): loramodel(
    (model): llamaforcausallm(
      (model): llamamodel(
        (embed_tokens): embedding(32000, 4096, padding_idx=31999)
        (layers): modulelist(
          (0-31): 32 x llamadecoderlayer(
            (self_attn): llamaattention(
              (q_proj): linear8bitlt(
                in_features=4096, out_features=4096, bias=false
                (lora_dropout): moduledict(
                  (default): dropout(p=0.05, inplace=false)
                )
                (lora_a): moduledict(
                  (default): linear(in_features=4096, out_features=16, bias=false)
                )
                (lora_b): moduledict(
                  (default): linear(in_features=16, out_features=4096, bias=false)
                )
                (lora_embedding_a): parameterdict()
                (lora_embedding_b): parameterdict()
              )
              (k_proj): linear8bitlt(
                in_features=4096, out_features=4096, bias=false
                (lora_dropout): moduledict(
                  (default): dropout(p=0.05, inplace=false)
                )
                (lora_a): moduledict(
                  (default): linear(in_features=4096, out_features=16, bias=false)
                )
                (lora_b): moduledict(
                  (default): linear(in_features=16, out_features=4096, bias=false)
                )
                (lora_embedding_a): parameterdict()
                (lora_embedding_b): parameterdict()
              )
              (v_proj): linear8bitlt(
                in_features=4096, out_features=4096, bias=false
                (lora_dropout): moduledict(
                  (default): dropout(p=0.05, inplace=false)
                )
                (lora_a): moduledict(
                  (default): linear(in_features=4096, out_features=16, bias=false)
                )
                (lora_b): moduledict(
                  (default): linear(in_features=16, out_features=4096, bias=false)
                )
                (lora_embedding_a): parameterdict()
                (lora_embedding_b): parameterdict()
              )
              (o_proj): linear8bitlt(
                in_features=4096, out_features=4096, bias=false
                (lora_dropout): moduledict(
                  (default): dropout(p=0.05, inplace=false)
                )
                (lora_a): moduledict(
                  (default): linear(in_features=4096, out_features=16, bias=false)
                )
                (lora_b): moduledict(
                  (default): linear(in_features=16, out_features=4096, bias=false)
                )
                (lora_embedding_a): parameterdict()
                (lora_embedding_b): parameterdict()
              )
              (rotary_emb): llamarotaryembedding()
            )
            (mlp): llamamlp(
              (gate_proj): linear8bitlt(in_features=4096, out_features=11008, bias=false)
              (down_proj): linear8bitlt(in_features=11008, out_features=4096, bias=false)
              (up_proj): linear8bitlt(in_features=4096, out_features=11008, bias=false)
              (act_fn): siluactivation()
            )
            (input_layernorm): llamarmsnorm()
            (post_attention_layernorm): llamarmsnorm()
          )
        )
        (norm): llamarmsnorm()
      )
      (lm_head): linear(in_features=4096, out_features=32000, bias=false)
    )
  )
)