三万字长文超详细解读LLama2！

自从Transformer架构问世以来，大型语言模型（Large Language Models, LLMs）以及AIGC技术的发展速度惊人，它们不仅在技术层面取得了重大突破，还在商业应用、社会影响等多个层面展现出巨大潜力。
随着ChatGPT的推出，这一技术日益走进大众视野，这也预示着一个由生成式AI塑造的未来正在加速到来。

与此同时，Meta AI Meta AI在2023年推出了LLama（Large Language Model Meta AI）系列大语言模型，这一模型初期是以较为封闭的形式面向特定研究人员开放。之后，又开源LLama系列模型LLama2。

什么是LLama2?

LLama2是Meta AI公司在2023年推出的一款半开源LLM(所谓半开源即为只有Inference没有Train过程)，它是Llama的下一代版本，训练数据集2万亿token，上下文长度由llama的2048扩展到4096，可以理解和生成更长的文本，包括7B、13B、70B三个模型，展现出了卓越的性能，使其迅速在基准测试中崭露头角，标志着生成式人工智能领域的一次重要进步。

LLama2模型的任务是在给定前n个单词的基础上预测句子中下一个单词。该模型的核心特点是其预测过程依赖于过去和当前的输入信息，而不考虑未来的信息。

该模型生成文本的过程中，每次迭代不仅需要提供当前待预测位置前n个单词作为输入，还需要将模型在前一次迭代中生成的单词作为新的输入的一部分。

例如，假设我们想要使用LLama2模型生成一句话，设定n=3，即模型每次基于前3个单词预测下一个单词。生成过程如下：

1. 初始输入：提供一个初始前缀，“今天天气”；模型接收到“今天天气”作为输入，预测下一个单词为“晴朗”。
2. 第二次迭代：将前一次的预测结果加入到输入序列，形成新的输入：“今天天气晴朗”；模型接收到“今天天气晴朗”作为输入，预测下一个单词为“，”。
3. 第三次迭代：将上一次预测的逗号“，”加入到输入序列中，形成新的输入：“今天天气晴朗，”；模型接收到“今天天气晴朗，”作为输入，预测下一个单词为“适合”。
4. 后续迭代：以此类推，每次模型预测出一个单词后，都将该单词添加到输入序列中，继续预测下一个单词，直到达到预设的终止条件（如生成一定长度的文本、遇到特定结束符等）。

如下图所示。

相比之下，CV模型在进行图像分类、目标检测等任务时，通常只需要一次性接收整个图像作为输入，然后经过一次推理过程就得出最终结果，无需像llama2这样的语言模型这样进行多次迭代和递归预测。

处理流程

在深入理解LLama2模型结构之前，我们先回顾一下LLM的一般处理流程：

输入

LLM的输入数据通常是一段或多段自然语言文本，可以是一个简单的句子或一段话。文本被表示成单词或字符的序列。

[岱宗夫如何？齐鲁青未了。造化钟神秀，阴阳割昏晓。]

tokenization

文本被切分为单词或字符，形成token序列。token序列进一步被序列化为列表或数组，并通过语料库进行索引化，将每个token映射到一个唯一的整数索引，便于模型内部计算。

序列化->[‘BOS’,‘岱’,‘宗’,‘夫’,‘如’,‘何’,‘？’,‘齐’,‘鲁’…‘阴’,‘阳’,‘割’,‘昏’,‘晓’,‘EOS’]
假设语料库索引化->[‘BOS’,‘10’,‘3’,‘67’,‘89’,‘21’,‘45’,‘55’,‘61’…‘7869’,‘9’,‘3452’,‘563’,‘56’,‘EOS’]

Embedding

tokenization之后的文本信息变为数字形式的token序列，然后通过Embedding层将数字token映射为一个实数向量Embeding Vector。其中，每个token对应的向量通常具有固定的维度d（如50、100、300、768等），向量中的每个元素（实数）表示token在特定语义空间中的某个属性或特征。

具体来说，Embedding Vector可以表示为一个二维数组或矩阵，其形状与token序列长度相同，每个元素是一个固定维度的向量。这里假设使用一个维度为d=10的Embedding向量，则经过Embedding层后得到的向量表示如下：

