Transformer 架构完全深度解析

核心论文：

Attention Is All You Need（Vaswani et al., 2017）arXiv:1706.03762

FlashAttention（Dao et al., 2022）arXiv:2205.14135

FlashAttention-2（Dao, 2023）arXiv:2307.08691

RoFormer: RoPE（Su et al., 2021）arXiv:2104.09864

GQA: Training Generalized Multi-Query Transformer Models（Ainslie et al., 2023）arXiv:2305.13245

LLaMA（Touvron et al., 2023）arXiv:2302.13971

一、为什么 Transformer 取代了 RNN
二、原始 Transformer 的整体结构
三、Token 嵌入（Token Embedding）
四、位置编码（Positional Encoding）
五、缩放点积注意力（Scaled Dot-Product Attention）
六、多头注意力（Multi-Head Attention，MHA）
七、RoPE：旋转位置编码（现代 LLM 标准方案）
八、GQA/MQA：注意力头分组优化
九、KV Cache：自回归解码的加速机制
十、前馈网络（FFN）及 SwiGLU 激活函数
十一、归一化层：RMSNorm 与 Pre-Norm
十二、残差连接（Residual Connection）
十三、FlashAttention：从数学等价的 IO 优化
十四、Transformer 层的完整数据流
十五、语言建模头（LM Head）与损失函数
十六、采样策略：从 Logits 到 Token
十七、现代 LLM 的典型超参数对比
十八、从 Encoder-Decoder 到 Decoder-only：为什么现代 LLM 都是 Decoder-only
十九、训练时的内存分析：激活值与梯度检查点
二十、总结：现代 Decoder-only LLM 的标准配置

一、为什么 Transformer 取代了 RNN

在 Transformer 出现之前（2017年之前），序列建模的主流方案是循环神经网络（RNN），具体形态包括 LSTM（Long Short-Term Memory）和 GRU（Gated Recurrent Unit）。

RNN 的工作原理是按时间步顺序处理序列，每一步的隐藏状态 $h_t$ 由前一步的隐藏状态 $h_{t-1}$ 和当前输入 $x_t$ 共同决定：

$$h_t = f(W_h h_{t-1} + W_x x_t + b)$$

这个设计有两个根本性的工程缺陷：

缺陷一：无法并行化。 时间步 $t$ 的计算必须等待时间步 $t-1$ 完成，因此一个长度为 $T$ 的序列必须串行完成 $T$ 步计算。在 GPU 大量并行计算能力面前，这是一个严重的浪费——GPU 上有数千个计算核心，但 RNN 在时间维度上只能利用其中一个。

缺陷二：梯度消失/爆炸问题。 信息从序列末尾传播回序列开头时，梯度需要经过 $T$ 次矩阵乘法。如果矩阵的最大奇异值大于 1，梯度会指数级增长（梯度爆炸）；如果小于 1，梯度会指数级衰减（梯度消失）。LSTM 通过门控机制部分缓解了这个问题，但没有根本解决。对于长度超过几百的序列，远距离依赖关系仍然难以学习。

Transformer 通过注意力机制（Attention Mechanism）彻底解决了这两个问题：

序列中任意两个位置的信息交互可以在一步内完成（不依赖路径长度），解决了梯度消失问题
所有位置的注意力计算可以并行进行，充分利用 GPU 的并行能力

二、原始 Transformer 的整体结构

原始 Transformer 论文（Vaswani et al., 2017）提出的是一个用于机器翻译的编码器-解码器（Encoder-Decoder）架构，如下：

输入序列（德语）        输出序列（英语）
    │                       │
  Embedding              Embedding + 右移一位
    │                       │
  位置编码               位置编码
    │                       │
┌───────────┐          ┌────────────┐
│  编码器    │          │   解码器    │
│ (N=6 层)  │          │  (N=6 层)  │
│           │          │            │
│ 自注意力  │──Cross───│ 掩码自注意力│
│    +      │  Attn    │     +      │
│  FFN     │          │ 交叉注意力  │
│           │          │     +      │
└───────────┘          │    FFN     │
                       └────────────┘
                              │
                         线性 + Softmax
                              │
                         预测下一个词

但现代 LLM（GPT 系列、Claude、LLaMA、DeepSeek）使用的都是仅解码器（Decoder-only）架构，原因是：对于语言建模（预测下一个 token），编码器是多余的——一个足够强大的解码器可以同时完成”理解”和”生成”两个功能。

