LLAMA2入门（一）-----预训练

Llama 2 是预训练和微调的LLM系列，Llama 2 和 Llama 2-Chat 模型的参数规模达到 70B。Llama 2-Chat 模型专门为对话场景进行了优化。

这是一个系列的文章，会分别从LLAMA2的预训练，微调，安全性等方面进行讲解。

1.数据来源

• 数据集：LLaMA 2 的训练数据来源于公开可用的数据集，并未使用来自 Meta 产品或服务的数据。为了避免隐私泄露，数据集中剔除了来自某些网站的高风险个人信息。
• 数据量：模型共预训练了 2 万亿个 token，比 LLaMA 1 增加了约 40%。同时，数据集上调了最具事实性的来源，以提升模型知识储备，减少幻觉（hallucination）的发生。
• 去重和清洗：数据经过了更严格的清洗和去重处理，旨在提供更干净且更具代表性的数据集。

2.训练细节

模型架构

LLaMA 2 使用了标准的 Transformer 架构（Vaswani et al., 2017），并基于 LLaMA 1 进行了优化，包括使用预归一化（pre-normalization）RMSNorm、SwiGLU 激活函数（Shazeer, 2020）以及旋转位置嵌入（RoPE, Su et al. 2022）。这些优化提升了模型处理长上下文的能力。

（1） 基础架构

• LLaMA 2 使用的是自回归语言模型（Auto-regressive Language Model），这意味着模型是通过预测每个词在给定上下文中的下一个词来进行训练的。
• 标准的 Transformer 架构：模型使用经典的 Transformer 架构 ，最初由 Vaswani 等人提出，这是一种基于自注意力机制的架构，能够处理长距离依赖关系并捕捉句子中的语义结构。
• 归一化层：LLaMA 2 采用了预归一化策略（Pre-normalization），具体使用了 RMSNorm 来替代传统的 LayerNorm。这使得模型的训练更为稳定，尤其是在处理深层网络时。

（2） 改进的激活函数

LLaMA 2 使用了 SwiGLU 激活函数，而不是传统的 ReLU 或 GeLU。SwiGLU 是一种更复杂的非线性激活函数，能够提高模型的表现，特别是在大规模模型中 。

为什么要使用SwiGLU激活函数？

SwiGLU 的公式如下：
$$SwiGLU(x)=(Linear(x)\cdot SiLU(Linear(x)))$$
其中，SiLU（Sigmoid Linear Unit）是一个常用的非线性激活函数，定义为：
$$SiLU(x)=x\cdot \sigma (x)$$
其中$\sigma (x)$ 是 Sigmoid 函数 $\frac{1}{1+e^{-x}}$

因此，SwiGLU 实际上是对输入进行了两次线性变换，并通过 SiLU 将两个结果结合起来，这种组合使得模型可以捕获更丰富的特征表示。

1. SwiGLU有更高的表达能力

   SwiGLU 通过双线性变换和 SiLU 非线性结合，提供了更强的表达能力。具体体现在：

   • 非线性表现更强：SwiGLU 在输入空间上提供了更复杂的非线性映射，能够捕获数据中更加复杂的模式。相比于传统的 ReLU，它避免了“梯度消失”或“梯度爆炸”的问题，同时也比 GeLU 更能灵活地处理不同输入。
   • 激活函数的平滑性：SiLU 是一种平滑的激活函数，这意味着它在小输入值附近表现平稳，输出变化相对缓和。与 ReLU 不同，SiLU 不会像 ReLU 那样在负数部分输出恒为零，这有助于模型更好地处理负值输入，从而提高训练稳定性。

2. 计算效率更高

   虽然 SwiGLU 包含了两次线性变换，但由于它与 SiLU 激活函数结合，SwiGLU 可以充分利用并行计算硬件（如 GPU 和 TPU）来高效执行。因此，它相比于一些复杂的非线性函数（如 ReLU 或 Tanh）在计算开销上并没有显著增加，同时提升了模型的性能。

3. 更好的梯度流动

   SwiGLU 结合了 SiLU 的平滑性和双线性变换的组合，使得它能够在梯度反向传播时提供更好的梯度流动，这尤其在深层网络中表现出色。在 LLaMA 2 这样的大规模 Transformer 模型中，模型层数很多，使用 SwiGLU 可以有效地缓解梯度消失问题，确保模型能够在训练过程中更好地学习到深层次的特征。

（3）位置嵌入

