从零手写推理模型：用NumPy实现Transformer核心模块

说实话，我入行AI工程这四年，最怕的不是模型训不出来，而是被一句话问住："你平时用的model.generate()，底层到底发生了什么？"我当年面试算法岗，简历上写着"熟练使用Transformer"，面试官让我手写一个attention的score计算，我写了三行就开始冒汗。后来我花了整整一个周末，用纯NumPy从零手写了一个小型语言模型，跑了训练，还让它生成出了几句人话。那次经历之后，我对整个AI工程的理解彻底变了——各种框架和工具包依然是工作的主力，但"从零实现一遍"给我建立的底层直觉，到现在都在帮我解决那些网上搜不到的玄学Bug。

这篇内容就是围绕"ai-engineering-from-scratch"这个主题，分享我手写推理模型的完整经验：从词表构建、注意力实现、反向传播，到训练调参、推理评估，再到后来迁移到生产环境的踩坑记录。适合那些已经会用深度学习框架、但想把底层原理真正吃透的工程师，也适合想自己动手造一个小型推理模型的学习者。你可以把它当成一条"从零复现"的路线图，每一个模块我都给出了我当时的选择逻辑和实际遇到的问题，照着走一遍，收获比看十篇框架教程都大。

1. 为什么我把"从零手写"当成AI工程的第一课

很多人觉得"从零手写"是学生时代的事，工作了就该用框架调包。但我在实际项目中越来越发现一个残酷的事实：越是依赖框架的封装，遇到问题越没有头绪。框架帮你省去了细节，同时也把排查问题的线索一起藏起来了。

1.1 "from scratch"到底指什么：边界与约定

"从零实现AI"这句话其实有个边界问题。我给自己定的规矩是：不用任何深度学习的自动求导框架，不调PyTorch/TensorFlow的nn.Transformer这类现成模块，所有核心计算用NumPy和纯Python手写。前提是——CPU上训练，模型规模控制在千万参数以内，数据几十MB级别，这样一台普通笔记本也能扛得住。

我选的实验对象是decoder-only的小型Transformer，因为它足够代表现代大语言模型的基础架构。训练目标是最简单的自回归语言建模：给定前文，预测下一个token。整个过程拆成六个模块：数据管线、tokenizer、embedding与位置编码、多头注意力、前馈网络与层归一化、训练与采样推理。每一块单独看都不复杂，但拼在一起，任何环节出错都会让最终结果变得不可控。

这个"约定"很重要。因为它决定了后面所有调试思路：既然没有框架帮我算梯度，那每一处矩阵乘法的shape变化必须自己盯死；既然没有自动求导，那反向传播的公式必须逐个验证。正因为约束严格，我从里面学到的东西才格外扎实。

1.2 手写一遍，到底能换来什么

我先说结论：手写一遍并不能让你写出比PyTorch更快、更稳的代码，但它能在三个层面上产生无法替代的价值。

第一，你会真正理解shape为什么这样变。比如多头注意力里(B, T, n_heads, head_dim)和(B, n_heads, T, head_dim)两种布局切换，框架里就是一行.transpose()的事，但当你手写时发现不转置就点乘不下去，你才会明白为什么要分成多个头、为什么每个头的维度要开根号缩放。

第二，你会建立"数值直觉"。第一次跑训练时，我遇到loss变成了NaN，第一反应是学习率太大。但我手写之后才知道还有另一个原因：softmax里减掉max之后，仍然可能因为过大的指数产生数值溢出；LayerNorm里给方差加的那个eps，小了会炸，大了会钝。这些经验在框架的nn.LayerNorm里根本感知不到。

第三，你会真正看懂框架API的设计意图。当我在PyTorch里看到nn.TransformerEncoderLayer那一长串参数时，脑子里自动映射到实现这些功能需要的几十行代码、十几个tensor操作的顺序和依赖关系，对照自己的实现，瞬间明白每个参数存在的必要。这种"框架向认知的迁移能力"，才是手写过一遍最大的红利。

