从零构建语言模型：AI工程的极简实践与避坑指南

如果只看现在的招聘 JD，你可能会觉得「AI 工程」是被大厂的 GPU 集群、算法团队和 MLOps 平台垄断的领域，个人开发者只能站在别人的模型后面调参数。但我决定反着来。两年前我开始了一个项目 ai-engineering-from-scratch，目标是在没有现成 transformers 库、没有预训练权重、没有 MLOps 工具的条件下，从数据、分词、Transformer 前向/反向传播到训练循环，把所有环节亲手实现一遍，最后得到一个小但能用的语言模型。这段经历对我的帮助，远比训练出一个高分模型更大。

这篇复盘适合三类人：想把大模型当黑箱用得明白的工程同学；准备转 AI 方向、从 sklearn 跨到 Transformer 的算法工程师；以及跟当初的我一样，会调框架但不懂原理的“调参侠”。你不需要很强的数学背景，但需要有点耐心。我的实践也参考了市面上那本很火的《Build a Large Language Model from Scratch》，但我不主张照抄它的代码，而是把它当成一份实验大纲，把每个思想拆开，再用自己的代码装回去。下面全部是我真正跑过、真正踩过坑之后留下来的记录。

1. 内容整体设计与思路拆解

1.1 为什么极简路线比堆料更接近 AI 工程

面对动辄几十万行代码的开源大模型库，很多人会陷入一种“看了等于会了”的错觉：库用得很顺，API 背得很熟，但真到模型输出胡言乱语、loss 不下降、显存不够用的时候，整个人是懵的。问题在于现代框架把太多细节封装在了抽象层下面，你看到的是接口，而不是原理。

from scratch 的思路恰好相反：每一个模块都短到能一眼看完，每一处乘法都能在白纸上推出来。比如我手写了softmax(q @ k.T) / sqrt(d_k)之后，才真正理解为什么 QK 点积结果要除以维度开根号——因为向量维度变大后，点积数值会跟着变大，softmax 会快速进入饱和区，梯度变得极小甚至消失。这种“知道为什么”的感觉，是直接调库永远给不了你的。

更重要的是，从零开始做一遍，等于给后续所有上层工作建了一张地图。以后再看到新论文里的 Mamba、MoE、MLA 这些变体时，你能瞬间判断它们改了哪一块、为什么要改，而不会像看天书一样。

1.2 项目的最终形态：一个能端到端跑通的最小仓库

我最终把这个项目做成了一个完整的仓库，里面不是几个孤零零的 .py 文件，而是一条看得见全貌的流水线：

ai-engineering-from-scratch/ ├── data/ # 原始语料 + 清洗脚本 ├── tokenizer/ # 自实现 BPE 分词与编码/解码 ├── model.py # decoder-only Transformer ├── train.py # 训练循环、日志、checkpoint ├── eval.py # 困惑度与生成评测 ├── generate.py # 交互式采样 └── experiments/ # 每次训练的实验配置与结果记录

这个小模型的配置并不豪华：理解它、训练它、修改它，全部控制在几百行代码以内。但结构和真实大模型项目是一样的：有数据管道、有词表训练、有模型定义、有训练策略、有验证和部署脚本。也就是说，我练的是“工程思维”，而不仅仅是“写个模型”。

1.3 选型背后的三条纪律

我给自己定了三条硬规矩，实践证明每一条都值得遵守。

第一，每个模块必须能在两百行以内读懂全貌。超过这个量级，说明设计有问题，或者你在提前复杂化。第二，一切必须能跑在一张消费级显卡甚至 CPU 上。因为只有在频繁实验的环境中，你才会真正珍惜每次 trial 的反馈速度，也才敢大胆调参数。第三，每次训练都要有版本记录和 baseline。没有 baseline 的实验就是盲人摸象，调了半天全靠感觉。

这三条纪律让这个项目没有走偏。它始终是一个“可复现、可诊断、可成长”的小型系统，而不是一个跑一次就丢的玩具脚本。

2. 核心细节解析与实操要点

2.1 模型结构：为什么是 decoder-only，而不是 encoder-decoder