以下我们重点解析现代 LLM 使用的 Decoder-only Transformer 的每个组件。

三、Token 嵌入（Token Embedding）

在进入 Transformer 之前，输入的文本必须首先被转化为向量。这个过程分两步：

步骤 1：分词（Tokenization）

现代 LLM 普遍使用 BPE（Byte Pair Encoding）算法将文本分割成 token。BPE 不是按字符也不是按单词分割，而是基于语料库中的频率统计，将高频字符组合合并为子词（subword）单元。

例如，”unbelievable” 可能被分割为 [“un”, “believ”, “able”]，而 “hello” 可能直接是一个 token。

GPT-2 使用 50,257 个 token 的词汇表；LLaMA 3 扩展到 128,256 个 token（大幅增加对多语言的覆盖）；DeepSeek-V3 使用 102,400 个 token。

步骤 2：嵌入查表（Embedding Lookup）

每个 token 对应词汇表中的一个整数 ID，通过嵌入矩阵 $E \in \mathbb{R}^{|V| \times d_{model}}$ 转化为 $d_{model}$ 维的稠密向量。

$$\text{token\_embed}(x_i) = E[x_i] \in \mathbb{R}^{d_{model}}$$

其中 $|V|$ 是词汇表大小，$d_{model}$ 是模型维度（GPT-3 是 12,288；LLaMA 3 70B 是 8,192）。

这个嵌入矩阵是可学习的参数，在预训练过程中与模型其他参数一起通过反向传播优化。注意：嵌入矩阵是模型参数中占比极大的一部分，对于词汇量 128K、维度 8192 的模型，仅嵌入矩阵就有约 10 亿参数。

四、位置编码（Positional Encoding）

Transformer 的注意力机制本身是排列不变的（permutation-invariant）——如果把输入序列的顺序打乱，注意力的计算结果（不考虑排列的影响）是一样的。但语言是有顺序的，”猫吃鱼”和”鱼吃猫”含义完全不同。

因此，必须将位置信息注入模型。原始 Transformer 使用正弦/余弦位置编码：

$$PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$
$$PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)$$

其中 $pos$ 是位置，$i$ 是维度索引，$d_{model}$ 是模型维度。

这种编码的位置向量被直接加到 token 嵌入上：

$$x_i = \text{token\_embed}(x_i) + PE_{(i)}$$

正弦位置编码的优点：

不需要学习额外参数
可以外推到训练时未见过的更长序列（因为是确定性函数）

正弦位置编码的缺点：

表示能力有限
在实践中发现，模型学到的相对位置关系不如专门设计的相对位置编码

这就引出了现代 LLM 普遍使用的 RoPE（Rotary Position Embedding），将在第七节详细介绍。

五、缩放点积注意力（Scaled Dot-Product Attention）

这是 Transformer 最核心的计算单元。给定输入矩阵 $X \in \mathbb{R}^{n \times d_{model}}$（$n$ 是序列长度），通过三个线性变换生成 Query、Key、Value 矩阵：

$$Q = X W^Q, \quad K = X W^K, \quad V = X W^V$$

其中 $W^Q, W^K \in \mathbb{R}^{d_{model} \times d_k}$，$W^V \in \mathbb{R}^{d_{model} \times d_v}$。

然后计算注意力输出：

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$$

逐步解析：

第一步：相似度计算 $QK^T$

$QK^T \in \mathbb{R}^{n \times n}$，其中元素 $(i, j)$ 表示位置 $i$ 的 Query 和位置 $j$ 的 Key 的点积，衡量位置 $i$ 对位置 $j$ 的”关注程度”。

第二步：缩放 $/ \sqrt{d_k}$

当 $d_k$ 较大时，点积的方差会随 $d_k$ 线性增长（因为每个维度都贡献了方差），导致 softmax 的输入进入极端区域，梯度极小。除以 $\sqrt{d_k}$ 将方差归一化到 1，使 softmax 数值更稳定。

直觉上：如果 $Q$ 和 $K$ 的每个维度都是均值 0、方差 1 的随机变量，那么它们点积的方差是 $d_k$，标准差是 $\sqrt{d_k}$，除以 $\sqrt{d_k}$ 后标准差变为 1。

第三步：因果掩码（Causal Mask，仅 Decoder 使用）

在语言模型的训练中，位置 $i$ 只能关注 $j \leq i$ 的位置（不能看到未来），这通过将 $j > i$ 的位置设置为 $-\infty$（softmax 后变为 0）来实现：

