LLM from Scratch · 从零手写大语言模型

从最朴素的 bigram 起步,手写到 GPT-2 124M 的预训练与推理,再做后训练、现代架构与推理效率, 最后补上分词器、注意力 kernel、scaling law 与 RLVR。每章配交互式可视化:先看懂直觉,再读代码。

PHASE 1 · nanoGPT

第一阶段 · nanoGPT

4 章
PHASE 1 · 124M

第一阶段 · 124M 模型

2 章
CHAPTER 05

第 5 章 · 跑起来:从 logits 到文字(采样与 KV-cache)

对应 phase1-124m/05_sample.py:把推理"从打分到吐字"拆成五步亲手拨——base 只续写不答题、temperature 捏尖/摊平分布、top_k 砍掉长尾、自回归逐 token 滚动、KV-cache 并排省掉重复计算(实测约 2.5–3.1×)。

CHAPTER 06

第 6 章 · 真训练:把玩具 GPT 喂真实数据

对应 phase1-124m/04_gpt2_124m.py:把 GPT-2 124M 预训练拆成五步亲手拨——配置从字符级玩具跃迁到 12/12/768、FineWeb-Edu 切 shard 顺序喂、梯度累积攒满 524288 token 再走一步、warmup + 余弦学习率实时画曲线、用两个真实端点(300M→3.65 / 10B→3.02)看数据量 33× 的收敛对照。

PHASE 2 · SFT / LoRA / DPO

第二阶段 · 后训练

4 章
CHAPTER 07

第 7 章 · 从"会续写"到"会听话":后训练地图

动手写代码前先对齐全景与术语:参数量 vs 数据量、任务 vs 手段(SFT/DPO × 全量/LoRA)、SFT 改了哪三处(数据/模板/loss mask)、偏好对齐里 RLHF 与 DPO 的关系。

CHAPTER 10

第 10 章 · 手搓 DPO:从会答到答得合人意

换任务:从"会答"到"答得合人意"。接在 SFT 之后,用偏好对(chosen/rejected)+ 冻结参考模型做 DPO——不训奖励模型、不走 RL。配套代码 08_dpo.py,真实曲线 loss↓ / margin↑ / 准确率 0→100%;加 --lora 可切 LoRA 手段。

PHASE 3 · MoE + LLaMA

第三阶段 · 现代架构

2 章
CHAPTER 11

第 11 章 · 手搓 MoE:参数翻倍,算力不变

把第 3 章 GPT 的 FFN 拆成 4 个专家 + 路由器:top-2 派单、分工自己涌现、专家塌缩与负载均衡 aux loss,再算清"总参数 ↑ 激活算力不变"这笔账。最后一步换上 DeepSeek 系的细粒度专家、共享专家与偏置均衡,在同一台模型上跑五组对照。全部在 1.35M 参数上跑通,大概算全球最小的 MoE 之一,机制却和 671B 的 DeepSeek 同款。配套代码 09_moe.py。

CHAPTER 12

第 12 章 · 把 GPT-2 升级成 LLaMA:四个零件

第 3 章那台 GPT 是 2019 年的样子。今天的开源模型换掉的不是 Transformer 本身,而是它身上四个零件:RoPE 把位置变成旋转角(注意力只看相对距离)、RMSNorm 少做一半、SwiGLU 给 FFN 装闸门、GQA 让 K/V 拼车(cache 减半)。每件都能单独开关,配套代码 10_llama.py,含六组真实对照训练。

PHASE 4 · 长上下文与推理效率

第四阶段 · 长上下文与推理效率

5 章
CHAPTER 13

第 13 章 · 线性注意力:把 KV cache 换成固定大小的记忆板

全注意力推理时每个 token 都要存 K/V,上下文越长 cache 越大。线性注意力每个头只留一块 32×32 的记忆板:写入是加一个外积,delta rule 先擦再写,遗忘门 α 按比例衰减,合起来就是 Qwen3-Next 用的 Gated DeltaNet。配套代码 11_linear_attn.py,9 组真实训练:测试长度拉到 512,全注意力 val loss 涨到 3.52,GDN 停在 1.53;键值检索 N=32 时 GDN 答对 99.9%,只加不减的 linear 16.6%。最后看工业界为什么用 3:1 混合。

CHAPTER 14

第 14 章 · 稀疏注意力:每个 query 只挑一部分历史去读

K/V 照样全存,每个 query 只挑一小部分去算。照 DeepSeek 的 NSA 手写三条支路:压缩给每块留摘要、挑选只细看 top-n 块、滑窗兜住最近的字,再用 sigmoid 门加权。配套代码 12_sparse_attn.py:256 上下文平均每个 query 读 76.8 个 key(全注意力 128.5);远处找钥匙任务上只滑窗 3.1%(乱猜),NSA 99.85%。再对照 DSA、DeepSeek-V4 的 CSA/HCA 与 Qwen3.8 的 QSA。

