ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

从零构建大模型:手写Transformer、BPE分词与推理模型训练实践

从零构建大模型:手写Transformer、BPE分词与推理模型训练实践

做AI工程这些年,我见过两类人:一类是每天早上打开API文档,复制几个接口就交差;另一类是连损失函数曲线都不会看,但非要号称自己“从零训练了大模型”。说实话,这两类人在真实项目里都会踩坑,只是坑的深度不同。今天想聊的是一个我最近完整走了一遍的项目——ai-engineering-from-scratch,用大白话讲就是“从零开始做AI工程”。

这个项目到底在做什么?简单说,就是不动任何现成的开源权重,不使用任何大模型API,从语料采集、分词器训练、Transformer模型搭建、预训练、指令微调,到最后把推理能力和评估体系全部手写一遍,目标是产出一个能回答数学推理问题的小模型。类似的玩法,大家如果看过《Build a Large Language Model from Scratch》或者最近很火的“build a reasoning model from scratch”系列,应该不陌生。

适合哪些人看?如果你刚入门NLP,天天被“自注意力”“位置编码”这些概念绕晕;或者你已经在调API做应用,却一直没搞懂模型内部到底怎么运转;又或者你纯粹想验证一下自己的工程能力,那这篇文章应该能帮到你。我会把整个过程里踩过的坑、验证过的方法、真正起作用的代码路径都摊开来说。

这个项目最终跑通的时候,我心里是很爽的:一个参数只有1亿左右的迷你模型,居然也能在简单的数学推理和自然语言理解任务上给出基本合理的答案。虽然和商用大模型差得远,但它让我对“AI工程”这四个字的理解,比调一年API都来得深。

1. 从零构建到底意味着什么:先想清楚再从零

动手之前,先把“from scratch”这个词拆清楚。很多人口口声声说从零,结果一查代码,用的是HuggingFace的transformers库,几行代码就把GPT-2的权重加载下来,这不算从零,这是组装。真正的从零,是让你亲手把每一块积木做出来,而不是把所有积木倒出来再拼回去。

1.1 调API和从零构建,差的不是代码量

我见过不少做应用的朋友,对模型的认知停留在“给它一段提示词,它就会回答”。这当然没错,但一旦遇到模型输出不稳定、回答逻辑不对、或者需要定制领域能力的时候,就会陷入一个尴尬的境地:你完全不知道去哪里调,只能靠改提示词碰运气。这就是调API的局限,你的控制边界停在“输入”和“输出”两端的中间地带,而对黑盒内部的运转逻辑一无所知。

从零构建模型则相反,你被迫把整条链路走一遍:数据怎么清洗、词表怎么定、序列怎么填充、注意力掩码怎么算、损失函数怎么设计、学习率怎么衰减、温度系数怎么影响采样。每走一步,都是在给“AI工程”四个字补课。代码量上,从零构建确实比调API多出好几个数量级,但这部分多出来的代码,恰恰是你对模型建立掌控感的过程。

还有一个很重要的原因是职业层面的。现在AI领域的岗位越来越卷,光会调API已经很难形成壁垒,但如果你能讲清楚自己从零训练过一个模型,哪怕只是一个小模型,面试官对你的认知也会完全不一样。毕竟真正做模型训练的团队,要的不是只会写Prompt的人,而是能看曲线、能调参、能定位bug的人。

1.2 “从零”到底指哪些环节:五层拆分

当我决定做这个项目时,我先把“从零”拆成了五个必须覆盖的环节,缺一个都算不上完整:

  • 数据层:语料采集、清洗、去重、切分。这一步常常被人忽略,但模型能学到什么,上限其实由数据决定。
  • 分词层:不直接沿用现成词表,而是自己训练一个BPE分词器,理解词表大小和压缩率之间的关系。
  • 模型层:手写Transformer架构,包含嵌入层、多头自注意力、前馈网络、层归一化和输出投影,还要自己实现位置编码。
  • 训练层:自己写训练循环,包括批次采样、前向传播、损失计算、反向传播、梯度累积、混合精度和动态学习率。
  • 推理层:实现自回归解码、温度采样、Top-k/top-p过滤,以及评估指标的计算逻辑。

这五个环节,对应着《Build a Large Language Model from Scratch》这本书的核心覆盖范围。说实话,书里有不少代码可以直接参考,但真正的价值在于你要亲手敲一遍,并且在自己的机器上把它跑起来。我建议任何人做这个项目时,都不要直接抄书上的代码,而是先自己写,卡住的时候再看书,这样印象深得多。