提示：如果你时间有限，不用把每个模块都手写。我的建议是至少手写attention、LayerNorm和训练循环这三个部分，它们是理解整个Transformer训练的钥匙。

2. 一个可运行的推理模型，核心模块究竟长什么样

下面进入正题，我把一个能训练的微型推理模型拆开讲。为方便说明，我用的是简化结构，完整的可运行代码逻辑会穿插给出关键片段，具体参数可以根据你的算力调整。

2.1 数据与词表：tokenizer不是调接口

很多人第一步就摔在这儿——直接用现成tokenizer库多舒服，为什么要手写？我的理由是：手写一个BPE并不复杂，而且能让你明白token和字符的关系。如果只想快速跑通，最简单的做法是先用字符级tokenizer，用中英文混合的训练集，词表只有几百个字符，对小型模型完全够用。

字符级tokenizer核心逻辑很简单：

# 构建字符词表 chars = sorted(set(text)) stoi = {ch: i for i, ch in enumerate(chars)} itos = {i: ch for i, ch in enumerate(chars)} def encode(s): return [stoi[c] for c in s] def decode(ids): return ''.join([itos[i] for i in ids])

这种极简tokenizer有个明显问题：序列太长。英文一句话可能几十个字符，一个token只覆盖一个字符，模型学到"词义"需要更多的层和上下文。所以我后来换成了字节对编码，用正则找出高频字符对合并，几轮合并之后词表能到几千，序列长度压缩了三四倍。遇到中文时，我先把字符串encode成UTF-8字节再训练BPE，这样能跨语言统一处理，而且不会出现中文分词不一致的问题。

在这个阶段我没花太多时间做数据清洗，只做了三件事：去重、过滤太短的行、统一换行符。因为训练目标是让模型学会"语言的统计规律"，而不是背书。

2.2 位置编码与embedding：把离散符号变成向量

Transformer没有循环结构，要让模型感知token的顺序，必须在输入里加入位置信息。最经典的是sinusoidal位置编码：

def sinusoidal_pos_embedding(seq_len, d_model): pos = np.arange(seq_len)[:, np.newaxis] # (T, 1) i = np.arange(d_model // 2)[np.newaxis, :] # (1, d_model//2) angle = pos / np.power(10000, 2 * i / d_model) pe = np.zeros((seq_len, d_model)) pe[:, 0::2] = np.sin(angle) # 偶数维度 pe[:, 1::2] = np.cos(angle) # 奇数维度 return pe

这个设计巧妙的点在于：相对位置可以表示成正弦函数的线性组合，模型更容易学到"两格之前出现某词"这类规律。embedding层就更简单了——一个vocab_size * d_model的矩阵，查表取值即可。但有一个细节新手容易忽略：embedding矩阵和输出投影矩阵的shape其实是同一个。很多现代做法会让这两者共享参数，能大幅减少参数量。我试过共享和分开两种，共享版在小型模型上收敛略慢，但最终效果差不太多，还省了一半内存。

2.3 多头注意力：从零实现最容易被忽略的细节

注意力机制是整个Transformer的心脏。我对它的理解可以概括为一句话：每个词去"问"所有词，收集它们"回答"的信息，再聚合成自己的新表示。实现上分三步。

第一步，把输入通过三组矩阵映射成Q、K、V：

# x: (B, T, d_model) -> q, k, v: (B, T, d_model) q = x @ Wq.T + bq k = x @ Wk.T + bk v = x @ Wv.T + bv

第二步，计算attention权重，关键是要除以sqrt(head_dim)：

scale = head_dim ** 0.5 scores = q @ k.transpose(-2, -1) / scale # (B, T, T) # 因果掩码：上三角置为负无穷 mask = np.triu(np.ones((T, T), dtype=bool), k=1) scores = np.where(mask, -1e9, scores) weights = softmax(scores, axis=-1) out = weights @ v

第三步，把每个头的结果拼接回去，再过一层输出投影。

网上很多教程到这里就停了，但实际手写时会遇到三个课本上不会提的陷阱：

