从零手写语言模型：拆解AI工程全链路与训练实操

后台时不时有人问我：想走AI工程这条路，第一步到底该怎么迈？我的答案一直很固定——从零手写一个语言模型，哪怕是玩具级别的。这个思路和《Build a Large Language Model From Scratch》那本书的核心主张一脉相承：不要上来就调现成的Transformer库，而是把数据管道、分词器、注意力机制、训练循环这些零件一个个亲手搭一遍。做完这件事，你会发现自己看大模型的眼光完全变了。以前觉得处处是黑盒，现在每一层、每一个参数、每一条loss曲线，都是可解释的工程决策。这篇文章就结合我自己的复现经历，把“从零开始AI工程”这条路的全链路拆开来讲，包括核心环节怎么设计、实操中怎么调参、以及那些教程里不会写的坑。

1. 为什么值得从零开始做AI工程

1.1 从零不是重复造轮子，而是拆解黑盒

先说一个经常遇到的灵魂拷问：现在轮子这么多，HuggingFace上一行代码就能加载模型，为什么还要自己从零搭？我的看法是，调包和从零复现，根本是两种能力。

调包解决的是“用”的问题，它让你快速交付业务结果；但从零复现解决的是“懂”的问题，它让你具备改造、调试、优化甚至发明新结构的能力。用做饭来类比：照着菜谱做一道菜，和每天切菜、调火候、观察食材变化练出手感，是完全不同的成长路径。前者能应付一顿饭，后者才让你有底气开一家自己的餐厅。

从零复现的成本其实没有想象中高。现在一块消费级显卡，哪怕是16G显存的入门款，就足够训练一个小规模的GPT模型。花上几个周末，你就能把一条完整的AI工程链路走通：数据清洗、分词器训练、模型搭建、训练循环、推理生成、性能优化。这个过程里的任何一个环节，都会在后续做项目时反复用到。在我看来，这就是性价比最高的AI工程入门方式。

1.2 从工程视角看LLM的知识地图

既然要从零走一遍，首先得把这条路上要踩的点位摸清楚。我按实际开发的顺序，把整条链路分成五层：

这张表就是你的行动清单。我在复现时就是照这个顺序一步步走的。每一层都花时间把原理搞清楚，再动手写代码，而不是一口气把代码抄完。这样虽然慢一点，但每一步踩实了，后面出问题才能定位到具体环节。

《Build a Large Language Model From Scratch》这本书的章节安排也基本沿着这条链路走：从数据处理讲起，逐步搭建注意力机制，再到训练、微调、多模态。它最大的价值不是给了多少代码，而是把“看起来神秘的AI”还原成“一步步可执行的工程”。你要是能自己把这张表上的每行都亲手做一遍，哪怕模型参数只有几百万，获得的经验也比单纯看十篇综述扎实。

2. 从零搭LLM的核心环节与技术要点

2.1 数据层：语料清洗与tokenizer训练

第一个挡在眼前的坎不是模型，是数据。我第一版数据管道几乎是把网上的语料直接塞进去训练，结果loss一直飘，生成出来的句子带着大量乱码。后来才意识到，数据清洗决定了一个模型能力的上限，模型架构只是在逼近这个上限。

清洗工作至少要覆盖三件事：第一，用规则把HTML标签、异常符号、重复段落过滤掉；第二，做文档级去重，不然模型会把高频句子背下来，泛化能力很差；第三，控制语料配比，不同来源的数据别让某一类占绝对主导，不然模型会说话带着一股“贴吧味”。这些工作不复杂，但很费时间，我第一次做的时候大约有一半的精力都花在写清洗规则和看样本上。

接下来是tokenizer。很多人贪方便直接用现成的GPT-2分词器，但如果你用的是中文语料，直接套GPT-2的BPE词表，效果会很差——很多常见汉字会被切成单字甚至字节碎片，序列长度暴涨，训练效率直线下降。更合理的做法是拿自己的语料重新训练一个BPE分词器，注意控制词表大小，中文场景我习惯用20k到50k之间，再检查一下训练语料的配比。

为什么BPE这么重要？因为它直接决定序列长度和词表大小，这两个参数又直接影响模型的内存占用和训练速度。一个好的分词器能让同样语义用更短的token序列表达，等于变相提升了模型的上下文长度。

2.2 模型层：注意力机制的实现细节

模型层是整个从零复现里最需要细抠的部分。Transformer里最核心的是自注意力机制，很多初学者把代码跑通就完事，但几个关键细节没搞懂，后面调优会不断踩坑。