我在实践过程中还加了一个第六层:指令微调和推理能力激发。因为光有预训练模型,它只会接着你的话继续写,不会真正“回答”问题。想让它变成一个能解决数学推理问题的模型,还得做后续的训练调整。

2. 核心技术点拆解:从数据准备到推理能力激发

“从零构建”不是一句口号,它具体体现在每一个技术决策上。下面我按数据工程、分词器、模型架构、训练目标、推理能力几个层次,把这套东西的核心逻辑说透。

2.1 数据工程:没有高质量语料,一切白搭

我一开始犯的最大错误,就是图省事,直接拿维基百科纯文本开练。结果练出来的模型什么都“知道”一点,但什么都不精,问它三加五等于几,它能跟你扯二十分钟哲学。原因很简单:维基百科是描述性文本,里面几乎没有“问答对”,模型根本没有机会学习“问题→答案”这种映射关系。

后来我把数据拆成了两部分:第一部分是通用语料,采用百科全书和小说的混合,目的是让模型掌握基础语法与世界知识;第二部分是指令与推理数据,采用数学计算题、常识推理题、多步逻辑题的问答对,格式统一为“问题 + 回答 [+ 推理过程]”。

数据清洗也是个大工程。HTML标签要剥掉、乱码字符要过滤、重复句子要去重、连续空格要归一化。我写了一个清洗管线,每一条文本进去,至少过五道关卡。清洗前后的文本总量从原始的八百兆文本缩到了两百多兆,其中将近一半都是重复和噪声。别看这一步不产生任何代码,它直接影响后面训练效率,脏数据进,垃圾模型出,这句话不是开玩笑。

2.2 分词器:BPE并不是玄学

分词器是很多人忽视的环节。你可能会想,直接用字符级分词不行吗?技术上可行,但序列长度会爆炸。假设一句话平均是50个字母,字符级分词后就是50个token,而用BPE分词后可能只有15个token。序列长度短,训练和推理都更快,而且模型能更好地处理“词根 + 前后缀”这类组合变化。

BPE(Byte Pair Encoding)的核心逻辑并不复杂:先以单个字符为基本单位,然后反复统计相邻单元的出现频率,把出现频率最高的一对单元合并成一个新单元,不断重复,直到词表达到预设大小。我训练了一个词表大小为12000的BPE分词器,合并次数大概是12000减256(256个字节初始单元)。这里的关键心得是:词表大小要和语料规模匹配。语料只有两百多兆,词表设到五万以上就会稀疏,很多token只出现几次,学不到可靠的词向量,反而拖慢训练。

分词器还有个细节容易踩坑:给模型输入时,序列要补齐到同一个长度。我个人用的方法是:超过最大长度(512)直接截断,不足的部分用特殊填充token补齐。而注意力掩码必须同步处理,把填充的部分遮住,让注意力机制不要被无意义的填充token干扰。

2.3 模型架构:Transformer的自注意力与位置编码

模型架构上,我选择了一个相对标准的Transformer解码器结构,参数大概是九千多万,接近一亿。具体配置如下:

超参数取值
层数(n_layers)8
嵌入维度(d_model)512
注意力头数(n_heads)8
前馈层维度(d_ff)2048
最大上下文长度512
参数量约9千万

自注意力是整套机制的核心。简单理解,它就是把输入序列里的每个token与序列里其他所有token做一次“相关性打分”,然后根据得分加权汇总信息。这样一来,模型在理解“它”指代什么的时候,能把整个上下文纳入考虑,而不是只看前后几个词。

位置编码则是为了解决“顺序”问题。自注意力本身是顺序无关的,你把“我爱你”和“你爱我”换成同样的词向量排列,在不加位置信息的情况下,模型看到的是完全一样的表示。我用的是可学习位置编码,也就是给每个位置分配一个可训练的向量,让模型自己学会位置间的相对关系。更进阶的做法是RoPE(旋转位置编码),它天然适合外推更长序列,我后面在部署阶段才替换成它。

一个容易被忽略但非常关键的细节是因果掩码,也叫自回归掩码。Transformer产生下一个token时,绝不能看到未来的信息。实现方式是在自注意力的得分矩阵上,对右上角区域加上一个极大的负数,经softmax后这些位置的权重趋近于零,模型就只能“向前看”。