$$\text{score}_{ij} = \begin{cases} \frac{q_i \cdot k_j}{\sqrt{d_k}} & \text{if } j \leq i \\ -\infty & \text{if } j > i \end{cases}$$

第四步：Softmax 归一化

对每行（每个 Query 位置）独立做 softmax，得到注意力权重矩阵 $A \in \mathbb{R}^{n \times n}$，每行之和为 1。

第五步：加权求和 $A \cdot V$

将注意力权重作为系数，对 Value 矩阵的行（每个位置的 Value 向量）做加权求和，得到最终的注意力输出。

六、多头注意力（Multi-Head Attention，MHA）

单个注意力头只能关注一种类型的关系（例如语法关系）。多头注意力让模型同时从多个”角度”关注信息。

具体做法是将模型维度 $d_{model}$ 分成 $h$ 个”头”，每个头有独立的 $W_i^Q, W_i^K, W_i^V$ 参数（每个头的维度为 $d_k = d_v = d_{model}/h$）：

$$\text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V)$$

将所有头的输出拼接后，通过输出投影矩阵 $W^O$ 混合：

$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) W^O$$

参数量计算（以 GPT-3 为例）：

$d_{model} = 12288$，$h = 96$ 个头，每头维度 $d_k = 128$
$W^Q$：$12288 \times 12288$
$W^K$：$12288 \times 12288$
$W^V$：$12288 \times 12288$
$W^O$：$12288 \times 12288$
每层注意力总参数：$4 \times 12288^2 \approx 603$M 参数

可以看出注意力层是模型参数的主要贡献者之一。

七、RoPE：旋转位置编码（现代 LLM 标准方案）

论文： RoFormer: Enhanced Transformer with Rotary Position Embedding（Su et al., 2021）arXiv:2104.09864

核心问题： 原始的绝对位置编码（加到 embedding 上）有一个缺点：模型很难直接学到两个 token 之间的相对位置。而语言理解中，相对位置（”这两个词相距多远”）往往比绝对位置（”这个词在序列的第几位”）更重要。

RoPE 的核心思想：

不把位置信息加到 embedding 上，而是在注意力计算中，用旋转矩阵对 Query 和 Key 进行变换，使得位置 $m$ 处的 Query 和位置 $n$ 处的 Key 的点积自然包含它们的相对位置信息 $(m-n)$。

数学形式：

对于位置 $m$ 处的 Query 向量 $q_m$ 和位置 $n$ 处的 Key 向量 $k_n$，RoPE 定义：

$$f_q(x_m, m) = R_{\Theta,m}^d W_q x_m$$
$$f_k(x_n, n) = R_{\Theta,n}^d W_k x_n$$

其中 $R_{\Theta,m}^d$ 是一个 $d \times d$ 的旋转矩阵，在每对维度 $(2i, 2i+1)$ 上按角度 $m \theta_i$ 旋转：

$$R_{\Theta,m}^d = \begin{pmatrix} \cos m\theta_1 & -\sin m\theta_1 & & \\ \sin m\theta_1 & \cos m\theta_1 & & \\ & & \ddots & \\ & & & \cos m\theta_{d/2} & -\sin m\theta_{d/2} \\ & & & \sin m\theta_{d/2} & \cos m\theta_{d/2} \end{pmatrix}$$

旋转角度 $\theta_i = 10000^{-2i/d}$（与原始正弦编码的频率设计类似，保证不同维度对应不同频率）。

关键性质： $f_q(x_m, m)^T f_k(x_n, n) = x_m^T W_q^T R_{\Theta,m-n}^d W_k x_n$

这意味着注意力得分只依赖于相对位置 $(m-n)$，而不是绝对位置 $m$ 和 $n$！

RoPE 的实际优势：

相对位置感知，外推性更好
不需要在 embedding 层额外加位置编码
在 KV Cache 中，已计算的 Key 向量天然包含位置信息，可以直接复用
通过调整 $\theta$（称为 theta base），可以扩展上下文长度

LLaMA 原版使用 $\theta = 10000$，LLaMA 3 将其扩展到 $\theta = 500000$ 以支持 128K 上下文（详见第 07 篇长上下文文档）。

八、GQA/MQA：注意力头分组优化

论文： GQA: Training Generalized Multi-Query Transformer Models（Ainslie et al., 2023）arXiv:2305.13245