​ 旋转位置嵌入（RoPE）：LLaMA 2 使用了旋转位置嵌入（Rotary Positional Embeddings, RoPE），这是一种改进的相对位置编码方式，能够处理更长的上下文信息，并提高模型在长文本上的性能 。

1. 位置编码的重要性

   Transformer 模型本质上是不具备顺序感的，因为它们并没有卷积或递归结构，而是通过自注意力机制（Self-Attention）在序列中的所有词之间计算注意力权重。因此，在处理语言任务时，必须显式地告诉模型每个词在序列中的位置，以便模型能够理解句子中的词序。

   传统的 Transformer 使用绝对位置编码（Absolute Positional Encoding），这意味着每个位置（token）的信息是通过一个固定的向量表示，该向量在训练开始时被注入到输入中。然而，绝对位置编码无法很好地建模词与词之间的相对位置关系，特别是当序列长度增加时，绝对位置编码的表现会下降。

2. RoPE的核心思想

   旋转位置嵌入（RoPE） 通过引入旋转操作来嵌入相对位置信息。它通过对输入的查询和键（Query 和 Key）向量进行逐元素的旋转变换，将位置信息嵌入到自注意力机制的计算中，从而保留了序列中词之间的相对位置。

   RoPE 的核心思想是将每个 token 的查询和键向量按照其位置进行旋转，而这种旋转是连续的，能够保持词与词之间的相对距离信息。这意味着，相比于绝对位置编码，RoPE 不需要依赖固定的位置编码表格，而是通过对向量进行旋转操作来隐式地表示相对位置。

   RoPE 的数学定义如下：

   旋转变换

   给定二维向量 $[x1,x2]$，我们可以定义它的旋转变换为：
   $$\left[\begin{array}{c}{x}_1^{\prime }\\ x_2^{\prime }\end{array}\right]=\left[\begin{array}{cc}\cos (\theta )& -\sin (\theta )\\ \sin (\theta )& \cos (\theta )\end{array}\right]\left[\begin{array}{c}{x}_1\\ x_2\end{array}\right]$$
   在 RoPE 中，这种旋转操作会应用到高维向量的每对相邻维度上（如第 1 维和第 2 维，第 3 维和第 4 维），每对维度的旋转角度由位置和频率确定。

   具体来说，RoPE 对每个位置 ppp 以逐维不同的频率来进行旋转变换，确保较低维度捕捉较短距离的依赖，较高维度捕捉较长距离的依赖。

   相对位置优势

   这种基于旋转的位置编码能够隐式地表示相对位置。RoPE 的一个关键优势在于，经过旋转编码后，两个位置 $p$ 和 $p+k$ 的旋转向量之间的差异只取决于相对位置 $k$，而不是绝对位置。这意味着即使文本序列长度增加，RoPE 也能保持相对位置的关系，使其更具泛化能力，特别是在长文本处理上表现优异。

3. RoPE的具体实现步骤

   RoPE 的实现涉及以下几个步骤：

   1. 输入处理：输入的查询向量 $q$ 和键向量 k 是没有位置信息的高维向量。
   2. 位置嵌入：对每个位置 $p$ 的查询和键向量，按照位置 $p$ 对查询 $q$ 和键 $k$ 的偶数维和奇数维进行不同的旋转嵌入。
   3. 自注意力机制计算：经过旋转嵌入后，查询和键向量会进入自注意力机制进行打分计算，这个过程中，自然地引入了位置相关性和相对位置偏置。
   4. 输出生成：自注意力机制根据注意力权重生成新的向量表示，传递到后续的 Transformer 层。

4. RoPE 相比传统位置编码的优势

   1. 处理长序列的能力

   RoPE 的设计使其可以扩展到更长的序列而不会丢失位置信息。传统的绝对位置编码在序列长度超过训练时的长度后，模型的表现会显著下降，因为位置编码无法很好地泛化。而 RoPE 通过旋转方式嵌入相对位置，具备良好的泛化能力，适用于长序列文本处理。

   2. 相对位置建模

   RoPE 的另一个显著优势是相对位置编码。对于许多 NLP 任务，词与词之间的相对距离往往比绝对位置更为重要。例如，在机器翻译或长文本生成任务中，模型必须理解词语之间的相对关系，而不仅仅是它们的固定位置。

   3. 计算效率

   RoPE 的旋转计算相对简单，能够高效地并行化计算。相较于需要额外层次处理的相对位置编码方案（如 Transformer-XL 中的相对位置编码），RoPE 在引入相对位置偏置的同时，保持了自注意力机制的计算效率。

