手撕Transformer：500行PyTorch实现可调试小型语言模型

1. 这不是又一篇“Transformer 入门科普”，而是一次亲手搭积木式的语言模型解剖实验

你有没有过这种感觉：读了十篇“图解 Transformer”、看了二十个动画演示、背熟了“Self-Attention 就是计算词与词之间的相关性”这句话，但一合上屏幕，脑子里还是雾蒙蒙的——QKV 矩阵到底长什么样？Positional Encoding 是加在哪儿的？为什么 LayerNorm 要放在残差连接之后而不是之前？更别提训练时 loss 怎么反向传播到 embedding 层，或者推理时那个“自回归生成”到底是怎么一词一词吐出来的。这不是你理解力的问题，而是绝大多数教程把 Transformer 当成一个黑箱来“讲”，而不是当成一个可触摸、可调试、可打断点的 工程实体 来“拆”。这篇内容，就是带你用不到 500 行 PyTorch 代码，从零构建一个真正能跑起来、能喂句子、能输出预测、能打印中间张量形状的极简语言模型——它没有 BERT 的预训练，没有 GPT 的千亿参数，但它有完整的 Encoder-Decoder 架构（我们聚焦 Decoder-only 路线，即 GPT 风格），有真实的 tokenization 流程，有可验证的 attention 权重热力图，有每一步 forward pass 的 shape 变化日志。核心关键词就三个： Transformer 解构、小型语言模型、PyTorch 实操 。它适合两类人：一类是刚学完《深度学习导论》、对矩阵运算不怵，但被 Attention 公式吓退的在校学生；另一类是工作三年、天天调 API 却想搞懂底层 why not 的工程师。它不承诺让你明天就复现 Llama，但它能确保你关掉这篇文章后，打开自己的 Jupyter Notebook，敲下 model = TinyLLM() ，然后 print(model) ，看到的不再是一串抽象的类名，而是一个由 Embedding 、 Linear 、 LayerNorm 这些你亲手定义的、有明确输入输出维度的砖块垒成的、呼吸着的模型。这才是“简单”的真意：不是降低复杂度，而是让复杂度变得 可见、可测、可干预 。

2. 整体设计思路：为什么必须“小”，以及“小”到什么程度才真正有用

2.1 “小”不是妥协，而是精准控制变量的科学实验

很多人一听到“小型语言模型”，第一反应是“这玩意儿能干啥？连‘Hello World’都续不出来吧？”——这恰恰暴露了对“小”的根本性误解。在科研和工程中，“小”从来不是目标，而是 最锋利的手术刀 。想象一下，你要研究汽车发动机的点火系统，是直接拆解一台 V8 涡轮增压发动机，还是先用一个单缸二冲程模型机？后者没有排气歧管、没有复杂的 ECU 控制逻辑、没有燃油喷射压力传感器，但它能把火花塞如何跳火、活塞如何被推动、曲轴如何旋转这一整套因果链，以最干净、最无干扰的方式呈现给你。我们的 TinyLLM 就是这个单缸模型机。它的“小”体现在三个硬性约束上，每一个都经过反复权衡：

1. 参数量级锁定在 10 万以内 ：最终模型总参数为 98,304。这个数字不是拍脑袋定的，而是基于一个关键计算：假设词表大小 vocab_size=256 （足够覆盖 ASCII 字符+基础标点），隐藏层维度 d_model=64 ，层数 n_layers=2 ，注意力头数 n_heads=4 。那么仅 nn.Embedding(vocab_size, d_model) 就占 256*64=16,384 参数；每个 nn.Linear(d_model, d_model*4) 的 FFN 层占 64*256=16,384 ，两层就是 32,768 ；两个 nn.MultiheadAttention 层，其 Q/K/V 投影权重合计 3 * d_model * d_model = 3*64*64=12,288 ，加上输出投影 d_model * d_model = 4,096 ，两层共 32,768 ；再加上所有 LayerNorm 的可学习参数（每个 2*d_model=128 ，共 4 个，总计 512 ）。加总： 16,384 + 32,768 + 32,768 + 512 = 82,432 ，再加一点 bias 和 final linear，稳稳落在 10 万内。这个量级，保证了模型能在任何一台现代笔记本的 CPU 上完成一次完整训练（约 2 分钟），且内存占用低于 500MB，让你可以毫无心理负担地 print(tensor.shape) 数十次。