标准多头注意力（MHA）中，每个注意力头都有独立的 Query、Key、Value 投影。这带来一个推理时的工程问题：KV Cache 占用过大。

KV Cache 问题：

在自回归解码（生成文本）时，为了避免重复计算，每一步会将已计算的 Key 和 Value 向量缓存起来。对于一个序列长度为 $n$、$h$ 个注意力头的模型，每层的 KV Cache 大小为：

$$\text{KV Cache Size per layer} = 2 \times n \times h \times d_k \times \text{bytes per element}$$

对于 LLaMA 3 70B（$h=64$，$d_k=128$，序列长度 $n=128000$）：

每层 KV Cache：$2 \times 128000 \times 64 \times 128 \times 2\text{ bytes（fp16）} \approx 4\text{ GB}$
80 层全部：约 320 GB

这显然无法在单张 GPU 上容纳。

三种方案的比较：

MHA（Multi-Head Attention）： 每个头有独立的 Q、K、V

参数最多，KV Cache 最大，性能最好

MQA（Multi-Query Attention，Shazeer 2019）： 所有 Query 头共享同一个 K 和 V

KV Cache 减少为 MHA 的 $1/h$
代价是性能下降，某些任务有明显差距

GQA（Grouped Query Attention）： $g$ 个 Query 头共享一个 K 和 V（介于 MHA 和 MQA 之间）

KV Cache 减少为 MHA 的 $1/g$（例如 $g=8$ 时减少 8 倍）
性能接近 MHA，损失较小

GQA 的 KV 数量：

$$\text{KV heads} = \frac{h}{g} \quad (\text{例如 }h=64, g=8 \Rightarrow \text{KV heads}=8)$$

实际使用情况：

LLaMA 3 8B：$h_Q = 32$，$h_{KV} = 8$（GQA，$g=4$）
LLaMA 3 70B：$h_Q = 64$，$h_{KV} = 8$（GQA，$g=8$）
DeepSeek-V2/V3：使用更激进的 MLA（Multi-head Latent Attention），将 KV 压缩为低秩向量，详见第 06 篇

九、KV Cache：自回归解码的加速机制

自回归解码的冗余计算

在自回归生成时，第 $t$ 步和第 $t+1$ 步的输入中，前 $t$ 个 token 完全相同。如果没有 KV Cache，每一步都需为所有历史 token 重新计算 Key 和 Value 向量——$O(n^2)$ 的重复计算。

KV Cache 原理

将已计算的 Key 和 Value 存储在显存中，新步骤只计算新 token 的 K 和 V，然后与缓存拼接：

第 t+1 步（有 KV Cache）：
  K_new = X[t+1] · W^K        ← 只计算一个新 Key
  V_new = X[t+1] · W^V        ← 只计算一个新 Value
  K_cache = concat(K_cache, K_new)
  V_cache = concat(V_cache, V_new)
  Attention(Q[t+1], K_cache, V_cache) → 输出

KV Cache 的显存增长

每层 KV Cache：$\text{KV Cache}{\text{layer}} = 2 \times n{\text{seq}} \times h_{\text{KV}} \times d_k \times \text{bytes}$

以 LLaMA 3 70B 为例（$h_{\text{KV}}=8$，$d_k=128$，bf16=2 bytes）：

序列长度	每层 KV Cache	80 层总计
4K tokens	8.4 MB	~672 MB
32K tokens	67 MB	~5.4 GB
128K tokens	268 MB	~21.5 GB

这就是 GQA（减少 $h_{\text{KV}}$）的根本动机：KV Cache 总大小与 KV 头数成正比。

Prefill vs Decode

Prefill（预填充）： 一次性并行处理整个 prompt。计算密集型（compute-bound），FLOPs 利用率高。

Decode（逐个生成）： 每次只处理一个新 token。内存带宽密集型（memory-bound）——瓶颈不在计算，而在从 HBM 读取不断增长的 KV Cache。典型 MFU（Model FLOPs Utilization）低于 5%，GPU 大部分时间在等待数据传输。

FlashAttention（本节十三）和 PagedAttention（第 08 篇）都是为解决 Decode 阶段的内存带宽瓶颈而设计的。

十、前馈网络（FFN）及 SwiGLU 激活函数

每个 Transformer 层除了注意力子层，还有一个前馈网络（Feed-Forward Network）子层。它独立作用于序列的每个位置（position-wise），不同位置之间没有信息交互：