5. 具体使用例子

   给定一个查询向量 qqq 或键向量 kkk，我们通过以下旋转公式将位置信息嵌入：
   $$\text{RoPE}(q_{2i},q_{2i+1})=\left[\begin{array}{cc}\cos ({\theta }_p)& -\sin ({\theta }_p)\\ \sin ({\theta }_p)& \cos ({\theta }_p)\end{array}\right]\left[\begin{array}{c}{q}_{2i}\\ q_{2i+1}\end{array}\right]$$
   这里：

   • $q_{2i}$和$q_{2i+1}$分别是查询向量 $q$ 的第 $2i$ 和第 $2i+1$ 维。
   • ${\theta }_p=p\cdot \theta$是与位置 $p$ 相关的旋转角度,通常设定为与维度相关的常数（如$\theta =10000^{\frac{-2i}{d}}$,其中$d$是查询向量的维度）
   • 这个变换将每对相邻维度 $(2i,2i+1)$旋转一个与位置 $p$ 相关的角度。

注意力机制的改进

引入了分组查询注意力机制（Grouped-Query Attention, GQA），特别是用于处理参数规模较大的模型（例如 34B 和 70B）。这一改进提升了推理的可扩展性，降低了注意力计算时的内存消耗。

查询注意力机制（Grouped-Query Attention, GQA） 是 LLaMA 2 模型中引入的一种优化注意力计算的机制，特别是在处理大规模模型（如 34B 和 70B 参数模型）时，GQA 大幅减少了注意力计算的内存开销，并提高了模型的推理效率。

为了理解 GQA，首先需要了解传统的多头自注意力机制（Multi-Head Self-Attention, MHA），然后我们会介绍 GQA 是如何改进 MHA 的。

回顾多头自注意力机制（MHA）

在传统的多头自注意力机制中，自注意力计算分为多个“头”（attention heads）。每个注意力头会为输入序列生成查询（Query, Q）、键（Key, K）和值（Value, V），并分别进行自注意力计算。具体流程如下：

1. 线性投影：每个输入向量（token embedding）会通过线性投影生成对应的 Q、K、V 向量。这意味着每个注意力头有自己独立的线性投影矩阵。

2. 注意力计算：每个头都会独立地进行注意力计算，计算公式如下：
   $$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$$
   其中 $Q$ 是查询，$K$ 是键，$V$ 是值，$d_k$ 是键向量的维度。

3. 多头组合:所有注意力头的输出会被拼接起来，经过一个线性变换后，生成最后的输出。

传统的 MHA 通过引入多个注意力头，允许模型在不同的子空间中学习不同的注意力模式。然而，MHA 的一个缺点是，当模型参数增加时（如 34B 或 70B 参数模型），每个头都有独立的 Q、K、V 线性投影矩阵，这会显著增加内存和计算开销。

传统多头注意力机制的代码：

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadSelfAttention(nn.Module):
    def __init__(self, embed_size, num_heads):
        super(MultiHeadSelfAttention, self).__init__()
        self.embed_size = embed_size
        self.num_heads = num_heads
        self.head_dim = embed_size // num_heads
        
        assert self.head_dim * num_heads == embed_size, "Embedding size needs to be divisible by heads"

        # 为每个头独立生成 Q、K、V
        self.query = nn.Linear(embed_size, embed_size, bias=False)
        self.key = nn.Linear(embed_size, embed_size, bias=False)
        self.value = nn.Linear(embed_size, embed_size, bias=False)

        # 最后的全连接层
        self.fc_out = nn.Linear(embed_size, embed_size)

    def forward(self, values, keys, query, mask):
        N = query.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

        # 对每个输入进行线性变换得到 Q、K、V
        queries = self.query(query)
        keys = self.key(keys)
        values = self.value(values)

        # 将维度 (N, seq_length, embed_size) -> (N, num_heads, seq_length, head_dim)
        queries = queries.view(N, query_len, self.num_heads, self.head_dim).transpose(1, 2)
        keys = keys.view(N, key_len, self.num_heads, self.head_dim).transpose(1, 2)
        values = values.view(N, value_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 注意力得分计算: QK^T / sqrt(d_k)
        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
        
        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))
        
        attention = torch.softmax(energy / (self.head_dim ** (1 / 2)), dim=3)

        # 使用注意力权重加权 V
        out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape(N, query_len, self.embed_size)

        return self.fc_out(out)

查询注意力机制（Grouped-Query Attention, GQA）