做文本生成任务，最直接的选择就是 decoder-only Transformer。Encoder-decoder 比如 T5 或 BART 在翻译、摘要这类“输入长度和输出长度差异大”的任务上确实有优势，但对于自回归生成来说，它的结构更复杂，需要在编码器和解码器之间做交叉注意力，反向传播路径更长，内存占用也更高。对于一个小型 from-scratch 项目，decoder-only 是最能聚焦核心问题的选择。

decoder-only 的核心机制是因果注意力，也叫 masked self-attention。意思是模型在预测第 t 个 token 时，只能看到前 t-1 个 token，绝对不能看到未来。这就像一个考生做题时必须把后面的试卷盖住，否则就变成开卷考试了。实现方式非常简单：把注意力分数矩阵的上三角部分替换成负无穷，softmax 之后这些位置的权重就会归零。

在真实代码里，我用了一个更工程化的写法：先预先构造下三角布尔矩阵，再用masked_fill把非下三角部分变成-inf。这也是训练稳定性的关键点之一。后面的翻车现场里我会专门讲，当 fp16 混合精度打开时，这个-inf如果被错误地传给 softmax 的指数运算，很可能直接变成 NaN。

2.2 数据与分词：BPE 的完整闭环

很多从零开始的项目把注意力全放在模型上，最后发现生成效果差得很，其实是分词器出了问题。我一开始甚至试过直接用字符级 token，因为中文字符单独编码看起来也能用，但效果不够好。后来还是老老实实用 BPE（Byte Pair Encoding）。

自实现 BPE 的步骤很简单，但每一步都有讲究：

1. 先把语料预分词成单词或子串，再转成 UTF-8 字节序列；
2. 统计相邻字节对的出现频率；
3. 每次都把出现频率最高的字节对合并成一个新 token；
4. 重复合并，直到词表达到目标大小，比如 8726。

中文用 BPE 有一个别的方案没有的好处：它能在字节层面自然地处理生僻字和混合语言，不会遇到“因为词表里没有这个词所以编码失败”的问题。至于那一堆“aaaabbbb”之类的字节对，它们只是中间产物，训练完成后词表会慢慢出现高频中文词、标点、英文单词碎片。

数据清洗是这里最繁琐但最值得做的工作。我跑完第一版模型后发现生成文本里有大量重复的“的的的”“了了了”，排查到最后发现是原始语料本身有很多重复句子。后来我写了一个简单的去重脚本，把相似度超过阈值的文本去掉，又丢掉所有包含 HTML 标签的行，再人工抽查 500 条，生成质量才明显改善。记住一个原则：模型能吃进多少数据不重要，数据干净程度才重要。

2.3 超参数与显存估算：怎么用数学而不是感觉定参

我第一次定超参数完全靠感觉，结果 loss 死活不降。后来我把纸笔拿出来算了一遍，才发现参数规模、上下文长度和训练代价之间的关系比想象中清晰得多。

我的最终配置如下：

这个参数量怎么来的？如果暂时忽略位置嵌入和 LayerNorm，计算如下：

• Token embedding 是vocab_size × d_model，约 8726 × 192 ≈ 1.68M；
• 每个 Transformer Block 中，注意力部分有 QKV 和输出投影四个矩阵，参数是4 × d_model²，即 4 × 192² ≈ 147K；
• MLP 部分通常是先升到 4 倍维度再降回来，参数是2 × d_model × (4 × d_model)，约 2 × 192 × 768 ≈ 295K；
• 一个 Block 合计约 442K，4 个 Block 就是约 1.77M；
• 如果让模型的最后一层 LM Head 与 token embedding 共享权重，总参数可以压到约 3.4M。

权重的内存只有3.4M × 4 字节 ≈ 13.6MB，加上梯度、AdamW 优化器状态和激活值，在现代显卡上非常轻松。很多人以为训练模型一定需要大显存，其实是把“大模型”和“训练大模型”强行绑定在一起了。小模型同样能训练，哪怕只想让 loss 降下去一点点，也需要经历完整的工程流程。

2.4 优化器、学习率与损失函数：训练稳定的最后一块拼图

模型结构对了，数据也干净了，训练还是会翻车，原因通常出在优化过程。我在这里吃过不少亏，所以直接给出结论。

损失函数用交叉熵，这没什么争议。需要注意的只有两个点：一是要把 logits 的形状从(B, T, vocab_size)展平成(B×T, vocab_size)，标签也同样展平；二是如果使用了 padding token，要记得在损失计算时把 padding 位置的标签设成ignore_index，否则模型会浪费时间学“填充符号也会被预测”这件事。

