从零构建推理模型：AI工程实战路线图

做 AI 工程的人这几年几乎都被同一个问题问过“现在 API 贵是贵但至少能跑你为什么还要自己写模型”我最初回答这个问题时也只会搬出“理解原理”这种略显空泛的理由。直到去年我认真动手做一个 from scratch 的推理模型实验——不依赖任何大模型框架的预训练权重从空白工程开始实现数据管线、tokenizer、Transformer 训练和推理服务才真正意识到所谓“AI engineering from scratch”价值根本不在“重新发明轮子”而在于它会把你的工程判断力逼到一个完全不同的层级。这也是为什么当我看到build a reasoning model from scratch这类话题和《Build a Large Language Model from Scratch》被越来越多人讨论时我并不意外——大家不是想要自己重新做一个 GPT而是想在黑盒时代重新夺回对模型的控制感。这篇文章我尽量不写成原理教科书而是按我自己动手时的真实路径来展开先想清楚“为什么值得从零写”再讲数据与分词器、Transformer 骨架、训练曲线、推理优化最后把我踩过的坑和一份可落地的路线图完整交给你。适合那些正准备写第一个 from-scratch 项目、或者已经在做但卡在某个环节的工程师。1. 为什么在满地 API 的时代我仍然坚持从零写一个模型1.1 从零开始到底值不值我的判断标准我给自己定的判断标准很简单如果项目的目标只是“把一个任务跑通”那直接用现成 API 和开源权重就够了一旦你的目标变成“理解任务为何跑通、为何跑不通、以及能否按自己的意图改造模型”从零开始就是唯一可靠的路径。拿我这次做的推理模型来说我需要模型在给定前提和问题后不只是输出一个“看起来合理”的答案还要能在中间生成可解释的推理轨迹。用 API 模型当然可以实现但我无法探测模型内部对每一步推理的置信度也无法控制它在某个中间环节是否真正“遵循了前提约束”。这些问题在推理 API 上几乎无解——你只能调整 prompt然后祈祷。而从零开始训练的小模型虽然绝对能力远不及商用 API但它的一切行为都是可拆解的我可以随时插入探针观察某一层 attention 的分布可以修改解码器让它在推理中途暂停并重试也可以直接干预采样过程来抑制跳跃性结论。这种“可操控性”是 API 永远给不了你的。1.2 除了“理解原理”还有三笔隐性收益第一笔收益是调试能力。使用黑盒 API 的时候遇到一个逻辑错误你大概率只能反复改 prompt但当你从零写过一个训练脚本你会自然建立起一套排查直觉loss 不降先怀疑数据还是模型重复输出先看采样参数还是训练语料这种底层判断力会在未来任何 AI 工程任务里持续发挥作用。第二笔收益是成本结构的变化。推理 API 按 token 计费做实验时一天几百万 token 的消耗是非常真实的成本压力量。自建一个 1 亿参数的小模型在单张消费级显卡上训练和推理主要开销是电费和显卡折旧。对需要大量试错的研究型工作来说这种“固定成本换边际成本趋近于零”的模式会极大解放你的实验空间。第三笔收益是架构自由。如果你需要尝试某种新机制——比如给 attention 加一个文档级 mask、自定义 KV cache 的淘汰策略、或者在 loss 里加入推理一致性约束——在别人的代码库上做深度改动往往比从零开始写一个新工程还要痛苦。因为成熟的代码库有很多隐形的工程假设你要先逆向推断它们的设计意图再小心翼翼地在不破坏其他模块的前提下改动。从零写所有边界都在你自己的掌控范围内。2. 万里长征第一步先搞定数据和小尺寸分词器2.1 别急着上大语料先建立一个干净的实验环境很多人一上来就注册去下载几百 GB 的大语料这其实是最容易掉进坑里的选择。T0 阶段的目标不是刷指标而是确保代码管线端到端可跑数据进入模型、loss 正常下降、生成一个通顺的句子。带着这个目标我的建议是先准备一份 5 万条以内的小型干净语料比如百科条目、文学作品段落或者代码文件。我自己的实验语料是大约 3 万条中文百科条目和 2 万条代码片段总大小不到 300MB训练一个 1 亿参数模型大约只需要 8 小时。这份小语料还有另一个作用它让你能快速做完一个完整的实验循环。数据过长时单个 epoch 就要跑几天你会发现自己的时间全都耗在等待和看日志上根本来不及迭代调参。而小语料允许你在一个工作日内完成“改代码-重训-看结果”的循环这对排查 bug 和积累经验来说太重要了。2.2 用 BPE 分词器给语言“切块”文本无法直接喂进神经网络得先切成整数 id。这一步的工程选择会直接影响后续训练效率和最终生成质量。我推荐自己训练一个 BPEByte Pair Encoding分词器而不是直接套用现成的开源分词器因为 from-scratch 项目讲究的是“每一个组件都可控”。BPE 的原理其实不复杂它从字符级开始反复统计语料中最高频的相邻 token 对把它们合并成一个新 token直到达到预设的词表大小。我在实现时把词表大小定在16,000 到 32,000之间。这个区间的理由很实际词表太小句子的 token 数会膨胀训练和推理速度都会受影响词表太大则嵌入矩阵参数量会急剧上升小模型根本学不过来。比如 32,000 词表配合 256 维嵌入仅 embedding 层就已经是 800 万参数占到一个 1 亿参数模型的 8%。词表大小平均每 token 覆盖字符数嵌入层参数量 (dim256)适合场景8,000约 1.4200 万极小模型快速验证16,000约 2.0410 万1 亿左右参数模型推荐32,000约 2.6820 万生成文本自由度要求较高时64,000约 3.11,640 万不建议参数量挤占 Transformer 层容量在工程实现上要特别注意 BPE 训练时的pre-tokenization 规则。中英文混合语料里一个常见问题是中文整字被拆成多个子词导致有效信息密度下降。我的经验是把中文字符视作独立的 pre-token让 BPE 在更“粗”的粒度上合并英文则按空格和标点先切词这样可以显著减少 token 数量训练速度能提升 30% 以上。2.3 数据质量的底线去重、过滤与采样窗口数据管线很容易被当作“不重要”的环节但实际训练时大量诡异现象都源于数据问题。我在第一版实验里脚本去重没写好结果语料里混进了几千条完全重复的记录直接导致模型对特定片段过拟合生成时反复输出一样的内容。我的数据预处理分了四步首先做精确去重对所有文本计算哈希删除完全重复的条目其次做模糊去重用字符 n-gram 集合的相似度过滤掉近似重复的内容然后做规则过滤去掉含有 HTML 标签、超长 URL、或连续特殊字符的无意义文本最后做语言检测确保语料标签与实际语言一致。这四步做完语料体积通常会缩小 15% 到 30%但训练稳定性的提升非常明显。训练样本的构造也要花心思。最朴素的做法是把每条文本单独 padding 到固定长度但这样做 token 利用率很低大量计算花在 padding token 上。更高效的方式是流式拼接把清洗后的文本按顺序拼成一个长流按固定窗口长度比如 512 或 1024 token切块并强制切块边界对齐到文档边界附近。这样每个训练样本几乎没有 padding有效 token 利用率可以提升到 95% 以上。同时我会保留一个“文档边界 mask”在注意力计算时禁止 token 跨文档相互关注防止模型学到莫名其妙的跨文本关联。3. 骨架设计把 Transformer 拆开我到底在构造什么3.1 把 Transformer 按网线接口的方式理解如果说数据处理是仓库管理那么 Transformer 骨架就是整个工程的承重墙。我不打算逐行贴代码因为这类代码在很多开源仓库里都能找到我更想说的是当你要从零实现一个可用的 Transformer 时哪些模块必须亲手写、哪些地方最容易写错。一个最小可用的 decoder-only Transformer 由以下几部分组成token embedding把整数 id 映射为向量、位置编码我用的 RoPE 旋转位置编码它对序列长度外推更友好、多层 decoder block每层包含 masked multi-head self-attention、feed-forward network、RMSNorm 和残差连接最后是输出投影层。这里没有 encoder因为自回归语言模型的训练和推理都只需要 decoder 结构。写代码时最容易糊弄过去的是causal mask 的实现方式。它的作用是保证第 i 个 token 在计算 attention 时只能看到前 i-1 个 token不能偷看后面的内容。我第一版错误地把 mask 加在了 softmax 之后的输出上导致模型在训练时能间接获取未来信息训练 loss 低得离谱但生成时完全崩坏。正确做法是把 mask 加到 softmax 之前的 logits 上将未来位置的注意力分数设为-inf这样 softmax 后那些位置的权重趋近于零。这个细节如果你自己手写一遍基本不会再忘。3.2 我在代码里保留的三个关键 Mask除了 causal mask我在代码里还额外保留了另外两个 mask它们都能直接影响模型行为。第一个是padding mask用于标记无效 token 位置让 loss 计算和 attention 都跳过这些位置。这个 mask 的设计会直接影响 token 利用率值得花时间写对。第二个是文档边界 mask。前面说到用流式拼接构造样本时如果文本跨文档切块模型可能会注意到“不同文档之间的奇怪拼接点”。我的做法是生成一个布尔矩阵记录哪些 token 对属于同一文档把跨文档的 attention 位置全部 mask 掉。实验下来这个细节对长文本生成的连贯性有明显帮助。第三个 mask 是我在尝试推理任务时加的“问题-前提混合 mask”。由于我的任务是给定一组前提和问题让模型生成推理轨迹我希望模型在编码前提时能看到全部前提 token而在生成答案时只能看到前面的推理过程。这种非对称的 mask 结构可以显著改善模型在 multi-step reasoning 任务上的表现。如果你只是做普通的对话模型不需要它但如果你想构建带推理能力的模型这会是绕不开的工程点。3.3 参数量和显卡容量的快速估算动手写代码之前最好先算一笔账我的显卡到底能装下多大的模型我用的估算公式来自实践经验训练总显存 ≈ 参数量 × (参数字节数 梯度字节数 优化器状态字节数) 激活值与临时缓冲。以 1 亿参数模型为例参数量是 1e8。混合精度训练fp16下参数占 2 字节梯度占 2 字节Adam 优化器需要保存一阶矩和二阶矩各 4 字节合计每参数约 12 字节即 1.2GB。激活值通常占训练峰值显存的 30% 到 50%所以 1 亿参数模型在 batch size 32、序列长度 512 时训练峰值大约需要 3 到 4GB 显存。如果你手里只有一张 8GB 显存的消费级显卡可以把模型规模压到 5000 万参数以下并把 batch size 调到 16基本就能跑起来。模型参数混合精度训练理论占用激活值估算推荐最小显存消费级显卡参考2000 万240MB300-500MB2GB笔记本老显卡可跑5000 万600MB1-1.5GB4GBRTX 3050 级别1 亿1.2GB2-3GB6-8GBRTX 3060 / 4060 级别3 亿3.6GB6-8GB12-16GBRTX 4070Ti / 4080 级别这个估算只是起步参考实际激活值会随 batch size 和序列长度线性增长跑之前最好用一个小 batch 做试算。4. 训练不是玄学loss 曲线里藏着所有答案4.1 warmup、衰减与 batch size 的联动训练稳定的第一根弦是学习率调度。我的第一版训练脚本直接使用了固定学习率 3e-4结果 loss 剧烈震荡甚至出现 NaN。后来我把调度改成warmup 余弦衰减前 1% 的步数把学习率从 0 线性升到峰值然后按余弦函数缓慢降到峰值学习率的 10%。这个做法的直觉是训练初期模型参数还很“脆弱”学习率过大会让梯度一步跨过最优区域等模型逐渐找到合理的方向后再加大步长加速收敛后期则用一个较小的学习率做精细打磨。batch size 对训练稳定性的影响经常被低估。我观察到在同样学习率下把 batch size 从 16 提高到 64模型的最终 loss 能降低 0.3 到 0.5。但 batch size 越大单步计算越贵显存也可能不够。一个折中方案是梯度累积将多个微批次的梯度累加后统一更新参数用时间换有效 batch size。我在实际项目中把有效 batch size 固定为 128即 4 个微批次累积一次训练曲线的平滑度比单批次训练好得多。4.2 loss 从 5 掉到 1 的典型过程与异常判断语言模型训练的经典曲线会经历几个阶段。最开始几轮步数loss 会从 8 以上快速掉到 5 左右这对应模型在疯狂学习高频词汇和基本语法结构这段往往只需要几百步。然后进入漫长的缓慢下降期loss 从 5 一路磨到 2 左右这时模型开始掌握复杂句式和上下文关联。再往下到 1 附近通常发生在训练后期或数据量很充分时对应模型能够非常有把握地预测绝大多数 token。我判断训练是否正常的关键指标有两个一是训练 loss 与验证 loss 的差值。差值小于 0.3 属于正常范围意味着模型还在持续泛化差值突然拉大到 1 以上基本可以确定过拟合需要增加数据量或加大 dropout。二是梯度范数。我每 50 步打印一次 grad norm正常时它应该稳定在一个较窄的区间如果突然飙升到之前数值的 5 倍以上一定要立刻降低学习率否则接下来几轮很容易引爆 NaN。4.3 梯度裁剪和日志监控训练稳定的最后防线即使做了 warmup 和 batch size 调整训练过程中仍然可能出现梯度爆炸。我的标准做法是设置max_grad_norm 1.0每步更新前对梯度做一次 L2 范数裁剪。这个操作几乎不增加成本却能把“偶尔一次异常梯度毁掉整个训练”的概率降到极低。日志监控方面我强烈建议把训练脚本做成“可观察”的而不只是在命令行打印几个数字。我会把每一步的 loss、token 级困惑度、学习率、grad norm、当前 token 吞吐量都记录到结构化日志里然后用一个简单的本地可视化页面实时查看。这个习惯在排查问题时相当高效——有一次训练曲线出现周期性的 loss 尖峰我正是靠日志发现数据加载阶段存在一个间歇性 shuffle bug。没有日志这种问题几乎不可能靠肉眼发现。5. 让模型开口说话解码策略、温度与 KV Cache 的性能账5.1 温度、top-k 与 top-p三个旋钮的直觉训练完成后模型输出的其实是下一个 token 的概率分布真正从分布中“抽”出一个 token 需要解码策略。这里最核心的旋钮是温度temperature、top-k 和 top-p。温度的直觉是温度越低概率分布越尖锐模型越倾向于选择概率最高的 token输出更保守、更确定温度越高分布越平坦模型更愿意冒险选择低概率 token输出更多样但也更容易跑偏。我在推理任务中经验值是0.1 到 0.4 之间太低会让输出陷入保守重复太高会让回答变得跳跃离谱。top-k 的作用是先把概率最高的 k 个 token 挑出来把其余全部踢出候选然后再做采样。top-p 则是选出累积概率超过 p 的最小子集。实际使用时两者通常会组合先 top-k 截断长尾再 top-p 做动态过滤。我给出的默认组合是temperature0.3, top_k40, top_p0.9大多数场景下都能拿到不错的平衡。5.2 KV Cache 为什么是“免费的午餐”如果每次生成一个新 token 都重新运行一次完整的 Transformer 前向计算代价会非常大——因为前文的所有 token 都要被重新计算一遍。KV Cache 的思路简单粗暴把每一层 attention 中已经算好的 Key 和 Value 向量缓存下来生成新 token 时只计算新增 token 的 K 和 V然后跟缓存拼接大幅减少重复计算。我在实验中用 1 亿参数模型做对比未开 KV Cache 时每秒只能生成 5 到 8 个 token开启后能到每秒 40 到 60 个 token几乎是数量级的差距而显存增加完全可以接受。实现 KV Cache 时最需要注意内存管理。当序列长度从 1 增长到 1024每层的 KV 缓存规模也在增长如果不对缓存做复用或及时释放长对话推理很快会撑满显存。我的做法是预分配一个固定的最大缓存容量并在序列长度变化时复用既有内存块尽量避免频繁申请和释放。5.3 显存不够时的推理妥协方案从零训练出的小模型虽然参数量不大但在低显存设备上依然可能推不动长序列。我的一个现实场景是在一台 4GB 显存的旧笔记本上跑 1 亿参数模型开启 KV Cache 后序列长度到 2048 就 OOM。我当时用了三个妥协方案分块推理把输入序列切块分别计算中间层表示再做拼接。这种方式增加了代码复杂度但能极大降低单次峰值显存。降低精度推理时用 fp16 已经是常态如果还紧张可以考虑 int8。我实测中 1 亿参数模型从 fp16 转 int8 后显存占用下降 40%生成质量损失在可接受范围内。限制上下文长度这是最简单直接的办法。很多任务其实 512 token 的上下文就够用强行用 2048 反而浪费时间。6. 从零路上交过的学费高频踩坑与完整排查链路6.1 症状一loss 完全不下降这是每个 from-scratch 新手最先遇到的“鬼打墙”。我当时的现象是训练循环跑了两千步日志里 loss 始终在 8.7 附近纹丝不动。排查链路我按顺序走了一遍第一步查数据是否正确喂入。我在某个 batch 后打印了几条 input_ids 和对应的文本发现一部分 token 被 padding 了而 loss 计算时没有应用 padding mask导致模型浪费大量精力去预测无意义的 padding token。修好后 loss 立刻有了下降趋势。第二步查学习率尝试从 3e-4 调到 1e-4loss 开始缓慢下降但速度仍不理想。第三步查初始化最终把残差层最后一部分的初始化比例调小loss 下降速度恢复正常。如果你也遇到 loss 卡死请按“数据 → mask → 学习率 → 初始化”的顺序排查大概率能定位到问题。6.2 症状二生成内容全是重复的“XXXX”模型训练完成后生成的前几个 token 还正常但是从某个位置开始就陷入循环一直输出同一个词。我排查时发现训练语料里存在一批连续重复度极高的记录导致模型学到了“当前 token 后面大概率还是同一个 token”的预测偏差。这个问题的修复分两层数据层我要做更严格的重复度过滤任何连续 n-gram 重复超过阈值的文本直接删除解码层我引入repetition penalty机制对已经出现过的 token 的概率乘以一个惩罚系数常见的比值是 1.1 到 1.3有效降低重复概率。两者结合后重复问题基本被解决。6.3 症状三OOM 与训练中断训练到第几百步直接 OOM 中断这是显存管理问题。我的排查结果是模型自带的激活值缓存没有做释放策略长时间运行后累积占用越来越大。解决方法是改用流式的梯度检查点gradient checkpointing计算激活值时只保留必要部分反向传播时再重算其余部分。这项技术显著增大了计算量但能把显存峰值降低一半以上。如果你的显存实在紧张另一个直接手段是把序列长度从 1024 降到 512token 利用率虽然下降但实验仍然能继续。6.4 症状四评估指标不错但生成质量很差在一次实验中模型的 next-token 预测准确率高达 45%但生成出的句子质量惨不忍睹。我起初怀疑是解码策略问题调了半天参数没有改善才意识到是数据级污染我的验证集和训练集高度重叠模型在验证集上的高精度只是“背诵”而非“理解”。之后我把训练集和验证集做了严格的文档级去重重新训练了一轮生成质量立刻好转。还有一个容易被忽略的原因是 beam search 等解码策略虽然能提升指标型评估却不一定适合开放式生成需要根据任务目标灵活选择。7. 给想动手的人一份可落地的工程路线图7.1 建议的最小闭环从 2000 万参数开始如果这是一次真正的 from-scratch 项目我强烈建议你不要从 10 亿参数起步。先做一个 2000 万参数的小模型用一份 2 万条左右的迷你语料目标设定为“能生成语法通顺、逻辑基本连贯的短文本”。2000 万参数在普通笔记本电脑上也能正常训练这能让你在半天内完整走一遍所有工程链路。等你对数据、训练、解码、评估都有了实感再逐步把参数量提升到 1 亿甚至 3 亿每一步都建立在已验证的基础上排查问题会容易得多。7.2 时间预算与三阶段路线图我做完这个完整闭环大概花了四周时间拆分为三个阶段第一阶段第 1 周构建数据管线和 tokenizer完成 Mini-Batch 训练的启动脚本。这一周的目标是让“data → model → loss → 反向传播”的链路跑通不需要追求好的 loss。第二阶段第 2-3 周调训练超参完善日志与评估模块把小模型训练到能够稳定生成通顺文本并加入简单的推理任务测试。这个阶段是知识密度最大的时期大量踩坑都发生在这里。第三阶段第 4 周推理优化与任务适配包括 KV Cache、采样策略调优、以及根据任务需求调整 mask 结构和 loss 设计。这三个阶段的时间分配可以根据你的基础调整但我不建议压缩第一阶段——数据管线没做扎实后面所有环节都会不断返工。7.3 参考资料与进阶路径如果你想系统性从零学习构建 LLM 的方法《Build a Large Language Model from Scratch》是一本很好的参考书籍它按数据准备、模型结构、预训练、微调到部署的顺序完整走了一遍。读过之后我的建议是立刻动手做一个自己的项目不需要试图复刻书里的完整体验只需要抓住“数据、架构、训练、推理”四条主线做出一个最小可用的模型即可。如果你有编程基础但没写过神经网络训练脚本可以先自己实现一个带反向传播的两层 MLP熟悉梯度下降和 loss 计算之后再进入 Transformer 语言模型。这一步绕不过去但也不是很难。我在整个从零过程中最深的体会是真正让你成长的并不是“知道 Transformer 有多少层”而是“当训练曲线异常、生成输出诡异、显存再次溢出时你能不借助任何黑盒工具靠一双手和日志把问题一层层挖出来”。这种能力在 AI 工程里会伴随你很久而且没有任何现成课程可以直接教。希望这篇文章能帮你把第一段路走得顺利一些。