Grouped-Query Attention (GQA) 通过对查询（Q）的投影进行分组优化，大幅减少了计算和内存的使用。具体来说，GQA 保留了多个注意力头的设计，但共享查询的投影矩阵，而不是为每个头都独立投影查询向量。GQA 的核心思想是将多个注意力头分成几组，每组内的所有注意力头共享同一个查询投影矩阵。

GQA 的主要改进点：

共享查询投影：在 GQA 中，多个注意力头会共享一个查询投影矩阵，而不再为每个头都单独分配一个查询投影矩阵。这样减少了查询投影的计算复杂度和内存需求。

• 例如，如果有 64 个注意力头，GQA 可能将这些头分为 4 组，每组 16 个头。组内的所有 16 个头共享同一个查询投影矩阵，而键和值仍然各自独立。

保持多头注意力的灵活性：尽管查询投影被分组共享，键（K）和值（V）的投影仍然是独立的。这意味着每个注意力头仍然可以在不同的子空间中捕捉不同的注意力模式，而不会丧失多头注意力机制的多样性。

GQA 的工作原理

查询投影分组：多个注意力头被分成几组（例如每组 8 或 16 个头），组内的所有注意力头共享同一个查询投影矩阵。这减少了模型中查询投影的总量。

键和值的独立投影：每个头仍然会独立生成键和值向量，从而保留每个头独立关注不同信息的能力。

自注意力计算：GQA 使用共享的查询投影矩阵来计算查询向量 QQQ，然后与独立的键和值一起进行标准的自注意力计算。尽管查询投影被共享，但每个头的键和值是不同的，因此每个头的注意力分数仍然可以有自己的模式。

多头组合：与传统 MHA 类似，所有注意力头的输出会被拼接起来，经过线性变换后生成最终的输出。

GQA 的优点

相比于传统的 MHA，GQA 在大规模模型中具有以下几个显著优点：

1. 减少内存占用

• 由于查询投影矩阵被分组共享，GQA 显著减少了查询投影矩阵的总量。这在处理大规模模型（如 LLaMA 2 的 34B 和 70B 参数模型）时，能有效降低内存占用。

2. 提高推理效率

• 通过减少查询投影的计算，GQA 可以加速推理过程，尤其是在大型模型中。传统的 MHA 随着注意力头的增加，计算量线性增长，而 GQA 通过共享查询投影减少了计算量，提升了推理效率。

3. 更好的可扩展性

• GQA 尤其适用于大规模模型。在大模型中，MHA 的计算和内存瓶颈会变得更加明显，而 GQA 通过减少查询的独立投影，可以更好地扩展模型参数，同时保持较高的计算效率。

4. 保持多头的灵活性

• 尽管 GQA 共享了查询的投影，但仍然保留了多头注意力机制的多样性。每个头的键和值投影仍然是独立的，因此 GQA 不会牺牲模型在不同子空间中学习不同模式的能力。

查询注意力机制（GQA）的代码

class GroupedQueryAttention(nn.Module):
    def __init__(self, embed_size, num_heads, num_groups):
        super(GroupedQueryAttention, self).__init__()
        self.embed_size = embed_size
        self.num_heads = num_heads
        self.num_groups = num_groups  # 分组数
        self.head_dim = embed_size // num_heads
        
        assert self.head_dim * num_heads == embed_size, "Embedding size needs to be divisible by heads"
        assert num_heads % num_groups == 0, "Number of heads must be divisible by the number of groups"

        # 共享查询的投影矩阵
        self.shared_query = nn.Linear(embed_size, embed_size // num_groups, bias=False)  # 每组共享

        # 为每个头单独生成 K 和 V
        self.key = nn.Linear(embed_size, embed_size, bias=False)
        self.value = nn.Linear(embed_size, embed_size, bias=False)

        # 最后的全连接层
        self.fc_out = nn.Linear(embed_size, embed_size)

    def forward(self, values, keys, query, mask):
        N = query.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]
        
        # 为每个头生成 K 和 V
        keys = self.key(keys)
        values = self.value(values)

        # 将每组内的所有头共享查询投影矩阵
        queries = []
        group_size = self.num_heads // self.num_groups
        for i in range(self.num_groups):
            # 对查询进行分组计算
            q = self.shared_query(query)
            q = q.repeat(1, 1, group_size)  # 共享矩阵复制给组内的所有头
            queries.append(q)
        
        queries = torch.cat(queries, dim=-1)  # 将分组后的 Q 拼接起来

        # 调整维度以适应注意力计算: (N, seq_length, embed_size) -> (N, num_heads, seq_length, head_dim)
        queries = queries.view(N, query_len, self.num_heads, self.head_dim).transpose(1, 2)
        keys = keys.view(N, key_len, self.num_heads, self.head_dim).transpose(1, 2)
        values = values.view(N, value_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 注意力得分计算: QK^T / sqrt(d_k)
        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])
        
        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        attention = torch.softmax(energy / (self.head_dim ** (1 / 2)), dim=3)

        # 使用注意力权重加权 V
        out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape(N, query_len, self.embed_size)

        return self.fc_out(out)

