写这篇文章之前，我想先把一个很常见的问题说清楚：市面上有大量教程教你“调用API”“加载现成模型”，但你一旦想自己动手构建一套AI推理系统，哪怕是做一个最小规模的演示模型，就会立刻发现信息断层。我最初也是从“调包侠”开始的，后来花了将近四个月时间，完整走了一遍从数据准备到训练再到推理部署的全链路，才敢说自己对“AI工程”这件事有了基本的体感。
这篇文章的核心，就是围绕“从零构建AI推理模型（reasoning model）”这条主线，把我在实际动手过程中踩过的坑、验证过的方法、以及那些书本上不会详细讲的细节全部摊开来。结构上会从整体设计讲起，接着是数据处理、核心模型实现、训练环节、推理链路，最后是问题排查清单。适合已经会用PyTorch做简单分类任务、但还没完整构建过语言模型的读者，也适合想要理解“模型内部到底发生了什么”的工程师。
1. 项目整体设计与思路拆解
1.1 “从零构建”到底构建的是什么
严格来说，“从零构建AI工程”并不是让你把CUDA内核从头写一遍，也不是拒绝使用任何现成库。更合理的定义是：不依赖现成的预训练模型和推理框架（比如直接load一个ChatGLM或Llama的权重来用），而是自己完成数据构造、模型架构实现、训练流程编写、推理逻辑设计这几个核心环节。
我给自己定的目标是做一个“能完成简单多步算术推理”的迷你推理模型。为什么选这个目标？因为算术推理足够简单——不需要大规模语料，不需要分布式训练，单张消费级显卡就能跑；同时它又足够复杂——模型必须学会理解问题格式、按步骤生成中间结果、最终输出答案，这恰好覆盖了推理模型中“链条生成”最核心的机制。
项目的技术栈最终确定为：Python + PyTorch + HuggingFace Tokenizers + 自研Transformer实现。没有使用transformers库的Trainer，因为我想亲手控制每一个训练细节；但分词器借用了HuggingFace的tokenizers库，因为从零实现一个BPE分词器虽然不难，却非常耗时，而且对核心学习目标帮助不大。这个取舍我想单独说一下：从零构建不是“什么都自己造”，而是“关键部分必须自己造，辅助部分可以借助成熟的轮子”。
1.2 为什么选择“最小可行模型”路线
很多人一上来就想着复现Llama 7B、复现GPT-3，这是完全错误的起点。我的经验是：先做一个参数量在10M到50M之间的小模型，把整个工程链路打通，然后再去扩展规模。
这里有一个非常实际的计算逻辑。假设我要做一个30M参数量的模型，用Adam优化器，参数量乘以4就是梯度、动量等额外显存占用，即120MB；加上中间激活值，一个batch的显存开销大概在200MB到500MB之间。这意味着单张24GB显存的显卡可以非常从容地训练。但如果你直接上手7B模型，光模型权重就要14GB（半精度），一个batch的激活值随便超过40GB，优化器状态更是直接爆显存。即便你有A100，排错成本也会高得让人崩溃。
所以“最小可行模型”的本质逻辑是：把工程复杂度控制在可处理范围内，让主要精力聚焦在正确性验证上。模型小，训练快，迭代周期从“天”缩短到“分钟”，这是任何从零构建项目最需要的加速器。
1.3 架构选型背后的取舍逻辑
我最终选择了标准的Decoder-only Transformer架构，没有做任何架构上的“创新”。原因很朴素：在从零构建的阶段，你追求的是“验证我理解原理”，而不是“发明新架构”。Decoder-only架构是最稳的，因果注意力天然适合逐token生成，推理逻辑也最简单——每次只预测下一个token。
配置上我参考了GPT-2 Small的相对比例并进行缩减：6层Transformer块，8个注意力头，嵌入维度256，前馈层维度1024。参数量大约在25M左右。为什么是8个头而不是4个或16个？因为256维的嵌入除以8个头，每个头分到32维——这个数值经过大量实验验证是合理的。头数太多，每个头的维度太小，注意力表达力不足；头数太少，多头机制带来的多样化特征提取效果就会打折扣。
需要强调的是，对于序列长度，我选择了128而不是更长。因为我的训练数据是短文本（算术题+推理过程），128的长度完全够用，而且更短的序列代表更小的注意力计算开销和显存占用。这也是一个常见误区：很多人一上来就用512甚至1024的序列长度，但在小模型阶段完全没有必要。
2. 数据构造与分词器：最容易被低估的环节
2.1 训练数据生成：不是“找数据”而是“造数据”
从零构建推理模型，第一道坎就是数据。网上能下载到的开源SFT数据大多是中文聊天对话或者英文通用指令，直接拿来训练算术推理模型效果很差，因为领域不匹配。我最终采用了“规则模板生成数据”的方式。
具体做法是：用Python脚本随机生成不同难度的四则运算题，包括两位数加减法、一位数乘两位数、多步混合运算等。每道题不仅生成问题和答案，还生成一步步的推理过程。比如问题“23 + 45 × 2等于多少”，推理过程就是“先计算45乘以2，得到90；然后计算23加90，得到113；所以答案是113”。
这一步看似简单，实际上决定了整个项目的上限。模型不会魔法般地学会它没见过的推理模式，训练数据的推理链质量直接决定了模型的推理能力。我花了整整一周时间调整数据生成逻辑，包括：统一术语（“先计算”vs“先算”会造成token分布分散）、控制数值范围（避免出现超出模型位数能力的长数字）、加入干扰项（让问题更接近真实场景）等。
2.2 分词器的构建与词表大小选择
分词器我用了BPE算法，训练语料就是上面生成的几十万条算术题和推理链。词表大小选择了2000，这个数字背后有一个权衡逻辑：词表太小，数字会被切得很碎（比如128被切成“1”“2”“8”三个token），模型学起来吃力；词表太大，嵌入层参数量会爆炸——以256维嵌入和20000词表计算，光嵌入层参数就是256乘以20000等于512万，占了整个25M模型参数的五分之一还多。
2000词表的好处是：常用数字（0到999）可以整体编码为单个token，运算符和中文连接词也都有独立token，嵌入层参数量只有51万，占比约2%，非常健康。
这里有一个实操经验想分享：训练BPE分词器时，要确保数字token不被切碎，可以在训练前对语料中的数字做保护处理——在数字前后加上特殊标记，BPE就不会把数字切开。这个细节我最初没有注意，导致模型生成“128”时经常输出“1”“2”“8”三个独立token，推理结果虽然语义上对，但格式完全错乱。
2.3 数据格式设计：让它看起来像一个“推理链”
我最终把每条训练样本的格式设计如下：
问题：23加45乘以2等于多少 思考：需要先计算乘法，45乘以2等于90。然后计算加法，23加90等于113。 答案：113这个格式的灵感来自一些开源reasoning model的训练格式。关键在于“思考”步骤要真正展示逐步推理的过程，而不是直接给答案。模型在训练中不是背答案，而是学习“遇到乘法优先计算”这类规则，并把规则外推到它没见过的新题目上。
我测试过一个有意思的现象：如果训练数据中只有“问题-答案”对，模型对简单加减法可以过拟合，但稍微变换一下数字顺序就完全崩溃；加入“思考”步骤后，模型的泛化能力明显提升，甚至能处理训练集里完全没有出现过的数字组合。这就是推理链数据的核心价值——它教会模型的不是“这个题等于113”，而是“这类题要按怎样的步骤来解”。
3. 模型架构实现：写一个属于你自己的Transformer
3.1 Transformer核心组件的代码实现要点
我用PyTorch从零实现了每个Transformer块，包括多头注意力、前馈网络、层归一化和残差连接。虽然代码量不算大，但有几个细节值得展开说。
先看多头注意力的实现。与单头注意力不同，多头注意力需要把输入投影到多组Q、K、V空间，分别计算注意力后再拼接。代码上不能用简单的for循环遍历每个头（虽然可以跑，但效率极低且反向传播开销大），而是要用矩阵重塑：把batch size乘以序列长度乘以嵌入维度，重塑成batch size乘以序列长度乘以头数乘以每头维度的四维张量，一次性完成所有头的计算。
因果掩码也是容易出错的地方。很多初学者会在注意力分数上加一个极大的负值（比如负1e9）来掩盖未来位置，但如果你用的是PyTorch的scaled_dot_product_attention，它会自动读取attn_mask参数，其中True的位置会被掩蔽。我在调试时发现的问题不是“有没有掩码”，而是“掩码维度对不对”——正确的维度是batch size乘以1乘以序列长度乘以序列长度，很多人在这一步少了一个维度，广播机制会自动补上，但补上之后语义就错了。
前馈网络部分用的是标准的“升维-激活-降维”结构：输入维度256升到1024，经过GELU激活函数，再降回256。GELU激活函数相比ReLU在小模型上效果略好，主要体现在梯度流动更平滑，训练初期的loss下降更稳定。这不算什么大发现，但实测下来确实少了很多“训练到一半突然loss爆掉”的情况。
3.2 参数初始化与优化器选择：这些细节决定成败
参数初始化是教科书里最少提、但实际影响最大的环节之一。我用的是均值为0、标准差为0.02的正态分布来初始化嵌入层和输出层；对于注意力内部的QKV投影，用的标准差更小，约0.02除以sqrt(每头维度)，这是为了控制初始注意力分数的方差，避免softmax过早饱和。
优化器方面，我选择了AdamW而不是原生Adam。两者的核心区别在于权重衰减的处理方式——Adam是把权重衰减和梯度动量混合在一起，AdamW则是把权重衰减单独作用在参数上，不参与动量计算。这个修正带来了明显的训练稳定性提升，在25M这个规模上也许感知不强，但如果你后续扩展到百M以上，AdamW几乎是必需品。
学习率调度我用的是warmup加余弦退火的组合。前1000步从0线性升到峰值学习率（1e-3），然后按照余弦曲线衰减到峰值的十分之一。为什么不直接用一个固定学习率？因为在训练初期，模型参数离最优解很远，太小的学习率会让训练极其缓慢，但此时梯度的方向波动又大，太大的学习率容易让loss发散。warmup就是为了解决这个“既要快又要稳”的矛盾。峰值学习率选1e-3也是基于小模型的常见经验值：模型越小，越可以使用较大的学习率；如果你用了一个千M级模型，1e-3大概率会直接loss发散。
3.3 损失函数：交叉熵背后的那些“小九九”
语言模型的损失函数是标准的交叉熵，但它有一个容易被忽略的实现细节：掩码。在训练数据中，几乎每条样本的长度都不完全一致，所以需要pad到统一长度（128）。在计算损失时，必须把这些pad token对应的位置遮蔽掉，否则模型会疯狂学习“预测空白”，实际的有效损失被稀释。
PyTorch的CrossEntropyLoss自带ignore_index参数，把pad token的索引传进去就行。但这又引出一个问题：pad token的索引是多少？这取决于你是如何构造batch的。如果使用DataLoader的collate_fn自定义padding，pad token的索引必须与loss计算中的ignore_index保持一致。听起来很基础，但我确实在这里吃过亏——第一次跑通时loss确实在下降，但生成结果全是乱码，排查了一整天才发现是pad token不一致导致的。
另外，标签是输入序列右移一位得到的。对于每一个位置，模型要预测的是“下一个token”，所以输入的第i个位置对应的标签是第i加1个位置的token。实现上，通常把输入序列作为模型输入，标签为输入序列从第1个位置开始的切片，并在末尾补一个eos token。这个“右移一位”的细节，决定了模型学到的到底是“预测下一个词”还是“复制当前词”，很多初学代码的人在这里前后端错位，导致模型学到了一个看似合理但完全无用的映射。
4. 训练循环与实验管理：从“能跑”到“跑得好”
4.1 训练循环的完整画布：一个batch的生命历程
一个完整的training step，从头到尾经历了这些阶段：从数据加载器中取出一个batch的input_ids和labels；把input_ids传入模型，得到logits；计算logits与labels之间的交叉熵损失；反向传播求梯度；梯度裁剪；优化器更新参数；学习率调度器更新当前学习率。
Dataloader是这里第一个性能瓶颈。我最初用的方式是每个epoch动态生成新的算术题，生成过程本身需要字符串拼接和分词器编码，耗时较长，导致GPU经常空闲等待。后来我改成预生成一个大文件（约50万条样本），分批读取，训练速度提升了三倍。这个优化虽然朴素，但效果非常直接。
另一个GPU利用率低的原因是没有开启batch的并行加载。PyTorch的DataLoader默认num_workers等于0，意味着数据加载在主进程中进行；设置为4或者8，数据加载在多进程中并行进行，GPU等待时间大幅缩短。但要注意，num_workers不是越大越好，过大会导致CPU成为瓶颈，甚至内存溢出。日常经验值是CPU核心数的一半左右。
梯度裁剪是我强烈建议加上的一步。把梯度的全局范数裁剪到1.0，作为所有上限。它不改变模型架构，不改变数据，但能有效防止个别异常样本导致的梯度爆炸。尤其是训练初期，模型尚未稳定，偶尔会出现loss突然跳到正常值十倍以上的情况，如果没有梯度裁剪，整个训练可能就此崩溃。
4.2 超参数调试：我的“黄金组合”和失败记录
经过多轮实验，我总结了一个在小模型上表现稳定的“黄金组合”：batch size 128，序列长度128，学习率峰值1e-3，warmup 1000步，训练10个epoch。在这个配置下，训练损失通常能从起始的7左右稳定下降到0.3以下，生成质量已经比较像样。
但这条路上我踩了很多坑。第一次训练时，我没有加warmup，直接以1e-3的学习率起步，结果开局300步内loss就从7冲到12以上，整个训练完全发散。我当时还以为是模型实现有问题，排查了两天，最后才意识到是学习率策略的问题。
还有一次，我把batch size从128调大到512，本意是提高训练稳定性、加速收敛，结果发现loss虽然降得更快，但生成结果出现了明显的重复和死循环现象。后来我意识到：大batch会降低梯度的噪声，在训练初期反而减少了探索能力，模型容易陷入局部最优。对小模型而言，128到256之间的batch size是性价比最高的区间。
另一个反直觉的经验是：训练轮数不是越多越好。在25M参数规模下，模型大约在第8个epoch之后就开始过拟合训练数据——具体表现是训练loss继续下降，但验证集（数据生成器按不同随机种子生成的新题）上的准确率开始持平甚至下降。所以我在项目里加了早停机制：验证loss连续3个epoch不下降就提前终止训练。最终模型在验证集上的推理正确率约为87%，对一个小模型来说已经相当理想。
4.3 训练监控：如何判断“模型正在变聪明”
判断训练是否健康，不能只看loss。Loss下降只是必要条件，并不充分——完全可能出现loss下降但模型学到的是“复读机”模式。我建立了一套多维度监控体系：
- 训练loss和验证loss曲线：两者同步下降是健康的信号；如果训练loss降但验证loss涨，就是过拟合前兆。
- 梯度范数：正常训练中梯度的全局范数应该围绕某个均值波动，如果频繁触达你设定的裁剪上限，说明优化过程很激进，可能需要降低学习率。
- 实际生成样本：每隔500步，让当前模型在几个固定问题上做推理，把输出打印出来。这是最直观的“模型到底学会了什么”的观测方式。早期模型的输出往往是乱码，中期开始出现“思考：”的前缀，后期能输出完整但偶尔算错的推理链。
- 注意力热力图：我在调试阶段周期性可视化某个特定head的注意力权重。合理状态下，模型在“答案：”这个token之后应该重点关注到问题中涉及的数字token；如果注意力权重完全分散，说明模型还没有建立起输入输出的对应关系。
5. 从语言模型到“推理模型”：推理链的工程化实现
5.1 让模型学会“思考”：数据构造与模型行为的联动
单纯的Decoder-only模型经过上面的训练，已经能够生成相对连贯的文本，但它本质上还是一个“文本续写器”，并不具备真正的推理能力。让“续写”变为“推理”的关键，藏在训练数据的格式设计和推理时机的使用策略上。
我的做法是在推理时显式地输出“思考：”前缀，然后用流式解码方式逐步生成，直到遇到“答案：”标记。这一设计让模型在生成推理链时处于一个“逐步思考”的状态——它已经学会看到“思考：”之后输出的是中间推导步骤，而不是直接跳到答案。
这里不得不提到一个关键概念：温度采样。在推理阶段，我使用temperature等于0.7的softmax采样，而不是贪心解码（temperature等于0）。为什么？贪心解码每次选择概率最高的token，结果稳定但容易陷入重复循环；少量随机性可以让模型偶尔探索其他路径，产出的推理链更多样。但温度太高也会引入错误——大于1.2时模型会频繁输出毫无意义的组合token。0.7到0.9是我在实际测试中效果最好的区间。
另外，我也在推理阶段加入了“最大步数限制”。因为模型在长推理链上可能陷入循环生成同一个token序列，此时需要硬性截断。我的经验值是最多生成256个token，超过则强制终止并输出“生成超时”。这道保险看起来简单，但如果没有它，某些异常输入会导致推理进程进入无限循环，卡死整个服务。
5.2 从监督微调到RLHF：推理模型的进阶路线
从“能生成推理链”到“能稳定生成正确的推理链”，中间还有一段路。监督微调（SFT）教会了模型推理的“格式”，但格式不等于质量。模型仍然经常在推理链中出现“中间步骤正确但最终答案错误”的情况，或者在第一步就算错但后续步骤严格执行了这个错误结果。
要解决这个问题，业界通常走两条路线：一条是拒绝采样，一条是强化学习（RLHF/GRPO）。拒绝采样的思路非常工程化——用当前模型生成多条推理链，只保留最终答案正确的那些样本，再用这些高质量样本做一轮SFT。我在实际测试中发现，这个简单的方法效果相当显著：一轮拒绝采样微调后，验证集正确率从87%提升到了91%。
至于强化学习，我在迷你模型上做过一次GRPO实验。GRPO的核心是用一组采样输出的相对优劣来构造奖励，替代传统PPO中需要额外Critic模型估计价值函数的方式。在简化实现中，我给每个正确样本奖励加1，错误样本加0，并用KL散度惩罚模型偏离初始SFT模型过远。训练感官上比SFT要微妙得多——奖励信号非常稀疏，即使正确率只有百分之十几，模型依然能缓慢改进。但这套体系在小模型上的投入产出比并不高，如果你的目标就是在小模型上验证工程链路，SFT加拒绝采样已经完全够用。
5.3 推理服务的工程化：从模型到可用产品
模型训练完成只是项目的一半，另一半是把它包装成一个可以被外部调用的服务。我使用FastAPI搭建了一个轻量级推理服务，对外暴露一个“请回答数学问题”的HTTP接口。
请求来了之后，推理服务内部执行这几个步骤：接收问题文本，用训练好的BPE分词器编码，送入模型逐步生成推理链和答案，最后把结果解析成结构化格式返回。最大的挑战是batch推理的效率。传统一次请求一个样本的推理方式，GPT的KV Cache利用率很低。我在服务端实现了简单的连续batch（continuous batching）机制：多个请求共享同一个GPU推理循环，某个请求生成完了就立即插入新请求。虽然实现起来增加了复杂度，但吞吐量提升了大约两倍，在演示场景下效果明显。
这里我强烈建议任何做从零构建项目的朋友，在模型训练之外至少留一周时间做服务化改造。因为本地跑notebook和面向真实请求提供服务之间有一条巨大的鸿沟——超时控制、异常输入处理、并发保护、结果校验，这些工程问题只会在服务化阶段暴露出来。
6. 常见问题与排查技巧实录
6.1 六个最容易踩的坑
| 问题现象 | 根本原因 | 解决方法 |
|---|---|---|
| 训练开始后loss快速发散 | 学习率过高或缺少warmup | 调低峰值学习率，加1000步warmup |
| 模型生成大量重复循环文本 | 训练数据推理链模式单一或温度过低 | 增加数据多样性，推理时temperature设为0.7以上 |
| 训练loss下降但生成完全乱码 | pad token与ignore_index不一致 | 核对collate_fn中的pad token索引与损失函数参数 |
| 数字经常被拆分生成 | 分词器没有保护数字完整性 | 训练BPE前对数字加保护标记 |
| 偶数epoch效果反而变差 | 过拟合，验证准确率出现“倒U型”曲线 | 加入早停机制，在验证loss最低点保存模型 |
| 推理服务偶发超时卡死 | 模型输出进入无限循环 | 设置最大生成步数，触发后强制截断 |
6.2 排查“模型不聪明”时的系统化思路
如果模型训练流程没有任何报错，但生成结果始终不满意，很多人会立刻怀疑模型实现有bug。我的经验是先做一轮“系统化体检”，而不是盲目调参。
第一步是过拟合测试。拿一小批训练样本（比如100条），在相同的batch下反复训练。如果loss在几个epoch内降到接近0，说明模型实现正确、容量足够、能记住数据；如果这个都做不到，那就不是过拟合或数据问题，而是模型实现本身的bug，应该回到代码层面排查。
第二步是检查生成格式。让模型输出一个固定前缀的续写，比如只给“思考：”三个字，看模型会接什么。如果模型能接出与训练数据风格一致的文本，说明预训练阶段基本成功；如果接的是乱码，那就回到分词和tokenizer的问题上。这个测试能快速定位问题是出在“数据链路”还是“模型生成策略”。
第三步是逐层检查梯度。如果loss不降，用hook打印每一层权重的梯度范数。正常情况是整个Transformer块的梯度范数依次递减（越靠近输入层越小，因为要经过更多层级传播）；如果发现某一层梯度范数为0或者突然暴涨，说明那一层有实现问题。这个方法在调试深度网络上比任何理论分析都管用。
6.3 我最后悔没早做的三件事
第一件是版本控制实验记录。最初我改了超参数只改代码里的变量，不做记录，导致同样配置下跑出的结果不一致时根本无法追溯是数据、代码还是环境变了。后来我强制自己用实验记录表，每次训练把数据版本、代码commit号、超参数、最终指标全部存下来，排查效率至少提高了一倍。
第二件是基准测试。我在做算术推理实验的同时，应该尽早建一个包含固定100道题的评测集，每次训练后都跑一遍，记录正确率。我一开始只靠印象判断模型好坏，后来才发现，没有基准测试凭感觉判断的“变好了”有很大主观成分，有时只是幸运抽到了简单题。
第三件是及早进行小规模模型验证。我不止一次犯同一个错误：在大模型上调整超参数浪费数小时，最后发现同样的logits在小模型上几分钟就能验证。正确做法是先在参数量缩减10倍的微型模型上验证所有可疑项，再上全量模型。这条经验在后续任何模型实验中都是最高效的路径。
结尾：一点个人的心得
整个项目做下来，我最大的感触是：从零构建不是为了“造一个更好的模型”，而是为了获得一种无法被替代的“工程体感”。在使用开源模型时，你看到的都是封装好的接口，遇到问题时的排查路径完全依赖别人的文档；自己构建过一次之后，对你而言这些概念变成了“我亲手写过的那个张量”“我调试过的那条梯度流向”。
最后分享一个特别实用的小技巧：如果你跟我一样在从零构建的路上，请务必保留一份“微型复现脚本”——把模型规模压缩到最小可训练状态，只保留最关键的数据逻辑和网络结构。当你改了某个数据字段或网络细节后，先跑这份微型脚本验证正确性，再上全量训练。这个习惯帮我节省的时间，保守估计足够再读完两本书。项目本身还有很多可以继续深化的方向，比如把推理链做长、引入工具调用、或者把SFT数据扩展到更多领域，但核心链路已经通了，后续的扩展都只是工程和时间问题。
