LLaMA (以LLaMA2为例，文末附加对比1 2 3 三个版本的变化）

补充背景：

关于Transformer和Llama架构的演进

一、背景

LLaMA 2 和 LLaMA2-Chat
参数规模：70亿、130亿和700亿
数据和训练规模：
上下文长度
训练资源
性能表现：

二、预训练 pretraining

1. 预训练数据

· 训练语料来自公开课用的数据源，不包括Meta的产品或服务数据
· 在2万亿个数据tokens上进行了训练
· 对真实的数据源进行上采样以提高只是并减少错误

2. 训练细节

2.1 标准的Transformer架构

2.2 RMSNorm归一化

2.3 SwiGLU激活函数

2.4 RoPE 旋转位置编码

import torch
import torch.nn as nn

class LlamaRotaryEmbedding(nn.Module):
    """
    计算 RoPE（旋转位置编码）所需的 cos(θ) 和 sin(θ) 值
    """
    def __init__(self, dim, base=10000):
        """
        初始化 LlamaRotaryEmbedding，计算逆频率 inv_freq
        :param dim: 需要旋转的位置编码维度（通常是 head_dim）
        :param base: 位置编码的基数（通常是 10000）
        """
        super().__init__()
        # 计算逆频率 inv_freq，维度为 [dim/2]
        # 公式：θ_i = 10000^{-2(i-1)/d}
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
        self.register_buffer("inv_freq", inv_freq)  # 不作为模型参数，但存入模型权重

    def forward(self, x, position_ids):
        """
        计算 cos(θ) 和 sin(θ) 值
        :param x: 输入 tensor（仅用于获取 batch 形状）
        :param position_ids: 位置索引 (batch_size, seq_len)
        :return: cos(θ), sin(θ)
        """
        # 扩展 inv_freq 以匹配 position_ids 的 batch 维度
        position_ids = position_ids.unsqueeze(-1)  # 形状变为 (batch_size, seq_len, 1)
        freqs = torch.einsum("bi,j->bij", position_ids.float(), self.inv_freq)  # 计算 mθ
        emb = torch.cat((freqs, freqs), dim=-1)  # 复制 freq，使其维度与 embedding 维度匹配

        # 计算 cos(θ) 和 sin(θ)，形状 (batch_size, seq_len, head_dim)
        cos = emb.cos()
        sin = emb.sin()

        return cos, sin


def rotate_half(x):
    """
    实现向量的旋转变换
    例如：输入 [x1, x2, x3, x4] -> 输出 [-x2, x1, -x4, x3]
    :param x: 输入 tensor 形状 (batch_size, seq_len, num_heads, head_dim)
    :return: 旋转后的 tensor
    """
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2:]  # 拆分 tensor
    return torch.cat((-x2, x1), dim=-1)  # 交换并加负号


def apply_rotary_pos_emb(q, k, cos, sin):
    """
    计算旋转位置编码后的 Q 和 K
    :param q: Query (batch_size, num_heads, seq_len, head_dim)
    :param k: Key (batch_size, num_heads, seq_len, head_dim)
    :param cos: cos(θ) 形状 (batch_size, seq_len, head_dim)
    :param sin: sin(θ) 形状 (batch_size, seq_len, head_dim)
    :return: 旋转编码后的 Q, K
    """
    # 进行旋转变换
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)

    return q_embed, k_embed


class LlamaSdpaAttention(nn.Module):
    """
    LLaMA 的注意力机制，包含 RoPE 旋转位置编码
    """
    def __init__(self, embed_dim, num_heads):
        """
        初始化注意力层
        :param embed_dim: 总 embedding 维度
        :param num_heads: 头数
        """
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads  # 每个头的维度

        # 线性变换层
        self.q_proj = nn.Linear(embed_dim, embed_dim)  # Query 投影
        self.k_proj = nn.Linear(embed_dim, embed_dim)  # Key 投影
        self.v_proj = nn.Linear(embed_dim, embed_dim)  # Value 投影

        # RoPE 旋转位置编码
        self.rotary_emb = LlamaRotaryEmbedding(self.head_dim)

    def forward(self, hidden_states, position_ids):
        """
        前向传播，计算注意力并应用 RoPE
        :param hidden_states: 输入 tensor (batch_size, seq_len, embed_dim)
        :param position_ids: 位置索引 (batch_size, seq_len)
        :return: 旋转编码后的 query_states, key_states, value_states
        """
        # 计算 Q, K, V
        query_states = self.q_proj(hidden_states)  # (batch_size, seq_len, embed_dim)
        key_states = self.k_proj(hidden_states)    # (batch_size, seq_len, embed_dim)
        value_states = self.v_proj(hidden_states)  # (batch_size, seq_len, embed_dim)

        # 重新调整形状以匹配多头注意力
        batch_size, seq_len, _ = hidden_states.shape
        query_states = query_states.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        key_states = key_states.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        value_states = value_states.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)

        # 计算 RoPE 旋转角度
        cos, sin = self.rotary_emb(hidden_states, position_ids)  # 计算 cos(θ), sin(θ)

        # 应用 RoPE 旋转编码
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        return query_states, key_states, value_states


# 测试代码
if __name__ == "__main__":
    batch_size = 2
    seq_len = 4
    embed_dim = 16
    num_heads = 4

    # 生成输入数据
    hidden_states = torch.rand(batch_size, seq_len, embed_dim)
    position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1)  # (batch_size, seq_len)

    # 实例化注意力层
    attention_layer = LlamaSdpaAttention(embed_dim, num_heads)

    # 计算 RoPE 旋转后的 Q, K, V
    query_states, key_states, value_states = attention_layer(hidden_states, position_ids)

    # 输出结果
    print("Query States:", query_states.shape)
    print("Key States:", key_states.shape)
    print("Value States:", value_states.shape)

2.5 GQA 分组查询注意力

2.6 Tokenizer分词器

三、微调 fine-tuning

1. 有监督微调 SFT

2. 基于人工反馈的强化学习 RLHF

3. 多轮对话中保持一致性的系统消息

四、LLaMA的前世今生（LLaMA1,2,3)

Llama1

动机：Meta认为推理成本更重要，所以提高数据量而不是模型大小，因为训练只需要一次，而推理是无数次的

具体行动：针对Transformer-decoder架构，做了以下修改：

和GPT-3一样将Normalization从每个子层的输出位置移动到了输入位置
将Layer Norm 改为 RMS Norm
动机：进行Norm时，对特征进行平移并不能改变特征的分布，所以可以去掉平移相关的部分
Note: 平移相关的部分指的是：
a. 输入特征-均值 $x - E[x]$
b. 对标准化后进行线性变化偏差的参数 $\beta$

LayerNorm（层归一化）

$ \text{LayerNorm}(x) = \frac{x - E[x]}{\sqrt{\text{Var}[x] + \epsilon}} * \gamma + \beta $

RMSNorm（均方根归一化）
$\ \text{RMSNorm}(x) = \frac{x}{\sqrt{\text{Mean}(x^2) + \epsilon}} * \gamma $

Note:
$\text{Var}(x) = E[x^2] - (E[x])^2$

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

    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)  # 乘以可训练参数