优化器我选了 AdamW，而不是普通 Adam。区别在于 weight decay 只作用于权重，不作用于偏置和 LayerNorm 参数。这对 Transformer 的稳定训练很重要，也是我调了一个晚上才明白的坑。

学习率必须配合 warmup 和 cosine 衰减。训练刚开始时梯度统计量还没热身，直接用大学习率容易让 loss 瞬间爆炸；训练后期用线性或 cosine 降到很低的学习率，则有助于收敛到更平缓的极小值区域。我最后用的 schedule 是 300 步 warmup，从 0 升至 6e-4，再按余弦曲线降到接近 0。

同时，梯度裁剪建议一定要加：torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)。这句话看着不起眼，但它拯救过我好几次。遇到 loss 突然跳变或者不收敛，先看是不是梯度范数爆了，再谈别的。

3. 实操过程与核心环节实现

3.1 先用手写 Numpy 把注意力想明白

正式切换到 PyTorch 之前，我先用纯 Numpy 实现了一个单序列的单头注意力。这一步极其推荐，因为 numpy 代码没有任何自动求导，你必须亲手写出每个矩阵的维度，才能真正理解注意力机制。

我写的最简版本是这样的：

import numpy as np def softmax(x): e = np.exp(x - x.max(axis=-1, keepdims=True)) return e / e.sum(axis=-1, keepdims=True) def self_attention(q, k, v, mask=None): # q, k, v 的形状都是 (T, d_k) d_k = q.shape[-1] scores = q @ k.T / np.sqrt(d_k) if mask is not None: scores = np.where(mask, scores, -np.inf) weights = softmax(scores) return weights @ v

第一次跑这个函数，我犯了一个特别蠢的错误：把q @ k.T写成了k @ q.T。结果输出当然全错，但也正因为这个错误，我意识到注意力矩阵的行索引对应查询 token，列索引对应键 token。这个感知在后来的所有调参中都帮了大忙。

Numpy 版本完整跑通之后，我才把它移植成 PyTorch 模块。这个“先痛一下再偷懒”的流程看起来很笨，实际回报非常高，因为在 PyTorch 里就算形状错了，自动广播也可能帮你把错误掩盖过去。

3.2 用 PyTorch 重写最小 GPT

PyTorch 版本不需要很长，核心就三块：因果自注意力、Transformer Block、GPT 主类。

import torch import torch.nn as nn import torch.nn.functional as F class CausalSelfAttention(nn.Module): def __init__(self, d_model, n_head): super().__init__() self.n_head = n_head self.d_head = d_model // n_head self.wq = nn.Linear(d_model, d_model, bias=False) self.wk = nn.Linear(d_model, d_model, bias=False) self.wv = nn.Linear(d_model, d_model, bias=False) self.wo = nn.Linear(d_model, d_model, bias=False) def forward(self, x): B, T, C = x.shape q = self.wq(x).view(B, T, self.n_head, self.d_head).transpose(1, 2) k = self.wk(x).view(B, T, self.n_head, self.d_head).transpose(1, 2) v = self.wv(x).view(B, T, self.n_head, self.d_head).transpose(1, 2) scores = (q @ k.transpose(-2, -1)) / (self.d_head ** 0.5) mask = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device)) scores = scores.masked_fill(~mask, float('-inf')) attn = F.softmax(scores, dim=-1) out = (attn @ v).transpose(1, 2).reshape(B, T, C) return self.wo(out) class Block(nn.Module): def __init__(self, d_model, n_head): super().__init__() self.ln1 = nn.LayerNorm(d_model) self.attn = CausalSelfAttention(d_model, n_head) self.ln2 = nn.LayerNorm(d_model) self.mlp = nn.Sequential( nn.Linear(d_model, 4 * d_model), nn.GELU(), nn.Linear(4 * d_model, d_model), ) def forward(self, x): x = x + self.attn(self.ln1(x)) x = x + self.mlp(self.ln2(x)) return x

这里我用了 Pre-Norm 结构，也就是先 LayerNorm 再进注意力或 MLP，而不是传统 Post-Norm。原因是 Post-Norm 在层数变多时容易出现梯度不稳定，Pre-Norm 在训练稳定性上更友好，小模型也适用。

