从零手搓AI工程:NumPy实现GPT与框架调包对比 1. 从零手搓AI工程为什么我不建议你直接调包第一次看到ai-engineering-from-scratch这个项目名的时候我脑子里蹦出来的画面是一个人坐在黑漆漆的终端前面从矩阵乘法开始一行一行把Transformer搭出来中间不碰任何现成的深度学习框架。后来我自己真去走了一遍这条路才发现它真正的价值不在于“造轮子”而在于把AI工程里那些被框架封装掉的决策点重新暴露出来。这个项目适合谁如果你已经能用PyTorch或TensorFlow跑通一个分类模型但被问到“为什么这里要除以根号d_k”“LayerNorm到底归一化了什么”“KV Cache省下的到底是哪部分计算”时答不上来那这个方向就是给你准备的。它解决的核心问题是把AI工程从“调包经验”变成“可推导的工程直觉”。你不再依赖文档告诉你某个参数该填多少而是能从计算量和数值稳定性两个角度自己算出来。我前后花了大概三周时间用纯NumPy加少量Python标准库从零实现了一个能跑通的小型GPT结构包括分词、嵌入、多头注意力、前馈网络、训练循环和推理采样。整个过程踩的坑比我想象的多但也正是这些坑让我对现在用的框架有了完全不同的理解。下面我把整个思路、关键细节和实操过程拆开讲你如果也想走一遍这条路可以直接照着抄作业。2. 整体设计与技术选型为什么是NumPy而不是别的2.1 核心思路把框架的“黑盒”拆成可验证的白盒从零做AI工程第一个决策就是用什么作为计算底座。可选方案无非三类纯Python列表、NumPy、或者自己写一个极简的张量库。我最终选了NumPy理由有三条。第一纯Python列表做矩阵乘法在规模稍大时慢到无法接受。一个(64, 128) (128, 256)的矩阵乘用嵌套循环要跑好几秒而NumPy底层是BLAS毫秒级完成。训练循环里这种操作成千上万次纯Python根本跑不动。第二自己写张量库虽然更能体现“从零”但会把大量时间花在内存管理和广播规则上偏离了AI工程本身的核心。我的目标是理解模型结构和训练逻辑不是重造一个数值计算库。第三NumPy的API足够底层np.dot、np.exp、np.max这些函数不会帮你做任何“智能”的事。比如softmax的数值稳定性NumPy不会自动帮你减最大值你得自己写。这正是我想要的效果——每一个数值细节都由我负责。提示如果你连NumPy的广播机制都不熟建议先花半天把np.newaxis、keepdims、axis参数搞清楚否则后面调试梯度时会非常痛苦。2.2 模型规模的选择为什么我选了一个“玩具级”配置真正动手前我给自己定了一个硬约束整个模型必须能在一台普通笔记本的CPU上几分钟内完成一次前向传播。这不是为了性能而是为了调试效率。最终配置是这样的词表大小5000嵌入维度128注意力头数4层数4上下文长度64前馈网络隐藏层维度512。参数量大概在200万左右。这个规模下一次前向传播在CPU上大约0.3秒一次完整训练迭代batch size 8大约2秒。我可以快速验证每个模块是否正确而不需要等半天看结果。如果你一上来就搞一个上亿参数的配置调试周期会拉长到无法忍受。从零实现的第一原则是先让一个小模型跑通再谈规模。这一点和用框架时的思路完全相反框架里你习惯先定大配置再调但从零实现时小配置才是你的朋友。2.3 模块划分六个必须独立验证的组件我把整个工程拆成六个独立模块每个模块都有明确的输入输出和验证方法分词器把文本转成整数ID序列再转回来。验证方法是往返一致性。嵌入层把ID序列映射成向量加上位置编码。验证方法是检查输出形状和位置编码的数值范围。单头注意力实现缩放点积注意力。验证方法是手动构造小矩阵对比手算结果。多头注意力把单头拆成多个头再拼接。验证方法是检查头维度是否正确分割。前馈网络与残差两层线性加激活配合残差连接和LayerNorm。验证方法是检查梯度是否消失。训练循环与采样交叉熵损失、反向传播、参数更新、自回归生成。验证方法是看损失是否下降、生成文本是否合理。每个模块单独测试通过后再组装成完整模型。这个顺序非常重要不要试图一次性把所有模块拼起来再调试那样出问题时你根本不知道是哪一层错了。3. 核心细节解析那些框架帮你做了但你必须自己写的东西3.1 数值稳定性softmax为什么要减最大值这是从零实现时第一个真正让我停下来思考的点。softmax的公式是exp(x_i) / sum(exp(x_j))看起来很简单。但如果你直接按公式写当x_i稍微大一点比如100exp(100)就溢出了结果是inf整个计算就废了。框架里你从来不用担心这个因为PyTorch的softmax内部自动做了最大值减法。但从零实现时你必须自己写def softmax(x, axis-1): x_max np.max(x, axisaxis, keepdimsTrue) exp_x np.exp(x - x_max) return exp_x / np.sum(exp_x, axisaxis, keepdimsTrue)减最大值不改变数学结果因为分子分母同时乘了exp(-x_max)但数值上把指数范围压到了[-inf, 0]exp的结果永远在(0, 1]之间不会溢出。这个细节在注意力机制里尤其关键。注意力分数是query和key的点积维度是d_k点积结果会随d_k增大而增大。这就是为什么注意力里要除以sqrt(d_k)——不只是为了梯度稳定也是为了把softmax的输入拉回一个数值安全的范围。注意如果你在实现时发现损失突然变成NaN第一个要检查的就是softmax有没有做最大值减法第二个是log里有没有加极小值。3.2 位置编码为什么用正弦函数而不是可学习嵌入位置编码有两种主流做法可学习的嵌入向量或者固定的正弦函数。我最终选了正弦函数原因有两个。第一正弦编码不需要训练参数在小模型上更不容易过拟合。我的训练数据量不大可学习的位置嵌入很容易记住训练集的位置模式泛化到稍长的序列就崩了。第二正弦编码的数值范围是固定的[-1, 1]和嵌入向量的量级匹配。可学习嵌入初始化时是随机的如果初始化方差没调好位置信息会淹没在词嵌入里。具体实现是这样的def positional_encoding(seq_len, d_model): pos np.arange(seq_len)[:, np.newaxis] i np.arange(d_model)[np.newaxis, :] angle pos / np.power(10000, (2 * (i // 2)) / d_model) angle[:, 0::2] np.sin(angle[:, 0::2]) angle[:, 1::2] np.cos(angle[:, 1::2]) return angle这里有个容易搞错的地方i // 2的整除操作。偶数维度用sin奇数维度用cos频率随维度指数衰减。我第一遍写的时候忘了整除导致sin和cos的频率对不上模型训练损失一直不降。后来打印出位置编码矩阵一看相邻维度的值几乎一样完全没有区分度。3.3 反向传播手写梯度比想象中更需要耐心从零实现最硬核的部分就是反向传播。框架的autograd帮你做了所有链式法则但从零实现时每一个操作的梯度你都得自己推、自己写。我以交叉熵损失为例。前向是loss -log(softmax(logits)[target])反向时梯度是softmax(logits) - one_hot(target)。这个结果很简洁但推导过程需要用到softmax的雅可比矩阵我建议你至少手推一遍否则很容易在实现时把维度搞错。矩阵乘法的反向传播更绕。如果前向是Y X W那么dX dY W.TdW X.T dY。这里的关键是维度必须对齐。我踩过的坑是当X是三维张量(batch, seq, dim)时X.T并不是我想要的转置。NumPy的.T默认反转所有轴而我只想交换最后两个轴。正确做法是np.swapaxes(X, -1, -2)。# 前向 Y X W # 反向 dX dY W.swapaxes(-1, -2) dW X.swapaxes(-1, -2) dY这个细节在批量训练时特别容易出错因为小批量测试时维度是二维的.T恰好正确一旦加上batch维度就错了。我的建议是所有矩阵操作都用swapaxes显式指定轴不要用.T。3.4 LayerNorm归一化的轴选择决定成败LayerNorm看起来简单减均值、除标准差、再缩放平移。但归一化的轴选错了整个模型就学不到东西。在Transformer里LayerNorm是对最后一维特征维度做归一化而不是对batch或序列维度。也就是说每个位置的每个样本独立归一化。我第一版实现时对batch维度做了归一化结果训练损失震荡得非常厉害因为不同样本的统计量被混在一起了。def layer_norm(x, gamma, beta, eps1e-5): mean np.mean(x, axis-1, keepdimsTrue) var np.var(x, axis-1, keepdimsTrue) x_norm (x - mean) / np.sqrt(var eps) return gamma * x_norm betaeps取1e-5是常见做法防止方差为0时除零。gamma和beta是可学习参数初始化为1和0这样训练初期LayerNorm近似恒等变换不会破坏预训练权重虽然我们从零训练但这个初始化习惯要保持。4. 实操过程从空文件到能生成文本的完整流程4.1 环境准备与数据管道环境很简单只需要NumPy和标准库。我用的Python 3.10NumPy 1.24。不需要GPUCPU足够跑这个小模型。pip install numpy数据方面我用了一段公开的英文文本作为训练语料大概5MB。分词器我选了最简单的字符级分词而不是BPE。原因很简单字符级分词不需要训练词表就是所有出现过的字符实现起来只有几十行代码。BPE虽然效果更好但训练分词器本身就是一个独立工程会分散注意力。字符级分词的实现class CharTokenizer: def __init__(self, text): self.chars sorted(list(set(text))) self.char_to_id {c: i for i, c in enumerate(self.chars)} self.id_to_char {i: c for c, i in self.char_to_id.items()} def encode(self, text): return [self.char_to_id[c] for c in text] def decode(self, ids): return .join([self.id_to_char[i] for i in ids])验证往返一致性tokenizer.decode(tokenizer.encode(hello)) hello。这一步必须通过否则后面所有训练都是错的。4.2 单头注意力的实现与验证单头注意力是整个模型的核心。我把它拆成四步线性投影得到Q、K、V计算注意力分数softmax归一化加权求和。def scaled_dot_product_attention(Q, K, V, maskNone): d_k Q.shape[-1] scores Q K.swapaxes(-1, -2) / np.sqrt(d_k) if mask is not None: scores np.where(mask, scores, -1e9) attn softmax(scores, axis-1) return attn V, attn验证方法构造一个2x2的Q和K手动算出注意力权重对比代码输出。我构造的例子是Q和K完全相同那么注意力权重应该是对角线占优。实测结果符合预期。mask的处理是另一个关键点。自回归生成时每个位置只能看到自己和之前的位置所以需要一个上三角为0的mask。我用np.triu生成mask np.tril(np.ones((seq_len, seq_len)))tril生成下三角矩阵1表示可见0表示不可见。在softmax之前把不可见位置设成-1e9softmax之后这些位置就变成0了。提示-1e9不要用-inf因为-inf参与减法会变成NaN。-1e9足够小softmax后是0又不会污染计算。4.3 多头注意力的拆分与拼接多头注意力的思路是把d_model维的Q、K、V拆成h个头每个头维度是d_model / h各自做注意力最后拼接再投影。拆分用reshape加transposedef split_heads(x, num_heads): batch, seq_len, d_model x.shape d_head d_model // num_heads x x.reshape(batch, seq_len, num_heads, d_head) return x.transpose(0, 2, 1, 3) # (batch, heads, seq, d_head)拼接是逆操作def combine_heads(x): batch, heads, seq_len, d_head x.shape x x.transpose(0, 2, 1, 3) # (batch, seq, heads, d_head) return x.reshape(batch, seq_len, heads * d_head)这里最容易错的是transpose的轴顺序。我第一版把(0, 2, 1, 3)写成了(0, 1, 2, 3)结果每个头看到的是错误的维度切片注意力完全失效。验证方法拆分后再拼接结果应该和原张量完全一致。4.4 训练循环损失下降只是第一步训练循环包括前向传播、损失计算、反向传播、参数更新。我用的是最简单的随机梯度下降学习率0.01没有动量。for epoch in range(num_epochs): for batch in get_batches(data, batch_size): logits model.forward(batch) loss cross_entropy(logits, batch) grads model.backward(batch) for param, grad in zip(model.params, grads): param - learning_rate * grad损失从初始的8.5左右log(5000)逐渐下降到2.0左右说明模型确实在学。但损失下降不代表模型有用还要看生成效果。我每隔500步做一次生成测试用模型自回归采样20个字符。初期生成的是乱码中期开始出现常见字母组合后期能生成类似单词的结构。这个过程非常直观比看损失曲线更有成就感。4.5 自回归采样温度参数的实际影响采样时有一个关键参数温度。温度小于1时概率分布更尖锐生成更确定但可能重复温度大于1时分布更平坦生成更多样但可能不连贯。def sample(model, start_ids, max_len, temperature1.0): ids start_ids.copy() for _ in range(max_len): logits model.forward(np.array([ids]))[0, -1] logits logits / temperature probs softmax(logits) next_id np.random.choice(len(probs), pprobs) ids.append(next_id) return ids我实测下来温度0.8左右生成质量最好既有一定多样性又不会太离谱。温度1.5以上基本就是随机字符了。5. 常见问题与排查技巧实录5.1 损失不下降的五个排查方向从零实现时损失不下降是最常见的问题。我整理了一个排查顺序按可能性从高到低排查项检查方法常见错误学习率打印梯度范数学习率太大导致震荡太小导致不收敛梯度计算数值梯度对比矩阵转置轴错误softmax梯度符号错误参数初始化检查初始损失初始损失应接近log(vocab_size)数据管道检查输入输出标签错位mask方向反了归一化检查LayerNorm轴对batch维度归一化而非特征维度我遇到最隐蔽的一个问题是标签错位。语言模型的目标是预测下一个字符所以输入是data[i:iseq_len]标签应该是data[i1:iseq_len1]。我第一版把标签写成了和输入一样模型学的是复制输入损失确实在降但生成时完全没用。5.2 梯度爆炸与裁剪从零实现时梯度爆炸比梯度消失更常见因为我没有用任何自适应优化器。一个简单的梯度裁剪就能解决def clip_gradients(grads, max_norm1.0): total_norm np.sqrt(sum(np.sum(g**2) for g in grads)) if total_norm max_norm: scale max_norm / (total_norm 1e-6) grads [g * scale for g in grads] return gradsmax_norm取1.0是常见值。裁剪后梯度方向不变只缩小模长能有效防止参数更新过大导致的NaN。5.3 内存与速度的平衡纯NumPy实现没有GPU加速内存和速度都需要注意。我的经验是batch size不要超过16否则内存占用会超过2GB。上下文长度不要超过128注意力矩阵是seq_len^2128已经产生16384个元素。前馈网络隐藏层维度不要超过4 * d_model再大收益递减但计算量线性增长。如果速度太慢可以用np.float32代替np.float64速度能提升约30%精度损失可以忽略。5.4 生成结果重复的解决思路自回归生成时模型容易陷入重复循环比如一直输出同一个字符。这是训练不充分和采样策略共同导致的。除了调温度还可以用top-k采样只从概率最高的k个字符里采样。def top_k_sample(probs, k10): top_k_ids np.argsort(probs)[-k:] top_k_probs probs[top_k_ids] top_k_probs top_k_probs / np.sum(top_k_probs) return np.random.choice(top_k_ids, ptop_k_probs)k取10到50之间效果比较好。太小会失去多样性太大和普通采样没区别。6. 从零实现之后我对AI工程的新理解走完这一遍之后我再看PyTorch的nn.MultiheadAttention感觉完全不一样了。以前我只知道它有三个投影矩阵现在我知道每个矩阵的形状、初始化方式、以及为什么d_model必须能被num_heads整除。以前我调temperature参数靠感觉现在我能从softmax的数值范围解释它为什么有效。这个项目最大的价值不是产出一个能用的模型——说实话这个玩具模型生成质量远不如任何现成的小型语言模型。它的价值在于把AI工程从“配置调参”拉回到“原理推导”。当你亲手写过一遍反向传播你就不会再对loss.backward()感到神秘当你亲手处理过softmax溢出你就不会再对log_softmax的存在感到困惑。如果你也想走这条路我的建议是不要追求一次写对要追求每次出错都知道为什么错。我前后重构了三次第一次是维度错误第二次是梯度错误第三次是数值稳定性问题。每一次重构都让我对模型的理解更深一层。这个过程中积累的调试直觉是任何框架文档都给不了的。最后分享一个我踩过的小坑在实现位置编码时np.power(10000, ...)里的10000不要写成1000或者100000这个值决定了不同维度频率的衰减速度改大了长序列的位置区分度会下降改小了短序列的相邻位置编码太相似。10000是经过验证的平衡点直接用就好。