从零构建 Transformer:毕业项目

从零构建 Transformer:毕业项目

本节摘要:十三节课,一个模型,没有捷径。你已经读了每一篇论文,实现了注意力、多头切分、位置编码、编码器和解码器块、BERT 和 GPT 损失、MoE、KV 缓存——现在让它们在一个真实任务上协同工作。毕业项目:端到端训一个小的 decoder-only Transformer,做字符级语言建模。它读莎士比亚,生成新的莎士比亚;小到笔记本上 10 分钟内能训完,正确到换上更大数据集和更长训练就能得到一个真正的语言模型。这是本课程的「nanoGPT」——不原创,Karpathy 2023 的 nanoGPT 教程是每个学生至少写一次的参考实现,我们沿用它的形状,按本系列覆盖的内容重新打磨。本节用 PyTorch 把第 04(位置)、05(RMSNorm/SwiGLU/pre-norm)、03(多头)、07(因果掩码、shift-by-one 损失)节的积木拼成一个能跑的 GPT,在 tinyshakespeare.txt 上训 2000 步,看训练损失从约 4.2(随机)降到约 1.5,采样输出是「莎士比亚形状」的——古体词、换行、ROMEO: 之类的专名涌现。读完本节,你拥有一个 80 万参数的 GPT,以及把它放大到 1.25 亿参数、生成流畅英语的清晰路径。

学习目标

阅读完本节,你应当能够:

1. 把前 13 节的积木拼成一个端到端可训练的 decoder-only GPT:嵌入+位置、pre-norm 块(RMSNorm + 因果多头注意力 + 残差 + SwiGLU FFN + 残差)、最终 RMSNorm、绑定 lm_head、shift-by-one 交叉熵。
2. 实现一个完整的训练循环:随机批量、前向、反向、AdamW、梯度裁剪、cosine 学习率。
3. 在 tinyshakespeare.txt 上训 2000 步,把验证损失压到 2.0 以下,采样出「莎士比亚形状」的文本。
4. 说清项目交付与省略:用了什么(RMSNorm、SwiGLU、pre-norm、绑定嵌入)、省略了什么(RoPE、KV 缓存、Flash Attention、MoE)以及为什么(简化),并知道如何把它们加回来。
5. 给出放大路径:换 BPE 分词器、上 OpenWebText/fineweb-edu、加 RoPE+KV 缓存+Flash Attention,得到一个 1.25 亿参数、生成流畅英语的 GPT。

一、问题与直觉

你已经实现了每一个组件,现在让它们协同。毕业项目:训一个小的 decoder-only Transformer,字符级语言建模,读莎士比亚、生成莎士比亚。

我们交付什么

• GPTConfig——一处配置所有超参。
• MultiHeadAttention——因果、批处理,可选 Flash 风格通路(PyTorch 的 scaled_dot_product_attention)。
• SwiGLUFFN——现代 FFN。
• Block——pre-norm、残差包裹的注意力 + FFN。
• GPT——嵌入、堆叠块、LM 头、generate()。
• 训练循环:AdamW、cosine LR、梯度裁剪。
• 莎士比亚文本的字符级分词器。

我们不交付什么(为简化)

• RoPE——第 04 节概念上实现过,这里用可学习位置嵌入(更简单),练习让你换上 RoPE。
• 生成时的 KV 缓存——每步重算全前缀注意力,慢但简单,练习让你加缓存。
• Flash Attention——PyTorch 2.0+ 在输入匹配时自动派发,我们用 F.scaled_dot_product_attention。
• MoE——每块单个 FFN,MoE 见第 11 节。

目标指标

Mac M2 笔记本上,4 层、4 头、d_model=128 的 GPT 在 tinyshakespeare.txt 上训 2000 步:

• 训练损失约 6 分钟内从约 4.2(随机)收敛到约 1.5。
• 采样输出是「莎士比亚形状」:古体词、换行、ROMEO: 之类的专名涌现。
• 验证损失(留出末尾 10%)紧贴训练损失,这个规模/预算下不过拟合。

💡 为什么选 Shakespeare:约 1.1 MB、65 个唯一字符、词汇表小到塞进 4 字节,没有 BPE 和分词器折腾,却能产生肉眼可辨的「像语言」输出——是字符级语言建模教学的标准玩具语料,2015 年以来每个 LM 教程都用它。

二、从零实现

本节用 PyTorch,装 torch(CPU 版即可)。完整代码见原课程 phases/07-transformers-deep-dive/14-build-a-transformer-capstone/code/main.py,脚本负责:缺失时下载 tinyshakespeare.txt、字节级字符分词器、90/10 训练验证划分、支持硬件上的 bf16 autocast、训练完成后采样。

Step 1:数据

text = open("tinyshakespeare.txt").read()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda xs: "".join(itos[x] for x in xs)
​

