“AI工程”这个词这两年算是被说烂了但真正上手时你会发现市面上绝大多数内容是教你调API、装框架真正讲清楚一个模型从数据到推理全链路怎么搭起来的却很少。我最近把项目标题定为“ai-engineering-from-scratch”核心就一句话不用任何现成大模型不靠黑盒推理接口从零开始做一个可以训练、可以推理、可以评估的微型模型。这篇文章就是整个项目的复盘包括数据、tokenizer、Transformer架构、训练策略以及最后用强化学习思路让一个小模型具备基础推理能力的过程。项目规模控制在20M参数左右单张消费级显卡就能跑通适合想彻底搞懂大模型原理的工程师、刚转行做算法的朋友以及不满足于“调包”的深度学习爱好者。很多人问我市面上已经有那么多开源模型了直接拿来做微调不香吗香但理解不了原理。就像你天天开车不一定知道发动机怎么点火、变速箱怎么换挡。真出了问题只能停在路边等救援。“从零构建”就是那个把你按在发动机旁边逼你拆开看一遍的过程。1. 项目定位为什么要做“从零构建”这件事1.1 从调包到懂原理之间差什么调API和做AI工程看起来都是“跟模型打交道”实际是完全不同的两个世界。调API时你关心的是提示词怎么组织、温度参数怎么设置模型内部对你来说就是个黑盒子。而AI工程关心的是数据分布对不对、tokenizer词表合不合理、模型参数量怎么分布、学习率曲线为什么是这个形状、梯度什么时候会爆炸。从零构建最大的价值是把整个链路的所有错误都亲手踩一遍。你会发现模型不work的原因千奇百怪数据里混了不可见字符、tokenizer出现了OOV、embedding初始化不合适、学习率太大导致loss直接飞到NaN、验证集和训练集分布没对齐……任何一环出错模型就废了。这些坑光看文档是看不出来的只有跑起来、崩溃、修复、再崩溃才会真正变成你自己的经验。我见过不少同学模型结构背得滚瓜烂熟能默写Transformer公式但一旦让他从一份txt文本训练出可用的生成模型就完全不知道从哪里下手。原因很简单知识是线性的工程是网状的。从零构建的价值不是让你重复造轮子而是让你拥有拆轮子的能力。1.2 目标边界20M参数能做什么做这种项目最怕的就是目标定得太离谱。我见过有人一上来就要复现GPT-31750亿参数单卡跑不起来就抱怨硬件最后项目烂尾。正确的做法是把目标压缩到可以“一个人、一张卡、几天内完成”的程度。我这边的项目边界是这样定的数据集TinyShakespeare大约1MB的莎士比亚文本切分后约100万token。公开可下载稳定句子结构清晰非常适合用来观察语言模型的生成效果。模型结构标准decoder-only Transformer6层、6个注意力头、d_model384、FFN中间维度1536整体约2000万参数。训练阶段先做自回归语言建模让模型学会“接话”再做一轮带奖励信号的微调让模型在简单符号推理任务上有稳定表现。硬件环境单张RTX 3090 24GB就够。没有3090的话用16GB显存也能跑只是batch要小一点。为什么选莎士比亚数据集两个原因。第一它是纯文本没有复杂的结构化标签词表相对封闭适合跑通全流程。第二它的语言有很强的文体特征训练几个小时就能看到模型生成出像模像样的“莎翁腔”成就感拉满。如果你手头有其他垂直领域文本比如代码、法律文书、病历也可以把后面这套流程原封不动搬过去。关于“reasoning model”的热度我想多说一句。最近大家都很关注OpenAI o1、DeepSeek R1这类推理模型但它们的核心突破其实不在模型结构而在训练策略——大规模强化学习让模型学会在推理时分配更多计算。所以这个项目第二部分我会在一个很小的自回归模型上尝试加入强化学习信号让模型完成类似“数字排序”这种简单推理任务。基座虽小但原理一条不落。2. 数据与Tokenizer模型的第一道关口2.1 数据清洗比想象中更关键做AI工程的第一课就是模型吃进去的是数据分布你喂进去的是垃圾后面所有努力都只是在给垃圾装修。数据清洗这件事听起来没什么技术含量但项目里80%的诡异问题追溯到底都是数据问题。TinyShakespeare原始文件本身算是干净的但我还是走了完整的数据预处理流程读入文本后统一把\r\n替换为\n消除Windows换行符差异。过滤掉所有控制字符只保留可打印字符和换行。检查文本中是否有重复的段落。虽然莎士比亚文本不会有但如果你换用自己的数据这一步必须做——重复数据会导致模型过拟合生成时无限复读。按9:0.5:0.5的比例切成训练集、验证集、测试集。注意切分时要按文本顺序切不要随机打乱后再切否则句子前后文关系会断裂。清洗完的纯文本还要转成模型能吃的东西。最朴素的做法是字符级编码但直接喂给模型不是最优解所以需要用分词器把文本切成更合理的token序列。这个东西太重要了我单独说。2.2 从零实现一个BPE分词器字符级tokenizer的实现很简单建一个字典映射每个字符到整数ID然后把文本逐字符编码。但字符级tokenizer有致命问题——上下文窗口变得非常“短”。比如上下文长度是128个token如果按字符切这128个token只覆盖了大约128个英文字符可能连一句完整的台词都覆盖不了模型能学到的东西极其有限。所以项目里我用了BPEByte Pair Encoding分词器。BPE的思路很朴素先从单个字符开始然后统计相邻字符对的出现频率把最高频的相邻对合并成一个新token重复这个过程直到词表达到预设大小。这样“the”这种常见词会被合并成一个token模型一眼就能看到整个单词效率提升非常明显。核心合并逻辑的代码其实很短你自己实现一遍就懂了def get_pair_stats(ids): stats {} for pair in zip(ids, ids[1:]): stats[pair] stats.get(pair, 0) 1 return stats def bpe_merge(ids, pair, new_idx): newids [] i 0 while i len(ids): if i len(ids) - 1 and ids[i] pair[0] and ids[i1] pair[1]: newids.append(new_idx) i 2 else: newids.append(ids[i]) i 1 return newids实际操作时先在训练集上统计字符对频率迭代执行merge操作得到一份merge记录。这份merge记录就是tokenizer的“词表”后续训练和推理都必须用同一份绝不能换。词表大小怎么选在TinyShakespeare上我试过512、1024、2048三个档位。512的词表偏小很多词根词缀拆不开2048词表在100万token的数据上有不少低频token几乎学不到最终选了1024。经验法则词表大小不要超过数据规模的0.1%~0.3%否则大量token会因为频率太低沦为噪声。这里必须提醒一个坑tokenizer和模型是强绑定的。训练时用了A词表推理时如果加载了B词表轻则出现一堆UNK重则生成内容彻底崩溃。我的做法是给每个tokenizer版本加一个hash值并在训练配置里记录保证任何一次实验都能回溯到正确的分词版本。3. 模型架构搭建从零手写Transformer3.1 自注意力机制的工程实现Transformer核心就是自注意力。理解它最直接的方式是把它看成“让序列里的每个token都向其他token询问我该关注谁”。具体来说每个token会生成三个向量Query我找什么、Key我是什么、Value我有什么。然后计算Query和所有Key的点积得到相似度分数经过softmax变成权重最后用权重加权求和所有Value得到这个token的新表示。在代码里缩放点积注意力就这几行def scaled_dot_product_attention(q, k, v, maskNone): d_k q.size(-1) scores q k.transpose(-2, -1) / math.sqrt(d_k) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) weights F.softmax(scores, dim-1) return weights v, weights注意那个除以sqrt(d_k)。为什么除因为当维度变大时点积的数值会变得很大softmax会迅速进入饱和区梯度接近于零。这就好比一群人同时大声说话你反而一个字都听不清。除以sqrt(d_k)相当于把音量压低到合适区间让注意力权重有区分度。另外decoder模型必须加因果掩码。因为自回归模型只能看到当前token之前的token不能“偷看”后面的答案。初始化时在注意力矩阵的上三角位置填负无穷softmax后这些位置会变成0模型就“看不见”未来了。3.2 多头注意力、残差、归一化与FFN多头注意力就是把注意力过程分成多个头并行做。每个头有不同的Q、K、V投影矩阵相当于开会时分了几个小组每个小组关注不同的维度一个组盯着句子的主谓结构另一个组盯着修辞关系最后把各组的结论拼在一起。这种“各看各的再汇总”的机制让模型表达能力大幅提升。一个标准的TransformerBlock包含四部分多头注意力、残差连接、层归一化、前馈网络。代码骨架如下class Block(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.attn nn.MultiheadAttention(d_model, n_heads, dropoutdropout, batch_firstTrue) self.ffn nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model), ) self.ln1 nn.LayerNorm(d_model) self.ln2 nn.LayerNorm(d_model) def forward(self, x, maskNone): x x self.attn(self.ln1(x), self.ln1(x), self.ln1(x), attn_maskmask)[0] x x self.ffn(self.ln2(x)) return x残差连接解决的是深层网络的梯度传播问题——让梯度可以“抄近路”直接流回浅层层归一化解决的是每层输出分布漂移的问题。这两个设计是训练深层Transformer的基石缺一个二十层以上的网络基本训不动。在小模型上它们仍然重要但作用没大模型那么“救命”。3.3 参数量手动计算“20M参数”是怎么来的我带你手算一遍。假设词表大小1024d_model384Token embedding矩阵1024×384 393,216约0.39M每层自注意力Q、K、V三个矩阵每个都是384×384 147,456三个就是442,368注意力输出投影矩阵又是147,456。加上bias每层注意力约0.59M每层FFN第一层升维384→1536参数384×1536 589,824第二层1536→384又是589,824。加上bias约1.18M6层加起来0.59 1.18×6 ≈ 10.6M最后的LM Head1024×384 ≈ 0.39M如果和embedding共享权重就不计再加上位置编码、层归一化参数总参数落在20M上下这里有一个参数分布陷阱如果词表太大embedding和LM Head会占据大量参数挤压Transformer主体的容量。所以我特别控制词表大小让模型参数尽量集中在“计算”的部分。这也是为什么很多专业模型都用subword而不是整个单词列表。另一个跟参数量直接相关的概念是显存占用。我们常说20M参数很小但实际训练时显存开销远远超过80MB参数本身。AdamW优化器要为每个参数保存一阶矩和二阶矩相当于三倍参数内存这就有240MB了。真正吃显存的是中间激活值——每层都会保存前向计算的中间结果用于反向传播。block_size越长、batch越大激活值占用越恐怖。所以显存不够时优先减batch_size而不是减模型参数。4. 训练工程让Loss真正掉下去的实操细节4.1 损失函数与优化器选型语言模型用的是交叉熵损失。为什么不能用MSE因为模型的输出是一个词表大小的概率分布交叉熵天然适合评估两个分布的差异而MSE是按回归任务设计的对概率分布不敏感。如果你用MSE去训练分类/生成任务梯度信号会变得很弱模型学得极慢。优化器我选AdamW。相比传统AdamAdamW把权重衰减从L2正则中解耦出来只对参数本身做衰减不对自适应动量做衰减。在Transformer这类模型上AdamW的效果明显更稳而且现在几乎成了标配。关键超参数学习率3e-4weight_decay0.1。optimizer AdamW(model.parameters(), lr3e-4, weight_decay0.1)这里需要提一个所有新手都会犯的错不设weight_decay或者设得太大。不设weight_decay模型在长训练下容易过拟合设成0.5以上模型权重被削得太狠loss下不去。0.1是Transformer里被广泛验证的安全区间。4.2 学习率调度与warmup学习率是整个训练过程中最敏感的超参数。两种典型失败模式学习率太大loss直接冲上NaN学习率太小loss降得极其缓慢几个小时后还在原地打转。项目里我用的是“warmup 余弦退火”策略。warmup阶段前500~1000步学习率从接近0线性升到峰值之后按余弦曲线逐渐下降到接近0。为什么需要warmup因为在训练的最初几步梯度方向极其不稳定相当于你在一片完全陌生的区域里跑如果一开始就冲刺很容易跑偏甚至掉沟里。先小步试探等模型大致找到方向了再加速效率反而更高。调度逻辑可以直接用PyTorch的LambdaLRdef lr_lambda(step): if step warmup_steps: return step / warmup_steps progress (step - warmup_steps) / (max_steps - warmup_steps) return 0.5 * (1 math.cos(math.pi * progress)) scheduler LambdaLR(optimizer, lr_lambdalr_lambda)配合峰值学习率6e-4在5000步的训练中loss能稳定地从初始的~8降到~1.5左右。如果是全量微调阶段学习率建议降到1e-4量级因为基座已经学好了更新太猛会把原有能力冲坏。4.3 训练循环的标准配置小模型训练最合适的单次batch大小是32个序列、每个序列128个token。但24GB显存可以一次喂64个序列为什么我还是坚持用32因为batch太小梯度噪声大batch太大又不利于泛化32是折中。如果显存不够就用梯度累积模拟大batchfor step in range(max_steps): x, y get_batch(train) logits, loss model(x, y) loss loss / accum_steps loss.backward() if (step 1) % accum_steps 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() optimizer.zero_grad()注意梯度裁剪。这个操作看似简单其实是训练稳定性的保障如果某个batch里出现极端样本导致梯度爆炸裁剪会把梯度的范数限制在1.0以内防止一次更新把模型权重推飞。我在项目里实测过不加梯度裁剪1000步内必出现一次loss尖峰加了之后曲线平滑很多。另外说下混合精度。20M参数的小模型理论上不需要AMP但在实验阶段我还是开了torch.cuda.amp.autocast因为后面的RL微调阶段需要更快的迭代速度。开了AMP之后要小心fp16的loss scale问题如果日志里出现大量NaN先关了AMP跑200步看看排除模型本身的问题再说。4.4 监控、日志与验证只看训练集loss是训练模型的大忌。我要求每次实验同时记录训练loss和验证loss。训练loss下降、验证loss也下降说明模型真的在学训练loss下降、验证loss不降甚至上涨说明过拟合了需要加dropout或提前停止。日志记录方面我用的是最朴素的方式每100步打印一次step、loss、lr、当前时间并定期保存checkpoint。有条件的可以接wandb或tensorboard但对于这个项目规模纯文本日志完全够用。关键是养成分阶段保存checkpoint的习惯每500步存一个这样万一后面某个时段模型训崩了可以从最近一个正常点恢复不用从头再来。5. 推理生成与推理能力增强5.1 自回归生成流程训练完的模型是个“预测下一个token”的机器。生成文本时把初始提示词编码成token序列让模型预测下一个token的概率分布采样一个token拼到序列末尾再把新序列喂回模型重复这个过程直到生成足够的token或遇到结束符。完整实现如下torch.no_grad() def generate(model, idx, max_new_tokens, temperature0.8, top_k50, top_p0.92): model.eval() for _ in range(max_new_tokens): idx_cond idx if idx.size(1) block_size else idx[:, -block_size:] logits, _ model(idx_cond) logits logits[:, -1, :] / temperature # top-k 过滤 if top_k is not None: v, _ torch.topk(logits, min(top_k, logits.size(-1))) logits[logits v[:, [-1]]] -float(Inf) # top-p 过滤 if top_p is not None: sorted_logits, sorted_indices torch.sort(logits, descendingTrue) cum_probs torch.cumsum(F.softmax(sorted_logits, dim-1), dim-1) sorted_mask cum_probs - F.softmax(sorted_logits, dim-1) top_p sorted_logits[sorted_mask] -float(Inf) logits torch.zeros_like(logits).scatter_(-1, sorted_indices, sorted_logits) probs F.softmax(logits, dim-1) idx_next torch.multinomial(probs, num_samples1) idx torch.cat([idx, idx_next], dim1) return idx这里有几个要点。第一自回归生成必须限制上下文长度只把最后block_size个token喂给模型否则序列无限变长显存爆炸。第二我用的是采样而非argmax因为argmax会退化成复读机随机性才能生成多样化的文本。5.2 采样策略温度、top-k、top-p生成这一步超参数对结果的影响比很多初学者想象的大得多。核心是温度T内部实现是softmax(logits / T)。T1保持原样T1时概率分布变平坦低概率token也有机会被选中生成更“发散”T1时分布变尖锐模型更倾向高概率token生成更“保守”。调温度的心得文学生成用T0.8代码/逻辑任务用T0.6甚至更低。温度太低会陷入重复循环生成内容死板温度太高模型会往句子里乱塞毫无关联的词看起来像病句集锦。top-k和top-p是两道过滤器。top-k只保留概率最高的k个token参与采样防止月度极低的“垃圾token”被选中top-p是动态截断保留累积概率达到0.92的那些token比top-k更灵活。两者一起用效果稳定实测在微型模型上top-k50、top-p0.92是比较好的组合。5.3 让小型模型具备基础推理能力一个RL微调扩展聊到“build a reasoning model from scratch”我们得把目光从生成流畅文本转向有目标的推理。直接用自回归训练出来的模型可以写出通顺的句子但让它解决“排序三个数字”这种简单任务大概率是瞎猜。原因很简单语言模型学的是“下一个字符是什么”不是“答案是什么”。推理能力需要额外的训练信号来引导。这里我用了一个极简的强化学习思路和大模型领域的GRPOGroup Relative Policy Optimization组相对策略优化思想一脉相承。任务是给模型输入[2, 1, 3] sort -期望输出[1, 2, 3]。训练时模型生成一段回答如果回答完全正确给正奖励错误给负奖励。然后用策略梯度方式更新模型增加正确回答的概率、降低错误回答的概率。伪代码如下# 生成一组候选回答 rollouts model.generate(prompts, n_samples4) # 按正确性打分 rewards [1.0 if is_correct(r) else 0.0 for r in rollouts] # 计算组内相对优势 baseline np.mean(rewards) advantages [r - baseline for r in rewards] # 对模型输出的logp计算policy gradient损失 loss -sum(advantage * logp for advantage, logp in zip(advantages, log_probs))这个思路看着简单但有几个必须注意的工程细节。第一基线设置很关键这里用的是组内平均奖励。直接把原始奖励当优势会导致即使全部回答正确模型还是强行调整概率训练不稳定。第二rollouts必须带dropout保持探索多样性否则模型很快就固化在单一答案模式。第三基座模型必须已经具备基本预测能力。你让一个loss还在5以上的模型去学推理它连“数字1后面接着什么”都不知道RL只会空转。我实测用这个方法让20M参数的小模型在“三个数字排序”上达到了相对稳定的正确率但说实话任务稍一复杂就崩。这恰恰说明了一个大模型领域反复被验证的事实推理能力不会凭空出现它需要基座足够的容错能力和表征空间。RL是放大器不是无中生有的魔法。6. 常见问题与排查技巧实录6.1 Loss不下降或直接变成NaN这是从零训练模型遇得最多的问题90%的原因出在三个地方学习率太大、数据顺序没打乱、混合精度溢出。排查顺序建议这样先把学习率降到1e-4关闭AMP确认数据加载时每个epoch都做了shuffle然后看前200步的loss曲线。如果loss在1e-4下还是掉不下去就检查数据——把输入文本打印前100个token确认不是乱码或者全空白。我踩过最冤的一次坑是数据文件里混入了一大段全角空格和零宽字符tokenizer把这些字符全分成了同一个token模型学到了“高频token等于空格”的假规律loss卡在某个点上不去怎么调超参都没用。最后打印输入样本才发现问题清洗完数据loss立刻恢复正常。6.2 生成结果反复重复模型生成的内容如果总是“I am ... I am ... I am ...”通常有四个可能温度太低、上下文太短、模型容量不够、训练步数不足。排查时先把温度升到0.9试试如果还重复再看训练日志——final loss如果远高于验证集上应该有的水平说明模型欠拟合需要更多训练步数如果训练loss已经很低但验证loss偏高那就是过拟合可以加dropout或者提前停止。还有一种情况context window设得太短120个token的上下文对学习莎士比亚这种长句子结构确实不够模型记不住前文只能靠当前几个词硬接。把block_size从128提到256通常有明显改善。6.3 显存溢出和训练崩溃OOMOut of Memory是训练标配不必慌。先算一笔账一个20M参数的模型纯参数占80MBAdamW优化器状态占240MB但中间激活值才是大头。如果batch32、seq_len128、6层Transformer激活值大概占2~4GB。如果超了优先减小batch_size然后考虑开启梯度累积替代、开启gradient checkpointing。20M参数用gradient checkpointing的开销不大可以接受。还有一个隐藏问题dataloader的num_workers设得太大磁盘IO跟不上训练循环卡在数据加载环节GPU利用率变成0。排查方法是看GPU利用率如果一直是0%先检查是不是数据加载瓶颈而不是模型问题。6.4 代码组织与实验可复现从零构建的项目最怕“昨天还能跑今天跑不了”。我强烈建议从第一天就做三件事配置用yaml管理、tokenizer版本绑定hash、每500步保存checkpoint并记录对应的训练配置。这样任何一次实验结果都可以精确复现。不要相信自己的记忆训练日志里要同时记录数据文件hash、模型配置、优化器参数、最终loss和生成样本。我用这套方法一次实验结束后三个月再回来仍然能100%重跑出当时的模型。最后说点实在的做“ai-engineering-from-scratch”这个项目最大的收获不是那个20M参数的模型本身而是我失去了一层“恐惧感”。以前看到别人分享的训练曲线、参数配置总觉得对方很神秘现在知道每个loss数字背后都有实实在在的物理含义每个超参数都有它的作用边界。模型报错的时候我不会再满头雾水地随机改参数碰运气而是能像看心电图一样从曲线形状判断出问题出在哪个环节。给想走这条路的朋友一个建议别去找什么现成的“从零构建大模型”的网盘资源也别抱着书看三遍不动手。直接打开编辑器下载一份莎士比亚文本从字符编码开始写一个BPE搭一个单层注意力模型看着loss曲线掉下去再亲手把梯度调崩一次。这个过程走一遍你学到的东西比看一百篇综述都扎实。AI工程说到底不是知识是手感。 
长河君 · 2026/10/2 5:31:51 · 阅读 1.2k 从零构建20M参数微型模型：数据、Tokenizer到Transformer全流程
“AI工程”这个词这两年算是被说烂了但真正上手时你会发现市面上绝大多数内容是教你调API、装框架真正讲清楚一个模型从数据到推理全链路怎么搭起来的却很少。我最近把项目标题定为“ai-engineering-from-scratch”核心就一句话不用任何现成大模型不靠黑盒推理接口从零开始做一个可以训练、可以推理、可以评估的微型模型。这篇文章就是整个项目的复盘包括数据、tokenizer、Transformer架构、训练策略以及最后用强化学习思路让一个小模型具备基础推理能力的过程。项目规模控制在20M参数左右单张消费级显卡就能跑通适合想彻底搞懂大模型原理的工程师、刚转行做算法的朋友以及不满足于“调包”的深度学习爱好者。很多人问我市面上已经有那么多开源模型了直接拿来做微调不香吗香但理解不了原理。就像你天天开车不一定知道发动机怎么点火、变速箱怎么换挡。真出了问题只能停在路边等救援。“从零构建”就是那个把你按在发动机旁边逼你拆开看一遍的过程。1. 项目定位为什么要做“从零构建”这件事1.1 从调包到懂原理之间差什么调API和做AI工程看起来都是“跟模型打交道”实际是完全不同的两个世界。调API时你关心的是提示词怎么组织、温度参数怎么设置模型内部对你来说就是个黑盒子。而AI工程关心的是数据分布对不对、tokenizer词表合不合理、模型参数量怎么分布、学习率曲线为什么是这个形状、梯度什么时候会爆炸。从零构建最大的价值是把整个链路的所有错误都亲手踩一遍。你会发现模型不work的原因千奇百怪数据里混了不可见字符、tokenizer出现了OOV、embedding初始化不合适、学习率太大导致loss直接飞到NaN、验证集和训练集分布没对齐……任何一环出错模型就废了。这些坑光看文档是看不出来的只有跑起来、崩溃、修复、再崩溃才会真正变成你自己的经验。我见过不少同学模型结构背得滚瓜烂熟能默写Transformer公式但一旦让他从一份txt文本训练出可用的生成模型就完全不知道从哪里下手。原因很简单知识是线性的工程是网状的。从零构建的价值不是让你重复造轮子而是让你拥有拆轮子的能力。1.2 目标边界20M参数能做什么做这种项目最怕的就是目标定得太离谱。我见过有人一上来就要复现GPT-31750亿参数单卡跑不起来就抱怨硬件最后项目烂尾。正确的做法是把目标压缩到可以“一个人、一张卡、几天内完成”的程度。我这边的项目边界是这样定的数据集TinyShakespeare大约1MB的莎士比亚文本切分后约100万token。公开可下载稳定句子结构清晰非常适合用来观察语言模型的生成效果。模型结构标准decoder-only Transformer6层、6个注意力头、d_model384、FFN中间维度1536整体约2000万参数。训练阶段先做自回归语言建模让模型学会“接话”再做一轮带奖励信号的微调让模型在简单符号推理任务上有稳定表现。硬件环境单张RTX 3090 24GB就够。没有3090的话用16GB显存也能跑只是batch要小一点。为什么选莎士比亚数据集两个原因。第一它是纯文本没有复杂的结构化标签词表相对封闭适合跑通全流程。第二它的语言有很强的文体特征训练几个小时就能看到模型生成出像模像样的“莎翁腔”成就感拉满。如果你手头有其他垂直领域文本比如代码、法律文书、病历也可以把后面这套流程原封不动搬过去。关于“reasoning model”的热度我想多说一句。最近大家都很关注OpenAI o1、DeepSeek R1这类推理模型但它们的核心突破其实不在模型结构而在训练策略——大规模强化学习让模型学会在推理时分配更多计算。所以这个项目第二部分我会在一个很小的自回归模型上尝试加入强化学习信号让模型完成类似“数字排序”这种简单推理任务。基座虽小但原理一条不落。2. 数据与Tokenizer模型的第一道关口2.1 数据清洗比想象中更关键做AI工程的第一课就是模型吃进去的是数据分布你喂进去的是垃圾后面所有努力都只是在给垃圾装修。数据清洗这件事听起来没什么技术含量但项目里80%的诡异问题追溯到底都是数据问题。TinyShakespeare原始文件本身算是干净的但我还是走了完整的数据预处理流程读入文本后统一把\r\n替换为\n消除Windows换行符差异。过滤掉所有控制字符只保留可打印字符和换行。检查文本中是否有重复的段落。虽然莎士比亚文本不会有但如果你换用自己的数据这一步必须做——重复数据会导致模型过拟合生成时无限复读。按9:0.5:0.5的比例切成训练集、验证集、测试集。注意切分时要按文本顺序切不要随机打乱后再切否则句子前后文关系会断裂。清洗完的纯文本还要转成模型能吃的东西。最朴素的做法是字符级编码但直接喂给模型不是最优解所以需要用分词器把文本切成更合理的token序列。这个东西太重要了我单独说。2.2 从零实现一个BPE分词器字符级tokenizer的实现很简单建一个字典映射每个字符到整数ID然后把文本逐字符编码。但字符级tokenizer有致命问题——上下文窗口变得非常“短”。比如上下文长度是128个token如果按字符切这128个token只覆盖了大约128个英文字符可能连一句完整的台词都覆盖不了模型能学到的东西极其有限。所以项目里我用了BPEByte Pair Encoding分词器。BPE的思路很朴素先从单个字符开始然后统计相邻字符对的出现频率把最高频的相邻对合并成一个新token重复这个过程直到词表达到预设大小。这样“the”这种常见词会被合并成一个token模型一眼就能看到整个单词效率提升非常明显。核心合并逻辑的代码其实很短你自己实现一遍就懂了def get_pair_stats(ids): stats {} for pair in zip(ids, ids[1:]): stats[pair] stats.get(pair, 0) 1 return stats def bpe_merge(ids, pair, new_idx): newids [] i 0 while i len(ids): if i len(ids) - 1 and ids[i] pair[0] and ids[i1] pair[1]: newids.append(new_idx) i 2 else: newids.append(ids[i]) i 1 return newids实际操作时先在训练集上统计字符对频率迭代执行merge操作得到一份merge记录。这份merge记录就是tokenizer的“词表”后续训练和推理都必须用同一份绝不能换。词表大小怎么选在TinyShakespeare上我试过512、1024、2048三个档位。512的词表偏小很多词根词缀拆不开2048词表在100万token的数据上有不少低频token几乎学不到最终选了1024。经验法则词表大小不要超过数据规模的0.1%~0.3%否则大量token会因为频率太低沦为噪声。这里必须提醒一个坑tokenizer和模型是强绑定的。训练时用了A词表推理时如果加载了B词表轻则出现一堆UNK重则生成内容彻底崩溃。我的做法是给每个tokenizer版本加一个hash值并在训练配置里记录保证任何一次实验都能回溯到正确的分词版本。3. 模型架构搭建从零手写Transformer3.1 自注意力机制的工程实现Transformer核心就是自注意力。理解它最直接的方式是把它看成“让序列里的每个token都向其他token询问我该关注谁”。具体来说每个token会生成三个向量Query我找什么、Key我是什么、Value我有什么。然后计算Query和所有Key的点积得到相似度分数经过softmax变成权重最后用权重加权求和所有Value得到这个token的新表示。在代码里缩放点积注意力就这几行def scaled_dot_product_attention(q, k, v, maskNone): d_k q.size(-1) scores q k.transpose(-2, -1) / math.sqrt(d_k) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) weights F.softmax(scores, dim-1) return weights v, weights注意那个除以sqrt(d_k)。为什么除因为当维度变大时点积的数值会变得很大softmax会迅速进入饱和区梯度接近于零。这就好比一群人同时大声说话你反而一个字都听不清。除以sqrt(d_k)相当于把音量压低到合适区间让注意力权重有区分度。另外decoder模型必须加因果掩码。因为自回归模型只能看到当前token之前的token不能“偷看”后面的答案。初始化时在注意力矩阵的上三角位置填负无穷softmax后这些位置会变成0模型就“看不见”未来了。3.2 多头注意力、残差、归一化与FFN多头注意力就是把注意力过程分成多个头并行做。每个头有不同的Q、K、V投影矩阵相当于开会时分了几个小组每个小组关注不同的维度一个组盯着句子的主谓结构另一个组盯着修辞关系最后把各组的结论拼在一起。这种“各看各的再汇总”的机制让模型表达能力大幅提升。一个标准的TransformerBlock包含四部分多头注意力、残差连接、层归一化、前馈网络。代码骨架如下class Block(nn.Module): def __init__(self, d_model, n_heads, d_ff, dropout0.1): super().__init__() self.attn nn.MultiheadAttention(d_model, n_heads, dropoutdropout, batch_firstTrue) self.ffn nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model), ) self.ln1 nn.LayerNorm(d_model) self.ln2 nn.LayerNorm(d_model) def forward(self, x, maskNone): x x self.attn(self.ln1(x), self.ln1(x), self.ln1(x), attn_maskmask)[0] x x self.ffn(self.ln2(x)) return x残差连接解决的是深层网络的梯度传播问题——让梯度可以“抄近路”直接流回浅层层归一化解决的是每层输出分布漂移的问题。这两个设计是训练深层Transformer的基石缺一个二十层以上的网络基本训不动。在小模型上它们仍然重要但作用没大模型那么“救命”。3.3 参数量手动计算“20M参数”是怎么来的我带你手算一遍。假设词表大小1024d_model384Token embedding矩阵1024×384 393,216约0.39M每层自注意力Q、K、V三个矩阵每个都是384×384 147,456三个就是442,368注意力输出投影矩阵又是147,456。加上bias每层注意力约0.59M每层FFN第一层升维384→1536参数384×1536 589,824第二层1536→384又是589,824。加上bias约1.18M6层加起来0.59 1.18×6 ≈ 10.6M最后的LM Head1024×384 ≈ 0.39M如果和embedding共享权重就不计再加上位置编码、层归一化参数总参数落在20M上下这里有一个参数分布陷阱如果词表太大embedding和LM Head会占据大量参数挤压Transformer主体的容量。所以我特别控制词表大小让模型参数尽量集中在“计算”的部分。这也是为什么很多专业模型都用subword而不是整个单词列表。另一个跟参数量直接相关的概念是显存占用。我们常说20M参数很小但实际训练时显存开销远远超过80MB参数本身。AdamW优化器要为每个参数保存一阶矩和二阶矩相当于三倍参数内存这就有240MB了。真正吃显存的是中间激活值——每层都会保存前向计算的中间结果用于反向传播。block_size越长、batch越大激活值占用越恐怖。所以显存不够时优先减batch_size而不是减模型参数。4. 训练工程让Loss真正掉下去的实操细节4.1 损失函数与优化器选型语言模型用的是交叉熵损失。为什么不能用MSE因为模型的输出是一个词表大小的概率分布交叉熵天然适合评估两个分布的差异而MSE是按回归任务设计的对概率分布不敏感。如果你用MSE去训练分类/生成任务梯度信号会变得很弱模型学得极慢。优化器我选AdamW。相比传统AdamAdamW把权重衰减从L2正则中解耦出来只对参数本身做衰减不对自适应动量做衰减。在Transformer这类模型上AdamW的效果明显更稳而且现在几乎成了标配。关键超参数学习率3e-4weight_decay0.1。optimizer AdamW(model.parameters(), lr3e-4, weight_decay0.1)这里需要提一个所有新手都会犯的错不设weight_decay或者设得太大。不设weight_decay模型在长训练下容易过拟合设成0.5以上模型权重被削得太狠loss下不去。0.1是Transformer里被广泛验证的安全区间。4.2 学习率调度与warmup学习率是整个训练过程中最敏感的超参数。两种典型失败模式学习率太大loss直接冲上NaN学习率太小loss降得极其缓慢几个小时后还在原地打转。项目里我用的是“warmup 余弦退火”策略。warmup阶段前500~1000步学习率从接近0线性升到峰值之后按余弦曲线逐渐下降到接近0。为什么需要warmup因为在训练的最初几步梯度方向极其不稳定相当于你在一片完全陌生的区域里跑如果一开始就冲刺很容易跑偏甚至掉沟里。先小步试探等模型大致找到方向了再加速效率反而更高。调度逻辑可以直接用PyTorch的LambdaLRdef lr_lambda(step): if step warmup_steps: return step / warmup_steps progress (step - warmup_steps) / (max_steps - warmup_steps) return 0.5 * (1 math.cos(math.pi * progress)) scheduler LambdaLR(optimizer, lr_lambdalr_lambda)配合峰值学习率6e-4在5000步的训练中loss能稳定地从初始的~8降到~1.5左右。如果是全量微调阶段学习率建议降到1e-4量级因为基座已经学好了更新太猛会把原有能力冲坏。4.3 训练循环的标准配置小模型训练最合适的单次batch大小是32个序列、每个序列128个token。但24GB显存可以一次喂64个序列为什么我还是坚持用32因为batch太小梯度噪声大batch太大又不利于泛化32是折中。如果显存不够就用梯度累积模拟大batchfor step in range(max_steps): x, y get_batch(train) logits, loss model(x, y) loss loss / accum_steps loss.backward() if (step 1) % accum_steps 0: torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() optimizer.zero_grad()注意梯度裁剪。这个操作看似简单其实是训练稳定性的保障如果某个batch里出现极端样本导致梯度爆炸裁剪会把梯度的范数限制在1.0以内防止一次更新把模型权重推飞。我在项目里实测过不加梯度裁剪1000步内必出现一次loss尖峰加了之后曲线平滑很多。另外说下混合精度。20M参数的小模型理论上不需要AMP但在实验阶段我还是开了torch.cuda.amp.autocast因为后面的RL微调阶段需要更快的迭代速度。开了AMP之后要小心fp16的loss scale问题如果日志里出现大量NaN先关了AMP跑200步看看排除模型本身的问题再说。4.4 监控、日志与验证只看训练集loss是训练模型的大忌。我要求每次实验同时记录训练loss和验证loss。训练loss下降、验证loss也下降说明模型真的在学训练loss下降、验证loss不降甚至上涨说明过拟合了需要加dropout或提前停止。日志记录方面我用的是最朴素的方式每100步打印一次step、loss、lr、当前时间并定期保存checkpoint。有条件的可以接wandb或tensorboard但对于这个项目规模纯文本日志完全够用。关键是养成分阶段保存checkpoint的习惯每500步存一个这样万一后面某个时段模型训崩了可以从最近一个正常点恢复不用从头再来。5. 推理生成与推理能力增强5.1 自回归生成流程训练完的模型是个“预测下一个token”的机器。生成文本时把初始提示词编码成token序列让模型预测下一个token的概率分布采样一个token拼到序列末尾再把新序列喂回模型重复这个过程直到生成足够的token或遇到结束符。完整实现如下torch.no_grad() def generate(model, idx, max_new_tokens, temperature0.8, top_k50, top_p0.92): model.eval() for _ in range(max_new_tokens): idx_cond idx if idx.size(1) block_size else idx[:, -block_size:] logits, _ model(idx_cond) logits logits[:, -1, :] / temperature # top-k 过滤 if top_k is not None: v, _ torch.topk(logits, min(top_k, logits.size(-1))) logits[logits v[:, [-1]]] -float(Inf) # top-p 过滤 if top_p is not None: sorted_logits, sorted_indices torch.sort(logits, descendingTrue) cum_probs torch.cumsum(F.softmax(sorted_logits, dim-1), dim-1) sorted_mask cum_probs - F.softmax(sorted_logits, dim-1) top_p sorted_logits[sorted_mask] -float(Inf) logits torch.zeros_like(logits).scatter_(-1, sorted_indices, sorted_logits) probs F.softmax(logits, dim-1) idx_next torch.multinomial(probs, num_samples1) idx torch.cat([idx, idx_next], dim1) return idx这里有几个要点。第一自回归生成必须限制上下文长度只把最后block_size个token喂给模型否则序列无限变长显存爆炸。第二我用的是采样而非argmax因为argmax会退化成复读机随机性才能生成多样化的文本。5.2 采样策略温度、top-k、top-p生成这一步超参数对结果的影响比很多初学者想象的大得多。核心是温度T内部实现是softmax(logits / T)。T1保持原样T1时概率分布变平坦低概率token也有机会被选中生成更“发散”T1时分布变尖锐模型更倾向高概率token生成更“保守”。调温度的心得文学生成用T0.8代码/逻辑任务用T0.6甚至更低。温度太低会陷入重复循环生成内容死板温度太高模型会往句子里乱塞毫无关联的词看起来像病句集锦。top-k和top-p是两道过滤器。top-k只保留概率最高的k个token参与采样防止月度极低的“垃圾token”被选中top-p是动态截断保留累积概率达到0.92的那些token比top-k更灵活。两者一起用效果稳定实测在微型模型上top-k50、top-p0.92是比较好的组合。5.3 让小型模型具备基础推理能力一个RL微调扩展聊到“build a reasoning model from scratch”我们得把目光从生成流畅文本转向有目标的推理。直接用自回归训练出来的模型可以写出通顺的句子但让它解决“排序三个数字”这种简单任务大概率是瞎猜。原因很简单语言模型学的是“下一个字符是什么”不是“答案是什么”。推理能力需要额外的训练信号来引导。这里我用了一个极简的强化学习思路和大模型领域的GRPOGroup Relative Policy Optimization组相对策略优化思想一脉相承。任务是给模型输入[2, 1, 3] sort -期望输出[1, 2, 3]。训练时模型生成一段回答如果回答完全正确给正奖励错误给负奖励。然后用策略梯度方式更新模型增加正确回答的概率、降低错误回答的概率。伪代码如下# 生成一组候选回答 rollouts model.generate(prompts, n_samples4) # 按正确性打分 rewards [1.0 if is_correct(r) else 0.0 for r in rollouts] # 计算组内相对优势 baseline np.mean(rewards) advantages [r - baseline for r in rewards] # 对模型输出的logp计算policy gradient损失 loss -sum(advantage * logp for advantage, logp in zip(advantages, log_probs))这个思路看着简单但有几个必须注意的工程细节。第一基线设置很关键这里用的是组内平均奖励。直接把原始奖励当优势会导致即使全部回答正确模型还是强行调整概率训练不稳定。第二rollouts必须带dropout保持探索多样性否则模型很快就固化在单一答案模式。第三基座模型必须已经具备基本预测能力。你让一个loss还在5以上的模型去学推理它连“数字1后面接着什么”都不知道RL只会空转。我实测用这个方法让20M参数的小模型在“三个数字排序”上达到了相对稳定的正确率但说实话任务稍一复杂就崩。这恰恰说明了一个大模型领域反复被验证的事实推理能力不会凭空出现它需要基座足够的容错能力和表征空间。RL是放大器不是无中生有的魔法。6. 常见问题与排查技巧实录6.1 Loss不下降或直接变成NaN这是从零训练模型遇得最多的问题90%的原因出在三个地方学习率太大、数据顺序没打乱、混合精度溢出。排查顺序建议这样先把学习率降到1e-4关闭AMP确认数据加载时每个epoch都做了shuffle然后看前200步的loss曲线。如果loss在1e-4下还是掉不下去就检查数据——把输入文本打印前100个token确认不是乱码或者全空白。我踩过最冤的一次坑是数据文件里混入了一大段全角空格和零宽字符tokenizer把这些字符全分成了同一个token模型学到了“高频token等于空格”的假规律loss卡在某个点上不去怎么调超参都没用。最后打印输入样本才发现问题清洗完数据loss立刻恢复正常。6.2 生成结果反复重复模型生成的内容如果总是“I am ... I am ... I am ...”通常有四个可能温度太低、上下文太短、模型容量不够、训练步数不足。排查时先把温度升到0.9试试如果还重复再看训练日志——final loss如果远高于验证集上应该有的水平说明模型欠拟合需要更多训练步数如果训练loss已经很低但验证loss偏高那就是过拟合可以加dropout或者提前停止。还有一种情况context window设得太短120个token的上下文对学习莎士比亚这种长句子结构确实不够模型记不住前文只能靠当前几个词硬接。把block_size从128提到256通常有明显改善。6.3 显存溢出和训练崩溃OOMOut of Memory是训练标配不必慌。先算一笔账一个20M参数的模型纯参数占80MBAdamW优化器状态占240MB但中间激活值才是大头。如果batch32、seq_len128、6层Transformer激活值大概占2~4GB。如果超了优先减小batch_size然后考虑开启梯度累积替代、开启gradient checkpointing。20M参数用gradient checkpointing的开销不大可以接受。还有一个隐藏问题dataloader的num_workers设得太大磁盘IO跟不上训练循环卡在数据加载环节GPU利用率变成0。排查方法是看GPU利用率如果一直是0%先检查是不是数据加载瓶颈而不是模型问题。6.4 代码组织与实验可复现从零构建的项目最怕“昨天还能跑今天跑不了”。我强烈建议从第一天就做三件事配置用yaml管理、tokenizer版本绑定hash、每500步保存checkpoint并记录对应的训练配置。这样任何一次实验结果都可以精确复现。不要相信自己的记忆训练日志里要同时记录数据文件hash、模型配置、优化器参数、最终loss和生成样本。我用这套方法一次实验结束后三个月再回来仍然能100%重跑出当时的模型。最后说点实在的做“ai-engineering-from-scratch”这个项目最大的收获不是那个20M参数的模型本身而是我失去了一层“恐惧感”。以前看到别人分享的训练曲线、参数配置总觉得对方很神秘现在知道每个loss数字背后都有实实在在的物理含义每个超参数都有它的作用边界。模型报错的时候我不会再满头雾水地随机改参数碰运气而是能像看心电图一样从曲线形状判断出问题出在哪个环节。给想走这条路的朋友一个建议别去找什么现成的“从零构建大模型”的网盘资源也别抱着书看三遍不动手。直接打开编辑器下载一份莎士比亚文本从字符编码开始写一个BPE搭一个单层注意力模型看着loss曲线掉下去再亲手把梯度调崩一次。这个过程走一遍你学到的东西比看一百篇综述都扎实。AI工程说到底不是知识是手感。