2.4 训练目标:从语言建模到推理能力的来源

预训练阶段,训练目标就是标准的语言建模:给定前面若干个token,预测下一个token。损失函数用交叉熵。这个目标看似简单,但它强迫模型学会语法、事实知识、上下文衔接的统计规律,构建出基础的世界表示。

但训练完语言模型之后,你会发现问题还是很大。你输入“三加五等于多少”,它大概率会输出“三加五等于多少”,因为它学到的只是“文本接龙”,不是“问题解答”。要让它变成真正的推理模型,就必须做指令微调。

指令微调本质上是把预训练得到的“续写能力”导向“问答能力”。做法是把训练样本组织成“用户指令 + 正确答案”的形式,继续用交叉熵训练。更进一步,还可以借鉴“build a reasoning model from scratch”的思路,加入思维链数据:不只要回答“等于八”,还要回答“设三为甲,加五为乙,合为八,首先取三,其次加五,得八”。让模型看到完整的推导步骤,它的内部表征才会慢慢建立起“逐步推理”的神经网络路径。

为什么要搞推理能力?拿数学题举例:模型如果能输出中间步骤,哪怕最后结果错了,你也知道它错在哪一步;而如果只输出一个最终答案,错了就是一片黑盒。可诊断性,就是推理能力的核心价值。

3. 实操记录:从零训练一个迷你推理模型

理论讲再多,不如亲手跑一遍。这一节我把实操过程完整记录下来,从环境准备到训练循环,再到推理评估,保证每一步都可以直接抄作业。

3.1 环境准备与工程骨架

我的开发机是单张RTX 4090(24GB显存),系统为Ubuntu 22.04,Python环境用的是3.10,深度学习框架用PyTorch 2.1,CUDA 12.1。这套组合的好处是稳定,网上资料多,遇到问题不容易卡住。

工程骨架我按模块划分,保持单文件内聚,避免在一个文件里写两千行代码:

ai-engineering-from-scratch/ ├── data/ │ ├── raw/ # 原始语料 │ └── processed/ # 清洗后的数据 ├── tokenizer/ │ ├── train_tokenizer.py # 训练BPE分词器 │ └── tokenizer.json # 保存的分词器 ├── model/ │ ├── attention.py # 自注意力实现 │ ├── transformer.py # 解码器核心 │ └── config.py # 超参数配置 ├── train.py # 训练主脚本 ├── generate.py # 推理脚本 ├── eval.py # 评估脚本 └── utils.py # 工具函数

这个结构不花哨,但够用。我强烈建议你从一开始就把训练、推理、评估拆成三个脚本,而不是混在一起,后面调参和复现成本会低很多。

3.2 构建最小数据集与训练循环

数据集我最终保留了约20万条训练样本,包含两部分:通用语料(约15万条)和推理问答对(约5万条)。每条样本根据最大序列长度512进行截断或填充,然后打包成batch=16,每个batch按最大长度动态padding,减少无效计算。

训练循环的核心部分如下:

import torch from torch.utils.data import DataLoader, Dataset from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR train_loader = DataLoader(dataset, batch_size=16, shuffle=True) model = build_gpt_model() model.train() optimizer = AdamW(model.parameters(), lr=1e-3, weight_decay=0.1) scheduler = CosineAnnealingLR(optimizer, T_max=5000) for step, batch in enumerate(train_loader): input_ids = batch["input_ids"].cuda() targets = batch["targets"].cuda() logits = model(input_ids) # [batch, seq_len, vocab_size] loss = cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() if step % 100 == 0: print(f"step {step}, loss {loss.item():.4f}, lr {scheduler.get_last_lr()[0]:.2e}")

这里有两个细节值得解释。第一,学习率为什么用1e-3而不是更小的值?原因是模型参数不到一亿,batch也不大,大学习率可以让前期收敛明显加快。但如果你的batch堆到64、128,我建议把初始学习率降到5e-4或3e-4,否则训练后期容易震荡。第二,为什么加梯度裁剪?因为Transformer的深层梯度范数容易出现极端值,裁剪到1.0可以防止单个异常batch把参数炸飞。

3.3 模型代码实现要点

我不打算贴完整的Transformer实现,那太长,只挑三个核心模块讲清楚。