原始 FFN（Transformer 论文）：

$$\text{FFN}(x) = \max(0, xW_1 + b_1) W_2 + b_2$$

即两层线性变换，中间用 ReLU 激活，内层维度 $d_{ff} = 4 d_{model}$。

现代 LLM 使用的 SwiGLU：

LLaMA 系列将激活函数改为 SwiGLU（来自 Noam Shazeer 2020 年的博客文章，arXiv:2002.05202）：

$$\text{SwiGLU}(x, W, V, W_2) = (x W \odot \text{Swish}(x V)) W_2$$

其中 $\odot$ 是逐元素乘法，$\text{Swish}(x) = x \cdot \sigma(x)$（$\sigma$ 是 sigmoid 函数）。

实际上，SwiGLU 是一种”门控线性单元（GLU）”的变体：

$xW$ 产生一个”内容”路径
$\text{Swish}(xV)$ 产生一个”门控”路径
两者逐元素相乘后，门控信号决定哪些内容通过

由于 SwiGLU 有三个权重矩阵（$W$、$V$、$W_2$）而不是两个，为保持参数量大致相同，内层维度改为 $d_{ff} = \frac{8}{3} d_{model}$（而不是 $4 d_{model}$）。

为什么 SwiGLU 更好？

在 PaLM 技术报告、LLaMA 的消融实验中都发现，SwiGLU 在相同参数量下，perplexity（困惑度，衡量语言模型质量的指标，越低越好）优于 ReLU 和 GeLU 约 0.5-1 PPL。从第一性原理来看，可能的原因是：

门控机制让 FFN 能”选择性”传递信息，类似 LSTM 的遗忘门
Swish 激活函数（相比 ReLU）在 0 附近是平滑的，梯度更友好
GLU 结构从理论上等价于更深的网络

FFN 的参数量（以 LLaMA 3 70B 为例）：

$d_{model} = 8192$，$d_{ff} = 28672$（$\approx 3.5 \times d_{model}$）
每层 FFN：$3 \times 8192 \times 28672 \approx 705\text{M}$ 参数
全模型 80 层：FFN 参数约 $56\text{B}$，占总参数的主要部分

十一、归一化层：RMSNorm 与 Pre-Norm

原始 Transformer 使用的 LayerNorm（Post-Norm）：

$$\text{LayerNorm}(x) = \frac{x - \mu}{\sigma} \gamma + \beta$$

其中 $\mu$ 是均值，$\sigma$ 是标准差，$\gamma$ 和 $\beta$ 是可学习的缩放和偏移参数。LayerNorm 同时归一化均值和方差。

现代 LLM 的两个改变：

改变一：Pre-Norm（归一化位置前移）

原始 Transformer 在残差连接之后做归一化（Post-Norm）：$\text{output} = \text{LayerNorm}(x + \text{SubLayer}(x))$

现代 LLM 在子层之前做归一化（Pre-Norm）：$\text{output} = x + \text{SubLayer}(\text{Norm}(x))$

Pre-Norm 使训练更稳定，即使不用学习率预热也能正常训练。

改变二：RMSNorm（去掉均值归一化）

论文：Root Mean Square Layer Normalization（Zhang & Sennrich, 2019）arXiv:1910.07467

$$\text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \gamma, \quad \text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2}$$

相比 LayerNorm，RMSNorm：

去掉了均值归一化（$\mu$ 项），只做方差归一化
去掉了偏移参数 $\beta$
计算量减少约 40%（省去了计算均值的操作）
实践中发现效果与 LayerNorm 相当

LLaMA 系列全部使用 RMSNorm。

十二、残差连接（Residual Connection）

每个子层（注意力和 FFN）都有残差连接：

$$\text{output} = x + \text{SubLayer}(\text{Norm}(x))$$

残差连接来自 ResNet（He et al., 2016），核心作用是让梯度可以直接从输出传播到输入，绕过子层。这解决了深层网络的梯度消失问题，使得训练几十、上百层的 Transformer 成为可能。

直觉理解：如果一个子层学到的变换是恒等映射（什么都不做），那么残差连接保证输出等于输入，网络不会因为加了这层而变差。子层只需要学习”残差”（对输入的增量改变），而不需要从零学习整个变换。

十三、FlashAttention：从数学等价的 IO 优化

论文： FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness（Dao et al., 2022）arXiv:2205.14135

FlashAttention 不改变注意力的数学结果，只优化其在 GPU 上的实现。理解它需要先了解 GPU 的内存层次结构。