'BOS'-> [p_{00},p_{01},p_{02},...,p_{09}]
'10' -> [p_{10},p_{11},p_{12},...,p_{09}]
'3'  -> [p_{20},p_{21},p_{22},...,p_{09}]
...
'EOS'-> [p_{n0},p_{n1},p_{n2},...,p_{09}]

位置编码

位置编码（Positional Encoding）用于标识每个token在序列中的位置。让模型在处理不同位置的token时，能够区分它们的相对位置，并为模型提供上下文关系信息。

对于每个位置i，预先计算一个固定的位置向量pe_i，其维度与Embedding相同。在输入模型前，将每个token的Embedding与对应位置的PE相加，得到包含位置信息的token表示：

token_i_with_pe=Embedding_i+pe_i

其中，Embedding_i是第i个token的Embedding，pe_i是第i个位置的位置编码向量。二者相加如下：

[p_{00},p_{01},p_{02},...,p_{09}]       [pe_{00},pe_{01},pe_{02},...,pe_{09}]
[p_{10},p_{11},p_{12},...,p_{09}]       [pe_{10},pe_{11},pe_{12},...,pe_{09}]
[p_{20},p_{21},p_{22},...,p_{09}]    +  [pe_{20},pe_{21},pe_{22},...,pe_{09}]
...                                       ...
[p_{n0},p_{n1},p_{n2},...,p_{09}]       [pe_{n0},pe_{n1},pe_{n2} ,...,pe_{09}]

transformer

目前大语言模型都是基于transformer结构。在生成任务中，如文本生成、对话响应生成、摘要生成等，模型(比如GPT、llama)通常只使用Transformer架构中的Decoder部分，也就是所谓的Decoder-Only结构。

自回归生成

在生成输出序列任务中，使用自回归（Autoregressive）方式，即每次只生成一个token，并且这个token的生成依赖于之前已经生成的所有token。例如下面的代码：

# 定义使用的LLaMA2模型
model = LLaMA2()

# 定义自回归生成函数
def generate(inputs, n_tokens_to_generate):

    # 自回归解码循环，迭代次数等于要生成的token数量
    for _ in range(n_tokens_to_generate):
        # 将当前输入传入模型进行前向传播
        output = model(inputs)

        # 使用贪婪采样（Greedy Sampling）策略，选取概率最高的token作为下一个预测结果
        next = np.argmax(output[-1])

        # 将预测的token添加到输入序列中，供下次迭代使用
        inputs.append(next)

    # 返回最后生成的n_tokens_to_generate个token
    return inputs[len(inputs) - n_tokens_to_generate :]

# 给定初始输入，包含特殊token 'BOS' 和两个汉字 '岱' '宗'
input = [p0, p1, p2]

# 请求生成3个新token
output_ids = generate(input, 3)  # 假设生成 ['p3','p4','p5']

# 将生成的token ID解码为实际字符
output_ids = decode(output_ids)  # 通过tokenization解码

# 将解码后的token ID转换为词汇表中的词汇（此处假设vocab是一个字典）
output_tokens = [vocab[i] for i in output_ids]  # 得到 "夫" "如" "何"

输出处理

生成的token序列通过一个输出层，将每个位置的概率分布转换为对应token的概率。根据概率，选择概率最高的token作为模型预测输出。