• 掩码位置必须用负无穷而不是零。如果用0，softmax后那些位置的权重不是零，是exp(0)/sum(exp(x))，信息照样泄漏。我第一次就是在这里埋了坑，训练出的模型loss降得很慢，生成时前后矛盾严重。
• QK^T之后如果做过mask，不能再对scores做任何加性操作，否则负无穷会被"污染"成普通数值，导致因果性失效。
• head_dim的选择最好是64或128的约数，同时和d_model整除。我一开始用d_model=128, n_heads=8, head_dim=16，效果正常；后来试过d_model=192, n_heads=7，发现整除不了，最后硬着头皮用n_heads=8，但每个head维度不均匀，代码里多了一堆reshape的破事。

多头注意力的意义在于：每个头可以关注不同的关系模式，比如一个头关注语法位置，另一个头关注指代关系。这些特性不是我刻意设计的，而是训练之后自然涌现的——这也是手写一遍之后你能直观体会到的奇妙之处。

2.4 前馈、LayerNorm与残差：为什么顺序不能乱

注意力之后是前馈网络和LayerNorm，它们的配合顺序是Transformer结构中最"细节敏感"的部分。我的实现采用的是post-norm结构，也是原始论文里的方案，顺序是：输入 -> Attention -> 残差相加 -> LayerNorm -> FFN -> 残差相加 -> LayerNorm。

LayerNorm的公式很简单，但它的作用我一直到调参时才真正理解：

def layer_norm(x, gamma, beta, eps=1e-5): mean = x.mean(axis=-1, keepdims=True) var = x.var(axis=-1, keepdims=True) return (x - mean) / np.sqrt(var + eps) * gamma + beta

为什么需要这个操作？因为深层网络中每一层的输入分布都可能漂移，LayerNorm把每个token的表示强行拉回标准分布，让梯度传播更稳定。我做过一个对照实验，去掉LayerNorm之后，6层Transformer在1500步时loss就开始震荡，3000步直接发散到NaN。

残差连接的意义则更直白：让梯度"抄小路"绕过非线性层，至少在反向传播时能给浅层传递得动信号。这两点是整个Transformer工程化的基石，没有它们，深度学习里的"深"字就无从谈起。

注意：如果你用pre-norm结构（每个子层前先norm再计算），训练会更稳定、更容易收敛，但最终性能略逊于post-norm。先固定一种结构把它调通，再换另一种体会差异。

3. 真正动手跑通训练：数学、内存与"玄学调参"

模块组装好之后，距离"能用的模型"还有一大段路——训练过程才是真正折磨人的地方。这一段我把手写反向传播、内存瓶颈和调参排查的经验一次性讲透。

3.1 手写反向传播的三种姿势

手写Transformer最硬核的部分是反向传播。我有三条路可走，从笨到聪明排开。

第一种，最原始也是最容易错的方法：对每一层手动推导梯度公式，逐个写出来。比如对attention里的scores = q @ k.T / scale，你要分别求出d_scores、d_q、d_k、d_v，每一处矩阵乘法都要记得转置关系和顺序。一旦某个维度写反了，梯度形状立马不对，排查起来很痛苦。

第二种，利用标量对矩阵求导的链式法则，把每层当黑盒，只写前向和反向"相邻两层之间的梯度传递"。这比第一种省事，因为你不用关心很远的上游梯度，只需要处理好当前层的输出梯度和输入梯度之间的关系。实现的时候我给自己定了个规矩：每写完两个连续的层，就用np.allclose对比一次"反向计算出梯度"和"用前向计算手动推导的梯度"是否一致。

第三种，我用到的取巧做法：用数值梯度校验。给定一个极小扰动h=1e-6，用(loss(x+h) - loss(x-h)) / (2h)近似梯度，跟手写的解析梯度对比。这个方法验证单个参数时非常高效，但全参数校验在模型稍大时慢得没法看。所以我的策略是：先校验单层模块（attention、LayerNorm各自正向逆向），再组装校验两层叠加，最后再全模型训练。每一步都确认无误再往下走。