第一个细节是为什么QK点积之后要除以根号d_k。如果不缩放，当维度变大时，点积的方差会跟着变大，softmax的结果会趋向one-hot，梯度几乎消失。除以根号d_k相当于把方差拉回到1附近，这是我后来训练变深时loss崩掉才真正体会到的地方。代码上只是多一步除法，但这一步背后是完整的数值稳定性设计。

第二个细节是mask。训练时要防止每个位置看到未来的token，所以需要构造一个上三角全为负无穷的注意力掩码。很多教程直接传一个布尔mask，但没有解释为什么需要这样，导致有人到推理阶段还以为模型“自动学会了”单向生成。

第三个细节是多头。多头不是简单把并行计算多来几份，而是让每个头关注不同的模式——有的头学语法关系，有的头学习语义相关性。头的数量通常是嵌入门大小的公倍数。我习惯用一个简单规则：d_model = n_head * head_dim，这样每个头的维度保持一致，计算效率最高。

我贴一段自己复现时简洁的自注意力代码，方便对照理解：

def self_attention(x, w_q, w_k, w_v, n_head, mask=None): B, T, C = x.shape q = x @ w_q # [B, T, C] k = x @ w_k v = x @ w_v head_dim = C // n_head q = q.view(B, T, n_head, head_dim).transpose(1, 2) k = k.view(B, T, n_head, head_dim).transpose(1, 2) v = v.view(B, T, n_head, head_dim).transpose(1, 2) scores = (q @ k.transpose(-2, -1)) / (head_dim ** 0.5) if mask is not None: scores = scores.masked_fill(mask == 0, float('-inf')) weights = scores.softmax(dim=-1) out = weights @ v return out.transpose(1, 2).contiguous().view(B, T, C)

这段代码虽然短，但把缩放、mask、多头三个关键点一次覆盖到了。建议自己从头写一遍，而不是复制粘贴，这样能注意到view和transpose之后维度变化的细节。

2.3 训练层：损失函数、优化器与学习率策略

模型搭好之后，训练环节是最让人头秃的。先说损失函数，语言模型用的都是交叉熵。为什么？因为自回归任务本质是一个token分类问题——在给定前文的情况下预测下一个token的概率分布，交叉熵恰好能衡量预测分布和真实分布之间的距离。如果你用的是PyTorch的CrossEntropyLoss，它会自动帮你做log_softmax和负对数似然的合并，数值上更稳定，不用手写。

优化器方面，我强烈建议用AdamW而不是Adam。区别就在权重衰减的处理方式上。Adam会把权重衰减混进梯度的一阶矩估计里，造成正则化效果不稳定；AdamW则是把权重衰减单独放在参数更新那一步，干净利落。对模型来说，正则化很重要，因为它能抑制过拟合，让生成文本更泛化。我在小模型上对比过，同样的超参数下，AdamW在验证集上的表现通常更稳。

学习率策略是训练能否收敛的分水岭。我踩过的大坑是没有加warmup就直接用大学习率开训，结果前几百步loss就飞了。为什么要warmup？因为训练刚开始时模型参数处于一个比较“离谱”的初始状态，梯度方向剧烈变化，这时候用大学习率很容易把参数推到不好的区域。先用小学习率走几百步，让参数慢慢进入一个较平滑的“地形”，再把学习率提到峰值，保险得多。

我的默认组合是这样的：

• 优化器：AdamW（beta1=0.9, beta2=0.95, eps=1e-8）
• 学习率峰值：3e-4（小模型规模）
• warmup步数：总训练步数的5%-10%
• 衰减策略：余弦衰减到峰值的10%

这个组合在从零训练的小模型上几乎不会出大问题，可以作为baseline。之后根据loss的下降速度再微调。

2.4 推理层：采样策略与KV Cache

训练跑通之后，你会进入一个很容易忽略的坑：模型学会了，但生成出来的文本质量很差。这不一定是你训练没到位，很可能是采样策略有问题。

自回归推理的核心是：每一步都根据历史token预测下一个token的概率分布，然后按某种策略从中选出下一个token。最简单的策略是贪心解码，每次都选概率最高的token。但这样做的结果是文本经常陷入重复循环。我一开始就是这么干的，输出两遍就听到模型开始“复读机”，语句来回倒。

更常用的方案是引入随机性。Temperature参数缩放概率分布的软硬程度，temperature大于1会让分布更平滑、输出更多样，小于1则更确定、更保守。Top-k是只从概率最高的k个token里采样，防止小概率垃圾token被选中。Top-p是从累积概率达到p的候选集合里采样，相当于自适应地调整候选数量。我实测下来，temperature设为0.8、top-p设为0.9这个组合，在续写任务上质量最好，既不太随机也不死板。

