chapter03.md — 架构细节：LLaMA 风格与 8k 扩展

开篇段落

本章将深入剖析现代大语言模型（LLM）的架构蓝图，以 LLaMA 及其后续演进为核心参照。我们将从“第一性原理”出发，探讨为何预归一化（Pre-normalization）、RMSNorm、SwiGLU 激活函数和旋转位置编码（RoPE）这些组件能够组合成一个如此稳定且高效的训练范式。学习本章后，您不仅能理解 3B/7B/13B 等不同模型规模下的超参数设计哲学与权衡，更将精通将模型上下文长度从标准 4k 扩展至 8k 乃至更长的核心技术——RoPE scaling，并能辨析 PI、NTK-aware 和 YaRN 各自的适用场景与内在机理。最后，我们将从算法与工程的交汇点，揭示 FlashAttention v2 和各类融合核函数（fused kernels）如何成为在 64x H100 这种规模下实现高吞吐、控制成本的关键，确保理论上的模型设计能够转化为经济上可行的训练项目。

文字论述

2.1 LLaMA 架构：简洁、稳定与高效的哲学

LLaMA 架构的成功并非源于革命性的全新模块，而是对现有成熟技术的精妙组合与优化，其设计哲学可以概括为：在保证模型表达能力的前提下，最大化训练的稳定性和计算效率。相较于早期如 T5（Encoder-Decoder）或 GPT-3（Post-normalization），LLaMA 采用的 Decoder-only 架构搭配一系列关键改进，已成为后续开源模型的事实标准。

  1. 预归一化与 RMSNorm (Pre-normalization with RMSNorm)

    • 背景：Pre-LN vs. Post-LN

      • 原始 Transformer 采用 Post-LN，即在残差连接之后进行层归一化。这种结构在网络较浅时表现良好，但随着层数加深，梯度范数在反向传播中容易逐层累积，导致训练初期不稳定，需要精细的 LR warmup 策略来“驾驭”。
      • Pre-LN 将归一化层置于子模块（多头注意力或 FFN）的输入端，并将残差连接“跨过”整个子模块。这相当于在每个模块的输入处都进行了一次“信号重置”，有效地控制了梯度范数，使得梯度在深层网络中传播更加平滑，极大地增强了训练稳定性，允许使用更大的学习率和更简单的 warmup 策略。

    • RMSNorm：更轻量的 LayerNorm

      • 标准的 LayerNorm 包含两个步骤：中心化（减去均值）和归一化（除以标准差），并辅以可学习的缩放（gain）和偏置（bias）参数。 $$ \text{LayerNorm}(x) = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} \cdot g + b $$ 其中 $\mu$ 和 $\sigma^2$ 分别是输入 $x$ 的均值和方差。

      • RMSNorm 的核心洞察是：在 LLM 这种大规模神经网络中，激活值的分布通常比较稳定，强制的中化（减去均值）可能是一个不必要的约束。RMSNorm 移除了均值中心化，仅通过输入的均方根进行缩放。 $$ \text{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{n}\sum_{i=1}^{n}x_i^2 + \epsilon}} \cdot g $$ 其中 $g$ 是可学习的增益参数，$\epsilon$ 是一个极小的常数（如 1e-6 或 1e-5）以防分母为零。

    • Rule-of-thumb：RMSNorm 在计算上比 LayerNorm 节省约 30-40% 的时间，因为它减少了一次遍历数据的操作（计算均值）。在实践中，其性能与 LayerNorm 相当甚至略优。对于 bf16 训练，将 epsilon 设置为 1e-5 可能比 1e-6 更能防止因分母过小导致的数值问题。

  2. SwiGLU 激活函数：门控增强表达力

    • 背景：从 ReLU 到 Gated Linear Unit (GLU)

      • FFN（前馈网络）是 Transformer 的核心“计算”单元。传统的 FFN 使用 ReLU 激活函数：$\text{FFN}(x) = \text{ReLU}(xW_1)W_2$。
      • GLU 引入了门控机制，其变体 SwiGLU 被证明在 LLM 中效果尤为出色。其 FFN 结构变为： $$ \text{FFN}_{\text{SwiGLU}}(x) = (\text{Swish}(xW_{\text{up}}) \otimes (xW_{\text{gate}})) W_{\text{down}} $$ 其中 $\text{Swish}(z) = z \cdot \text{sigmoid}(z)$，$\otimes$ 表示逐元素相乘。$W_{\text{up}}$ 和 $W_{\text{gate}}$ 将输入 $x$ 投影到中间维度，而 $W_{\text{down}}$ 将其投影回 d_model。

    • 门控机制的价值：$xW_{\text{gate}}$ 部分经过 Sigmoid 函数后，其输出值在 (0, 1) 之间，可以看作一个动态的“门”。这个门决定了 $\text{Swish}(xW_{\text{up}})$ 中每个元素有多少信息能够通过。这种数据依赖的过滤机制，使得 FFN 能够更灵活地路由和处理信息，从而增强了模型的表达能力。

    • 参数与计算的权衡：为了使 SwiGLU FFN 的总参数量与标准 ReLU FFN（中间层维度为 4*d_model）大致相当，LLaMA 将其 FFN 中间层维度 d_ffn（intermediate_size）设置为约 (2/3) * 4 * d_model。这是一种经过精心设计的权衡，用略微增加的计算换取了显著的模型性能提升。

  3. 旋转位置编码 (Rotary Positional Embedding, RoPE)

    • 核心思想：RoPE 是一种将相对位置信息注入自注意力机制的精妙方法。它没有采用可学习的位置嵌入或固定的正弦位置嵌入，而是在计算 Query 和 Key 之前，对它们的向量进行“旋转”。

    • 数学直觉：将 head_dim 维的向量 $v$ 看作 $\frac{\text{head_dim}}{2}$ 个复数（或二维向量）。对于位置为 $m$ 的 token，其 Query 向量 $q_m$ 中的每一对特征 $(q_m^{(2i)}, q_m^{(2i+1)})$ 会被旋转一个角度 $m\theta_i$。同样，位置为 $n$ 的 token 的 Key 向量 $k_n$ 也会被旋转角度 $n\theta_i$。 $$ f(v, m)_i = v_i \cdot e^{im\theta_i} \quad (\text{复数形式})$$ 当计算 $q_m$ 和 $k_n$ 的点积时，由于复共轭的性质，其结果只与相对位置 $(m-n)$ 有关： $$ \langle R(q, m), R(k, n) \rangle = \text{Re} \left( \sum_{i} (q_i e^{im\theta_i}) (k_i e^{-in\theta_i}) \right) = \sum_{i} \text{Re}(q_i \bar{k}_i e^{i(m-n)\theta_i}) $$ 这使得注意力分数天然地具备了对相对位置的敏感性，且随着距离增大，相关性自然衰减。

    • 优势：

      1. 无参数：不增加任何可学习的参数。
      2. 相对编码：直接编码相对位置信息。
      3. 外推性：理论上具有一定的长度外推能力，为 RoPE scaling 提供了基础。