这里有一个经验教训：永远不要直接写完整模型的梯度再一次性跑。除非你天赋异禀，否则99%的概率会出现某一处维度错、符号反、或者mask没有参与梯度计算的问题。分模块验证，最多多花半天时间，但能救你一周的Debug。

3.2 显存/内存的真实天花板

纯NumPy在CPU上训练，最大的瓶颈不是计算速度，而是中间张量的内存占用。假设batch_size=32, seq_len=64, d_model=128，一个batch的前向过程中，Q/K/V分别是32*64*128，注意力权重是32*8*64*64，乘以每个头的独立计算，这一层的激活值就有几十MB，堆到6层之后几GB是常态。

我实际测量过一组数据：

所以如果你用笔记本跑，建议直接从第三套配置起步。降d_model比降层数更省内存，因为注意力矩阵的大小直接和head_dim线性相关，而d_model决定所有线性层的宽度。

内存不够还有一个工程化的解决思路：梯度累积。因为我的batch不能开太大，可以把一个batch拆成4个微批，每个微批独立前向反向并累积梯度，然后统一做参数更新。这个技巧在后来的生产环境里也被广泛用在显存受限的卡上，思路完全一致。

3.3 学习率、初始化与"训练不收敛"的排查清单

我自己实际跑训练时遇到的第一个大坑是：模型训练了2000步，loss从6.5降到4.2，然后开始原地踏步。这时候我用的学习率是固定3e-4，看了不少教程说这个值是Transformer的黄金默认参数，但在我这个小模型上其实略大了——前期降得快，后期卡在局部震荡里。

后来换成warmup + cosine退火的学习率策略：前500步从1e-5线性升到3e-4，之后按余弦曲线衰减到1e-5。loss曲线的形态立刻不一样了，最终比固定学习率低了不少。为什么有效？因为transformer的深层网络对学习率极其敏感，一开始就用大学习率容易让梯度方向飘忽不定，warmup让参数先在平坦区域站稳，再进入高速更新阶段。

初始化也是一个容易被忽视的大坑。Xavier和He初始化是针对tanh/ReLU设计的，对于Transformer这种带残差和LayerNorm的结构，一个小规模但有效的做法是：在所有线性层中采用标准差为0.02的随机初始化，并将输出投影层初始化为零偏置。这个数值来自GPT系列论文的经验值，直接抄作业在千万参数级别以下非常稳。如果初始化标准太大，第一轮前向的注意力分布就会接近one-hot，loss下不去；太小则梯度消失，学习几乎停滞。

如果模型还是训不动，我建议按这个顺序排查：

1. 先看输入数据是否normalize过（token id是否从0开始连续编号，如果中间有空位，embedding层会学到无意义的向量）。
2. 打印前向时attention权重的均值——正常情况下应该近似均匀分布（1/T），如果出现某个token占据了80%以上的注意力，就是初始化或scale出了问题。
3. 检查梯度范数：前5步梯度范数如果超过10，说明网络已经开始爆炸，先降学习率到1e-4再观察；如果前5步梯度范数接近0，说明信号传不进去，检查残差和LayerNorm的顺序。
4. 最后再看loss曲线形状——正常应该是平滑下降、没有锯齿波动；如果剧烈震荡，大概率是batch太小或者学习率过大。

这套排查清单给我省了无数时间。后来我去帮同事排查他用PyTorch训的一个模型出现NaN的问题，第一反应就是"看看你初始化标准差设的多少"，结果真是默认初始化的锅。

4. 评估推理能力：别只看loss，要看它怎么"想"

模型训练到一定程度，loss变成一个钝化的指标——它还在下降，但你无法感知模型"能力"是否真的有本质提升。这时候需要一个更系统的评估视角。

4.1 从"生成流畅"到"真的会推理"

训练了一个约600万参数的小模型之后，我让它生成一段话：

输入："我感到很疲惫，因为今天" 输出示例："我感到很疲惫，因为今天我跑了很多代码，解决了很多问题，但是没有时间休息。"