2. 架构极度精简，只保留不可删减的核心组件 ：我们砍掉了所有“锦上添花”的模块。没有 Dropout （在如此小的模型上，它带来的正则化效果远不如数据噪声本身，反而会掩盖梯度流动的真实路径）；没有 Learned Positional Embedding （改用经典的 sin/cos 固定编码，因为它不引入额外参数，且其数学形式能让你一眼看穿“位置信息是如何被注入到向量空间”的）；没有 Tied Weights （即 embedding 和 final linear 不共享权重，虽然大模型常用，但在这里会混淆“输入表示”和“输出 logits”的概念边界）；甚至 MultiheadAttention 也手动实现，而非调用 torch.nn.MultiheadAttention ，因为后者内部封装了 in_proj_weight 等细节，会模糊 Q/K/V 的独立计算过程。每一行代码，都对应一个清晰、单一、可解释的数学操作。

3. 任务场景聚焦于“字符级”而非“词级”建模 ：我们不使用 tiktoken 或 SentencePiece ，而是将整个文本视为一个字符序列。这意味着 vocab_size=256 是硬编码的 ASCII 表。好处是：tokenization 过程完全透明（ ord('A')=65 ， chr(65)='A' ），没有任何神秘的 subword 切分；模型学到的“语言规律”必然是最底层的——比如“q”后面大概率跟“u”，“.”后面大概率是空格或换行。这种粒度，让 loss 曲线的每一次波动、每一个预测概率的分布，都能被你用肉眼关联到具体的字符模式上，这是词级模型永远无法提供的“显微镜视角”。

提示：选择字符级建模，并非因为它更“先进”，而是因为它把“语言建模”这个宏大命题，压缩到了一个你可以用纸笔推演的尺度。当你看到模型在第 100 个 epoch 后，对 'th' 后接 'e' 的预测概率从 0.12 稳步上升到 0.87，那种“我亲眼见证了它学会了一个英语规则”的实感，是任何大模型 demo 都无法替代的。

2.2 为什么坚持“Decoder-only”？这是理解生成式 AI 的唯一入口

Transformer 架构有 Encoder、Decoder、Encoder-Decoder 三种经典范式。BERT 是 Encoder，它擅长理解；T5 是 Encoder-Decoder，它擅长翻译；而 GPT 系列是 Decoder-only，它擅长生成。本项目坚定选择 Decoder-only，原因只有一个： 生成，是语言模型最本质、最不可绕开的能力，而 Decoder 的结构，正是为生成而生的精密设计 。Encoder 的核心是双向上下文聚合，它像一个全知的裁判，能看到一句话的所有词，然后给出一个综合评分。但人类的语言产出（说话、写作）是单向的、自回归的——你不可能在说出“我今天吃”之前，就决定好最后一个字是“饭”还是“面”。Decoder 的 Masked Self-Attention，就是这个单向约束的数学实现。它的 attention mask 是一个下三角矩阵，强制每个位置 i 只能“看到”位置 0 到 i-1 的信息，完全模拟了人类思考的时序性。如果你跳过 Decoder，直接去学 Encoder，你就永远无法理解：为什么 ChatGPT 不能“编辑”一段已有的文字？为什么它每次回答都是从头开始“编”？为什么它的“思考”永远是向前的、不可逆的？在 TinyLLM 中，我们将亲手实现这个 causal_mask ，并可视化它——你会看到，当输入序列长度为 8 时，mask 是一个 8x8 的矩阵，主对角线及右上方全是 -inf （代表被屏蔽），只有左下方是 0 （代表允许计算）。这个矩阵，就是生成式 AI 的“时间之箭”，是它区别于