2.2 模型规模与超参设计空间

为 3B/7B/13B 等规模选择超参数是在有限的参数预算内，对模型宽度、深度和注意力复杂度的综合权衡。下表展示了 LLaMA 系列的典型配置，代表了该领域一条经过验证的有效路径。

| 模型规模 | 参数量 (N_params) | 隐藏层维度 (d_model) | 层数 (n_layers) | 注意力头数 (n_heads) | 注意力头维度 (head_dim) | FFN 中间层维度 (d_ffn) | 词表大小 (vocab_size) |

模型规模 参数量 (N_params) 隐藏层维度 (d_model) 层数 (n_layers) 注意力头数 (n_heads) 注意力头维度 (head_dim) FFN 中间层维度 (d_ffn) 词表大小 (vocab_size)
~3B 2.7B - 3.5B 3200 32 32 100 8640 32000 (padded)
~7B 6.7B - 7.5B 4096 32 32 128 11008 32000 (padded)
~13B 12.5B - 13.5B 5120 40 40 128 13824 32000 (padded)

Rule-of-thumb：

  • 深宽比：随着模型规模增大，优先增加深度（n_layers）而非无限增加宽度（d_model）。更深的模型能够学习到更层次化的特征抽象。
  • 头维度 (head_dim)：128 是一个常见的选择，它在性能和硬件利用率（尤其是在 A100/H100 上）之间取得了很好的平衡。
  • FFN 维度 (d_ffn)：LLaMA 使用 int((2/3 * 4 * d_model) / 256) * 256 的方式计算，这里的 256 是为了硬件对齐（Tensor Core），确保计算效率。
  • 词表大小 (vocab_size)：通常会向上填充到 64 或 128 的倍数，以优化 embedding 查找和 softmax 计算的硬件效率。

2.3 上下文扩展：RoPE Scaling 深入解析