多头自注意力的正确实现顺序是:先做线性投影得到Q、K、V,再按头数拆分,计算缩放点积注意力,最后合并头并做输出投影。这里特别提醒一点,Q和K做完点积之后,一定要除以sqrt(d_k),即每个头的维度开根号。如果不缩放,点积结果随维度增大而增大,进入softmax后梯度会变得极小,训练基本原地踏步。

前馈网络内部通常包含两次线性变换和一次激活函数,我用的是GELU。这一步占据了模型最多的参数量,贡献了非线性映射能力。FFN的层数扩张,是容量提升最直接的途径之一。

层归一化的位置有讲究。原始Transformer是“后归一化”,也就是把LayerNorm放在残差连接之后;后来的GPT系列改成“前归一化”,把LayerNorm放在子层之前。前归一化训练更稳,深层数时不容易梯度爆炸。我直接用了前归一化,实测收敛更平滑。

3.4 推理与评估流程

预训练加指令微调总共跑了大概3000步,耗时约8小时。训练结束时的loss降到了2.1左右。这个值看着不小,但考虑到12000的词表和混合数据分布,属于正常范围。

推理阶段,我用贪心解码先做冒烟测试:

prompt = "Q: 小明有3个苹果,妈妈又给了他5个,他一共有几个苹果?\nA:" tokens = tokenizer.encode(prompt) input_ids = torch.tensor([tokens]).cuda() with torch.no_grad(): generated = input_ids for _ in range(64): logits = model(generated)[:, -1, :] / temperature probs = torch.softmax(logits, dim=-1) next_token = torch.argmax(probs, dim=-1, keepdim=True) generated = torch.cat([generated, next_token], dim=1) answer_ids = generated[0, input_ids.size(1):].tolist() print(tokenizer.decode(answer_ids))

temperature参数我把默认值设为0.1,这个值下输出比较保守、稳定。如果设到1.0,输出多样性高,但更容易跑偏,数学推理场景不适合太高的温度。

评估方面,我构造了一个50道题的数学问答测试集,包含个位数加减法、乘法、两步综合题。最终准确率约68%。说实话,这个成绩放到大模型面前不够看,但作为从零训练的迷你模型,我心里是满意的。更重要的是,模型输出的推理过程出现了明显的“中间步骤”,这说明指令微调和思维链数据确实把推理路径建出来了。

4. 踩坑实录:那些课本上不会写的成长经历

从零训练模型,最大的确定性就是充满不确定性。这一节把我的典型失败经历整理出来,每一个都是用真金白银的显卡时间换来的教训,建议收藏。

4.1 显存溢出:你以为的内存是不够,其实是批量和序列长度的问题

我第一次跑训练,batch设的是32,序列长度是512,结果刚跑两步就报torch.cuda.OutOfMemoryError。最开始我以为是显存不够,准备换更大的显卡,后来一算账才明白,24GB显存的瓶颈主要在激活值和中间梯度,而不是模型参数。一亿参数模型的权重才400MB左右,但一个batch所有token的中间激活值能轻松超过8GB。

解决方案很简单:把batch减半,从32降到16,再开启梯度累积。梯度累积的意思是,每16个样本算一次梯度但不立刻更新参数,累积4次,等价于batch=64的更新效果。显存占用不变,训练稳定性反而提升了。

4.2 损失不下降:最隐蔽的杀手是学习率过大和过小

训练过程中最让人头大的事就是loss卡着不动。第一次遇到这个问题时,我的loss在2.5附近震荡,前后500步都没变化。排查半天,发现不是模型结构问题,而是学习率设到1e-3之后,后期更新步长太大,参数在最优区间附近反复横跳。

正确做法是配合学习率调度器,让学习率按余弦曲线衰减到接近0。还有一个思路是重新加载一个中间检查点,把学习率调小一半继续训练。我在项目后期测试过:loss不动的时候,把学习率降到原来的一半,loss立刻开始缓慢下降。这一招在从业者圈子里叫“学习率重启”,朴素但极其有效。

4.3 复读机问题:推理生成陷入死循环

推理阶段遇到一个特别搞笑的事:模型从第4个token开始,就一直在重复同一个词,比如“八八八八八八”。这在大模型里也偶尔出现,尤其是小模型和低数据量的组合,更容易复读。