CHAPTER 15

第 15 章 · N-gram 查表:给模型挂一本短语词典

给第 12 章的小模型挂一张哈希 n-gram 表:每个 token 按最近两三个 token 算出 8 个行号,只读 8 行 × 16 维,再由一个门决定加多少进残差流。表涨到主干的 4.6 倍,每 token 算力不变;行号在 forward 之前就能算出,表可以留在主机内存。配套代码 13_engram.py,FineWeb-Edu 上 5000 步 val loss 降约 0.06(单种子),tiny shakespeare 字符级反例里表把训练集背了下来,另测了表在 GPU / CPU 的吞吐。

CHAPTER 16

第 16 章 · 投机解码与 MTP:先便宜地猜,再一次验收

自回归生成每吐一个字都要跑一次主模型。投机解码让草稿员先猜 γ 个字,主模型一次 forward 算出这些位置的 p,以 min(1, p/q) 接受、拒绝就从 max(0, p−q) 重抽,输出分布不变。配套代码 14_mtp.py,在 RTX 4090 上实测:1 层小草稿模型接受率 0.667–0.745,γ=6 时每次主 forward 吐 3.09 个字;DeepSeek-V3 式 MTP 头接受率 0.899,墙钟 1.27×;20 万次模拟后与 p 的总变差距离 0.0005。最后算清玩具规模下省了 forward 为什么没怎么省时间。

CHAPTER 17

第 17 章 · 量化:把 124M 从 32 位压到 4 位

把第 6 章训出来的 GPT-2 124M 量化到 8 / 4 / 3 / 2 位:一个 scale 管多少权重决定误差,一个 10σ 的离群值就能让整矩阵 int4 崩到 val loss 9.93,留出 0.1% 离群值又回到 3.26。再拆开 Q4_K、NVFP4、MXFP4 三种 4 位格式怎么存 scale。配套代码 15_quant.py,15 组真实评测:4.5 位的 NVFP4 从 497.9 MB 压到 126.9 MB,val loss 多 0.04。Phase 4 收尾。

PHASE 5 · 补齐全链路

第五阶段 · 补齐全链路

4 章
CHAPTER 18

第 18 章 · 手写 BPE:自己训一个分词器

第 6 章起一直调 tiktoken,这一关自己写:从 256 个字节出发,数相邻对、合并出现最多的一对,在 8 MB FineWeb-Edu 上训出 16,384 词表,编码再解码逐字还原。配套代码 16_bpe.py,纯 CPU 约 2 分钟:每 token 4.25 字节(GPT-2 的 50,257 词表 4.62);只按空格预切分时,词表里 2,279 个 token 是「词 + 标点」的变体,GPT-2 正则下只有 7 个;英文语料训出的词表一个汉字都没合并。

CHAPTER 19

第 19 章 · 手写 FlashAttention:拆开 SDPA 这个黑盒

朴素注意力要把 T×T 的分数矩阵整个写进显存,T = 4096 时约 7.5 GB。online softmax 边读边更新最大值与分母,分块之后这个矩阵不再落到显存。配套代码 17_flash_attn.py,手写一个 Triton 前向 kernel:RTX 4090 上 T = 4096 用时 0.77 ms、占 24 MB,PyTorch 的 SDPA 0.76 ms;朴素写法 31.8 ms,T = 8192 起显存不够。换进第 6 章的 124M,val loss 与 SDPA 相差 0.0002。

CHAPTER 20

第 20 章 · Scaling Law:用小模型预测 124M 的 loss

第 6 章把 124M 训了一整天。能不能先训一批小模型,算出它会训到多少?训 23 次 0.6M–25M 参数的小模型(5000 万到 16 亿 token),拟合 L(N, D) = E + A/N^α + B/D^β,外推到 124M 训 10B token:预测 3.36,实测 3.03。换三组拟合点对照,外推越远误差越大,主拟合的 bootstrap 区间也没盖住实测。配套代码 18_scaling.py,顺带排查两次训练发散:第 2 层注意力分数涨到约 1000 万,加 qk-norm 后同配置 val loss 从 4.95 降到 4.19。

CHAPTER 21

第 21 章 · RLVR 与 GRPO:只给判分器,不给答案

只给一个会对答案的判分器,不给标准答案:同一道题采样 8 次,组内比较得出优势,按优势更新。配套代码 19_grpo.py,把 124M 教会两位数加法,并在同一热身起点、同一评测、同样 4096 道题下和 SFT、DPO 同台比,学习率各自按 dev 挑。热身 1000 道的起点上 GRPO 从 41.0% 到 84.8%,有标准答案的 SFT 到 93.0%,离线 DPO 47.8%(学习率 3e-8,再大就掉到 warm 以下);热身 200 道的起点上 GRPO 的组大多 8 次全错,几乎学不到。Phase 5 收尾。