Base LLM | 从 NLP 到 LLM 的算法全栈教程 第六天

Llama2架构

与原始 Transformer 解码器相比，Llama2 及其同类模型进行了一系列改进，以提升性能和训练稳定性。它的数据流可以概括为：

（1）输入嵌入 ：将 token_ids 转换为词向量。

（2）N x Transformer 层堆叠：数据依次通过 N 个相同的 Transformer Block。

• 预归一化 ：在进入子层之前，先进行一次 RMSNorm。
• 注意力子系统 ：包含旋转位置编码 、分组查询注意力（GQA） 和 KV 缓存机制。
• 前馈网络子系统 ：采用 SwiGLU 激活函数。

（3）最终归一化与输出：在所有层之后，进行最后一次 RMSNorm，并通过一个线性层将特征映射到词汇表 logits。

预归一化

不再均值中心化，在dim维度进行均方根缩放，保留可学习的参数weight。

# code/C6/llama2/src/norm.py
class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim)) # 对应公式中的 gamma

    def _norm(self, x: torch.Tensor) -> torch.Tensor:
        # 核心计算：x * (x^2的均值 + eps)的平方根的倒数
        # torch.rsqrt平方根的倒数
        # pow的幂计算是按元素操作，输入输出的shape完全一致
        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        out = self._norm(x.float()).type_as(x)
        return out * self.weight

旋转位置编码RoPE

方案	原理	痛点
绝对位置编码	给每个位置一个固定向量，直接加在 token embedding 上	外推极差！训练时用了 4k 长度，推理时用 8k 就废了 ------ 因为模型从没见过第 5000 位的位置向量
传统相对位置编码	在注意力分数里，手动加一个和 "两个 token 距离" 有关的偏置	实现复杂，不支持线性注意力，长序列下计算量暴增

（1）**precompute_freqs_cis**: 预计算一个包含旋转角度信息的复数张量 freqs_cis，这个张量在模型初始化时计算一次即可。它的输入包括 head 的维度 dim、序列最大长度 end 以及一个用于控制频率范围的超参数 theta，最终输出一个形状为 [end, dim / 2] 的复数张量。

def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> torch.Tensor:
    # 1. 计算频率：1 / (theta^(2i/dim))
    # 我们要把 dim 维的向量，两两分组做二维旋转，拿到每个二维组的起始索引。
    # [: (dim // 2)]是为了解决奇数偶数的问题
    # /dim是为了做归一化，将索引缩放到0-1之间
    # freqs 取倒数，使得前面变化频率大，后面频率小，频率大容易走入循环，主要用于看短期信息
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    # 2. 生成位置序列 t = [0, 1, ..., end-1]
    t = torch.arange(end, device=freqs.device)
    # 3. 计算相位：t 和 freqs 的外积，位置*旋转角度
    freqs = torch.outer(t, freqs).float()
    # 4. 转换为复数形式 (cos(theta) + i*sin(theta))极坐标转欧拉公式
    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
    return freqs_cis

**apply_rotary_emb**: 将预计算的 freqs_cis 应用于输入的 Query 和 Key 向量。输入端需要传入形状为 [batch_size, seq_len, n_heads, head_dim] 的 Query 向量 xq、形状为 [batch_size, seq_len, n_kv_heads, head_dim] 的 Key 向量 xk，以及预计算的旋转矩阵切片 freqs_cis，它会输出旋转后的 xq 和 xk，且两者的形状保持不变。

def apply_rotary_emb(
    xq: torch.Tensor,
    xk: torch.Tensor,
    freqs_cis: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    # 将 Q/K 向量视为复数
    # （batch、seq_len、n_head）都保留，*xq.shape[:-1] 就是把这些维度原样展开
    # 把最后一维（head_dim），拆成 [d//2, 2]，也就是我们之前说的两两分组
    # 将最后两个维度转成复数的实部虚部，二维组 [q0, q1]，转完就变成复数 q0 + q1*i,
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    
    # 准备广播将freqs_cis的shape变为[1, seq_len, 1, d//2]
    freqs_q = reshape_for_broadcast(freqs_cis, xq_)  # 针对 Q 的广播视图
    
    # 复数乘法即为旋转
    xq_out = torch.view_as_real(xq_ * freqs_q).flatten(3)
    
    # K 向量可能与 Q 向量有不同的头数（GQA），所以需单独生成广播视图
    freqs_k = reshape_for_broadcast(freqs_cis, xk_)
    # 两维实数变成复数，在复数层级的四维，转化成实数变为三维，用flatten合并
    xk_out = torch.view_as_real(xk_ * freqs_k).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xq)

分组查询注意力