GPT 主类负责把 token embedding、位置 embedding 和所有 Block 拼起来：

class GPT(nn.Module): def __init__(self, vocab_size, d_model, n_layer, n_head, context_len): super().__init__() self.tok_emb = nn.Embedding(vocab_size, d_model) self.pos_emb = nn.Parameter(torch.zeros(1, context_len, d_model)) self.drop = nn.Dropout(0.1) self.blocks = nn.ModuleList([ Block(d_model, n_head) for _ in range(n_layer) ]) self.ln_f = nn.LayerNorm(d_model) self.lm_head = nn.Linear(d_model, vocab_size) for p in self.parameters(): if p.ndim > 1: torch.nn.init.normal_(p, mean=0.0, std=0.02) def forward(self, idx): B, T = idx.shape h = self.drop(self.tok_emb(idx) + self.pos_emb[:, :T, :]) for block in self.blocks: h = block(h) logits = self.lm_head(self.ln_f(h)) return logits

位置 embedding 我一开始用的是torch.randn初始化，后来改成torch.zeros。原因是随机初始化会在训练早期引入较大的噪声，而零初始化让模型先从纯 token embedding 学起，位置信息通过梯度慢慢注入，训练更平稳。

3.3 训练循环：从代码到可观察的日志

模型定义完了，训练循环反而没那么神秘。但细节决定成败，尤其是批数据构造、学习率调度和梯度裁剪的顺序。

我的核心训练代码大约长这样：

x, y = next(batch_iter) logits = model(x) loss = F.cross_entropy( logits.view(-1, vocab_size), y.view(-1) ) opt.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() lr_scheduler.step()

这里有一个容易忽略的点：y应该是x的下一 token 序列。也就是说，对同一个输入批次，模型看到[t1, t2, t3, ...]，预测目标却是[t2, t3, t4, ...]。这种滑动窗口式的构造方法，保证了每个 token 都既当过输入也当过预测目标。

实验记录也很关键。我每 500 步跑一次验证集困惑度，把 train loss、val loss、learning rate 全部写进一张 CSV 表。这些日志后来帮我看清了很多问题，比如”loss 在训练集上很好看，验证集上却不降”，基本都能从日志里直接定位。

3.4 采样与生成：temperature、top-k、top-p 怎么选

模型训练完，下一步是让它生成文本。生成的核心是自回归采样：一次只预测下一个 token，把它拼到输入后面继续预测。看似简单，但采样策略直接决定了生成结果的可读性。

我常用的生成参数是temperature=0.8, top_k=40, top_p=0.9。temperature 控制概率分布的尖锐程度，值越小越保守，值越大越天马行空；top_k 只从概率最高的 40 个 token 里采样；top_p 则从累计概率达到 0.9 的最小 token 集合里采样。几个策略一起用，是为了既保证多样性，又避免选到那些明显不合理的低概率 token。

值得一提的是，如果你的模型很小、训练数据量也不大，生成结果会频繁出现重复循环，比如“太阳太阳太阳”这种。这时不要急着加更大模型，先把 top_k 调小一点、temperature 调低一点，效果立竿见影。

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

4.1 训练日志里四个危险信号

从零开始的路上，我记录了很多训练日志，慢慢总结出四个危险信号。

第一个是 loss 直接变 NaN。原因通常有三个：学习率太大、梯度爆炸、注意力 mask 里的-inf在 fp16 下“漏电”。第二个是 train loss 持续很低但 val loss 很高，这是过拟合的典型标志。第三个是 loss 在某一步突然跳高，但随后又恢复正常，这是数据里有异常样本或学习率调度碰到了尖锐极值。第四个是生成文本出现大段重复，说明模型对上下文利用不足，或者数据本身多样性差。

光看一个信号不能定位问题，最好每次都记录两条曲线：train loss 和 val loss。我见过不少人只盯着训练集 loss 看，训练结束才发现验证集一塌糊涂，这等于考试看答案答题。

4.2 三次真实翻车现场与修复过程

第一次翻车是 NaN。我打开了torch.autocast混合精度，训练到第 300 步 loss 变成 NaN。排查了半天，终于发现是注意力分数里的-inf在低精度下变成了无效值。解决办法是暂时关闭 attention 部分的 autocast，或把masked_fill里的负无穷改成一个足够大的负数，比如 -1e9，在 fp16 下更安全。