65 个唯一字符,小词表,塞进 4 字节 vocab_size,无 BPE、无分词器折腾。

Step 2:模型

块就是第 05 节的教科书版——pre-norm、RMSNorm、SwiGLU、因果 MHA。4/4/128 配置约 80 万参数。

Step 3:训练循环

取随机批量(长 256 的 token 窗口)→ 前向 → shift-by-one 交叉熵 → 反向 → AdamW 步 → 记录 → 重复:

for step in range(max_steps):
    x, y = get_batch("train")
    logits = model(x)
    loss = F.cross_entropy(logits.view(-1, vocab_size), y.view(-1))
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # 防爆
    opt.step()
    opt.zero_grad()
​

Step 4:采样

给定提示,反复前向、从 top-p logits 采样、追加、继续,500 token 后停。

Step 5:读输出

2000 步后:

ROMEO:
Away and mild will not thy friend, that thou shalt wit:
The chief that well shame and hath been his friends,
...
​

不是莎士比亚,但是莎士比亚形状的。对 80 万参数、笔记本 6 分钟而言,这是清晰的胜利。

设计要点:绑定嵌入(lm_head 权重 = token 嵌入矩阵的转置)既省一半参数,又让「读词」和「写词」共享同一套语义空间——这是 GPT-2 以来的标准技巧,在第 07 节已讲过,这里直接用。梯度裁剪到范数 1.0 是防止训练爆炸的廉价保险,几乎所有生产训练都开。

三、框架对比:把它变成真的

这个毕业项目是参考架构。把它变成真东西的三个扩展:

1. 换分词器。用 BPE(如 tiktoken.get_encoding("cl100k_base")),词表从 65 跳到约 5 万,模型容量要相应 scale。
2. 上更 corpora。用 OpenWebText 或 fineweb-edu(HuggingFace),单张 A100 上 1.25 亿参数 GPT 训 100 亿 token 约 24 小时。
3. 加 RoPE + KV 缓存 + Flash Attention。练习带你逐个加。

最终得到一个 1.25 亿参数、生成流畅英语的 GPT。不是前沿模型,但同一条代码路径——只是更大——就是 Karpathy、EleutherAI、Allen Institute 在 2026 年训研究检查点用的。

四、可复用产物

原课程产出 outputs/skill-transformer-review.md:一个审查 Skill,对照前 13 节审查一个「从零实现 Transformer」实现的正确性。

五、练习

1. (Easy) 跑 code/main.py,验证训练模型末步验证损失低于 2.0。把 max_steps 从 2000 改到 5000,验证损失还在提升吗?
2. (Medium) 把可学习位置嵌入换成 RoPE,在 MultiHeadAttention 内对 Q 和 K 施加旋转,训练并验证损失至少一样低。
3. (Medium) 在采样循环里实现 KV 缓存,有/无缓存各生成 500 token,笔记本上墙钟应改善 5~20 倍。
4. (Hard) 给模型加第二个头,预测「下一个的下一个」token(MTP——DeepSeek-V3 的多 token 预测),联合训练,有帮助吗?
5. (Hard) 把每块单个 FFN 换成 4 专家 MoE(路由器 + top-2 路由),看等激活参数下验证损失如何变化。

本节要点回顾

1. 十三节积木拼成一个 GPT:嵌入+位置、pre-norm 块(RMSNorm+因果 MHA+残差+SwiGLU+残差)、最终 RMSNorm、绑定 lm_head、shift-by-one 交叉熵。
2. 交付:GPTConfig、因果批处理 MHA、SwiGLU FFN、pre-norm Block、GPT(含 generate)、AdamW+cosine LR+梯度裁剪训练循环、字符分词器。
3. 省略以简化:RoPE(用可学习位置)、KV 缓存(每步重算)、Flash Attention(用 F.scaled_dot_product_attention)、MoE(单 FFN)——练习带你逐个加回。
4. Shakespeare 是教学标准:1.1 MB、65 字符、词汇表小、输出肉眼可辨「像语言」。
5. 绑定嵌入:lm_head = token 嵌入转置,省一半参数,读写共享语义空间。
6. 梯度裁剪范数 1.0:防爆的廉价保险,生产训练几乎全开。
7. 目标指标:4/4/128、2000 步、约 6 分钟,训练损失 4.2→1.5,验证损失 <2.0,输出莎士比亚形状。
8. 放大路径:BPE 分词器 → OpenWebText/fineweb-edu → RoPE+KV 缓存+Flash Attention,得到 1.25 亿参数、流畅英语的 GPT。

下一节,我们回头系统梳理注意力变体——滑动窗口、线性注意力、稀疏注意力、Mamba/SSM 等,看清 softmax(QKᵀ)V 之外还有哪些选择,以及它们各自换掉了 O(N²) 的哪一部分。