KV Cache则是一个性能优化点。自回归生成时，历史token的Key和Value其实没有变化，每次推理都重新算一遍，纯属浪费算力。KV Cache的核心思想是只算新token的K和V，把旧的缓存下来直接拼接，生成速度能提升好几倍。我在玩具模型上就做过对比，开启KV Cache之后，同样生成100个token，耗时差不多降了一半。这个优化到上线阶段几乎是必须的。

3. 实操过程：mini GPT从零到生成

3.1 超参数选择与参数规模估算

理论讲了一堆，现在进入实际操作环节。我建议的第一只“小白鼠”配置是这样的：

这套配置的参数量大概在1000万到2000万之间。怎么估算？Transformer参数量基本可以用这个公式粗算：参数量 ≈ L × (4×d_model² + 4×d_model×d_ff) + 词表大小×d_model。把数字代入，4层就是 4×(4×128² + 4×128×512) + 30000×128，大约等于 4×(65536 + 262144) + 3840000 = 445万左右，加上没细算的Embedding和输出层，粗估也就千万级别。这个规模在单卡上完全跑得动。

为什么选这个规模而不是更大？因为我要验证的是“从零走通全链路”这件事，而不是追求一个能跟别人比拼的结果。模型越小，迭代周期越短，越方便调参。我把数据量控制在1亿token左右，在消费级显卡上大约几小时就能完成一轮完整训练。等这个规模跑通了，再往大了推，公式和踩坑经验都还能用。

3.2 搭一个可监控的训练循环

训练循环是AI工程里的“日常作业”，它的好坏直接影响你的调试效率。我自己的训练循环长这样：

for step, batch in enumerate(train_loader): x, y = batch logits = model(x) loss = criterion(logits.view(-1, vocab_size), y.view(-1)) optimizer.zero_grad() loss.backward() grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step()

这段代码虽然简单，但里面藏着一个我后来才学会的重要动作：梯度裁剪。Transformer训练特别容易在某个batch上出现极端梯度，如果不裁剪，一次大的梯度波动就可能让loss彻底爆掉。clip_grad_norm_把整个模型的梯度范数限制在1.0以内，相当于给每次更新加了个安全带。我后来把梯度范数也打进日志里，只要发现它经常触顶，就说明学习率或数据批次里有异常值。

监控指标方面，我强烈建议记录这几样：step loss、学习率当前值、梯度范数、每秒处理的token数。前两个用来判断收敛状态，第三个用来判断是否快崩了，第四个用来评估训练效率。如果你用的是W&B，直接把这些指标打进去就能看实时曲线；不想引入额外工具的话，用CSV记录也完全够用。

这里还有一个省显存的技巧：gradient accumulation。当batch size大到单卡放不下的时候，可以把一个大的batch拆成几个小batch，分别算梯度，累加起来，每攒够一定步数再做一次参数更新。效果上等价于大batch训练，但显存压力小得多。我第一次训练时batch size只敢设32，用梯度累积之后等效batch size能达到128，训练稳定性明显改善。

3.3 判断模型真的收敛了：loss曲线与生成测试

训练跑完，怎么判断模型是真的学会了，还是只是没崩？我习惯双管齐下：看loss曲线，也做生成测试。

先看loss曲线。正常收敛的曲线是前几百步快速下降，然后逐渐放缓，最终在一个平台期附近波动。如果loss曲线一直在震荡但总体向下，这个属于正常，不用担心；如果直接变成一条水平线，可能是学习率太低或数据量不够；如果出现突然的尖峰然后恢复，大概率出现过大的梯度。

再看生成测试。这是AI工程里最朴素的验收手段：给模型一个开头，让它续写，看文本是否流畅、是否连贯。我在测试时会给模型几个不同类型的prompt，比如日常对话、短新闻、故事开头，观察它能不能保持主题不跑偏。

举个例子，我当时用2亿token的数据训练了一个小模型，给它的prompt是“秋天的公园里，银杏叶落了一地”，模型续写出来的是“阳光透过稀疏的树枝，洒在铺满金黄落叶的小路上，几个孩子追逐着翻滚的叶子，笑声在空气中回荡”。虽然称不上惊艳，但读起来通顺、有场景感，这就说明模型已经学到了基本的语法和语义关联。如果生成出来的句子出现大量重复、成分残缺或者驴唇不对马嘴，那就说明训练还不够，或者数据质量有问题。