预训练好的 RoPE 模型直接处理超出训练长度的序列时性能会崩溃，因为位置编码的“频率”超出了其见过的范围。RoPE scaling 技术通过修改位置编码的计算方式，优雅地解决了这个问题。

  1. 位置插值 (Position Interpolation, PI)

    • 核心思想：将新的、更长的位置索引 "线性压缩" 到原始训练长度的范围内。假设原始训练长度为 $L_{\text{orig}}$，目标长度为 $L_{\text{target}}$，则新的位置索引 $m'$ 为： $$ m' = m \cdot \frac{L_{\text{orig}}}{L_{\text{target}}} $$ 然后用 $m'$ 计算 RoPE。这就像把一把 8000 厘米的尺子上的刻度，重新标记到一把 4000 厘米的尺子上。

    • 优点：实现极其简单，在 CPT 或微调阶段应用时非常稳定，是扩展上下文的可靠基线。

    • 缺点：压缩操作使得相邻 token 之间的相对位置差异变小，这会损失模型对高频（精细）位置信息的感知能力。对于需要精确位置信息的任务，性能可能会下降。

  2. NTK-aware Scaling

    • 核心思想：PI 改变位置 $m$，而 NTK-aware scaling 改变 RoPE 的旋转基底（base $\theta$）。它受到神经正切核（NTK）理论的启发，认为高频信息对于长距离建模不那么重要。通过修改旋转频率，使得长距离的旋转“更慢”，从而为长上下文保留更多的分辨能力。 $$ \theta'_i = \theta_i \cdot \alpha^{(d/d-2)} \quad \text{where} \quad \alpha = L_{\text{target}}/L_{\text{orig}} $$ 其中 $d$ 是 head_dim。这有效地降低了所有位置的旋转频率。

    • 优点：相比 PI，能更好地保留长距离依赖和局部高频信息，理论外推性能更强。

    • 缺点：单纯修改基底有时会在微调中引入不稳定性，或产生一些非自然的注意力模式。

  3. YaRN (Yet another RoPE extensioN)

    • 核心思想：YaRN 可以看作是 PI 和 NTK-aware 的“集大成者”，并修正了它们各自的缺陷。它包含三个关键部分：

      1. 融合 PI 与 NTK：同时对位置索引进行插值和对旋转基底进行修改，取两者之长。

      2. 注意力温度缩放：YaRN 的作者发现，插值会使相对距离变小，导致 softmax 函数的输入方差减小，输出分布变得更“尖锐”（低熵）。这会使模型过于自信。YaRN 引入一个温度参数 $t$ 来缩放 QK 点积，将其拉回到原始分布： $$ \text{softmax}\left(\frac{QK^T / \sqrt{d}}{t}\right), \quad t = 0.1 \cdot \ln(\alpha) + 1 $$

      3. 非均匀插值：对不同频率的维度采用不同的插值策略，进一步保留高频信息。

        • 优点：在长上下文任务上取得了 SOTA 级别的性能，是目前扩展上下文的首选方案，因为它系统性地解决了插值带来的副作用。
        • Rule-of-thumb：对于从 4k 扩展到 8k 的 CPT 任务，YaRN 是首选。如果需要一个快速、简单的实现，PI 也是一个非常稳健的选项。NTK-aware 单独使用的情况较少，但其思想是 YaRN 的重要组成部分。

2.4 训练吞吐与数值稳定性优化