'''
从给定的概率分布中采样一个token，采用top-p策略
probs: 表示给定的概率分布
p: 表示概率阈值，在采样过程中，只保留累积概率小于p的部分
'''
def sample_top_p(probs, p):
    # 1.概率降序排序：对输入的 probs 张量按最后一个维度（即每个概率向量内部）进行降序排序
    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True) #给定的概率降序排序
    # 2. 计算概率的累计和：对排序后的概率向量进行累计求和，得到一个新的张量，表示每个概率向量的累计概率
    probs_sum = torch.cumsum(probs_sort, dim=-1)
    # 3. 计算累积概率减去当前概率值是否大于p
    # 生成一个布尔型张量 mask，其中 True 表示该位置的累积概率减去当前概率值大于p，False则反之
    mask = probs_sum - probs_sort > p
    # 4. 在 probs_sort 张量中，将 mask 为 True 的位置（累积概率超过p 的部分）的值置为0。
    # 这样就实现了仅保留累积概率小于p的部分
    probs_sort[mask] = 0.0
    # 5. 归一化处理：对经过截断处理后的probs_sort张量进行归一化，使其概率总和为1
    # 使用 sum(dim=-1, keepdim=True) 计算每个概率向量的总和，并保持维度不变。然后进行元素级除法操作，使每个概率向量成为一个合法的概率分布。
    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
    # 6. 随机采样：进行一次随机抽样，得到一个形状为 (batch_size, 1) 的张量，表示每个批次数据采样到的 token 索引
    next_token = torch.multinomial(probs_sort, num_samples=1)
    # 7. 还原原始索引：根据next_token中的索引，从probs_idx中提取对应的原始索引
    next_token = torch.gather(probs_idx, -1, next_token)
    return next_token

模型结构

目前主流的LLM模型大多都是基于Transformer构建，llama2也不例外。LLM是根据给定输入文本序列的上下文信息预测下一个token，因此通常只需要Transformer Decoder部分。
而Decoder与Encoder的本质区别就是在计算Q*V时引入Mask以确保当前位置只关注前面已经生成的内容。

llama2的模型结构与Transformer Decoder部分基本一致，主要由32个Transformer Block组成。

同时在Transformer Decoder基础上做了如下改进：

1. 前置归一化，使用RMSNorm
2. Q与K相乘之前，使用旋转位置编码ROPE(Rotary Position Embeddings)
3. KV Cache，并采用Group Multi-Query Attention
4. Feed Forward SwiGLU

前置归一化-RMSNorm

为什么要进行归一化？

归一化(Normalization)是指将数据按照比例缩放，使其落入一个小的特定区间，通常是0-1，这样做有助于加快模型训练速度，提高模型性能。比如某模型训练场合，一些特征数值是远大于其他特征值的，像人的身高(eg:180cm)与体重(eg:75kg)，这样直接训练，可能会影响损失函数梯度，导致优化过程困难，同时容易造成梯度消失或爆炸。

前置归一化是指在每一层神经网络计算之前先进行归一化操作，与传统后置归一化(即在每一层计算之后进行归一化)相对。具体实现如下：

1. 第一层归一化：先对输入进行归一化，再送入多头注意力层（Multi-Head Attention, MHA）进行计算。
2. 第二层归一化：先对从MHA输出的特征进行归一化，然后再输入到全连接前馈神经网络（Feedforward Neural Network, FNN）进行计算。
3. 残差连接：多头注意力层的原始输出会与经过第一层归一化后的输入相加，然后再输入到全连接层。

Transformer中的Normalization层一般都是采用LayerNorm来对Tensor进行归一化，LayerNorm的公式如下：

$$y=\frac{x-E[x]}{\sqrt{Var[x]}+ϵ}\ast \gamma +\beta$$

其中，

$$E[x]=\frac{1}{N}\sum _{i=1}^N{x}_i$$

$$Var[x]=\frac{1}{N}\sum _{i=1}^N(x_i-E[x]{)}^2$$

llama的前置归一化采用RMSNorm(Root Mean Square Normalization)，省去了求均值的过程，也没有了偏置$\beta$。

$$y=\frac{x}{\sqrt{Mean(x^2)}+ϵ}\ast \gamma$$

其中，$\gamma ,\beta$是可学习的参数，且：

$$Mean(x^2)=\frac{1}{N}\sum _{i=1}^N{x}_i^2$$

# RMSNorm
class RMSNorm(torch.nn.Module):
    '''
    RMSNorm：归一化
    :param dim: int，待归一化的特征维度（通常是通道数）
    :param eps: float，默认值为 1e-6，用于防止除法运算时分母过小导致数值不稳定
    '''
    def __init__(self, dim: int, eps: float = 1e-6):
        # 初始化父类（即 torch.nn.Module）的属性和方法
        super().__init__()
        self.eps = eps # ε
        # 初始值为全1向量
        self.weight = nn.Parameter(torch.ones(dim)) #可学习参数γ