第二次是分词器引起的“灾难”。我在语料里加了<|endoftext|>标记，但训练代码里的特殊 token ID 和分词器训练时的不一致，导致模型经常生成乱码。最后我把vocab_size统一成“BPE 词表大小 + 特殊 token 数”，并强制禁止模型生成几个保留 token，问题才解决。

第三次是过拟合。我的模型只有 3.4M 参数，但训练数据只有几万行小故事，跑 5 个 epoch 之后 val loss 开始回升。我没有粗暴地增加模型容量，而是先加了 Dropout、增大数据多样性、提前早停。事实证明在小模型上，增加数据远比增加参数有用。

4.3 避坑速查表：一个问题的症状与对策

我把踩过的坑整理成一张表，每次实验前扫一眼都能省很多时间。

这张表的价值在于“先检查什么”。很多人一看到 loss 变差就马上改模型结构，其实 80% 的问题都出在更基础的地方。

5. 工具选型与进阶方向

5.1 PyTorch、JAX、Numpy：为什么这么混着用

有时候会收到这类问题：“既然有 PyTorch，为什么还要先手写 Numpy？”我的回答是：它们解决的问题不一样。

Numpy 版本是教学工具，让每个矩阵乘法都暴露在眼皮底下。PyTorch 版本是真正的训练引擎，它帮你处理自动求导和 GPU 加速。JAX 则提供了函数式编程风格，在大规模并行训练和论文复现中更受欢迎，但对调试不太友好，新手容易陷入“哪里出错都不知道”的困境。TensorFlow 当然也能用，只是我的个人体验是它在动态图和生态上的学习成本偏高。

我的建议很直接：动手项目从 PyTorch 开始，遇到需要验证数学原理的地方切回 Numpy 手推一遍，等对整个流程很熟了再考虑 JAX。工具是服务理解的，没必要为了“潮流”牺牲调试体验。

5.2 从语言模型到推理模型：尝试 build a reasoning model from scratch

很多人问“练完一个语言模型之后还能做什么”。最近热门的答案之一是：接着搭一个 reasoning model。这跟你训练出来的普通语言模型最大的区别在于，普通模型只会“接下文”，推理模型要会“多走几步再下结论”。

从一个完型的语言模型往推理方向扩展，我亲测可行的一条路是做 SFT，也就是用推理轨迹微调。我构造了一批小学数学题，每条数据都包含思考过程和最终答案，格式类似：

问题：一袋糖果有 34 颗，小明吃了 7 颗，还剩多少颗？ 思考：34 减去 7，先算 30 减 7 得 23，再加 4，结果是 27。 答案：27

训练时让模型学会先输出思考过程，再输出最终答案。评测时只对“答案”部分做精确匹配。这个思路其实就是现在很多推理模型的基础形态，只是规模小得多。

更进阶的路是引入可验证奖励的强化学习。简单说，就是让模型对同一个问题采样多个答案，用规则自动判断答案是否正确，再根据“正确答案的样本概率被提升、错误答案的概率被压低”来更新模型。在小规模场景下，不需要搬出复杂的 PPO，用带 baseline 的 REINFORCE 就能写出一个雏形。这个方向很适合作为下一个 from-scratch 项目，难度比训练语言模型高，但价值也更明显。

5.3 部署与评测：让最小模型完成工程闭环

训练和评测都通过之后，我把它拆成了一个能调用的服务。没有上 Kubernetes，也没有搞模型网格，因为小模型根本不需要这些东西。我只写了一个 FastAPI 接口，把generate.py里的采样逻辑包一层，然后用一个简单的请求队列做并发控制。这是最朴素的工程闭环，也是最有教学意义的。

评测部分我分了两层：一是困惑度，衡量模型在验证集上的整体表现；二是任务级准确率，比如回答数学题时是否输出了正确数字，生成文本时是否包含指定关键词。后者更能反映业务价值，也更容易让人理解模型有没有进步。

如果你也想开一个类似的项目，我的建议是：从 10 分钟能跑完的训练开始，先做小，再做大。先接受 loss 为 4，再去追求 2。每一步都记录为什么，别急着跨过中间过程。那些中间过程不是阻碍，它们才是这个项目的真正收获。