单看语法和语义，它已经学会了基础的搭配和常识路径。但这不叫"推理"，充其量是"统计层面的模仿"。真正的推理模型要求模型在给定前提时，能一步步推出结论，而不是直接蹦出训练数据里的高频搭配。这是从"语言模型"到"推理模型"的门槛，也是标题里那个build a reasoning model from scratch让我最感兴趣的方向。

我当时做了一个很小的"推理实验"：给模型输入几道简单的算术题，比如"23 + 47 ="，模型输出的结果不稳定——有时给出70，有时73，有时干脆是乱码。这说明它没有真正计算能力，只是在模仿数字符号的分布。要让它具备推理能力，光靠加大数据量和参数是不行的，必须改变训练范式。

4.2 推理模型的最小实现：链式思考与回退机制

受OpenAI o1和各类reasoning model的启发，我开始尝试在小模型上复现"链式思考"（Chain-of-Thought）的简化版。原理不复杂：不直接给模型<问题, 答案>这类训练样本，而是给<问题, 逐步推理过程, 答案>三段式样本，让模型学会"先思考再回答"的内部模式。

我在自己收集的数学题数据集上做了三组对照实验，每组都训了相同步数：

数据非常有意思：显式CoT让模型生成长度暴增6倍，最终答案准确率提升了近3倍。这说明哪怕是一个600万参数的小模型，只要训练数据里包含"推理步骤"，它也能在token之间建立起一个临时的"工作记忆"，通过一步步生成来降低最终答案的难度。就像人做数学题一样，"先列式，再计算"比"心算两步跳答案"要可靠得多。

但这套方案最大的局限性是：模型生成的推理步骤可能是错的，但答案却恰好对了，或者反过来。这正是推理模型需要"回退机制"的原因——它不是一次性输出，而是生成几步之后评估一下当前假设是否合理，不合理就回退重来。我在自己的实验里加了一个很粗糙的版本：在采样过程中连续计算每一步生成token的置信度，如果连续5步的置信度都低于阈值，就回退到上一个"决策点"重新采样。效果谈不上惊艳，但确实让最终答案的随机性降低了。

4.3 评估集设计与人工检查

评估一个推理模型，我的经验是不能只用一个自动指标。Perplexity衡量的是"模型对数据的拟合程度"，Rouge/BLEU衡量的是"生成文本和参考答案的字符重叠度"，两者都测不出"推导过程是否合理"。

我给自己设计了三层评估方案：

• 第一层是symbolic evaluation，用程序化的方法检查答案是否精确匹配，适合数学题、代码生成这类有确定性答案的场景。
• 第二层是semantic evaluation，用一个小型BERT模型算生成文本与参考答案的语义相似度，适合开放问答。
• 第三层是人工抽检，每100条生成结果随机抽10条，直接看模型输出的完整推理链，评判标准不是最终答案对错，而是"推理链内部是否连贯、有没有跳步、有没有幻觉"。

这个第三层最耗费时间，但往往最能暴露问题。有一次我检查生成结果时发现模型输出了一段"推理"，里面的数字跟题目完全无关，但结论却和标准答案一样。后来我去看训练数据，发现有十几条样本的推理链写得非常敷衍，模型学到的就是"先抄一段，再猜个答案"的投机策略。数据质量对推理模型的影响，比我之前以为的大得多。

5. 从NumPy到生产：我在迁移和迭代中踩过的坑

手写版跑通只是第一步。我评估完模型的推理能力后，开始把它往真实业务场景迁移，这段经历里踩过的坑比写代码本身还要多。总结起来有三类问题最值得注意。

5.1 数值精度：float32 vs float64的差别是致命的

手写时我全程用的np.float64，因为调试阶段精度高、梯度校验方便。但到了迁移阶段，我在PyTorch里默认用的是float32，结果训练稳定性和最终效果出现了肉眼可见的下降。具体表现是：loss下到2.8附近就开始震荡，而且反向传播的梯度在某些层出现了"数值消失"——attention里softmax的指数计算在float32下的舍入误差被放大，深层梯度传到第3层时几乎为0。

我不是说float32不能训，而是要认清它和float64的边界：