标准的多头注意力（Multi-Head Attention, MHA） 为每个 Query 头都配备了一组独立的 Key 和 Value 头。这意味着 K 和 V 投影矩阵的尺寸以及推理时 KV 缓存的大小都与总头数 n_heads 成正比，当模型规模增大时，这部分开销变得非常显著。而分组查询注意力（Grouped-Query Attention, GQA） 4 就是对此的核心优化，它的思路是允许多个 Query 头共享同一组 Key 和 Value 头 。具体来说，MHA 中每个 Q 头都有自己的 K/V 头（即 n_heads 与 n_kv_heads 相等），而 GQA 则是让每组 Q 头共享一组 K/V 头（此时 n_heads 大于 n_kv_heads）。还有一种特殊情况是多查询注意力（MQA） ，所有 Q 头共享唯一的一组 K/V 头（n_kv_heads 等于 1），可以被视为 GQA 的特例。

# code/C6/llama2/src/rope.py
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    batch_size, seq_len, n_kv_heads, head_dim = x.shape
    if n_rep == 1:
        return x
    # 现创建一个维度然后复制n_rep次，然后合并
    return (
        x[:, :, :, None, :]
        .expand(batch_size, seq_len, n_kv_heads, n_rep, head_dim)
        .reshape(batch_size, seq_len, n_kv_heads * n_rep, head_dim)
    )