GPU 内存层次结构（以 A100 为例）：

SRAM（片上内存，共享内存）：192 KB per SM，带宽 19 TB/s，延迟极低
  ↕（速度快）
HBM（高带宽内存，显存）：40-80 GB，带宽 2 TB/s
  ↕（速度慢，约10倍差距）
DRAM（CPU 内存）：几百 GB，带宽 ~50 GB/s

标准注意力的 IO 瓶颈：

标准注意力计算的步骤（在 HBM 和 SRAM 之间的数据移动）：

1. 从 HBM 读取 Q, K → SRAM，计算 S = QK^T，写回 HBM      [读 Q+K，写 S]
2. 从 HBM 读取 S → SRAM，计算 softmax(S) = P，写回 HBM    [读 S，写 P]
3. 从 HBM 读取 P, V → SRAM，计算 O = PV，写回 HBM          [读 P+V，写 O]

对于序列长度 $n$，中间结果 $S \in \mathbb{R}^{n \times n}$，HBM 读写量是 $O(n^2)$。

FlashAttention 的关键洞察： 注意力矩阵 $n \times n$ 太大，无法整体放入 SRAM，但可以分块（tiling）计算。

FlashAttention 的核心技术：

将 Q、K、V 分成若干块，在 SRAM 内逐块计算注意力：

将 Q 分成 $T_r = \lceil n/B_r \rceil$ 块，K 和 V 分成 $T_c = \lceil n/B_c \rceil$ 块
对每对块 $(Q_i, K_j, V_j)$，在 SRAM 内完成计算，并增量更新输出 $O_i$
关键技巧：使用在线 Softmax（online softmax），允许分块计算并维护正确的归一化

在线 Softmax 原理：

$\text{softmax}(x_1, x_2) = \frac{e^{x_i}}{e^{x_1} + e^{x_2}}$ 通常需要两遍扫描（先找最大值，再计算指数和）。

在线 Softmax 通过维护当前见到的最大值 $m$ 和归一化分母 $\ell$，可以在一遍扫描中增量更新：

$$m^{new} = \max(m^{old}, x_i), \quad \ell^{new} = e^{m^{old} - m^{new}} \ell^{old} + e^{x_i - m^{new}}$$

这样每个块的计算结果可以正确合并，而无需存储完整的 $n \times n$ 矩阵。

FlashAttention 的效果：

内存：从 $O(n^2)$ 降低到 $O(n)$（不需要存储完整的注意力矩阵）
速度：在 A100 上，序列长度 2K 时比标准注意力快 6.7 倍
数学结果：完全等价（exact attention，不是近似）

FlashAttention-2（Dao, 2023）arXiv:2307.08691 进一步优化：

减少非矩阵乘法（non-matmul）FLOPs 的数量（这些操作在 GPU 上效率更低）
改善线程并行策略（在序列长度维度上并行而不是只在批次和头上并行）
在 A100 上达到约 72% 的理论最大吞吐量（相比 FlashAttention-1 的约 35%）

FlashAttention-2 是目前几乎所有顶级模型训练框架的标准配置。

十四、Transformer 层的完整数据流

以一个现代 Decoder-only Transformer 层为例，数据流如下：

输入 x ∈ ℝ^{n × d_model}
  │
  ├─── 残差路径（x，直通）
  │
  ▼
RMSNorm(x) → x_norm
  │
  ▼
Self-Attention（因果掩码）
  ├─ Q = x_norm W^Q ∈ ℝ^{n × d_model}
  ├─ K = x_norm W^K ∈ ℝ^{n × d_k_heads}  （GQA：少于 Q 的头数）
  ├─ V = x_norm W^V ∈ ℝ^{n × d_k_heads}
  ├─ 对 Q, K 应用 RoPE
  ├─ score = QK^T / sqrt(d_k) + mask  → softmax → A
  ├─ attn_out = A V W^O
  └─ attn_out ∈ ℝ^{n × d_model}
  │
  ▼
x = x + attn_out   （残差连接）
  │
  ├─── 残差路径（x，直通）
  │
  ▼
RMSNorm(x) → x_norm2
  │
  ▼
SwiGLU FFN
  ├─ gate = Swish(x_norm2 W_gate)
  ├─ content = x_norm2 W_up
  ├─ ffn_out = (gate ⊙ content) W_down
  └─ ffn_out ∈ ℝ^{n × d_model}
  │
  ▼