3.4 用推理接口把模型变成可用产品

模型训练好之后，工程化最后一步是把推理封装成可用接口。我会写一个简单的生成函数，把温度、top-p、max_new_tokens这些参数暴露出来，方便调用方按需求调整：

def generate(model, input_ids, max_new_tokens=100, temperature=0.8, top_p=0.9): model.eval() for _ in range(max_new_tokens): logits = model(input_ids)[0, -1, :] / temperature probs = torch.softmax(logits, dim=-1) # 用nucleus sampling过滤低概率token sorted_probs, sorted_indices = torch.sort(probs, descending=True) cum_probs = torch.cumsum(sorted_probs, dim=0) remove_mask = cum_probs > top_p sorted_probs[remove_mask] = 0.0 sorted_probs /= sorted_probs.sum() next_id = sorted_indices[torch.multinomial(sorted_probs, 1)] input_ids = torch.cat([input_ids, next_id.unsqueeze(0)], dim=1) if next_id == eos_token_id: break return input_ids

这个函数看起来简单，但覆盖了常见场景的需求。把它包一层HTTP服务，用FastAPI两三行就能起一个最小可用的GPT应用。我在这个阶段通常还会顺手加一个最大长度限制，防止生成失控耗光显存。

4. 常见问题与排查技巧实录

4.1 loss变成NaN：先查数值稳定性

从零训练最吓人的事情就是loss突然变成NaN。我第一次遇到的时候整个人是懵的，后来总结出一个排查顺序，屡试不爽：

第一步，检查学习率是不是太高。这是最大的元凶，因为学习率高会让参数更新幅度过大，数值溢出。我建议先降低学习率到1e-4以下试试。

第二步，检查数据里有没有异常值。语料里偶尔会出现极长的连续符号或者全角空格之类的脏数据，这些会在计算时产生极端输入，导致激活值爆炸。清洗规则里要加一道长度过滤和特殊字符过滤。

第三步，检查梯度里有没有NaN。可以在backward之后打印每层参数的梯度范数，找到是哪一层爆的。通常注意力层的概率值softmax如果出现极端情况，也会传染给梯度。

第四步，如果用了混合精度训练，确认梯度缩放和溢出检测是否正常工作。fp16的表示范围比fp32小很多，很容易下溢或上溢，建议直接开启动态loss scaling，或者暂时关掉混合精度排查问题。

4.2 显存不足：梯度累积与检查点技术

显存不够是玩本地模型的老大难。除了前面提到的梯度累积，还有两个常用技巧。

第一个是gradient checkpointing，也叫激活值重计算。它把前向传播时保存的中间激活值丢弃一部分，反向传播需要时再重新计算一遍。代价是训练时间增加大约30%，但显存占用可以大幅下降，有些场景能省掉一大半显存。原理上就是用算力换显存，对单卡用户来说非常划算。

第二个是混合精度训练。把模型参数和梯度用fp16存储，只在需要累加的时候用fp32。AMP是PyTorch内置的支持，开启之后显存占用能下降近一半，训练速度还有提升。需要注意的有两点：一是某些层对精度敏感，比如LayerNorm，需要保留fp32；二是训练中要密切观察loss，如果出现异常NaN，优先排查混合精度相关的数值问题。

4.3 生成质量差：重复循环与空洞内容

训练正常但生成质量差，是最让人焦虑的问题。常见的症状有两个：一是输出陷入重复循环，二是内容空洞无语义。

重复循环的根源，在于自回归模型天然倾向于重复近期出现的高频token。应对方法除了我前面提到的采样参数调整，还有一个实用的工具：repetition penalty。它的做法是在计算下一个token的logits时，把已经出现过的token的分数压低，让模型更不愿意重复。我常用的是1.2左右的惩罚系数，太大会让文本变得不自然，太小又压不住重复。

内容空洞的问题，通常指向数据质量和模型容量。如果你给模型喂的语料都是几条新闻反复洗出来的，它当然只能产出套路化文本。解决办法是加高质量、多样性的语料，尤其是那些信息密度高、逻辑性强的长文本。模型容量不够的话，就得考虑扩大d_model和层数，但要注意别步子迈太大，每次加参数量之后重新调学习率。

4.4 过拟合与小数据陷阱

从零训练小模型还有一个隐蔽的陷阱：过拟合。因为从头训练时语料规模往往有限，模型很容易在训练集上背下来，而不是学出泛化规律。判断方法很简单：如果训练loss一直在降，但验证loss升了，那就是过拟合的典型信号。