# code/C6/llama2/src/attention.py
class GroupedQueryAttention(nn.Module):
    def __init__(self, dim: int, n_heads: int, n_kv_heads: int | None = None, ...):
        ...
        self.n_local_heads = n_heads
        self.n_local_kv_heads = n_kv_heads
        self.n_rep = self.n_local_heads // self.n_local_kv_heads # Q头与KV头的重复比
        ...
        self.wq = nn.Linear(dim, n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(dim, n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(dim, n_kv_heads * self.head_dim, bias=False)
        ...

    def forward(self, x, start_pos, freqs_cis, mask):
        xq = self.wq(x).view(batch_size, seq_len, self.n_local_heads, self.head_dim)
        xk = self.wk(x).view(batch_size, seq_len, self.n_local_kv_heads, self.head_dim)
        xv = self.wv(x).view(batch_size, seq_len, self.n_local_kv_heads, self.head_dim)

        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
        
        # ... KV Cache 逻辑 ...

        keys = repeat_kv(keys, self.n_rep)   # <-- 关键步骤
        values = repeat_kv(values, self.n_rep) # <-- 关键步骤

        scores = torch.matmul(xq.transpose(1, 2), keys.transpose(1, 2).transpose(2, 3)) / ...
        ...

SwiGLU

# code/C6/llama2/src/ffn.py
class FeedForward(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, multiple_of: int, ...):
        super().__init__()
        # hidden_dim 计算，并用 multiple_of 对齐以提高硬件效率
        hidden_dim = int(2 * hidden_dim / 3)
        ...
        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)

        self.w1 = nn.Linear(dim, hidden_dim, bias=False) # 对应 W
        self.w2 = nn.Linear(hidden_dim, dim, bias=False) # 对应 W2
        self.w3 = nn.Linear(dim, hidden_dim, bias=False) # 对应 V

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # F.silu(self.w1(x)) 实现了 swish(xW)
        # * self.w3(x) 实现了门控机制
        return self.w2(torch.nn.functional.silu(self.w1(x)) * self.w3(x))

整体代码

# code/C6/llama2/src/transformer.py
class LlamaTransformer(nn.Module):
    def __init__(self, vocab_size: int, ...):
        ...
        self.tok_embeddings = nn.Embedding(vocab_size, dim)
        self.layers = nn.ModuleList([TransformerBlock(...) for i in range(n_layers)])
        self.norm = RMSNorm(dim, eps=norm_eps)
        self.output = nn.Linear(dim, vocab_size, bias=False)
        self.register_buffer("freqs_cis", precompute_freqs_cis(...))

    def forward(self, tokens: torch.Tensor, start_pos: int) -> torch.Tensor:
        h = self.tok_embeddings(tokens)
        
        # 1. 准备 RoPE 旋转矩阵
        freqs_cis = self.freqs_cis[start_pos : start_pos + seq_len]

        # 2. 准备因果掩码 (Causal Mask)
        mask = None
        if seq_len > 1:
            mask = torch.full((seq_len, seq_len), float("-inf"), device=tokens.device)
            mask = torch.triu(mask, diagonal=1)
            # 考虑 KV Cache 的偏移
            mask = torch.hstack([torch.zeros((seq_len, start_pos), ...), mask]).type_as(h)

        # 3. 循环通过所有 TransformerBlock
        for layer in self.layers:
            h = layer(h, start_pos, freqs_cis, mask)
        
        h = self.norm(h)
        logits = self.output(h).float()
        return logits

MoE架构

稠密架构 ：对于每一个输入的 Token，模型中所有的 参数（从第一层到最后一层）都会参与计算，代表是 Llama2、GPT-3。混合专家模型（Mixture of Experts, MoE） 。MoE 技术通过一种 "稀疏激活" 的机制，兼具了大规模参数的知识容量与极低的推理成本。

如果我们尝试让一个网络同时学习多个截然不同的子任务（例如既学做菜又学修车），往往会出现**"强干扰效应（Strong Interference Effects）"**。这是因为网络的所有权重都参与了所有任务的计算，当网络调整参数以适应任务 A 时，可能会破坏它在任务 B 上已经学到的特征表示。

为了解决这个问题，论文提出了一种基于**"分治（Divide and Conquer）"**策略的系统架构：

• 专家网络：系统包含多个独立的神经网络（可以是简单的前馈网络）。每个专家不再需要处理全局任务，只需专注于输入空间中的一个局部区域或一类特定的子任务。
• 门控网络 ：充当协调者的角色。它接收与专家相同的输入 xx，并输出一组**混合比例（Mixing Proportions）**pipi，即选择每个专家的概率。它就像一个软性的随机开关，决定当前的输入案例应该由哪位专家来主导处理。

**如何稀疏选择：**在门控输入中加入可训练的高斯噪声，计算后仅保留权重最大的 k 个专家（例如 k=4），将其余所有专家的权重强制置为 −∞（即概率为 0）。随着稀疏专家的增多，效果显著提升。

如何防止专家崩塌： 总损失函数中加入了额外的辅助损失（Auxiliary Loss） ，包含 Importance Loss 和 Load Loss。这些损失函数并不直接服务于预测准确率，而是专门用来惩罚"分配不均"的现象，强制门控网络"雨露均沾"，确保所有专家都能接收到大致相等的样本量，从而得到充分的训练。

Gshard

数据并行与模型并行：

• 非 MoE 层（如 Attention） ：采用**复制（Replicated）**策略。所有设备持有相同的副本，进行标准的数据并行训练。
• MoE 层 ：采用**分片（Sharded）**策略。专家网络被切分并分布在不同设备上（例如 2048 个专家分布在 2048 个 TPU 核上）。
• 模型参数量增加 16 倍（从 37.5B 到 600B），训练算力成本仅增加了不到 4 倍。

Switch Transformer

单专家路由

TotalTokens =Batch Size*序列长度 (Sequence Length)

专家容量 = 平均分配token数 * 容量因子

超出专家容量的token被丢弃，直接通过残差连结到下一层，很多优化：用路由前全局 shuffle token，其实丢弃的token也比较少。

大规模稀疏模型训练极易不稳定，Switch Transformer 提出了一系列改进方案：

• Router z-loss： 为了提高训练稳定性，Switch Transformer 引入了 z-loss 来惩罚门控网络中过大的 logit 值。这有助于减少数值溢出问题，使训练过程更加平稳。

• 选择性精度（Selective Precision） ： 在混合精度训练（通常用 bfloat16）中，路由器的 Softmax 计算容易导致数值不稳定。Switch Transformer 创新地在局部路由计算部分使用 float32，而在其他部分保持 bfloat16。这既保证了稳定性，又没有增加昂贵的 float32 通信成本。

• 更小的初始化方差： 将权重初始化的高斯分布标准差缩减为原来的 1/101/10（例如 s=0.1s=0.1 而非 1.01.0），显著提升了训练初期的稳定性。

• 专家正则化（Expert Regularization） ： 在微调阶段，为了防止过拟合（特别是专家层参数量巨大），模型对专家层内部采用了更高的 Dropout 比率（如 0.4），而非专家层保持较低比率（如 0.1）。

GLaM

• 隔层稀疏：类似于 GShard，GLaM 采用隔层替换策略，将每隔一个 Transformer 层中的前馈网络（FFN）替换为 MoE 层。
• Top-2 路由 ：每个 MoE 层包含 64 个专家，对于每个输入 Token，门控网络会选择权重最高的 2 个专家进行处理。
• 活跃参数 ：尽管总参数量高达 1.2T，但对于每个 Token，仅激活 966 亿（96.6B） 参数（约占总量的 8%）。这意味着在推理时，GLaM 的计算量（FLOPs）仅为 GPT-3（175B 全激活）的约一半。

DeepSeekMoE

上下文学习与提示词技术

**零样本学习（Zero-Shot Learning）**是指在提示词中仅提供任务指令和待处理的输入内容，而不提供任何预期输出的参考示例。

如果我们需要模型严格遵循某种特定的输出格式，或者处理一些边界模糊、容易混淆的复杂情况，仅仅依靠指令说明可能无法得到稳定可靠的结果，模型很多时候还会像零样本例子中那样，输出多余的解释性自然语言或产生判断偏差。这就必须引入具备"观摩"作用的少样本学习（Few-Shot Learning），在上下文中补充几个规范的输入输出示例，向大模型清晰地演示预期效果。

上下文内在机制

为了解释为什么在不更新任何权重的情况下，大语言模型依然能在提示词上下文中"学会"新任务，学术界近年来提出了几种主流的假设与验证机制：

（1）感应头机制 2：这是从模型内部注意力层面的机械可解释性角度提出的。研究发现，Transformer 模型在预训练的特定阶段会激发出一种特殊的注意力头------"感应头"。它的核心行为模式是"匹配并复制"，当模型在上下文中发现当前输入 A 此前出现过，它会将注意力回溯到前一个 A，并倾向于直接预测其后跟随的 Token B 作为当前的输出。这种底层的复制机制是少样本学习中模式匹配的重要微观基础之一。

（2）隐式学习动力学（隐式权重更新） 3：有研究从"前向计算本身就是一种学习过程"的角度解释上下文学习。在 Transformer block 内，自注意力会把上下文中的示例信息写入激活，再与后续的 MLP 组合，可近似理解为产生一种对后续计算起作用的低秩"权重更新/适配"效应。这个过程无需显式反向传播或持久化的权重改动，但会以瞬时激活的形式影响后续 token 的 logits，从而动态改变输出分布。

（3）（近似）贝叶斯视角与其检验 45：有一类工作把上下文学习理解为"在上下文中对潜在任务/潜变量做（近似）推断"的过程，用先验---后验来解释零样本与少样本的差异；也有研究提出可操作的统计检验，并在其实验设置下观察到 LLM 的上下文学习会偏离严格的贝叶斯性质。总体而言，贝叶斯更像是一类有启发性的解释框架，而非已被普遍证明的严格等价。

这三大理论从微观执行（注意力单元复制）、中观适配（隐式权重更新）到宏观统计（概率推断），共同揭示了现代大语言模型强大泛化能力的底层本质。

思维链与思维树

思维链（Chain-of-Thought, CoT） 技术 6 鼓励模型在给出最终答案之前，先显式地输出中间的推理步骤。这种方法不仅客观上增加了生成过程的计算步数，让模型获得了更多的"思考时间"，而且将一个复杂的大问题拆解成了多个简单连贯的小逻辑节点。在应用层面上，这项技术经历了一条清晰且快速的演进路线。最初的研究提出了需要手动编写详尽推理示例的多样本思维链（Few-Shot CoT） 。随后，研究人员观察到在一些模型与任务上，在提示词末尾添加一句"Let's think step by step"也可能显著提升推理过程的显式展开，也就是常说的零样本思维链（Zero-Shot CoT） 。后续，学术界又提出了尝试让大模型自动化构建推理示例的自动思维链（Auto-CoT） 7。对于基础大模型，我们可以用一个非常直观的零样本测试样例去感受思维链被激活的过程。例如向模型抛出一道暗含生活常识与逻辑陷阱的问题：

我想洗车，洗车店离我家 50 米。我应该开车去还是走路去？Let's think step by step

不加Let's think step by step就会回答错误

在线性、单向的思维链基础上，研究人员进一步提出了思维树（Tree of Thoughts, ToT） 框架 8。对于更具挑战性、需要全局规划或容易陷入逻辑死胡同的任务（如算 24 点、填字游戏或规划调度），单向单线的 CoT 很可能有去无回------中间某一步哪怕只犯了微小的评估错误，就会使得推理陷入局部最优 ，导致最终结果全盘皆输。如图 6-42 所示，思维树将线性的推理链条扩展成了具有多个分支的树状拓扑结构。它的核心理念可以总结为"系统性多路径探索 + 智能评估 + 回溯机制"。它允许模型在推理过程中像树枝一样展开多个可能的探索分支，并交替执行以下环节：

• 生成候选：在当前步骤生成多种可能的下一步解法方向。
• 状态评估：由模型自身作为"智能裁判"，对各个候选路径的后续成功潜力进行打分排查。
• 选择与回溯 ：借助经典的搜索算法（如广度优先搜索 BFS 或深度优先搜索 DFS），根据评分选择最优路径继续深入。一旦发现某条路径走进了"死胡同"，立刻向后回溯到上一个安全节点并尝试其他选项。

专业推理模型

在基础预训练之后，模型不再单纯依赖传统的人类偏好对齐，还会通过大规模的强化学习提升逻辑推导能力。在数学求解或代码生成等具有明确结果反馈的场景中，系统基于规则的奖励机制引导模型自我探索，促使模型更愿意延长推导链条以换取更高奖励，形成更强的自我校验与纠错倾向。