📑 本页目录(点开跳转)
02b · 现代 LLM 架构的四处改动
⏱ 88 分钟 | ⭐⭐ 读完这章,你能看懂任何一份 LLaMA 类模型的代码
🎯 一句话
2017 年的原始 Transformer 和 2026 年的 LLaMA / Qwen,骨架完全一样,但有四个零件被换掉了:LayerNorm 挪了位置、LayerNorm 换成了 RMSNorm、FFN 换成了 SwiGLU、绝对位置编码换成了 RoPE。 这四处改动不是审美,每一处都在修原始设计的一个具体毛病。
🚪 为什么要单独讲这一章
第 2 章讲完,你脑子里的一层 Transformer 是这样的:
注意力 → +残差 → LayerNorm → FFN → +残差 → LayerNorm
然后你打开一份真实的 LLaMA 实现,第一屏就会看到:
class RMSNorm(nn.Module): ← 这是什么?
self.rotary_emb = ... ← 这又是什么?
F.silu(self.gate_proj(x)) * self.up_proj(x) ← FFN 怎么有三个矩阵?
hidden = self.input_layernorm(x) ← 怎么 norm 跑到注意力前面去了?
第 2 章那张图上,这四个东西一个都没有。 这是学习路径上的一个硬断点:你以为自己懂了 Transformer,结果代码看不懂。
第 2 章自己在「深挖的话要学什么」里把这四项列成了必修——这一章就是那一行的展开。
🔑 读法建议:四处改动彼此独立,可以只挑你眼下需要的看。 只想看懂代码 → 二、三、四节;想搞清长上下文为什么能撑到 100 万 → 直接跳第五节 RoPE。
🗺️ 一、先看全景:两张图的差别
一张对照表:
| 零件 | 原始 Transformer(2017) | GPT-2 / GPT-3 | LLaMA / Qwen / Mistral |
|---|---|---|---|
| 归一化位置 | Post-LN(残差之后) | Pre-LN | Pre-LN |
| 归一化种类 | LayerNorm | LayerNorm | RMSNorm |
| FFN 激活 | ReLU | GELU | SwiGLU |
| FFN 中间维度 | 4d | 4d | 约 8/3 · d(对齐取整) |
| 位置编码 | 正弦(加在输入) | 可学习绝对位置(加在输入) | RoPE(乘在每层的 Q/K 上) |
| 归一化的 bias | 有 | 有 | 没有(连线性层的 bias 也普遍去掉了) |
⭐ 注意中间那一列:这四处改动不是一次换完的。Pre-LN 在 GPT-2 时代就换了,RMSNorm/SwiGLU/RoPE 是 2021-2023 陆续换的。 所以你看不同年代的代码会看到不同的组合,这很正常。
🧩 二、改动①:Post-LN → Pre-LN
原始设计是什么
原始 Transformer 的一个子层写成:
Post-LN: x_out = LayerNorm( x + Sublayer(x) )
↑ 归一化套在最外面,把残差通路也一起归一化了
它的毛病:残差通路被截断了
回想第 12 章 ResNet 里残差为什么有用——因为它给梯度开了一条加法通路:y = x + f(x) 求导得到 1 + f'(x),那个 1 保证梯度不会被连乘吃光。
Post-LN 把这个 1 弄没了。 因为 x 不是直接输出的,它先进了 LayerNorm。梯度往回传时,每穿过一层就要过一次 LayerNorm 的雅可比矩阵——不再是「乘 1」,而是「乘一个跟当前激活值尺度有关的量」。
6 层(原论文的规模):连乘 6 次,还行
96 层(GPT-3 的规模):连乘 96 次
→ 靠近输入的层,梯度要么被压到近乎为 0,要么被放大到爆炸
→ 而且不同层的尺度不一致,同一个学习率对某些层太大、对另一些层太小 💀
warmup 是个补丁,不是解法
这解释了第 9 章那句「大模型必配 warmup」的根源之一:
Post-LN 训练的头几百步,各层的激活尺度还很乱
→ 梯度的尺度也乱
→ 这时候如果学习率已经是满值,一步就能把某几层的权重推到再也回不来的地方
→ 所以要先用极小的学习率"暖机",等尺度稳下来再放开
⚠️ 补丁的问题在于它掩盖了病因:warmup 的长度成了一个必须精心调的超参。调短了发散,调长了浪费算力,而且换个模型深度就要重调。
Pre-LN:把归一化挪进残差分支里
Pre-LN: x_out = x + Sublayer( LayerNorm(x) )
↑ 主干上的 x 一路裸奔到顶,谁也没碰它
把 L 层展开:
x_L = x_0 + f_1(LN(x_0)) + f_2(LN(x_1)) + … + f_L(LN(x_{L-1}))
↑
⭐ 从输出到输入有一条【纯加法】的路,中间没有任何归一化
梯度沿这条路回到第 0 层时,是完整的、没被缩放过的。这就是「梯度直通」。
后果:
| 现象 | Post-LN | Pre-LN |
|---|---|---|
| 能堆多少层 | 十几层就开始难训 | 上百层没问题 |
| 需要 warmup 吗 | 必须,且长度敏感 | 可以短,甚至可以不要 |
| 学习率能开多大 | 小 | 大得多 |
| 同深度下的最终效果 | 略好 ⭐ | 略差 |
⚠️ 这是个权衡,不是纯赢
最后一行别跳过。学界的共识(也是面试常问的那个反直觉点)是:同等深度下,Post-LN 训得起来的话,效果通常比 Pre-LN 好一点。
一个直觉解释:Pre-LN 里每一层往主干上加的东西,相对于已经累积的主干越来越小(主干在长大,分支的输出被归一化后尺度固定)。到了深层,一层的贡献可能只占主干的百分之几——看着是 96 层,实际"有效深度"没那么深,行为上更接近一个更浅但更宽的模型。
⭐ 但在千亿参数尺度,这个权衡根本不用想: Post-LN 更好的前提是「它训得起来」。在几千张卡上跑几个月,一次发散可能要回滚几天的进度。 稳定性的价值,远大于那一点点效果。 这就是 Pre-LN 全面胜出的真实原因。
💀 真实代价的一个参照:Meta 训 OPT-175B 时公开了完整的训练日志,里面记录了大量因 loss 发散和硬件故障导致的中断、回滚与超参调整。它用的是 Post-LN 系的配置,配合 992 张 A100 训练。"训练过程本身需要人 7×24 盯着"是那个年代的常态——而这正是后来所有人都换 Pre-LN(并加一堆稳定性技巧)的动机。🔗 《AI基础设施》第 21 章整章在讲这件事。
💡 还有一个必须记住的细节:Pre-LN 的最后一层输出没有经过归一化(它是主干直接出来的)。 所以所有 Pre-LN 模型在整个 stack 的最后都会补一个 final norm,然后才接输出层。看代码时你会在
model.norm或ln_f这个名字下找到它。
中间路线也有人做:Sandwich-LN(分支前后各来一个)、DeepNorm(在残差相加时给 x 乘一个跟深度有关的常数 α,让 Post-LN 也能堆到 1000 层)。知道有这些就行,主流仍是 Pre-LN。
📏 三、改动②:LayerNorm → RMSNorm
LayerNorm 做了两件事
① re-centering(再居中):减去均值 μ → 让这一行的均值变成 0
② re-scaling (再缩放) :除以标准差 σ → 让这一行的方差变成 1
③ 然后乘一个可学的 g,加一个可学的 b
(如果对「一行」是哪一行还有疑问:LayerNorm 是横着归一化,对一个 token 的所有特征维度做,不跨样本。🔗 第 10 章有那张 4×3 的对照图。)
RMSNorm 砍掉了第一步
$$\text{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{d}\sum_i x_i^2 + \epsilon}} \cdot g$$
变化只有两处:不减均值了,也不加 bias 了。分母从「标准差」变成了「均方根 RMS」——两者的唯一区别就是标准差要先减均值。
为什么去掉也没事
RMSNorm 论文(Zhang & Sennrich, 2019)的核心实验发现是:LayerNorm 真正起作用的是 re-scaling,不是 re-centering。 把减均值那步去掉,效果基本不掉。
一个直觉:
归一化在这里的职责是【控制数值尺度】,防止层层往上漂移
→ 这件事是【除以 σ】干的
→ 【减 μ】管的是"平移",而后面紧跟的线性层本来就有平移能力
(原本 LayerNorm 的 bias 也是干这个的,一并去掉了)
⚠️ 别把这个结论推广得太远。它是在 Transformer 这个特定结构里、后面紧跟线性层的经验结论,不是「所有归一化都可以不减均值」。
省了多少
从 FLOPs 看,归一化在整个模型里占比极小(远不到 1%),砍掉一次减法看起来毫无意义。但这么想就错了。
⭐ 关键洞察:归一化是逐元素算子,瓶颈在显存带宽不在算力。 它要把整个激活张量从显存读出来、算一下、再写回去。GPU 的算力远远快于显存带宽, 所以这类算子的实际耗时,跟它读写了多少字节成正比,跟它算了多少次乘法几乎无关。 🔗 这正是《AI基础设施》第 4 章 Roofline讲的 memory-bound 区间。
在这个视角下 RMSNorm 的收益就清楚了:
LayerNorm:要先遍历一遍求 μ,再遍历一遍求 σ(或者用一趟算法但要多存中间量)
RMSNorm :一趟就够 —— 只需要平方和
→ 少一次对张量的遍历,少一批中间量的读写
→ RMSNorm 论文报告的端到端加速在 7%~64% 之间(取决于模型和实现)
⚠️ 那个上限 64% 是特定小模型上的数字,别拿它到处引用。在现代大模型里,换 RMSNorm 的实际收益是个位数百分比——但它是白捡的(效果不掉),所以没有理由不换。
代码:两者放一起看
import torch
import torch.nn as nn
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)) # ⭐ 只有缩放 g,没有 bias
def forward(self, x: torch.Tensor) -> torch.Tensor:
dtype = x.dtype
x32 = x.float() # ⭐ 低精度下必须先升到 fp32 再算
rms = torch.rsqrt(x32.pow(2).mean(-1, keepdim=True) + self.eps) # ⭐ eps 在开方【里面】
return (x32 * rms).to(dtype) * self.weight
# 对照:nn.LayerNorm 比它多了"减均值"和"加 bias"
ln = nn.LayerNorm(8)
rn = RMSNorm(8)
h = torch.randn(2, 5, 8)
print(ln(h).shape, rn(h).shape) # 形状完全一样,可以直接替换
print(ln.weight.shape, ln.bias.shape) # LayerNorm 有 bias
print(rn.weight.shape) # RMSNorm 没有
⚠️ 两个真实的坑(都在上面的代码里标了):
| 坑 | 后果 |
|---|---|
| eps 加在开方外面 | 当整行接近 0 时分母炸掉,出 NaN。必须写成 sqrt(mean(x²) + eps) 而不是 sqrt(mean(x²)) + eps |
| fp16/bf16 下直接算平方和 | x² 很容易上溢/下溢。所有正经实现都会先 .float() 再算,最后转回去 |
🛑 第一个歇脚处 —— 归一化那两处改动讲完了(约 28 分钟)。 你已经拿到:Pre-LN(梯度直通,稳定性压倒一点点效果) 和 RMSNorm(去掉减均值,省的是带宽不是算力)。 后面还有:SwiGLU 门控 · RoPE 旋转位置编码 · 长度外推的 PI/NTK/YaRN · 拿到陌生 config 该看什么。 回来的时候不用重读,直接从下一节接着看就行。
🚪 四、改动③:FFN → SwiGLU
先看原来的 FFN
$$\text{FFN}(x) = W_2 \cdot \text{ReLU}(W_1 x)$$
W₁ 把 d 维放大到 4d,激活一下,W₂ 压回 d 维。参数量:
W₁: d × 4d = 4d²
W₂: 4d × d = 4d²
合计 8d² ← 记住这个数,下面要用
中间那步:ReLU → GELU
GPT-2/GPT-3/BERT 都用 GELU(Gaussian Error Linear Unit):
$$\text{GELU}(x) = x \cdot \Phi(x)$$
Φ 是标准正态分布的累积分布函数。人话:ReLU 的"软"版本——ReLU 是「大于 0 就全放过,小于 0 就一刀切掉」,GELU 是「按这个值有多大,按概率放过一部分」。
ReLU: x = -0.1 → 输出 0 (切得干脆,但这里【梯度也是 0】)
GELU: x = -0.1 → 输出约 -0.046 (留了一点,梯度不为 0)
⭐ 好处:负半轴不再是死区,不会出现"死 ReLU"(神经元被推到负区后永远不再更新)
而且它处处可导,优化更平滑
⭐ 代价:要算 erf 或用 tanh 近似,比 ReLU 贵
(🔗 激活函数为什么必须是非线性的、异或那个实验,在第 7 章。)
GLU 门控:加一路"闸门"
GLU(Gated Linear Unit)的想法:不要只走一路。走两路,一路当内容,一路当闸门,然后逐元素相乘。
普通 FFN(一路)
x → [W₁] → 激活 → 一个向量
GLU-FFN(两路)
① 内容路:x → [W_up] → 内容向量(这一路【不过激活】)
② 闸门路:x → [W_gate] → 激活 → 闸门向量
③ 把两路【逐元素相乘】⊗ → 结果
💡 为什么这有用:普通 FFN 里,每个输出通道拿到的是输入的加权和再过激活——本质是加法组合。门控引入了乘法交互:闸门那一路可以根据输入内容,逐通道地决定「这个维度放多少过去」。
闸门值接近 1 → 这个通道完整放行
闸门值接近 0 → 这个通道被关掉 ⭐ 相当于对每个 token 动态地选一组子通道
这是纯加权和做不到的事——它让 FFN 的行为变成「输入相关的」,而不是对所有 token 用同一套固定变换。
SwiGLU 就是把闸门那一路的激活选成 Swish(也叫 SiLU):$\text{Swish}(x) = x \cdot \sigma(x)$
$$\text{SwiGLU-FFN}(x) = W_{\text{down}} \Big( \text{Swish}(W_{\text{gate}}\,x) \odot (W_{\text{up}}\,x) \Big)$$
(⊙ 是逐元素相乘。注意 W_up 那一路没有激活函数——它是纯线性的"内容"。)
⭐ 关键的一步:为什么中间维度要乘 2/3
数一下矩阵:SwiGLU 有三个(W_gate、W_up、W_down),原来只有两个。
如果中间维度还用 4d:
3 × (d × 4d) = 12d² ← 比原来的 8d² 多了 50% 参数!
那它效果好,可能只是因为它更大,这个对比就不公平了
所以要把中间维度 h 定成让参数量持平:
3 × d × h = 8d² → h = 8d/3 = (2/3) × 4d ⭐
这就是那个 2/3 的来历:不是什么神秘常数,就是「三个矩阵要装进两个矩阵的预算里」。
LLaMA-7B 的真实数字:
d = 4096
4d = 16384
× 2/3 = 10922.67
向上取整到 256 的倍数 → 11008 ⭐ 这就是 LLaMA-7B config 里的 intermediate_size
import torch
import torch.nn as nn
import torch.nn.functional as F
class SwiGLU(nn.Module):
"""LLaMA 式 FFN:三个矩阵,不是两个。"""
def __init__(self, dim: int, hidden: int):
super().__init__()
self.w_gate = nn.Linear(dim, hidden, bias=False) # ⭐ 闸门那一路
self.w_up = nn.Linear(dim, hidden, bias=False) # ⭐ 内容那一路(无激活)
self.w_down = nn.Linear(hidden, dim, bias=False)
def forward(self, x):
return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x)) # ⭐ 逐元素相乘 = 门控
def llama_hidden(dim: int, multiple_of: int = 256) -> int:
h = int(2 * (4 * dim) / 3) # ⭐ 4d 先乘 2/3
return multiple_of * ((h + multiple_of - 1) // multiple_of) # 再向上对齐
print(llama_hidden(4096)) # 11008 —— LLaMA-7B 的真实值
ffn = SwiGLU(4096, llama_hidden(4096))
print(sum(p.numel() for p in ffn.parameters()) / 1e6, "M 参数")
# 对照:原始 FFN 是 2 × 4096 × 16384 = 134.2M
它到底为什么更好
诚实的答案是:没人说得清。
提出这个改动的论文(Noam Shazeer, GLU Variants Improve Transformer, 2020)在做完一堆对照实验、确认 SwiGLU/GeGLU 在同参数量下确实更好之后,结尾写了一句在圈内很出名的话——大意是:我们对这些架构为什么 work 不提供解释,把它的成功归于神的恩赐。
⭐ 这句话值得记住,不是因为它幽默,而是因为它诚实。 深度学习里有相当一部分架构选择就是这样:大量对照实验证明它好,但没有令人满意的理论解释。 遇到这种情况,正确的态度是「知道它在同参数量下经过了公平比较」,而不是硬编一个说得通的故事。
代价也要说清楚:
| 代价 | 说明 |
|---|---|
| 多一次矩阵乘 | 三次 GEMM 而不是两次。参数量持平,但 kernel 启动次数变多 |
| 维度不好对齐 | 8d/3 基本不会是漂亮的数,必须手动向上取整到 128/256 的倍数(否则 GPU 上 tensor core 利用率掉一截) |
| 激活值显存变多 | 训练时要为 gate 和 up 两路各存一份中间激活,峰值显存比两矩阵 FFN 高 🔗 显存优化 |
🛑 第二个歇脚处 —— 三处改动讲完了(约 45 分钟),刚好到全章一半。 最后一节是 RoPE,它是四处改动里最独立、也最值得单独花一次时间的一节——它直接决定了模型能不能撑起长上下文。 这一节可以独立读,回来直接从下一节接着看就行。
🌀 五、改动④:绝对位置编码 → RoPE
这是四处改动里影响最深远的一处——它直接决定了模型能不能撑起长上下文。
原来的做法和它的两个毛病
原始 Transformer:算一堆不同频率的 sin/cos,加到输入 embedding 上,只在第 0 层加一次。GPT-2 换成了可学习的位置向量,做法一样也是加在输入。
毛病一:加法会被稀释
位置信息和词的语义信息挤在同一组数字里,往上传几十层,
经过几十次注意力混合和 FFN 变换,还剩多少是位置?没人说得清
毛病二:换长度就废 💀
可学习位置编码:训练时只见过 0~511 号位置 → 第 512 号位置的向量【根本不存在】
正弦位置编码:理论上能外推,实际一超出训练长度就崩
RoPE 的想法:不加,改成"转"
核心一句话:把位置编码成"旋转的角度",而不是"加上去的向量"。
具体做法:把每个注意力头的 d 维向量两两配对,看成 d/2 个二维平面上的点。位置 m 的向量,在第 i 个平面上旋转 m·θᵢ 弧度。
位置 0 的词:转 0 度 (原地不动)
位置 1 的词:转 1×θ
位置 2 的词:转 2×θ
位置 m 的词:转 m×θ ⭐ 转多少完全由位置决定,和词是什么无关
⚠️ 注意作用位置:RoPE 不作用在输入 embedding 上,而是作用在每一层注意力的 Q 和 K 上(V 不转)。所以它每层都在起作用,不会被稀释——这直接解决了毛病一。
⭐ 为什么相对位置"自动"就出现了
这是 RoPE 最漂亮的地方,也是面试最爱问的一句。
用一句公式说清楚:
把 q 转 m·θ,把 k 转 n·θ,然后算内积
→ 内积 = |q||k| · cos( 原夹角 + (m−n)·θ )
↑
⭐ m 和 n 只以【差】的形式出现,绝对位置消失了
🔑 这就是那句标准答案的完整含义: RoPE 用绝对位置的方式实现(每个 token 独立地转一个只跟自己位置有关的角度),却得到相对位置的效果(注意力分数只依赖两者之差)。 而且它不用像 T5 那样为每个相对距离额外学一张偏置表——相对性是几何送的,不花参数。
🔗 为什么相对位置比绝对位置更符合语言的规律(形容词紧跟名词、代词指代多远),《ML 基础》第 14 章讲过。RoPE 是那一段的具体实现。
θᵢ 的设计:一组转速不同的指针
d/2 个平面,每个平面转速不同:
$$\theta_i = \text{base}^{-2i/d}, \quad \text{base 通常取 } 10000$$
i = 0 (最前面的维度):θ 最大 → 转得最快 → 高频
i = d/2−1 (最后面的维度):θ 极小 → 转得最慢 → 低频
💡 一个好用的类比:钟表。
秒针(高频):转得快,能分辨"差了几秒",但转一圈就绕回来了 → 分不清远距离
时针(低频):转得慢,一天才半圈,能区分"上午还是下午" → 但分不清秒
⭐ 一起看,才能唯一确定一个时刻
RoPE 就是这样:高频维度负责精细的邻近关系,低频维度负责粗糙的长距离关系,所有维度合起来才唯一地表示一个位置。
import torch
def build_rope_cache(seq_len: int, head_dim: int, base: float = 10000.0):
# ⭐ θ_i = base^(−2i/d):i 越大,转得越慢
inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
pos = torch.arange(seq_len).float()
angles = torch.outer(pos, inv_freq) # (seq_len, head_dim/2)
return angles.cos(), angles.sin()
def apply_rope(x, cos, sin):
"""x: (batch, heads, seq, head_dim)。两两配对,在每个平面上转一下。"""
x1, x2 = x[..., 0::2], x[..., 1::2]
cos, sin = cos[None, None], sin[None, None]
out1 = x1 * cos - x2 * sin # ⭐ 就是二维旋转矩阵,没有别的
out2 = x1 * sin + x2 * cos
return torch.stack((out1, out2), dim=-1).flatten(-2)
cos, sin = build_rope_cache(seq_len=8, head_dim=16)
q = torch.randn(1, 2, 8, 16) # 1 个样本、2 个头、8 个 token、头维度 16
print(apply_rope(q, cos, sin).shape) # ⭐ 形状不变 —— 旋转不改变维度
print(cos.shape, sin.shape) # (8, 8):8 个位置 × 8 个平面
⚠️ 两个实现细节:① RoPE 只作用在 Q 和 K,V 不转(V 是要被取回来的内容,转了没意义)。② 真实实现里常用「前半段和后半段配对」而不是「相邻两个配对」,因为前者的张量操作更快——两种配对方式不通用,权重不能混用,这是加载别人权重时的经典报错来源。
长度外推:PI / NTK / YaRN 到底在动什么
RoPE 不是天生就能外推。训练时最长见过 4096,你直接喂 32768 会怎样?
高频那些维度(转得最快的)在 4096 个位置里已经转了很多圈
→ 推到 32768,它们会转到训练时【从没见过的角度组合】
→ 注意力分数变成噪声,模型输出崩坏 💀
三种解法,它们动的其实是同一个东西:每个频率的"指针"该转慢多少。
| 方法 | 做法 | 直觉 | 代价 |
|---|---|---|---|
| PI(位置插值) | 把位置索引统一除以 s(如 8) | 所有指针一律转慢 8 倍,32768 个位置被压回 4096 的角度范围内 | 邻近位置被挤在一起,高频的精细分辨率变差 |
| NTK-aware | 不均匀缩放:高频少动、低频多缩(等价于把 base 从 10000 调大) | 「精细分辨率归高频管,别动它;长距离归低频管,缩它」 | 比 PI 温和,但仍需要少量微调才最优 |
| YaRN | NTK 的精细版:分频段处理 + 同时调整注意力的温度 | 把维度按"在训练长度内转了几圈"分成三档区别对待 | 实现最复杂,效果最好 |
⭐ 实践里最简单粗暴的一招:直接把 base 从 10000 改成 1000000,然后在长数据上继续训一小段。很多长上下文模型(LLaMA 3、Qwen 的长文版本)都做了这件事——base 变大等于所有 θ 变小等于所有指针转慢,本质就是 NTK 的思路。
🔗 第 3 章说的「上下文窗口是怎么撑到 100 万的」,位置编码这一侧的答案就在这里。而那一章同时提醒你:窗口撑大了不等于有效——大海捞针能过,不代表长文推理能过。
🧰 六、实用:拿到一份陌生模型,看什么
打开任何 HuggingFace 模型的 config.json,这几个字段基本就把架构说清了:
| 字段 | 它在说什么 |
|---|---|
hidden_size |
就是 d。4096 ≈ 7B 级,8192 ≈ 70B 级 |
intermediate_size |
FFN 中间维度。如果它约等于 8/3 × hidden_size → 是 SwiGLU;如果约等于 4× → 是老式两矩阵 FFN ⭐ |
num_attention_heads / num_key_value_heads |
两者不等就是 GQA(🔗 KV Cache) |
rope_theta |
RoPE 的 base。10000 是原版,500000 / 1000000 说明做过长上下文扩展 ⭐ |
rope_scaling |
有这个字段 = 用了 PI / NTK / YaRN 之一,里面会写 type 和 factor |
rms_norm_eps |
光是这个字段名出现,就说明用的是 RMSNorm 不是 LayerNorm |
max_position_embeddings |
宣称的上下文长度。⚠️ 这是宣称值,不是有效值 |
num_hidden_layers |
层数 |
💡 一个能立刻用的小推断:
intermediate_size / hidden_size这个比值, 约 2.67 → SwiGLU 家族;约 4.0 → GELU 两矩阵 FFN。 一眼就能判断这个模型是哪个年代的设计。
🔗 这一章连到哪里
| 去哪 | 为什么 |
|---|---|
| 02 · Transformer 原理 | 本章是它的下游。如果你对 QKV、残差、FFN 分工还不熟,四处改动会看不出「改的是什么」 |
| 02c · 为什么都是 Decoder-only | 这一章讲的是「块内部换了什么零件」,那一章讲「整个模型为什么长成这个形状」——两个问题不同层次 |
| 02d · MoE 混合专家 | 第五处改动。本章的 SwiGLU 讲的是 FFN 怎么变强,MoE 讲的是 FFN 怎么变多 |
| 03 · Tokenizer 与上下文 | RoPE 的外推是「上下文窗口怎么撑到 100 万」的一半答案,另一半(有效性、中间迷失)在那一章 |
| 《ML 基础》14 · 通往 Transformer | 那里有位置编码四选项对照表和「相对位置为什么更好」的直觉,本章的 RoPE 节是它的展开 |
| 《ML 基础》10 · 正则化全家桶 | LayerNorm vs BatchNorm、为什么 Transformer 必须用 LN。读 RMSNorm 之前最好先有这个底 |
| 《ML 基础》09 · 优化器与学习率 | 「大模型必配 warmup」的另一半答案就在本章第二节:warmup 是 Post-LN 的补丁 |
| 《ML 基础》12 · CNN | ResNet 的 +I 梯度高速公路。Pre-LN 之所以成立,靠的就是这条通路不被打断 |
| 《ML 基础》07 · 从线性到神经网络 | ReLU / GELU / 激活函数为什么必须非线性,异或那个实验在这里 |
| 《AI基础设施》04 · Roofline 与 MFU | 「RMSNorm 省的是带宽不是算力」这句话的完整依据——memory-bound 到底是什么意思 |
| 《AI基础设施》21 · 训练稳定性与故障恢复 | Pre-LN 想解决的「训练发散」在千卡规模到底有多贵,那一章有完整的工程视角 |
| 《AI基础设施》16 · KV Cache | 看 config 时判断 GQA 用的 num_key_value_heads,机制在那一章 |
| 《AI基础设施》08 · FlashAttention | 本章讲注意力「算什么」,那一章讲同样的东西「怎么算得快」 |
✅ 检查点
- Post-LN 和 Pre-LN 的公式分别怎么写?为什么说 Pre-LN 让「梯度直通」?
- warmup 为什么被说成是「补丁而不是解法」?它掩盖的病因是什么?
- Pre-LN 是纯赢吗?它的代价是什么?为什么大模型还是全都选它?
- RMSNorm 相比 LayerNorm 去掉了哪两样东西?为什么去掉也没事?
- 归一化在整个模型里 FLOPs 占比不到 1%,为什么换成 RMSNorm 还是有意义?
- SwiGLU 有几个矩阵?为什么中间维度要乘 2/3?LLaMA-7B 的
intermediate_size是多少,怎么算出来的? - GLU 门控引入了什么普通 FFN 做不到的东西?
- RoPE 作用在哪几个张量上?为什么「相对位置自动出现」?
- RoPE 里 θᵢ 为什么要设计成随维度递减?用钟表类比说一遍。
- PI、NTK-aware、YaRN 三者动的是同一个东西,是什么?
👀 答案
- Post-LN:
x_out = LayerNorm(x + Sublayer(x));Pre-LN:x_out = x + Sublayer(LayerNorm(x))。Pre-LN 把 L 层展开后是x_L = x_0 + f_1(LN(x_0)) + … + f_L(LN(x_{L-1})),从输出到输入有一条纯加法的路,中间没有任何归一化,梯度回传时不被缩放。Post-LN 则每穿过一层就要过一次 LayerNorm 的雅可比,96 层连乘会消失或爆炸。 - 因为病因是 Post-LN 把残差通路截断了、各层梯度尺度不一致。warmup 只是用「前期学习率极小」躲过最危险的阶段,没修好结构。代价是 warmup 长度成了敏感超参:调短了发散、调长了浪费算力,换个深度就要重调。
- 不是纯赢。同等深度下 Post-LN 训得起来的话效果通常略好——因为 Pre-LN 里深层往主干加的东西相对越来越小,「有效深度」不如标称深度。但大模型全选 Pre-LN,因为 Post-LN 更好的前提是「它训得起来」,几千张卡跑几个月时一次发散可能回滚几天,稳定性的价值远大于那一点效果。
- 去掉了 减均值(re-centering) 和 bias。因为 LayerNorm 起作用的主要是 re-scaling(控制数值尺度),而平移这件事后面紧跟的线性层本来就能做。⚠️ 这是 Transformer 这个特定结构里的经验结论,别过度推广。
- 因为归一化是逐元素算子,瓶颈在显存带宽不在算力——耗时跟读写多少字节成正比,跟算多少次乘法几乎无关。LayerNorm 要遍历张量两次(先 μ 后 σ),RMSNorm 一趟就够。论文报告 7%~64% 的加速(上限是特定小模型的数字),现代大模型上是个位数百分比,但效果不掉所以是白捡的。
- 三个(
W_gate、W_up、W_down),原来只有两个。原始 FFN 参数量2 × d × 4d = 8d²,要让三矩阵持平就得3 × d × h = 8d²→ h = 8d/3 = (2/3)×4d。LLaMA-7B:d=4096 → 4d=16384 → ×2/3 = 10922.67 → 向上对齐到 256 的倍数 = 11008。 - 乘法交互。普通 FFN 每个输出通道拿到的是输入的加权和(加法组合);门控让另一路根据输入内容逐通道地决定放多少过去,闸门接近 0 就等于关掉这个通道——相当于对每个 token 动态选一组子通道,纯加权和做不到。⚠️ 但为什么这样更好,提出它的论文自己承认给不出解释。
- 只作用在每一层注意力的 Q 和 K 上,V 不转(而不是像旧方案那样加在输入 embedding 上,所以不会被深层稀释)。相对位置自动出现是因为:q 转 m·θ、k 转 n·θ 之后,内积 = |q||k|·cos(原夹角 + (m−n)θ)——m 和 n 只以差的形式出现,绝对位置抵消了。而且这个相对性是几何送的,不额外花参数。
θᵢ = base^(−2i/d),base 通常 10000。钟表类比:秒针(高频、i 小)转得快能分辨"差了几秒"但转一圈就绕回来分不清远距离;时针(低频、i 大)转得慢能区分上午下午但分不清秒。高频维度管精细的邻近关系,低频维度管粗糙的长距离关系,合起来才唯一确定一个位置。- 都在动「每个频率的指针该转慢多少」。PI 一律转慢 s 倍(代价:高频的精细分辨率变差);NTK-aware 不均匀缩放,高频少动低频多缩(等价于把 base 调大);YaRN 分频段处理 + 调注意力温度。实践里最简单的一招是直接把
rope_theta从 10000 改成 1000000 再在长数据上续训一段。
🛑 可以停在这里
⚡ 走神救援
原始 Transformer(2017) 和 LLaMA 骨架一样,只换了四个零件。①Post-LN → Pre-LN:Post-LN 是
LN(x + f(x)),归一化套在最外面把残差通路截断了,96 层连乘梯度就消失/爆炸——warmup 只是补丁(前期小学习率躲过危险期),病因没修,且换深度就要重调。Pre-LN 是x + f(LN(x)),展开后x_L = x_0 + Σf_i(LN(x_i)),从输出到输入有一条纯加法的路,梯度直通,能堆上百层。⚠️但不是纯赢:同深度下 Post-LN 效果通常略好(Pre-LN 深层贡献被稀释,有效深度不足)——大模型仍全选 Pre-LN,因为"效果好"的前提是"训得起来",千卡上一次发散回滚几天。Pre-LN 的最后一层没归一化,所以 stack 末尾要补一个 final norm。②LayerNorm → RMSNorm:去掉减均值和 bias,只除均方根。因为起作用的是 re-scaling 不是 re-centering。⭐省的是带宽不是算力——归一化是逐元素算子、memory-bound,LN 要遍历两次 RMSNorm 一次,论文报 7%~64%(现代大模型上个位数,但效果不掉是白捡)。⚠️ 两个坑:eps 必须在开方里面、fp16 下必须先 .float() 再算平方和。③FFN → SwiGLU:GELU 是 ReLU 的软版(负半轴不再是死区);GLU 门控走两路——一路内容一路闸门,逐元素相乘,引入了纯加权和做不到的乘法交互。⭐三个矩阵不是两个,所以中间维度要乘 2/3 才保持参数量:3×d×h = 8d²→ h = 8d/3。LLaMA-7B:4096 → 16384 → ×2/3 = 10922.67 → 对齐 256 得 11008。为什么更好?提出它的论文自己说给不出解释——这份诚实值得学。④绝对位置 → RoPE:旧方案把位置加在输入上(深层被稀释、换长度就废);RoPE 把 d 维两两配对成 d/2 个平面,位置 m 转 m·θᵢ,只作用在每层的 Q 和 K 上,V 不转。⭐相对位置自动出现:内积 = |q||k|·cos(原夹角 + (m−n)θ),m、n 只以差的形式出现,绝对位置抵消,而且不花参数。θᵢ 随维度递减像钟表的秒针/时针:高频管精细邻近、低频管粗糙长距。外推的 PI / NTK / YaRN 动的是同一件事——每个频率该转慢多少;实践里最简单是把rope_theta从 10000 改成 1000000 再续训。看 config 一眼判年代:intermediate_size / hidden_size≈ 2.67 是 SwiGLU,≈ 4.0 是老式 FFN。
下一节 👉 02c-为什么都是Decoder-only.md