原因有两个方向:一个方向是采样温度太低,导致模型只走最确定的那条路径,一旦某个token概率异常高,后续就锁定住;另一个方向是模型容量有限,没能学会合理的停止条件,一到生成长尾就直接放弃思考。

我的应对方法是三管齐下。调高temperature到0.3;在采样时启用top-p(累积概率超过0.9就不采样);再加一条硬规则,连续出现同一个token超过6次,强制停止。组合拳打下来,复读率从30%降到了5%以下。

4.4 训练秒崩:fp16的隐藏雷区

为了提速,我一开始直接上了混合精度训练,也就是fp16。结果训练不到100步,loss直接变成NaN,然后秒崩。查来查去,原因不是代码逻辑错,而是部分参数在fp16下动态范围太小,梯度下溢成了0。

解决办法是启用动态损失缩放,PyTorch的GradScaler就干这个事:每步先检查梯度是否溢出,溢出了就把损失缩放因子减半,重新反传;正常的话,每2000步再把缩放因子增大一点。开了这个之后,整个训练过程再没出现NaN。另外还有一个省钱经验:对Layer Norm和最后的输出层,尽量保留fp32精度,它们对数值稳定性极其敏感。

提示:训练中途突然NaN,先别怀疑模型结构,90%的情况是数值精度问题,优先查学习率、fp16缩放、数据里有没有NaN值。

5. 从最小Demo到真实工程化:没有捷径但有路径

一个能跑通的小模型只是起点。如果你真的想在AI工程这条路上走下去,下面几个方向是必须补齐的,它们也是“build a reasoning model from scratch”这类进阶项目最喜欢展开的部分。

5.1 分布式训练:单卡练不了大模型

一亿参数在单卡上没问题,但想体验百亿参数的训练,就必须上分布式。核心概念是三件套:数据并行(把数据分给多张卡,每张卡持有完整模型副本)、模型并行(把模型切成多段,每张卡持有部分模型)、张量并行(把单个矩阵运算拆到多张卡上并行计算)。

工程上最省心的做法是直接用PyTorch自带的模型并行封装,一套代码改几个装饰器就能跑多卡。但有一个细节容易被忽略:不同显卡之间的梯度通信是全量传输的,网络一慢,整体训练效率就断崖式下跌。要么保证节点内部走NVLINK或PCIe直连,要么把模型并行度控制在一定阈值内,不要盲目堆显卡。

5.2 对齐与推理模型:RLHF和GRPO到底在干什么

预训练学到的是“语料里的分布”,指令微调学到的是“用户想要的形态”,但想让模型真正确认“什么答案好”,还得靠强化学习。RLHF(人类反馈强化学习)是OpenAI当年的成名招数:先训练一个奖励模型给回答打分,再用强化学习(PPO)让模型输出尽量拿到高分。

现在更流行的是以GRPO为代表的轻量强化学习方案,它不再需要独立的奖励模型,而是利用群体采样,通过多个回答之间的相对优劣直接计算优势值。如果你关注“build a reasoning model from scratch”这个热门主题,会发现很多开源项目就是用GRPO来激发模型的推理能力,比如让模型回答数学题,对答案正确的样本给正强化,对错误的给负强化。我在小模型上试过简化版的GRPO,两轮迭代后,测试集准确率从68%提升到了73%。强化学习对推理能力的激发,确实不是玄学。

5.3 工程化部署与评测:训练完了才是真正的开始

最后这个大模型训练完,并不等于工程完成。部署阶段要处理的事情一箩筐:模型量化(把fp16的权重转成int8或int4,显存占用直接缩到四分之一到八分之一,代价是精度略降)、推理优化(KV Cache缓存历史计算状态,避免每个新token都重新计算全部注意力)、在线评测(把模型输出逐条和标准答案比对,记录准确率、困惑度、响应延迟)。

我个人的建议是,把评估脚本提前到训练还没结束的时候就写起来。我见过太多人训练一结束才发现没有评估工具,临时写,然后发现测试集格式和训练集不一致,又要回头补数据。这些环节在真实的AI工程里,工作量占比经常超过模型训练本身。

最后再分享一个经验总结:从零做AI工程,最大的收获不是模型本身的表现,而是每次卡住时排查问题的思路。后面我再做类似项目时,遇到奇怪的现象,第一反应已经不再是问别人,而是自己对照数据、参数和代码,一步步定位。这个从“会用”到“能造”的跨越,才是这种项目真正值钱的地方。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进