← 返回首页
12_llama_viz.html
把 GPT-2 升级成 LLaMA:四个零件
第 3 章那台 GPT 是 2019 年的样子。今天的开源模型(LLaMA / Qwen / Mistral / DeepSeek)长得不太一样了 ——
但换掉的不是 Transformer 本身,而是它身上四个零件:
RoPE、
RMSNorm、
SwiGLU、
GQA。
注意力、残差、训练循环一行没动。本关一次换一个,每个都能单独开关,最后跑一组
六选一的真实对照。配套代码 phase3-moe/10_llama.py。
这一关
n_embd 128 主干维度
n_head 4 query 头数
head_size 32 =128/4
n_kv_head 2 K/V 头数(GQA)
block_size 64 上下文
底座 = 第 3 章那台 4 层 GPT · 配套 10_llama.py
① 老架构的四处别扭
第 3 章那台 GPT 能跑、也训得动。但把它放到 2023 年之后的标准看,有四个地方会被换掉。
点下面任意一处,看它在第 3 章代码里是哪一行、代价是什么、今天换成了什么:
四处别扭里,只有 GQA 是纯粹为了省(省推理显存),另外三个都是"既省一点、也更好训一点"。
这一关按 位置 → 归一化 → FFN → 注意力 的顺序一个一个换。
换了这四件,它还算 Transformer 吗?
算,而且骨架一动没动:逐 token 嵌入 → N 层(注意力子层 + FFN 子层,各带残差与归一化)→ 映射回词表,
这条主干和 2017 年那篇论文一模一样。四个零件换的是"每个子层内部用什么实现":
位置怎么进来、归一化怎么算、FFN 用什么非线性、K/V 存几份。
所以你在第 3 章学的那套结构,今天读 LLaMA 的代码依然直接够用。
为什么这四件是一起出现的?它们之间有关系吗?
没有强耦合,四件互相独立,任何一件都能单独用(所以本关每件都做成了单独的开关)。
它们凑在一起,是因为都在 2021–2023 年被反复验证"便宜且不亏",于是被 LLaMA 一次性打包进同一套配方,
后来的开源模型就照着抄。真正有点关系的是 RoPE 和 GQA:两个都作用在注意力里,
而且都对"推理时怎么处理长上下文"有影响。
↳ 下一步:先修最要命的那个 —— 位置表是"背下来"的,一超长度就抓瞎。
② RoPE:把位置从"查表"变成"转角度"
第 3 章的位置信息来自一张 nn.Embedding(block_size, n_embd) 表:第几个位置就查第几行,加到 token 向量上。
RoPE 换了个思路:不加任何东西,而是把 q 和 k 这两个向量按"你在第几个位置"旋转一个角度。
向量两两配对当成平面上的点,位置越靠后转得越多。拖滑块看它怎么转:
0
四对维度 = 四个不同的转速。靠前的转得快(管邻近关系),靠后的转得慢(几十个位置才转一圈,管远距离)。
真实每个头有 16 对(head_size=32),这里只画前 4 对。
关键转完之后,注意力只看得见"相对距离"
这是 RoPE 真正的好处,不是省参数。把 q 放在位置 m、k 放在位置 n,
算出来的注意力分数 q·k 只跟 m−n 有关,跟它俩具体在第几个位置无关。
拖两个滑块试试:一起往右拖,距离不变,分数纹丝不动。
↳ 代码:10_llama.py 的 build_rope_cache / apply_rope · 开关 --rope 1
"只依赖相对距离"为什么值钱?
因为语言里真正重要的是"这个词和那个词隔多远",不是"它们在第几个字符"。
同一句话出现在文章开头还是第 900 个 token,内部关系应该一样。
可学习的位置表做不到这点:它给每个绝对位置一行独立的参数,第 900 行和第 3 行毫无关系,
而且训练时没见过的位置(超过 block_size)根本没有那一行 —— 这就是老架构"外推不了"的根源。
RoPE 把位置变成角度之后,超出训练长度也照样能转(效果会下降,但不是无中生有),
所以后来那些"把上下文从 4K 拉到 128K"的做法(NTK / YaRN 之类),动的都是 RoPE 的频率。
为什么是"两两配对"地转?
旋转是平面上的操作,得有两个坐标才能定义"转多少度"。
所以 RoPE 把 head_size=32 的向量拆成 16 对,每一对当成一个平面上的点,各自按自己的频率转。
每对配一个不同的转速,是为了让模型同时拥有"看近"和"看远"两种尺度 ——
转得快的那几对,隔几个位置角度就差很多,适合分辨近处;转得慢的那几对,几十个位置才转一点,
用来编码长距离。这和傅里叶那套"用不同频率叠出复杂信号"的思路是一个味道。
它对 v 也转吗?
不转,只转 q 和 k。因为位置信息是要影响"谁该看谁"这个打分过程,
而打分只用到 q·k;v 是"被搬运的内容",不参与打分。代码里也就是
apply_rope(q, …) 和 apply_rope(k, …) 两行,没有 v 的份。
↳ 下一步:位置解决了。再看每一层都要做两次、但其实做多了的那个操作 —— 归一化。
③ RMSNorm:同一件事,少做一半
LayerNorm
做两件事:①减均值(把这根向量挪到 0 附近)、②除以标准差(把尺度拉回 1 附近)。
RMSNorm 发现第 ① 步基本没用,干脆只做第 ②:直接除以均方根。
点下面切换,看同一根向量在两种归一化下的处理步骤:
↳ 代码:10_llama.py 的 RMSNorm · 开关 --rmsnorm 1
减均值这一步,凭什么说它"基本没用"?
这是个经验结论,不是数学定理。RMSNorm 那篇论文(2019)做的就是拆开来做消融:
把 LayerNorm 的"减均值"去掉,在多个任务上训练稳定性和最终效果都没有明显损失,
但省掉了一次求均值、一次减法,以及每层的 bias 参数。
后来 LLaMA 等一批模型跟进采用,等于用大规模训练又验证了一遍。
本关的玩具规模也能看到同样的现象:val loss 基本打平(见 STEP 6 的对照),参数还少了一点。
省下的那点算力单看很小,但归一化在每层要做两次、几十上百层累起来就不小了。
↳ 下一步:归一化省完了,轮到占参数大头的 FFN —— 那个一刀切负数的 ReLU。
④ SwiGLU:给 FFN 装一道闸门
第 3 章的 FFN 是 放大 → ReLU → 压回。ReLU 很粗暴:负的全归零,正的原样通过。
SwiGLU 换成门控:同一个输入走两路,一路过
SiLU
当闸门,一路当内容,两路逐元素相乘——
闸门决定每个通道放多少内容过去。拖滑块看两种做法怎么处理同一个值:
-1.20
算账三个矩阵,为什么参数没变多?
↳ 代码:10_llama.py 的 SwiGLU · 开关 --swiglu 1
门控到底强在哪?ReLU 不也能"关掉"通道吗?
差别在谁来决定关不关、关多少。
ReLU 的开关是它自己:这个通道的值为负就关死、为正就全开,没有中间态,也没有第二个信号参与。
门控的开关来自另一路独立的线性变换:内容是一路算出来的,闸门是另一路算出来的,
于是模型可以学出"这个通道的内容有用,但在当前上下文里只放三成过去"。
多出来的那一路矩阵,买到的就是这种逐通道、连续可调的控制力。
实测上它通常带来一点点稳定收益,不是革命性的 —— 属于"便宜且不亏"那一类改动。
SiLU、Swish、GELU 是什么关系?
SiLU 和 Swish 是同一个函数的两个名字(x·sigmoid(x)),分别在两篇论文里被独立提出。
GELU 是另一个形状很接近的平滑激活(GPT-2 用的就是它)。
三者共同点是:都比 ReLU 平滑,负半轴不硬归零而是缓慢趋近 0,处处可导。
实践里换来换去差别很小,真正带来变化的是"门控"这个结构,而不是选了哪个平滑激活 ——
所以名字里那个 GLU(Gated Linear Unit)才是重点。
↳ 下一步:前三件都在省算力和参数。最后一件省的是另一样东西 —— 推理时的显存。
⑤ GQA:让几个 query 头拼一份 K/V
第 5 章讲过 KV-cache:
推理时把每个 token 的 K/V 存下来复用。问题是它会一直变大 ——
上下文每多一个 token,就要多存一份 K/V,长对话里它能吃掉比模型权重还多的显存。
GQA 的办法很直接:query 头保持 4 个,但 K/V 只留 2 份,每 2 个 query 头共用一份。
拖滑块看连线和账本:
2
↳ 代码:10_llama.py 的 Attention(k_proj/v_proj 只投 n_kv_head 份 + repeat_interleave)· 开关 --n-kv-head 2
共享了 K/V,注意力还算"多头"吗?
还算。每个 query 头依然有自己独立的 Wq,依然算出自己的一套注意力权重、聚合出自己的一根输出向量。
共享的只是"被查询的那一侧"。打个比方:4 个人各自带着自己的问题(query)去查资料,
原本一人一套资料(K/V),现在两人合看一套 —— 问的问题还是各问各的,得到的答案也各不相同,
只是可查的资料副本少了。表达力确实略有损失,所以是 4→2 这种温和的共享,
而不是一步到位全共享。
为什么不干脆压到 1 份(MQA)?
可以,那就是 MQA(Multi-Query Attention),省得最狠(本关配置下 KV cache 只剩 1/4),
代码里 --n-kv-head 1 就能跑。代价是表达力损失更明显、训练也更容易不稳。
GQA 正是被当作 MHA 与 MQA 之间的折中提出来的:
大部分显存收益拿到手,质量损失又小到可以忽略。所以现在主流模型多用 GQA(比如 64 个 query 头配 8 份 K/V),
而不是走极端。
↳ 下一步:四件都换完了。回到第 3 章那台 GPT 边上,把六组真实训练摆一起算总账。
⑥ 四件齐上:六组真实对照
同一份数据、同一套超参、同样 5000 步,只换零件,一共跑了六组:
老架构基线、四个零件各自单开、以及四件齐上。点柱子看每一组的数:
诚实交代:哪些差异是真的,哪些只是噪声
六组都是单种子、各跑一次,所以差多少要分开看:
RMSNorm(−0.001)和 GQA(−0.005)这个量级基本就是噪声,
换个随机种子完全可能翻过来 —— 别据此说"RMSNorm 更好"。
SwiGLU 单开 −0.039、四件齐上 −0.054 幅度大得多,更可能是真的,
但严格讲也该多跑几个种子才敢下结论。
更要紧的是:这四件在真实模型上被采用,靠的本来就不主要是 loss,而是
① 大规模下的训练稳定性;② 长上下文的外推能力(RoPE);③ 推理显存与吞吐(GQA)。
这三样在 0.8M 参数、64 上下文的玩具规模上根本显不出来。
本关能确实测到、且不依赖规模的只有一条:KV cache 实打实少了一半。
那我什么时候能看到这四件的真实收益?
把三个条件里的任意一个放大就能看到:模型更大(几亿参数往上,训练稳定性开始成为瓶颈)、
上下文更长(几千 token 往上,RoPE 的外推和 KV cache 的显存都成为硬约束)、
真做推理服务(并发几十上百条对话,KV cache 直接决定你一张卡能塞多少用户)。
本项目停在玩具规模,是因为再往上就不是"看懂机制"的问题,而是要真金白银的算力了。
达标本关实测(六次训练 · 各 5000 步 · tiny shakespeare)
Q&A常见疑问
今天的模型只改了这四处吗?
这四件是覆盖面最广、几乎人人都用的一组,但不是全部。常见的还有:
权重绑定与词表大小的取舍、QK-Norm(对 q/k 再做一次归一化,稳大模型训练)、
MoE(第 11 章那关)、更长上下文的位置插值(NTK / YaRN)、
以及注意力实现层面的工程优化(FlashAttention、PagedAttention)。
再加上 MoE,你手上这套已经能拼出一个相当接近现代开源模型的骨架了。
第 11 章的 MoE 和这一关能一起用吗?
能,而且真实模型就是这么干的。DeepSeek-V3、Qwen 的 MoE 版本,都是
RoPE + RMSNorm + SwiGLU + GQA 打底,再把 FFN 换成 MoE。
因为两边动的是不同零件:这一关换的是位置编码、归一化、FFN 内部结构和 K/V 份数,
第 11 章换的是"FFN 有几份、每个 token 走哪几份"。
本项目把它们拆成两关,就是为了让你能一件一件看清楚,而不是一上来面对一整套配方。
学完这一关,我离读懂真实开源模型的代码还差多远?
结构上基本不差了。去翻一份 LLaMA 或 Qwen 的建模文件,你会看到的正是:
RMSNorm、rotary_emb、gate_proj/up_proj/down_proj、
num_key_value_heads —— 这一关这四个零件的名字。
剩下的差距主要在工程侧:分布式训练怎么切、推理怎么批处理和调度、量化怎么做。
那些是"把它跑到生产上"的问题,不再是"它是怎么算的"的问题。
↳ 跑法:python 10_llama.py(四件齐上)· --preset gpt2(老架构对照)· --preset rope 等单开一件
🎉 把 GPT-2 升级成 LLaMA · 通关
你已经亲手把位置表换成旋转、把 LayerNorm 换成 RMSNorm、给 FFN 装上闸门、让 K/V 拼车。再加上第 11 章的 MoE,现代开源模型那套配方,你四件零件都手搓过一遍了。