在 64x H100 集群上，每一秒的训练时间都成本高昂。最大化 tokens/s（每秒处理的 token 数）是项目成功的关键。

  1. FlashAttention v2

    • 核心问题：标准自注意力的计算瓶颈在于 GPU 的内存带宽（HBM I/O），而非计算能力（FLOPs）。$O(L_{\text{ctx}}^2)$ 的注意力矩阵需要反复从 HBM 读写，这极大地拖慢了速度。
    • 解决方案：FlashAttention v2 是一种 I/O 感知的注意力算法。它将输入 Q, K, V 分块（tiling），在 GPU 高速但容量小的 SRAM 中完成一小块注意力矩阵的计算、softmax 和与 V 的乘积，然后才将最终结果写回 HBM。这个过程中，它通过在线重计算的方式避免了存储巨大的中间注意力矩阵。
    • 影响：
      • 吞吐：对训练速度有 2-4 倍的提升，上下文越长，效果越明显。
      • 内存：将注意力的显存占用从 $O(L_{\text{ctx}}^2)$ 优化到 $O(L_{\text{ctx}})$，使得单卡训练 8k 甚至更长上下文成为现实。
      • Rule-of-thumb：在任何涉及 Transformer 的大规模训练中，FlashAttention v2 是非选项，而是必选项。确保你的环境（PyTorch, CUDA版本）正确支持它。

  2. 融合核函数 (Fused Kernels)

    • 是什么：将多个连续的、元素级的 GPU 操作（如加法、乘法、激活函数）合并成一个单一的 CUDA kernel。例如，一个 fused_add_rmsnorm kernel 可以一次性完成残差连接的加法和后续的 RMSNorm。
    • 影响：
      • 减少 Kernel Launch 开销：每次调用 CUDA kernel 都有微秒级的 CPU-GPU 通信开销。将 3 个操作融合成 1 个，就减少了 2/3 的开销。
      • 减少 HBM 读写：数据可以在 GPU 寄存器或 SRAM 中停留更久，无需在每个操作后都写回 HBM 再读出。
      • 典型融合点：RMSNorm、SwiGLU 激活部分、RoPE 的应用、优化器更新步骤（Fused AdamW）。
    • Rule-of-thumb：优先使用 PyTorch 2.x 的 torch.compile，它可以自动地进行大量的算子融合。对于无法自动融合的关键路径，可以考虑使用如 apex 或 xformers 提供的预编译融合核函数。

  3. 其他稳定性考量

    • Dropout：在从零预训练（1T tokens 级别）中，dropout 通常被设置为 0。海量、多样化的数据本身就是最强的正则化器。关闭 dropout 不仅能略微提升吞吐，还能消除一个随机性来源，使得实验更具确定性。
    • 初始化 (Initialization)：权重初始化对训练初期的稳定性至关重要。LLaMA 遵循 GPT-2 的方案，采用均值为 0，标准差为 0.02 的正态分布初始化大部分权重，而对残差连接路径上的层（如 W_down 和 embedding）则采用更小的标准差（例如 $0.02 / \sqrt{2 \cdot N_{\text{layers}}}$），以防止残差累积过快。

本章小结

  • LLaMA 架构精髓：以 Pre-LN + RMSNorm 保证训练稳定性，SwiGLU 提升模型表达力，RoPE 优雅地编码相对位置，共同构成了一个强大而高效的基座。
  • 模型超参设计：模型缩放（scaling）是一门艺术，通常倾向于优先加深网络，并选择硬件友好的度（如 head_dim=128，维度对齐 256）。
  • 上下文扩展：RoPE scaling 是突破预训练长度限制的关键。YaRN 通过融合 PI、NTK 和温度缩放，提供了当前最先进的性能。PI 则是简单、可靠的备选方案。
  • 性能优化铁律：FlashAttention v2 是长上下文训练的非 negotiable 项，它能同时解决速度和显存两大瓶颈。尽可能利用 Fused Kernels（通过 torch.compile 或专用库）来压榨硬件的每一分性能。
  • 训练稳定性细节：大规模预训练中关闭 Dropout，并采用恰当的权重初始化策略，是确保训练过程平稳、可复现的重要保障。

常见陷阱与错误 (Gotchas)

  1. RoPE Scaling 的训练-推理不一致：最常见的错误之一是在微调或 CPT 时使用了某种 RoPE scaling（如 YaRN），但在部署推理服务时忘记应用或使用了错误的参数。这会导致模型在处理长文本时输出完全无意义的内容。必须确保推理时的位置编码逻辑与训练时完全一致。
  2. 数值溢出与 Fused Kernels：某些版本的融合核函数可能在 bf16 精度下，其内部计算顺序与标准 PyTorch 实现略有不同，可能导致在极端值下出现数值问题（NaN）。如果怀疑是 Fused Kernels 的问题，可以尝试临时禁用它们进行调试，或者调整 RMSNorm 的 epsilon 值。
  3. 上下文扩展后的“中间遗忘” (Lost in the Middle)：成功将上下文扩展到 8k 并不意味着模型能同等关注上下文中的所有信息。许多研究发现，模型对开头和结尾的 token 最敏感，而对中间部分的信息容易“遗忘”。这并非架构缺陷，而是长上下文模型的普遍挑战。需要通过特定的评估基准（如 Needle-in-a-Haystack 测试）来诊断，并通过调整数据和训练策略来缓解。
  4. FlashAttention 安装与环境兼容性问题：FlashAttention 对 CUDA 版本、PyTorch 版本和 GPU 架构（如 Hopper、Ampere）有严格要求。环境配置不当会导致编译失败或运行时错误。在开始训练前，务必在一个最小化的脚本中验证 FlashAttention 是否能被正确调用和运行。
  5. 不经微调直接外推的诱惑：虽然 RoPE scaling 技术提供了外推的能力，但未经任何长文本微调就直接将 4k 模型用于 8k 推理（zero-shot extrapolation）的效果通常不理想，困惑度会显著上升。可靠的做法是通过 CPT 或微调，让模型在目标长度的数据上“适应”新的位置编码机制。