这两个月我干了一件特别“找虐”的事：给自己定了个项目代号ai-engineering-from-scratch，目标是从零开始构建一个能跑通数据、模型、训练、推理、部署全流程的AI工程，而不是调用现成的 Transformers 工具包。这个项目本质上包含两条主线：一条是build a reasoning model from scratch，也就是手写一个具备一定推理能力的模型；另一条是build a large language model from scratch，把 GPT 风格的 Decoder-only 大语言模型从矩阵乘法开始搭起来。做这件事的最大收获不是得到一个还不错的模型，而是把从前只会调包时那些模模糊糊的“直觉”，变成了可以逐行解释的“工程决策”。如果你也想真正吃透大模型，而不是停留在跑通例子的水平，这篇文章值得你看完。
1. 为什么非要从零开始：AI工程的“拆轮子”思维
1.1 调包侠无法回答的三个问题
在动手之前，我列了一张“拷问清单”，全是过去用别人框架时答不上来的问题。第一个问题是：nn.Linear里的参数矩阵如果初始化成零，模型会怎么样？第二个问题是：训练时为什么要按 batch 计算 loss，而不是把整个语料拼成一条超长序列？第三个问题是：推理时每次生成一个 token，为什么显存占用反而比训练还高？
这三个问题看着简单，但如果不读源码、不自己重写一遍，很难给出有底气的答案。我见过不少同学能熟练跑完微调脚本，但对Forward里的矩阵形状变换、mask的作用、KV cache的生命周期一无所知。从零构建的意义就在于此：它逼你把每个组件都过一遍，知其然也知其所以然。
1.2 从零构建的工程价值与适用场景
这并不代表所有场景都该从零造轮子。如果你的目标是快速验证业务想法，直接用 HuggingFace 是正确选择。但你的目标如果是以下三类，从零构建的性价比极高：第一，你要在资源受限的端侧设备上部署模型，必须知道哪一层可以剪、哪一层可以量化；第二，你要修改模型结构，比如加入某种自定义注意力机制，不熟悉底层实现根本改不动；第三，你要深入做性能优化，需要理解算子融合、显存分配、梯度累积这些概念。
我把项目划分为五个层次：数据工程、模型实现、训练循环、推理优化、部署评测。每一层都使用最朴素的 Python + PyTorch 基础算子完成（某些地方甚至手写了 NumPy 版本），不引入任何高级模型库自带的抽象。下面按我当时动手的顺序逐一展开。
2. 地基先行：数据准备与分词器全解析
2.1 数据集选型和处理流水线
所有大语言模型的故事都从数据开始。我最初以为随便找几个开源语料就能训练，结果发现数据质量直接影响收敛速度。我最终选择的是清洗后的纯英文文本混合数据集：维基百科的子集、代码片段、数学题库，以及一部分自制的“推理链”语料。处理流水线非常朴素：去掉重复段落、过滤超短碎片、按文档边界切分、统一为 UTF-8 后存储成 JSONL。
这一步的关键细节是“去重”。我第一版没做去重，导致模型在验证集上反复出现同一批句子，loss 降得非常快但泛化很差。后来用 MinHash 做了近似去重，把重复比例从 12% 压到了 1% 以下，训练曲线立刻健康了很多。建议大家在动手写模型之前，先把数据管线跑通并生成一个几万条的小样本集，这样后续迭代模型时能秒级验证，不用每次都等全量数据。
2.2 从零手写BPE分词器
我这次没有直接用tokenizers库，而是手写了一个BPE（Byte Pair Encoding）分词器。BPE 的核心思路并不复杂：先把文本拆成单字符，然后反复统计相邻字符对的出现频率，把频率最高的对合并成一个新符号。例如"low"和"lower"中频繁出现的"low"会被合并成独立 token，这样模型可以更高效地学习常见子词单元。
实现 BPE 时最容易踩的坑是“合并顺序的不可逆性”。训练词表时，你必须记录每次合并的规则，之后切分新文本时按同样顺序应用，否则同一个词在不同地方可能被切成不同结果。我的实现采用了预正则化：先将所有文本转换为小写，再用正则按字符类别预切分，这样数字、标点不会被错误地拼进单词。词表大小我定在 8000，加上 256 个字节级 fallback 和几个特殊 token，总词表 8260。对于一个小型演示模型，2000 万 token 的训练数据用 8000 词表比较经济；如果数据量超过 1 亿 token，建议把词表加大到 16000 以上。
2.3 词表大小与特殊token的权衡
这里展开讲一下我踩过的另一个坑。一开始我天真地用了 32000 词表，觉得越大越“像样”，结果在小数据集上大量 token 出现次数极少，嵌入层几乎是在随机初始化。更严重的是，词表太大导致 embedding 矩阵占了模型参数的一半以上，训练速度被矩阵乘法拖慢。后来我把词表缩小到 8000，同样数据量下每个 token 的平均出现次数翻了四倍，模型困惑度明显下降。
特殊 token 的处理同样重要。除了<|endoftext|>之外，我加入了[BOS]、[EOS]、[PAD]和一个用于指令数据的[SEP]。特别要提醒的是，特殊 token 在分词器中的索引必须固定下来，不能因为训练语料里自然出现的配对关系而把它们合并掉。我在手写 BPE 时特意在词表最后预留了特殊 token 的槽位，并在合并规则中禁止任何导致特殊 token 被拆分的操作。这一点听起来细枝末节，但一旦忽略，训练时会出现大量非法 id，报错还特别隐晦。
3. 核心模型：从注意力机制到大语言模型
3.1 注意力机制从矩阵乘法开始
重头戏来了。我从scaled dot-product attention开始写，定义非常简单：输入Q、K、V三个矩阵，先计算Q和K的点积除以sqrt(d_k)，再经过 softmax 得到权重，最后与V相乘。这里最关键的问题是为什么除以sqrt(d_k)：因为当d_k变大时，点积结果的方差会随之增大，导致 softmax 后的梯度趋近于零。除以一个缩放因子能让梯度保持在比较合理的区间。
我在实现多头注意力时，为了直观，没有直接用einops，而是手动 reshape 和 transpose。这过程非常磨人，但让我彻底记住了(batch, seq_len, num_heads, head_dim)的顺序。建议新手第一次实现时，先用随机张量把每个中间步骤的 shape 打印出来，对照公式检查。我当时犯过一个低级错误：在reshape时把num_heads和head_dim的顺序搞反，结果模型 loss 一直不降，整整查了两天。
3.2 位置编码的选择：RoPE还是正弦
Transformer 是“置换等变”的，如果不加位置信息，模型看到的句子只是词袋。我对比了两种方案：传统正弦位置编码和RoPE（旋转位置编码）。从零构建时，正弦编码实现简单，但它的外推性较差，超过训练长度后注意力分数会变得混乱。RoPE 虽然计算复杂一点，但通过旋转 Q 和 K 的向量来注入相对位置信息，在长文本上有更好的外推效果。
我最终选了 RoPE，因为它更贴近现代开源模型，也方便后续对比其他项目。RoPE 的本质是将每个 token 的d维向量分成若干对，每一对按频率做二维旋转。频率是个关于位置索引的几何函数，通常形式为theta_i = 10000^(-2i/d)。实现时只需要预计算一个形状为(max_seq_len, d//2, 2)的旋转矩阵，然后在 forward 里对 Q 和 K 做旋转。我最初写的是逐 token 循环，慢得离谱，后来改成预计算矩阵并用torch.gather批量索引，速度提升了近 50 倍。
3.3 残差、归一化与FFN：把decoder block拼起来
一个标准的 Decoder-only block 由三个子层组成：多头注意力、前馈网络（FFN）、以及各自后的残差连接和层归一化。我有一个偏执：所有归一化都使用Pre-Norm结构，也就是在每个子层之前做归一化。原因很简单：Pre-Norm 的梯度路径更干净，训练更稳定，尤其对从零训练的小模型友好；Post-Norm 在某些实现里效果更好，但对学习率和初始化非常敏感，容易让我这种野路子选手陷入不收敛的泥潭。
FFN 部分我使用了近似 GPT 的配置：先线性升维 4 倍，再经过 GELU 激活，最后线性降维回模型维度。这里有个容易忽略的细节：第一次线性层之后的 GELU 要用近似形式还是精确形式。我起初直接调用了F.gelu，后来为了深入理解，手写了一个 tanh 近似版本，两边的微调结果几乎一致。如果你想完全从零，建议从近似 GELU 开始，代码只有两行，但能让你明白激活函数的可替代性。
3.4 参数量与内存估算：用公式算清家底
写模型前建议先纸上估算参数量，避免配错显存。以我的最终模型为例：hidden_size = 256，num_layers = 6，num_heads = 8，head_dim = 32，词表大小8260，最大序列长度1024。嵌入层参数量 = 8260 × 256 ≈ 211 万。每个 Transformer block 中，注意力部分的 QKV 和输出投影都是256 × 256的矩阵，加上 bias 总共约 26 万参数；FFN 两个线性层约 52 万参数；层归一化每层约 512 参数。单 block 大约 78.6 万参数，6 层约 471 万。模型总参数约 682 万，单精度浮点占内存约 27MB。
这样大小的模型在单张消费级显卡上跑起来很轻松，但你是不是以为显存只需 27MB？错了。训练时还需要保存梯度、优化器状态、激活值。我实际开训练时，batch size 为 8、序列长度 512，动量优化器 + Adam 状态让显存占用直接飙到 1.8GB。所以要做任何训练之前，先用公式算好显存预算：总显存大约等于参数内存 × (1 + 梯度倍率 + 优化器倍率) + 激活内存。这句话建议刻在脑子里。
4. 训练与调优：让模型从“复读机”变成“推理者”
4.1 训练目标与loss设计
大语言模型的基础训练目标就是自回归语言建模：给定前文预测下一个 token。实现起来就一行CrossEntropyLoss，但要把 label 左移一位。比如输入序列是[a, b, c]，模型看到[a, b]时预测[b, c]。这里面有一个关键决策：如何处理批次内的 padding。我选择不对 padding token 计算 loss，方法是把 label 中 padding 位置的数值设成-100，PyTorch 的交叉熵会自动忽略它们。
但单靠语言建模 loss 并不足以让模型“会推理”。我的做法是在训练语料中加入约 10% 的“推理链”数据：把数学题的逐步解答拆成带思考步骤的文本，然后把每一步都暴露给自回归目标。这样模型不再只是预测单个结果，而是学会模仿“先说什么、再说什么”的中间过程。这种思路非常接近现在流行的 reasoning model 训练范式，只是我的规模很小，只能算一个玩具级别的验证。
4.2 学习率策略与优化器选择
从零训练时，我强烈建议用AdamW + warmup + cosine decay。warmup 是很多人容易忽略的一步：刚开始优化器对参数的更新步长需要较小，否则模型输出会快速飞向某个极端，后面很难拉回来。我试过直接全量学习率，loss 在前几百步剧烈震荡，一度冲到 10 以上，后来改成前 2000 步线性 warmup 到峰值，再把学习率按余弦降到峰值的 10%，训练瞬间稳定了。
峰值的选取也有讲究。隐藏层 256、batch size 8 的小模型，我用3e-4效果不错；但如果把 batch size 调大，学习率需要相应微调。一个实用的经验法则是：线性缩放规则——batch size 翻倍，学习率可以按约 1.4 倍调整。但这条规则对超大 batch 不适用，只适合小规模试错。还有，千万不要把 AdamW 的weight_decay设成 0，我一开始为了省事设置成 0，结果过拟合速度极快，验证集 loss 和测试集表现差距越来越大。
4.3 数据配比：通用语料+推理链路数据
数据配比是我想重点强调的训练杠杆。自行实现模型时，很多人只用一份语料从头跑到尾，这样很容易让模型学到“表层统计”，而不是通用能力。我给自己的数据配比设成了：70% 通用百科和博客文本、20% 代码文本、5% 数学题、5% 手工书写的推理链数据。这个比例让模型在验证集上的几个任务表现更均衡，尤其是经过逻辑推导的简单问答，准确率从几乎为 0 提升到约 35%。
每个 epoch 结束时，我专门抽出 200 条“推理题”作为验证集，而不是随机采样的通用文本。这样做的原因是，语言模型 loss 是一个平滑的平均值，某类任务的表现会被其余数据淹没。如果你关心“模型是不是真的在学推理”，必须单独测这一类样本，否则你看到的往往只是“loss 下降了 0.1”，但回答质量毫无变化。
4.4 监控指标：除了loss还要看什么
训练过程中我定了四个监控项：训练 loss、验证 loss、token 准确率、以及“最长可生成文本的重复率”。最后一项听起来很怪，但它能暴露模型的退化现象。我曾遇到一个情况：训练 loss 一路下降，验证 loss 也正常，但生成结果全是同一个句子的循环。后来查日志发现，模型在小数据集上把某些高频 token 的 embedding 学成了一个“吸铁石”，采样时反复落到那少数几个 token 上。
解决办法是加入重复惩罚（repetition penalty）和多样性采样，但更根本的做法是把数据去重做好，并在训练时随机遮盖掉部分高频 token 的 loss 权重。如果你用现成框架，可能永远遇不到这种问题，因为它背后已经有大量工程保护。但从零构建时，这些保护措施都需要你亲手补上，而这恰恰是这个项目最珍贵的学习过程。
5. 推理阶段工程化：从KV缓存到采样策略
5.1 KV缓存的实现细节
模型训练完成后，推理环节同样有无数的“隐形工作”。第一件要做的是KV cache：在每一步生成时，不需要对所有历史 token 重新计算 key 和 value，只需要缓存上一层的 K、V 矩阵，让新 token 只和缓存做注意力计算。这个优化在小模型上可能只带来 1.5 倍加速，但在长序列上差距是指数级的。
实现 KV cache 时我踩了一个坑：没有区分训练和推理模式。训练时 K、V 需要完整计算，不能使用缓存；推理时则要一边生成一边更新。我的解决方案是在模块里加一个use_cache标志，推理模式把 K、V 存储到一个列表，每次生成时 concat。另外要注意剔除 KV cache 对反向传播的影响，推理模式下不需要保存计算图，所以要用torch.no_grad()包住整个循环，不然显存会越跑越大。
5.2 温度、top-p、top-k采样
从推理模型输出一个 token，本质是在词表上得到一个概率分布。最简单的是贪心解码，直接取概率最大的 token，但这会导致生成文本过于枯燥且容易陷入重复。我实现了温度采样、top-k 和 top-p 三种策略的组合：先用温度对 logits 进行缩放，温度越高，分布越平滑，结果越随机；然后截断到概率累计值 top-p，比如 p=0.9，只保留累计概率达到 90% 的最少 token 集合。
我的经验是：做数学推理题时，温度要控制在 0.3 以下，top-p 用 0.8；做开放式写作时，温度可以提到 0.8 到 1.0，top-p 用 0.9。这个经验并不是拍脑袋，而是因为推理任务对“确切性”要求高，采样时的随机性很容易导致中间步骤出错。另一个关键点是，对 EOS token 的概率要特殊处理。如果某种采样策略会把 EOS token 过滤掉，模型可能永远停不下来，必须设置最大生成长度作为兜底。
5.3 按需加载与轻量部署
我最后把训练好的模型导出一份纯 NumPy 的推理脚本，不依赖 PyTorch 也能跑。过程很简单：把模型权重state_dict里的每个张量用numpy.save存为.npy文件，然后在纯 Python 环境里手写 forward 和采样逻辑。这种“退步式”部署反而让我对权重存储格式、字节序、精度转换有了深刻认识。
如果你也想把模型部署到云端函数或者移动端，建议在导出时把精度从 FP32 转成 FP16。我试过直接转换，发现推理速度提升了约 40%，但某些层的损失相对明显，尤其是 layer norm 的方差计算。后来我换成只量化线性层，保留归一化层为 FP32，效果和速度取得了很好的平衡。对一个 680 万参数的小模型来说，FP16 后权重文件只有 13MB，放进轻量级服务非常舒适。
6. 踩坑实录：训练不收敛、显存溢出、推理太慢
6.1 loss卡住不变？检查数据与初始化
我遇到过印象最深的一次 bug：模型 loss 在 9.5 附近卡了整整 1 个 epoch，纹丝不动。排查步骤不是先调代码，而是先检查数据：打印第一条 batch 的输入 ID 和 label，我立刻就发现 label 里有一半是 -100，因为 padding token 被填到了序列中间。正确地做法是使用attention_mask配合 padding，让模型忽略 padding 位置，而非靠 -100 硬屏蔽。
如果数据和 mask 都没问题，概率最大的嫌疑就是初始化。我之前看到某些教程直接把 attention 的 bias 全部初始化为 0，这种设定会让梯度在注意力层消失。标准做法是线性层的权重用均匀分布或正态分布初始化，偏置设为 0，而最后一层输出 LM Head 权重可以和 embedding 层共享，这样既能缩小参数，也能加快收敛。我的经验是，从零写模型时最好加一个“初始化检查脚本”：随机输入一条短文本，跑一次 forward 和 backward，确认没有 nan 且梯度范数在合理范围内，再开始完整训练。
6.2 显存OOM：梯度累积与混合精度
当我把 batch size 从 8 调到 16 时，显存直接爆掉。我不愿意缩小序列长度，于是采用了梯度累积：累积 4 个 step 的梯度再更新一次参数，等效 batch size 变成 64。这个办法不需要改动数据加载逻辑，只需要在 optimizer 更新前判断累积步数。要注意的是，梯度累积时 loss 要除以累积步数，否则等效 batch 变大但 loss 尺度没变，会导致学习率需要重调。
另一个极大缓解显存的手段是混合精度训练。PyTorch 自带torch.cuda.amp，但我为了搞懂原理，自己实现了简单的 FP16 前向和 FP32 梯度回传：权重复制一份 FP16 用于前向，反向时用 FP32 主权重更新。这个工程量不小，但做完之后我终于明白为什么推荐用 AMP 而不是全程 FP16——因为某些中间梯度的数值范围差距过大，FP16 无法兼顾，需要损失缩放。如果你不想重复造轮子，直接用官方 AMP 就好，但心里得知道它帮你在后台做了哪些事。
6.3 推理速度慢：batch vs streaming
把模型跑通之后，我一度嫌它每生成一个 token 要 200ms，太慢了。首先我检查了是否在循环里反复创建了 attention mask 和 position ids，如果是，应该把它们提到循环外面只创建一次。其次，我将单条文本生成改成多请求 batch 生成：一次前向处理多个独立的 prompt，每个请求分别维护自己的 KV cache slice，这样 GPU 利用率显著上升，整体吞吐提高了约 3 倍。
如果你也需要实时流式输出，还需要处理一个隐藏问题：当 batch 中某个序列已经生成 EOS 时，不能把它从 batch 中剔除，否则会破坏别的序列的缓存索引。标准做法是把对应位置的finish标志设为 True，后续生成时仍然占位但不采样。这块逻辑很绕，我建议专门写一个SequenceGenerator类，把“是否结束”作为状态机管理，不要和模型前向逻辑混在一起。
6.4 长文本输出退化：重复惩罚与信息密度
最后一个让我挠头的现象是：模型生成长文本时，总是重复最近出现过的短语，到了 200 token 之后几乎变成了复读机。哪怕我加了 top-p 采样也没完全解决。后来我加了repetition penalty：在采样前，遍历已经生成的 token，对它们的 logits 统一减去一个惩罚值（我用 1.1），这样模型在选择下一个 token 时会倾向于避开刚说过的话。
但重复惩罚不能加得过大，否则模型会开始说一些完全不连贯的内容，因为它在每个位置都强行避免历史 token，导致逻辑链断裂。我反复测试后，取了一个比较温和的值 1.05，配合 top-p 0.85，效果最佳。老实说，这种小模型的长文生成能力天花板很低，如果想真正提升，应该去优化数据质量和模型规模，而不是一味调采样参数。不过，从零项目里能亲手调节惩罚值的体验，让我对线上大模型生成的“神秘感”少了很多。
顺带一提：这次从零构建，我最满意的一点
整个项目下来，我最大的感受是：“重新发明轮子”其实是一种极其高效的深度学习方式。我不否认现成框架的价值，但如果没有亲手搭一遍，你不会知道一个看似简单的torch.nn.Transformer里藏着多少工程权衡。我的建议是，如果你想做类似的项目，不要一开始就追求完美，先用最小数据集跑通一个 2 层、128 维的模型，然后再逐步加数据、加层数。
还有一个小技巧想分享给动手实践的人：每次修改代码或数据之后，固定跑同一个“冒烟测试”prompt，比如让模型做一道简单的两位数加法，并记录输出。这相当于给模型训练加上了一个“行为回归测试”。我就是靠这条冒烟测试，发现了好几次“loss 正常但回答质量下降”的问题，比只看 loss 曲线可靠得多。
ai-engineering-from-scratch 这个项目，对我来说就像按着图纸重新组装了一遍发动机。现在再看到各种大模型开源项目时，我能准确判断哪些地方是核心、哪些地方只是锦上添花。如果你也想真正进入 AI 工程的大门，别再满足于跑通别人的 demo 了，从矩阵乘法开始，手写一个只属于你自己的大语言模型吧。