​    '''
    _norm 是一个私有方法，仅在 RMSNorm 类内部使用
    '''
    def _norm(self, x):
        # RMSNorm
        '''
        :param x: 待归一化的tensor，通常为模型某一层的输出，形状为 (batch_size, ..., dim)
        '''
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
​    # 前向传播
    # 返回经过标度的归一化张量作为RMSNorm层的输出
    def forward(self, x):
        # 将输入张量 x 转换为浮点类型（float），确保计算精度
        # 调用私有方法 _norm 对 x 进行根均方归一化，得到归一化后的张量 output
        # 将 output 的数据类型恢复为与输入 x 相同，以保持数据类型一致性
        output = self._norm(x.float()).type_as(x)
        # 将归一化后的张量 output 与可学习权重参数 self.weight 相乘，对归一化结果进行标度
        return output * self.weight

其中，计算RMSNorm的具体步骤如下：

1. 对输入x按照最后一个维度（即dim维）计算元素平方x.pow(2)。
2. 在最后一个维度上计算平均值mean(-1)，保留维度大小（keepdim=True），得到形状为 (batch_size, …, 1) 的均方值向量。
3. 向均方值向量中添加 eps 平滑项，避免除以过小的数值。
4. 使用torch.rsqrt计算均方值向量的逆平方根，得到归一化因子。
5. 将归一化因子逐元素与输入x相乘，完成根均方归一化RMSNorm计算。

RoPE-rotary positional embeddings

Transformer模型通常在输入序列经过Embedding层后只做一次位置编码，而lamma2模型选择在每个Attention层中分别对Query(Q)和Key(K)进行旋转位置编码（Rotary Positional Embedding, RoPE），即每次计算Attention时，都需要对当前层的Q和K进行位置编码。

RoPE是为了解决什么问题？用提出者苏大神的话来说，“就是‘通过绝对位置编码的方式实现相对位置编码’，这样做既有理论上的优雅之处，也有实践上的实用之处，比如它可以拓展到线性Attention中就是主要因为这一点。”。

假设通过下述运算来给q,k添加绝对位置信息：

$${\bar{q}}_m=f_q(x_m,m),{\bar{k}}_n=f_k(x_n,n)$$

上述函数处理后，使得${\bar{q}}_m、{\bar{k}}_n$是带有位置m、n的绝对位置信息。

Attention的核心运算是内积，所以我们希望的内积的结果带有相对位置信息，因此假设存在恒等关系：

$$<f_q(x_m,m),f_k(x_n,n)>=g(x_m,x_n,m-n)$$

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

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

$$f_q(x_m,m)=(W_q{x}_m)e^{in\theta }$$

$$f_k(x_n,n)=(W_k{x}_n)e^{in\theta }$$

$$g(x_m,x_n,,m-n)=Re[(W_q{x}_m)(W_k{x}_n{)}^{\ast }{e}^{i(m-n)\theta }]$$

乍一看挺复杂哈，怎么理解呢？

首先我们先回顾一下复数基础。f和g公式中都有一个指数函数$e^{ix}$，这不是欧拉公式么，x表示任意实数，e是自然对数的底数，i是复数中的虚数单位，则根据欧拉公式可知：

$$e^{ix}=cos(x)+i\ast sin(x)$$

因而，上述指数函数可以表示为实部为cos(x)，虚部为sin(x)的一个复数，从而建立了指数函数、三角函数和复数之间的桥梁。则：

$$e^{im\theta }=cos(m\theta )+i\ast sin(m\theta )$$

$$e^{in\theta }=cos(n\theta )+i\ast sin(n\theta )$$

$$e^{i(m-n)\theta }=cos((m-n)\theta )+i\ast sin((m-n)\theta )$$

此时再回看：

$$f_q(x_m,m)=(W_q{x}_m)e^{im\theta }$$

其中，$W_q$是一个二维矩阵，$x_m$是二维向量，二者相乘也是二维向量，用$q_m$表示如下：

$$q_m=(\begin{array}{c}{q}_m^{}\end{array}$$