应对策略我按优先级排序：第一，加数据量，这是最根本的解法；第二，加正则化，weight decay可以适当调大一点；第三，把模型容量降一点，强迫模型学更通用的特征；第四，训练中间做early stopping，因为小语料上继续训练只会加重记忆，不如适时收手。

我还发现一个有意思的现象：小模型在小数据上“背课文”特别快，但你给它一个没见过的prompt，输出立刻变形。这其实说明模型没有真正学到语言的统计规律。真正有用的做法是不断用新鲜、多样的语料喂它，别让它闲着。

5. 从复现到工程化：把玩具模型变成可用系统

5.1 从单卡到多卡：分布式训练的必经之路

玩具模型跑通之后，如果想把规模放大，很快会遇到单卡瓶颈。分布式训练是AI工程绕不开的一站，但它不是简单装个库就能跑。核心要理解两个概念：数据并行和模型并行。

数据并行最简单，每张卡拿一份模型副本，分一批不同的数据各自计算梯度，然后互相同步梯度，取平均后统一更新。PyTorch里的分布式数据并行就是干这个的，配置好init_process_group之后就能用。模型并行则针对模型太大、单卡装不下的情况，把模型的层拆分到多张卡上，数据要按顺序在卡间流动，通信开销大，但对超大模型是刚需。

从零走一遍之后，你会发现分布式训练的关键不只是代码，还有数据加载和通信效率。数据的shuffle和prefetch如果不做好，多卡训练往往会因为GPU空等而效率反降。我建议从DDP开始，先跑通小规模多卡训练，再考虑更复杂的模型并行方案。

5.2 量化与推理优化：把小模型做快做稳

模型要上线做服务，光能推理还不够，还得快、稳、省资源。模型的量化是个绕不开的话题。

量化就是把模型参数从fp32或fp16压缩到int8甚至int4，换来更小的内存占用和更快的推理速度。你可以直接用现成的量化库，把模型精度降到8bit，模型占据的内存和推理延迟都能下降不少。我用过一个6.7B的模型做量化对比，fp16占13.4G，int8占6.7G，速度大约提升近一倍。注意量化会对精度有轻微损失，对生成类任务往往影响不大，但有些任务，比如结构化输出，会变得不稳定，建议上线前做一轮质量评估。

推理加速还有一招叫批量推理。很多业务场景是多个请求并发，如果把多个请求的token合并成同一个batch再交给模型，GPU的利用率会大幅提升。vLLM这类推理框架在调度和显存管理上做了大量优化，从零复现完之后再去用这些东西，你会看得懂它们到底在优化什么，而不是只会调参。

5.3 从LLM到reasoning model：下一步怎么走

最近“build a reasoning model from scratch”这个话题很火，很多人关心怎么从普通语言模型进化到有推理能力的模型。我的理解是，在完成从零复现LLM之后，道路其实是清晰的：推理能力不是凭空出现的，它来自“训练方式”的升级。

第一步是学会思维链数据。把训练语料中增加带推理过程的长文本，比如数学题分步解答、逻辑推理论述，让模型在自回归生成时习惯“先推理再给结论”的格式。第二步是用强化学习循环，比如GRPO这类方法，让模型在推理任务上反复试错并优化策略。这个方向已经把“从零训练一个LLM”扩展成了“从零训练一个reasoning model”，但它的底座，仍然是你从零复现时掌握的那些东西：数据管道、模型架构、训练循环、推理采样。

所以我建议已经走通从零训练的读者，下一步别急着追新模型，先把思维链数据做出来，再尝试用强化学习在小模型上微调推理能力。你会惊讶地发现，原来“推理”也可以通过一套清晰的工程流程来逼近。

做完整轮从零复现之后，我最深的感触是：AI工程不是看会的技能，是做会的技能。你读十篇讲注意力机制的文章，都不如自己亲手把那个mask写错一次然后debug两小时来得深刻。现在再遇到一个新模型，我看它的眼光完全不一样了——先问数据怎么洗的，再问训练策略怎么配的，最后才轮到模型结构本身。这也是为什么我一直建议想入行的人从from scratch开始：它给你的不是一堆API调用经验，而是一条完整的、属于自己的工程直觉。

最后分享两个小技巧。第一，尽量保存每个阶段的checkpoint，包括中间状态。某次训练失败后，你想回退到几天前的参数再调，没有checkpoint就只能整个重来。第二，训练日志一定要记全，尤其是学习率和梯度范数。它们在你排查问题时，是最可靠的“案发现场”线索。