• float64适合小规模、高精度要求的实验，比如梯度检查、数值稳定性分析。
• float32适合正常训练，但要保证模型的scale设计合理，比如给attention内部加上scale参数，以及用torch.nn.utils.clip_grad_norm_做梯度裁剪。
• 如果量化到float16甚至int8，那就得配合混合精度训练的各种trick，否则loss发散几乎是必然的。

我在迁移时犯的一个具体错误是：初始化的时候直接拷贝了NumPy版本里float64的权重，然后转成float32，导致前几轮输出分布和预期不符。后来改用PyTorch自己的默认初始化，发现两者最终收敛效果差不多，根本不需要费劲做权重迁移。

5.2 迁移到PyTorch时保留的"校验锚点"

从手写版迁移到框架版，最容易忽视的问题是你不知道自己的实现什么时候开始和原版"分叉"。我的做法是给整个过程设置了一个校验锚点：把同一个随机种子、同一个batch的输入分别喂给NumPy版和PyTorch版，前向传播中各层的输出必须满足allclose，至少前5轮迭代loss的差值要保持在1e-3级别以内。

这个校验过程让我发现了不少隐性bug。比如PyTorch的MultiheadAttention默认会把batch_first设成False，而我的数据是按(B, T, d_model)排列的，如果不指定batch_first=True，运行时能通过但结果完全不同。再比如nn.LayerNorm的elementwise_affine默认是True，但如果你从外部传入了从NumPy导出的gamma和beta，就会发生参数冲突。这些细节不校验根本发现不了。

校验通过之后，性能优化才有了可靠的基础。我的经验是：先确保行为一致，再谈速度提升。如果你在迁移的同时还想顺手改模型结构、优化训练策略，那出了问题根本不可能定位——变量太多，没法排查。

5.3 推理性能优化：KV cache、批处理与并发

迁移到生产之后，纯NumPy实现完全扛不住线上推理的QPS要求。这个阶段我把注意力转向了工程优化，最有价值的一个改造是给模型的解码过程加上了KV cache。

没有KV cache的时候，每生成一个新token，模型都要把之前所有token的Key和Value重新计算一遍。生成100个token的时间复杂度是O(T^2)，token一长，延迟直线上升。加上KV cache之后，只有新token的K和V需要计算，之前的结果直接存起来复用，生成200个token的耗时直接降到了原来的1/3左右。

除了KV cache，我还对推理服务做了批量化改造——把多个用户请求拼成一个batch喂给模型，而不是一个用户一次请求。因为GPU或CPU的矩阵运算在小batch下计算效率很低，batch从1涨到8，总吞吐量能提升5倍以上。但batch size也不能无限大，受限于显存和延迟上限。

最后一个工程细节是并发控制。Python的GIL会让多线程推理变成"假并发"，我采用进程池 + 每个进程独占模型副本的方案，在8核机器上跑出了接近线性的加速比。这部分经验放到今天依然是部署推理服务的基础功课。

提示：KV cache虽然好用，但显存占用会随上下文长度线性增长。如果你的服务面向超长文本场景，建议先用分段式KV cache或者对历史KV做量化压缩，不然很容易出现显存耗尽的问题。

写在最后

从零手写一个推理模型，再到把它的能力一点点逼出来，整个过程花费的大块时间远超我最初的预期，但现在回头看，这笔投入的性价比太高了。我最大的收获不是"我会手写Transformer了"这个标签，而是建立了一套从数据到渲染再到推理部署的全局视野——之后不管用哪套框架，遇到奇怪的问题第一反应不是查文档，而是回到"模型内部到底在算什么"这个层面去推理。

最后分享一个小技巧：调试小型模型时，我最喜欢打印attention的可视化热力图。当你看到在不同训练阶段，某些attention头开始稳定聚焦到特定的位置或词性时，那种"模型开始学习规律了"的实感，比任何指标都来得震撼。如果你也在从零实现某个AI项目，欢迎在评论区分享你卡住的地方，我踩过的坑里也许正好有你要的答案。