x = x + ffn_out    （残差连接）
  │
输出 x ∈ ℝ^{n × d_model}

十五、语言建模头（LM Head）与损失函数

最后一个 Transformer 层的输出经过最终的 RMSNorm 后，通过语言建模头（LM Head）投影到词汇表维度：

$$\text{logits} = \text{LayerNorm}(h_L) W_{vocab}$$

其中 $W_{vocab} \in \mathbb{R}^{d_{model} \times |V|}$ 通常与词嵌入矩阵 $E$ 共享权重（weight tying），这样可以减少参数量并提高效果。

然后对 logits 做 softmax 得到下一个 token 的概率分布，训练时最小化交叉熵损失：

$$\mathcal{L} = -\frac{1}{T} \sum_{t=1}^T \log P(x_t | x_{

这就是下一个 token 预测（Next Token Prediction）的损失，也称为自回归语言建模损失。

十六、采样策略：从 Logits 到 Token

语言建模头输出 logits 后，采样策略决定选择哪个 token——深刻影响生成文本的流畅性、多样性和事实性。

贪心解码（Greedy Decoding）

$$y_t = \arg\max_i P(y_t = i | y_{

确定性最强，但缺乏多样性，易陷入重复循环（”I am a a a a…”）。

温度采样（Temperature Sampling）

对 logits 除以温度 $T$ 后再 softmax：

$$P(y_t = i) = \frac{\exp(z_i / T)}{\sum_j \exp(z_j / T)}$$

$T \to 0$：趋向贪心（保守、确定性）
$T = 1$：标准 softmax（训练分布）
$T > 1$：分布更均匀（更多样、更”有创造力”，但可能胡言乱语）

实践参考值： 代码生成 $T=0.7$，创意写作 $T=0.9-1.0$，需要事实准确性的任务 $T=0.1-0.3$。

Top-K 采样（Fan et al., 2018）

仅从概率最高的 $k$ 个 token 中采样。缺点： 固定 $k$ 不适应不同上下文——模型很确定时 top-k 可能遗漏好候选，不确定时 top-k 可能包含大量无意义 token。

Top-P / Nucleus 采样（Holtzman et al., 2020）arXiv:1904.09751

累加概率直到达到阈值 $p$，仅从该”核心集”中采样：

$$S_p = \arg\min_{S} |S| \quad \text{s.t.} \quad \sum_{i \in S} P(y_t = i) \geq p$$

优点： 动态适应分布形状——模型确定时核心集小，不确定时自然扩大。Claude 和 GPT-4 均使用类似机制，典型值 $p=0.9$ 或 $0.95$。

现代实践：温度 + Top-P 组合

当前顶级 LLM 的采样管线通常为：

logits → logits / T → softmax → top-p 过滤 → 多项式采样 → 下一个 token

Claude API 中 temperature 和 top_p 是用户可调参数。一般建议：调其中一个，不同时调二者。

停止条件

采样循环在以下任一条件满足时终止：

生成特殊 EOS（End of Sequence）token
达到最大生成长度（max_tokens）
触发自定义停止序列（如 \n\nHuman: 或 }）

十七、现代 LLM 的典型超参数对比

模型	$d_{model}$	层数	$h_Q$	$h_{KV}$	$d_{ff}$	参数量	词汇表
LLaMA 3 8B	4096	32	32	8	14336	8B	128K
LLaMA 3 70B	8192	80	64	8	28672	70B	128K
LLaMA 3 405B	16384	126	128	8	53248	405B	128K
DeepSeek-V3	7168	61	128	128(MLA)	MoE	671B	102.4K

（数据来源：各模型官方技术报告）

十八、从 Encoder-Decoder 到 Decoder-only：为什么现代 LLM 都是 Decoder-only

原始 Transformer 的 Encoder-Decoder 设计适合”输入-输出对”明确的任务（如机器翻译：德语→英语）。

但 GPT（Generative Pre-trained Transformer）系列的成功证明：

统一的语言建模目标更强大：Decoder-only 模型用单一目标（预测下一个 token）就能学到丰富的语言知识，而不需要专门为编码和解码设计不同的子任务。
In-context Learning 从 Decoder-only 天然涌现：模型在输入序列中同时看到”示例”和”问题”，用同一个因果注意力处理，这种统一性使得 few-shot 学习更自然。
RLHF/指令微调更简单：对于生成式任务，只需要 Decoder 即可。编码器的存在增加了架构复杂度但不必然带来收益。
规模定律在 Decoder-only 上验证最充分：OpenAI/DeepMind/Meta 的大量实验都是基于 Decoder-only 架构，工程优化（FlashAttention、GQA 等）也主要针对此架构设计。

十九、训练时的内存分析：激活值与梯度检查点

理解 Transformer 的训练内存需求，对于理解模型规模的经济可行性至关重要。

训练内存的三个组成部分

$$\text{Total Memory} = M_{\text{params}} + M_{\text{optimizer}} + M_{\text{activations}}$$

1. 模型参数 $M_{\text{params}}$： LLaMA 3 70B（bf16）= $70 \times 10^9 \times 2$ bytes ≈ 140 GB

2. 优化器状态 $M_{\text{optimizer}}$： AdamW 为每个参数维护一阶矩 $m$ 和二阶矩 $v$（fp32），加 fp32 主权重副本：

$$M_{\text{optimizer}} = M_{\text{params}} \times (2 \times \frac{4}{2} + \frac{4}{2}) = 70 \times 10^9 \times 6 \text{ bytes} \approx 420 \text{ GB}$$

仅参数 + 优化器就约 560 GB，远超单张 H100（80GB）。这是模型并行（张量并行 + 流水线并行 + 数据并行）存在的根本原因（详见第 02 篇）。

3. 激活值 $M_{\text{activations}}$： 前向传播的中间结果——注意力矩阵、Q/K/V、FFN 中间值（SwiGLU gate/content 乘法）、RMSNorm 统计量。

对于 batch $b$、序列长度 $n$、hidden 维度 $d$、层数 $L$：

$$M_{\text{activations}} \approx b \times n \times d \times L \times c$$

其中 $c$ 约 15-20（取决于注意力头数和 FFN 内层维度）。以 LLaMA 3 70B（$n=4096$，$d=8192$，$b=4$，$L=80$）：

$$M_{\text{activations}} \approx 4 \times 4096 \times 8192 \times 80 \times 15 \text{ bytes} \approx 160 \text{ GB}$$

梯度检查点（Gradient Checkpointing）

论文：Training Deep Nets with Sublinear Memory Cost（Chen et al., 2016）arXiv:1604.06174

核心思想：用时间换空间。 不在前向传播中保存所有中间激活值，在反向传播需要时从最近的”检查点”重新计算子图来恢复激活值。

标准做法：
  前向：计算并保存所有激活值
  反向：直接使用已保存的激活值求梯度

梯度检查点：
  前向：仅保存"检查点"层的激活值
  反向：从检查点重新运行子图，恢复中间激活值

典型策略：每个 Transformer 层边界设检查点。激活值内存从 $O(L)$ 降到约 $O(\sqrt{L})$。

代价： ~33% 额外前向计算。LLaMA 3 405B 预训练大量依赖此技术使 16K 序列训练可行。

与 FlashAttention 的协同： FlashAttention 不存储完整 $n \times n$ 注意力矩阵（分块 + 在线 softmax），进一步降低每个检查点需保存的信息量。两者叠加使长序列训练真正可行。

二十、总结：现代 Decoder-only LLM 的标准配置

综合来看，以 LLaMA 3 为代表的现代 LLM 相较于 2017 年的原始 Transformer，主要做了以下改进：

组件	原始 Transformer	现代 LLM 标准配置	改进原因
架构	Encoder-Decoder	Decoder-only	语言建模更统一
位置编码	正弦绝对编码（加法）	RoPE（乘法，作用于 Q/K）	相对位置感知，更好外推
归一化位置	Post-Norm	Pre-Norm	训练更稳定
归一化方法	LayerNorm（均值+方差）	RMSNorm（仅方差）	计算更高效，效果相当
激活函数	ReLU	SwiGLU	更低 perplexity
注意力头	MHA（全独立 KV）	GQA（共享 KV）	减少 KV Cache
注意力实现	朴素实现	FlashAttention	GPU IO 优化，显存降 $O(n^2)→O(n)$
偏置项	有 bias	无 bias	参数效率和稳定性
权重绑定	否	词嵌入与 LM Head 共享	减少参数量

这套组合（Pre-Norm + RMSNorm + SwiGLU + RoPE + GQA + FlashAttention）是目前开源 LLM 的”黄金标准配置”，理解这些组件的每个设计决策，就理解了现代 LLM 架构的全貌。