说实话我入行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), dtypebool), k1) scores np.where(mask, -1e9, scores) weights softmax(scores, axis-1) out weights v第三步把每个头的结果拼接回去再过一层输出投影。网上很多教程到这里就停了但实际手写时会遇到三个课本上不会提的陷阱掩码位置必须用负无穷而不是零。如果用0softmax后那些位置的权重不是零是exp(0)/sum(exp(x))信息照样泄漏。我第一次就是在这里埋了坑训练出的模型loss降得很慢生成时前后矛盾严重。QK^T之后如果做过mask不能再对scores做任何加性操作否则负无穷会被污染成普通数值导致因果性失效。head_dim的选择最好是64或128的约数同时和d_model整除。我一开始用d_model128, n_heads8, head_dim16效果正常后来试过d_model192, n_heads7发现整除不了最后硬着头皮用n_heads8但每个head维度不均匀代码里多了一堆reshape的破事。多头注意力的意义在于每个头可以关注不同的关系模式比如一个头关注语法位置另一个头关注指代关系。这些特性不是我刻意设计的而是训练之后自然涌现的——这也是手写一遍之后你能直观体会到的奇妙之处。2.4 前馈、LayerNorm与残差为什么顺序不能乱注意力之后是前馈网络和LayerNorm它们的配合顺序是Transformer结构中最细节敏感的部分。我的实现采用的是post-norm结构也是原始论文里的方案顺序是输入 - Attention - 残差相加 - LayerNorm - FFN - 残差相加 - LayerNorm。LayerNorm的公式很简单但它的作用我一直到调参时才真正理解def layer_norm(x, gamma, beta, eps1e-5): mean x.mean(axis-1, keepdimsTrue) var x.var(axis-1, keepdimsTrue) 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对比一次反向计算出梯度和用前向计算手动推导的梯度是否一致。第三种我用到的取巧做法用数值梯度校验。给定一个极小扰动h1e-6用(loss(xh) - loss(x-h)) / (2h)近似梯度跟手写的解析梯度对比。这个方法验证单个参数时非常高效但全参数校验在模型稍大时慢得没法看。所以我的策略是先校验单层模块attention、LayerNorm各自正向逆向再组装校验两层叠加最后再全模型训练。每一步都确认无误再往下走。这里有一个经验教训永远不要直接写完整模型的梯度再一次性跑。除非你天赋异禀否则99%的概率会出现某一处维度错、符号反、或者mask没有参与梯度计算的问题。分模块验证最多多花半天时间但能救你一周的Debug。3.2 显存/内存的真实天花板纯NumPy在CPU上训练最大的瓶颈不是计算速度而是中间张量的内存占用。假设batch_size32, seq_len64, d_model128一个batch的前向过程中Q/K/V分别是32*64*128注意力权重是32*8*64*64乘以每个头的独立计算这一层的激活值就有几十MB堆到6层之后几GB是常态。我实际测量过一组数据配置参数量前向中间激活峰值训练总时长(5000步)6层, d128, 8头约1100万约1.2GB4.5小时4层, d128, 8头约730万约0.8GB2.8小时6层, d64, 4头约600万约0.4GB2.1小时所以如果你用笔记本跑建议直接从第三套配置起步。降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-hotloss下不去太小则梯度消失学习几乎停滞。如果模型还是训不动我建议按这个顺序排查先看输入数据是否normalize过token id是否从0开始连续编号如果中间有空位embedding层会学到无意义的向量。打印前向时attention权重的均值——正常情况下应该近似均匀分布1/T如果出现某个token占据了80%以上的注意力就是初始化或scale出了问题。检查梯度范数前5步梯度范数如果超过10说明网络已经开始爆炸先降学习率到1e-4再观察如果前5步梯度范数接近0说明信号传不进去检查残差和LayerNorm的顺序。最后再看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的简化版。原理不复杂不直接给模型问题, 答案这类训练样本而是给问题, 逐步推理过程, 答案三段式样本让模型学会先思考再回答的内部模式。我在自己收集的数学题数据集上做了三组对照实验每组都训了相同步数训练方式直接回答准确率链式思考准确率生成Tokens/问题直接问答12%10%5隐式CoT只给答案不显式要求推理过程15%14%6显式CoT要求先输出推理步骤再给答案11%31%32数据非常有意思显式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_firstTrue运行时能通过但结果完全不同。再比如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项目欢迎在评论区分享你卡住的地方我踩过的坑里也许正好有你要的答案。