代码解释：

energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])

该行代码使用了 torch.einsum 来计算注意力得分矩阵，其中 queries 和 keys 是多头自注意力机制中的查询和键张量。

queries 和 keys 的形状为 (N, S, num_heads, head_dim)，经过变换，当前它们的形状为：

queries：(N, query_len, num_heads, head_dim)，即nqhd：

N：批次大小（batch size）

q：查询序列长度（query length）

h：注意力头的数量（number of heads）

d：每个头的嵌入维度（embedding size per head）

keys`：`(N, key_len, num_heads, head_dim)`，即 `nkhd

• N：批次大小
• k：键的序列长度（key length）
• h：注意力头的数量
• d：每个头的嵌入维度

爱因斯坦求和约定

torch.einsum 使用字符串形式的爱因斯坦求和约定，"nqhd,nkhd->nhqk"，用来描述张量的运算规则。这个规则可以分为几个部分：

• 输入部分

  "nqhd,nkhd"
  

  描述了输入张量的维度：

  • queries 的维度是 nqhd，即 (N, query_len, num_heads, head_dim)。
  • keys 的维度是 nkhd，即 (N, key_len, num_heads, head_dim)。

• 输出部分

  "nhqk"
  

  描述了输出张量的维度：

  • n：批次大小。
  • h：注意力头的数量。
  • q：查询序列长度。
  • k：键序列长度。

在这个表达式中，d（head_dim，嵌入维度）被省略掉了，因为 d 是被求和维度。torch.einsum 将在 d 维上进行内积操作，然后输出 (N, num_heads, query_len, key_len) 形状的张量。

注意力得分矩阵

这一步的结果 energy 是一个形状为 (N, num_heads, query_len, key_len) 的张量，它包含每个注意力头中每个查询位置与所有键位置之间的点积结果。该点积值表示查询和键之间的相似度，也就是注意力机制中的原始得分。

接下来，通常会对这个 energy 张量进行 softmax 归一化，以计算注意力权重，从而决定每个查询应该重点关注哪些键值。

传统多头自注意力机制（MHA）与 GQA 的区别

特性	传统多头自注意力机制（MHA）	查询注意力机制（GQA）
查询（Q）投影矩阵	每个注意力头都有独立的查询投影矩阵	查询投影矩阵在分组内的所有头中共享，减少了查询投影的计算量
键（K）和值（V）投影矩阵	每个头有独立的键和值投影矩阵	每个头仍然有独立的键和值投影矩阵，保持注意力头的灵活性
内存需求	随着注意力头的数量线性增加	通过共享查询投影减少内存占用
计算效率	当注意力头数较多时，计算复杂度和内存开销大	通过共享查询投影，计算效率更高，尤其是在大模型中
灵活性	每个头都有独立的查询投影，因此具有较大的灵活性	查询投影共享在一定程度上减少了灵活性，但仍保持键和值的独立性

上下文长度

LLaMA 2 将上下文长度从 2048 token 增加到 4096 token，这使模型能够处理更长的文本段落，特别是在对话场景和长文本总结中表现更好。

3.训练超参数

• 使用了 AdamW 优化器，参数为 β1=0.9、β2=0.95，学习率使用余弦衰减调度，初始预热 2000 步，最终学习率衰减至峰值的 10%。模型使用了权重衰减（weight decay）0.1 和梯度裁剪（gradient clipping）1.0。
• 训练时长和资源：为了减少训练时间，LLaMA 2 采用了 Meta 的研究超级计算集群（RSC），并使用 NVIDIA A100 GPUs。每个 GPU 的功率限制为 350-400W，总共花费 330 万 GPU 小时 进行预训练，碳排放总计 539 吨 CO2，100% 通过 Meta 的可持续性计划进行了抵消 。

4.优化点

更强的泛化能力：LLaMA 2 的设计在处理长上下文、减少幻觉、提升模型推理速度和内存效率等方面做了大量优化。这些优化使得模型在对话生成、文本总结等任务中表现优异，特别是在需要保持上下文一致性的场景下效果更好。