从零实现 Mini-GPT，深入理解 Attention 数学本质、RoPE、KV-Cache、Flash Attention，完整预训练/SFT/DPO 流程。
本章面向有基础深度学习经验的工程师，特别是已经在使用 Transformer 做感知任务的 CNN 自动驾驶算法工程师。我们不仅讲解"怎么做"，更重要的是解释"为什么这样做"。
学完本章你将获得：
- 对 Transformer/LLM 每个组件的深度理解（数学原理 + 工程实践）
- 一套完整可运行的 Mini-GPT 代码（~1500 行，包含详细注释）
- 预训练 / SFT / DPO 三阶段训练的完整经验
- LoRA 参数高效微调的实践能力
在 Transformer 之前，处理序列数据的主流方案是 RNN/LSTM。它们有两个根本问题：
1. 顺序计算瓶颈：必须按时间步逐个处理，无法并行
2. 长程依赖问题：经过多层传播后，远距离 token 之间的信号变得极弱
Self-Attention 的解决方案是：让每个 token 直接与所有 token 交互。
注意力机制的核心公式：
Attention(Q, K, V) = softmax(QK^T / √d_k) V
让我们逐步分析：
Step 1: QK^T —— 计算"相关性"
Q (Query) 和 K (Key) 的内积衡量了两个 token 之间的"相关性"或"匹配度"。如果 token i "想了解" token j（即它们语义相关），那么 Q_i · K_j 会比较大。
Step 2: 除以 √d_k —— 方差控制
假设 Q 和 K 的每个分量独立同分布，均值为 0，方差为 1。那么 Q_i · K_j 的方差是多少？
Var(Q·K) = Σ Var(q_m · k_m) = d_k · 1 = d_k
当 d_k = 64 时，标准差为 8；d_k = 128 时，标准差为 ~11。大的内积值会导致 softmax 输出极端（接近 0 或 1），使得梯度几乎为零（softmax 饱和区）。
除以 √d_k 将方差控制回 1，保持梯度在合理范围。
Step 3: Softmax —— 归一化为概率
Softmax 将注意力分数转化为概率分布：每行和为 1，表示"这个 query 对各 key 的关注度"。
Step 4: 乘以 V —— 加权聚合
最后，用注意力概率对所有 Value 做加权求和，得到该 query 的输出表示。
为什么要多头？一个头只能学习一种"关注模式"。多头允许模型同时关注：
- 语法关系（主语-谓语、动词-宾语）
- 语义关系（同义词、指代消解）
- 位置关系（相邻词、远距离依赖）
- 特定模式（代码缩进、格式化）
每个头有独立的 Q、K、V 投影矩阵，工作在低维子空间（d_model / n_head）。
工程实现技巧：
GPT-2 将 Q、K、V 的投影合并为一个大矩阵乘法，减少 kernel launch 开销：
# 而非分别计算：
# q = self.w_q(x); k = self.w_k(x); v = self.w_v(x)
# GPT-2 的做法（更高效）：
qkv = self.c_attn(x) # 一次性投影
q, k, v = qkv.split(n_embd, dim=-1)
自回归语言模型在训练时使用因果掩码（Causal Mask），确保 token i 只能看到 token 0, 1, ..., i（不能看到未来的 token）。这在概念上与驾驶中的因果性一致：规划只能基于过去和当前的感知，不能基于未来。
Self-Attention 对位置不敏感：交换两个 token 的位置，它们的注意力输出也会相应交换。没有位置信息，模型无法区分"狗咬人"和"人咬狗"。
Sinusoidal (原始 Transformer)
- 优点：可外推，无需学习
- 缺点：表达能力有限
- 适合：教学理解、特殊外推需求
Learned (GPT-2, BERT)
- 优点：灵活，模型可学习最优模式
- 缺点：无法外推，增加参数
- 适合：固定长度的任务
RoPE (LLaMA, Mistral, Qwen — 当前主流)
- 优点：天然编码相对位置，可外推
- 缺点：计算稍复杂（但可忽略）
- 核心公式：q_m @ R(θ, m) · (k_n @ R(θ, n))^T = q_m @ R(θ, m-n) @ k_n
- 旋转编码使得内积只依赖于相对位置 (m-n)
ALiBi (BLOOM)
- 优点：极简，强外推，零参数
- 缺点：绝对位置信息弱
- 方法：直接在注意力分数上加线性偏置
选择建议（2025年）： 新项目首推 RoPE，它是当前几乎所有主流 LLM 的选择。
RoPE 的核心思想：通过对 Q 和 K 向量施加位置相关的旋转，使得注意力计算自然地编码了相对位置。
对每个维度对 (x_{2i}, x_{2i+1}) 施加 2D 旋转：
x_{2i}' = x_{2i} · cos(m·θ_i) - x_{2i+1} · sin(m·θ_i)
x_{2i+1}' = x_{2i} · sin(m·θ_i) + x_{2i+1} · cos(m·θ_i)
其中 θ_i = 1 / (base_theta)^{2i/d}，不同维度有不同的旋转频率。
这是 Transformer 训练稳定性的关键选择。
Post-Norm (原始方案)：
x = LayerNorm(x + Attention(x))
x = LayerNorm(x + FFN(x))
问题：随着层数增加，残差路径上的信号可能不断放大，导致训练不稳定。
Pre-Norm (现代方案，GPT-2 起)：
x = x + Attention(LayerNorm(x))
x = x + FFN(LayerNorm(x))
LayerNorm 放在子层"之前"而非"之后"，有效控制每一层输入的方差，允许训练更深的网络（100+ 层）。
你的代码中： 默认使用 Pre-Norm，这是生产环境的推荐选择。
FFN 占 Transformer 参数量的 ~2/3，负责 token 级的非线性变换：
FFN(x) = W_2 · GELU(W_1 · x)
其中 W_1 将维度从 d_model 扩展到 4×d_model，W_2 再压缩回去。
GELU (Gaussian Error Linear Unit) 是 ReLU 的平滑版本，在 0 附近提供非零梯度，对训练有帮助。
Token IDs → Token Embedding + Position Embedding
→ Dropout
→ [TransformerBlock × N]
→ Final LayerNorm
→ LM Head (Linear)
→ Logits
对于序列 [t_0, t_1, ..., t_n]，目标是预测 [t_1, t_2, ..., t_{n+1}]。这是自监督学习，不需要人工标注。
输入的 Token Embedding 和输出的 LM Head 共享权重矩阵。这减少了 ~30% 的参数量（对于大 vocab），并提供了自然的正则化。
自回归生成时，每生成一个新 token，需要将所有历史 token 重新输入模型计算。不做优化的话：
- 每步都重新算所有历史 K、V
- 第 n 步的计算量是 O(n²)
- n 步总计是 O(n³)
缓存之前计算好的 K 和 V。新 token 只需要：
1. 计算自己的 Q、K、V
2. 将新 K、V 追加到缓存
3. Q 与全部 K 做注意力
这样每步的计算量是 O(n)，总计 O(n²)。
进一步优化：将 KV-Cache 分页管理，按需分配，减少显存碎片。vLLM 借此实现了 2-4x 的推理吞吐提升。
标准注意力的中间结果 S = QK^T 是 O(N²) 的，必须写入 GPU HBM：
- N=2048: ~16MB / 头 (fp16)
- N=8192: ~256MB / 头
- 当 N 很大时，HBM 带宽成为瓶颈
不在 HBM 中存储完整的 S 和 P 矩阵，而是：
1. Tiling: 将 Q/K/V 分块，每次只在 SRAM 中处理一块
2. Online Softmax: 增量计算 softmax（运行中维护 max 和 sum）
3. Recomputation: 反向传播时重新计算 S（而不是存储）
结果：HBM 读写量大幅减少，在 A100 上可达 2-4x 加速。
PyTorch 2.0+ 内置了 Flash Attention 支持：
output = F.scaled_dot_product_attention(q, k, v, is_causal=True)
这会自动选择最优实现，无需额外配置。
Byte Pair Encoding 反复合最高频的相邻 token 对，直到达到目标词汇量。它平衡了字符级（太细）和词级（词汇太大）的缺点。
将对话格式化为模型可理解的文本。不同模型有不同的模板：
| 模型 | 模板格式 | ||||
|---|---|---|---|---|---|
| LLaMA 2 | [INST] ... [/INST] |
||||
| ChatML | `< | im_start | >role\nmsg< | im_end | >` |
| Vicuna | USER: ... ASSISTANT: |
正确的 Chat Template 对指令微调的效果至关重要。
目标： 学习语言的统计规律和世界知识
数据： 海量无标注文本（几百 GB 到几 TB）
Loss： 所有 token 的交叉熵
关键技巧：
- 大 batch size + 梯度累积
- Warmup + Cosine LR 调度
- 梯度裁剪 (max_norm=1.0)
- 混合精度训练 (bf16)
目标： 让模型学会遵循指令格式
数据： (指令, 回复) 对（几千到几十万条）
Loss： 只在回复部分计算（prompt 部分 labels = -100）
关键技巧：
- 从预训练 checkpoint 开始
- 更小的学习率 (1e-5 ~ 2e-5)
- 更少的训练步数
目标： 让模型的回复更符合人类偏好
数据： (prompt, chosen, rejected) 三元组
Loss： DPO loss
L = -log σ(β(log Pθ(chosen)/Pref(chosen) - log Pθ(rejected)/Pref(rejected)))
直觉：
- 增加 chosen 回复的概率
- 减少 rejected 回复的概率
- β 控制偏离 reference 的程度
DPO vs RLHF：
- RLHF 需要训练独立的 Reward Model + PPO
- DPO 直接从偏好数据优化策略
- DPO 更简单、更稳定、训练成本更低
预训练权重 W 不更新，而是学习低秩增量：
W' = W + (α/r) · B · A
其中 B ∈ R^{d×r}, A ∈ R^{r×k}, r << min(d,k)。
微调时权重的更新 ΔW 具有低的"内在秩"（通常 r=4~64 就足够）。大模型的微调只需要在少数方向上调整参数。
- 可训练参数减少 99%+
- 训练速度快 2-3x
- 显存占用大幅降低（可以微调更大的模型）
- 可以快速切换多个 LoRA 适配器
- 合并后推理零开销
| LLM 能力 | 自动驾驶应用 |
|---|---|
| 文本生成 | 驾驶决策描述、场景报告 |
| 推理能力 | 复杂交通场景的推理和决策 |
| 少样本学习 | 快速适应新的驾驶场景 |
| 多模态融合 | 将视觉信息与语言指令结合 |
本章：LLM 基础 (GPT, Attention, Training)
↓
第2-3章：视觉 + 多模态 (ViT, CLIP, LLaVA)
↓
第4-5章：BEV + 感知 (LSS, BEVFormer)
↓
第14-15章：VLA (RT-2, OpenVLA, 端到端驾驶)
你在本章学到的 Transformer、训练 pipeline、LoRA、DPO 等知识，是后续所有章节的基础。
Transformer 有三种变体：
- Encoder-only (BERT): 双向注意力，适合理解任务
- Encoder-Decoder (T5): 适合序列到序列任务
- Decoder-only (GPT): 单向注意力 + 自回归生成
Decoder-only 在语言建模和生成任务上表现最好，且架构更简单，更容易规模扩展。
Pre-Norm 确保每层输入都经过归一化，信号不会逐层放大。这使得：
- 训练更稳定（可以用更大的学习率）
- warmup 要求降低
- 可以训练更深的网络
对于 LLaMA-7B (32层, 32头, head_dim=128)：
- 单 token: 32 × 2 × 32 × 128 × 2 bytes = 0.5 MB
- 2048 tokens: ~1 GB
- 因此推理时需要预留足够的显存给 KV-Cache
2025年的主流选择是 RoPE（LLaMA、Mistral、Qwen、Gemma 都在用）。
ALiBi 简单但绝对位置表达力弱。RoPE 是当前事实标准。
完成本章后，建议：
1. 运行所有演示脚本，确保代码可以正常工作
2. 尝试在小数据集上完成一次完整的预训练
3. 实验不同的超参数（学习率、batch size、模型大小）
4. 阅读 GPT-2 和 LLaMA 的原始论文，对比你的实现
5. 进入第2章：视觉基础模型
本章项目地址： vla_learn/chapter01_transformer_llm/
完整代码行数： ~2000+ 行 Python
预计学习时间： 2-4 天（含代码阅读和实验）