从零手写大模型：三个月复现语言模型与推理模型全记录

这两年AI engineering领域最热的方向绕不开一个词from scratch。从《Build a Large Language Model (From Scratch)》被技术社区反复推荐到越来越多人尝试build a reasoning model from scratch复现类似R1那样的长思维链推理大家的共识越来越清楚——天天调API、套框架换来的是对模型内部机制越来越薄的直觉。这篇文章是我用三个月时间从零构建自己的小语言模型和推理模型的完整记录。不用Transformers库不加载任何预训练权重从tokenizer到训练循环全部手写最后在千万级参数的小模型上跑通了一条从语言建模、到CoT微调、再到强化学习推理训练的完整链路。如果你已经会写PyTorch、用过HuggingFace却总觉得大模型是个黑盒或者你对推理模型到底是怎么被训练出来的这件事充满好奇那这份记录就是我踩完坑之后最想递给你的地图。1. 我为什么坚持从零手写而不是继续用框架调参先说这段经历的起点。之前几年我一直做模型微调和部署HuggingFace生态用得飞起fine-tune、量化、推理加速都上手过自认为对Transformer相当熟悉。直到有一天同事问我这段attention的耗时为什么和序列长度呈二次方增长能不能在不引入FlashAttention的情况下自己写一个融合版本我愣了一下发现自己连QKV投影的shape变化都要现翻文档。那次之后我做了个决定把手头项目放一放花三个月时间不开任何现成模型把一个大语言模型从空白搭起来。参考书我主要用《Build a Large Language Model (From Scratch)》Sebastian Raschka著配合几个开源课程的代码对照读。项目名就叫ai-engineering-from-scratch很直白——就是要搞清楚每一步到底发生了什么。1.1 一个关于注意力层的追问暴露了我的底牌我当时的真实水平是能跑通训练脚本能调超参能部署服务但模型内部每一步算的是什么基本靠猜。transformers库把一切封装得太好了embedding层、attention mask、KV cache、layer norm全是一行调用的事。出了问题我看错误日志的时间比看论文的时间还多。从scratch这件事最直接的价值就是逼你把手伸进泥里。当你自己写causal mask、自己处理softmax的数值溢出、自己调试首次生成的一堆乱码时你才会真正理解那些框架层背后的设计决策。build a logic model from scratch我这里特指逻辑推理模型的热潮里很多人都以为难点在模型结构真正上手才发现推理模型的难点在采样循环、奖励设计和数据格式上——这些恰恰是框架藏得最深的部分。1.2 我给自己划定的从零边界从零构建如果不划边界很容易变成某种偏执。我做之前先把范围写清楚了手写范围BPE tokenizer、词嵌入、多头自注意力、前馈网络、LayerNorm、训练循环、采样器、GRPO更新逻辑。不手写范围autograd自动微分、CUDA kernel、底层矩阵乘法。AdamW优化器我单独手写了一遍纯粹为了检验自己对梯度更新的理解。目标模型参数规模10M到60M之间训练数据自己采集并清洗最后在一套小的数学/逻辑多选题数据集上做CoT SFT GRPO。这里要特别说清楚我不是反框架派。生产环境该用vLLM用vLLM该用HuggingFace pipeline就用pipeline。从零是教学工具目的是把直觉补上。等你真正手写过一遍再回去看框架的源码你会发现自己第一次能读懂那些抽象层到底在帮你做什么。2. 动手前必须拿下的三项底层能力否则后面全是翻车现场很多人一上来就照教程写Transformer结果卡在第一个loss就是NaN然后就开始怀疑人生。我踩过的顺序是先补底层能力再写模型。三项里缺任何一项后面都会花双倍时间补课。2.1 数学基础矩阵形状、softmax数值稳定和梯度直觉做这个项目你不需要把《深度学习》里的公式全部推导一遍但有三样东西必须过手第一每一层的shape变化。输入是(B, T, C)经过embedding变成(B, T, d_model)attention之后还是(B, T, d_model)但Q/K/V的维度怎么拆分、多头怎么reshape和permute这些必须闭着眼睛写得出来。我写代码期间一半的bug都是shape写错。第二softmax的数值稳定。这个坑我踩得特别深第一版cross entropy在fp16精度下直接溢出loss变成inf。正确做法是softmax前先减去最大值scores - scores.max(dim-1, keepdimTrue)。原理很简单指数函数在负数区间不会溢出。第三梯度从loss反向流到每一层的shape直觉。你不需要手算每个偏导但至少要能回答这层的梯度为什么是那个shape。我建议正式动手前先干一件小事用纯numpy写一个两层神经网络手动实现一次反向传播哪怕不看答案抄一遍也行。这个半天工程能帮你省掉后面两周的排错时间。2.2 数据管线清洗时间远超写模型结构的时间这是我觉得整个项目里最反直觉的一点写模型结构用了大概一周清洗数据用了将近两周。很多人跟着教程搭完模型就开训结果模型生成的全是碎片第一反应是改结构、调超参其实问题根源在数据。我的数据管线长这样原始收集了大约2.4GB英文文本加少量中文经过重复行过滤、超短行过滤、语言检测、文档切分之后剩下1.7GB左右。然后用其中大概5万个代表性句子训练BPE tokenizer词表大小选8192。这里有个容易被忽视的教训如果数据集里有一条新闻被重复了几百次小模型会直接把这条新闻背下来而不是学语言规律。另外验证集一定要提前切好否则训练集里泄漏了验证样本你看到的loss全是虚的。2.3 算力预算与实验纪律我实际用的是一张RTX 4090 24GB再加一些云GPU按需实例。刚开始千万不要规划什么130B大模型老老实实从千万级参数起步。我给自己排了三组实验实验线参数量序列长度每步batch预计时长baseline10M256321~2小时主力模型33M512328~12小时压力测试60M51216 梯度检查点1~2天实验纪律也很重要固定随机种子、每个配置一个独立文件夹、每100步记录一次loss、每个实验结束必须导出生成样例。这不是强迫症——三周之后你根本回想不起哪个配置跑出了最好的loss那时候你就知道日志和文件夹命名多值钱了。3. 最小可用LLM的完整拆解从BPE到能输出像样句子这一节是整个项目最核心的部分。我按实际动手顺序拆开讲每个模块都附上关键代码和当时的判断逻辑。3.1 Tokenizer为什么我坚持自己训练BPE很多人觉得tokenizer是现成的东西直接pip install tiktoken不就行了。但从scratch的角度看加载GPT-2的tokenizer只是调用不是理解。自己训练BPE会让你明白vocab_size、未登录词、special token到底意味着什么。我用的BPE核心逻辑是这个整体思路就是反复找最高频的相邻token对、合并、编码成新tokenimport collections from typing import List def get_stats(ids: List[int]): counts collections.defaultdict(int) for pair in zip(ids[:-1], ids[1:]): counts[pair] 1 return counts def merge(ids: List[int], pair, new_idx: int): new_ids [] i 0 while i len(ids): if i len(ids) - 1 and ids[i] pair[0] and ids[i 1] pair[1]: new_ids.append(new_idx) i 2 else: new_ids.append(ids[i]) i 1 return new_ids实际训练前需要先做预分词把文本切成字母数字片段和标点然后在预分词后的整数序列上反复执行上面两步。我保留了|begin_of_reason|和|end_of_reason|两个占位token专门为后面的推理阶段用这在后面会提到。这里有个选型的坑vocab过大的话embedding层会吃掉小模型一半以上的参数。对小项目来说8192到16000都是合适的区间。词表大小本身不需要卷。3.2 自注意力与Transformer从公式到可运行代码接下来是重头戏。我强烈建议先写单头的自注意力确认理解后再扩成多头。核心是这几件事Q、K、V从哪里投影来、为什么要除以sqrt(d_k)、因果掩码怎么加。import torch import torch.nn.functional as F def causal_self_attention(x, w_q, w_k, w_v, w_o): B, T, C x.shape d w_q.shape[1] # 每个head的维度 q x w_q # (B, T, d) k x w_k v x w_v scores q k.transpose(-1, -2) / (d ** 0.5) mask torch.tril(torch.ones(T, T, devicex.device)).view(1, T, T) scores scores.masked_fill(mask 0, float(-inf)) weights F.softmax(scores, dim-1) out weights v return out w_o注意除以sqrt(d_k)的原因如果不做这个缩放当d_k变大时attention分数的方差会跟着变大softmax会过早饱和梯度就消失了。这是transformer原论文里刻意设计的细节自己写一次才会真正记住。整体结构我选了GPT式的pre-norm每个block先LayerNorm再进attention/FFN残差连接包在外面。实测下来pre-norm在小规模训练里比post-norm稳定得多学习率稍微大一点也不太容易炸。位置编码用的是可学习参数最大序列长度设了1024。写到这里有个体会你在框架里点一个model.generate()觉得云淡风轻但自己写gen循环时会第一次意识到每生成一个token都要重新跑一遍前面的上下文KV cache的空间换时间意义就是这样暴露出来的。这个直觉在后面看推理框架源码时价值连城。3.3 训练循环、损失曲线与能不能说人话的评判标准模型结构写完训练循环反而简单cross entropy损失、AdamW优化器、warmup 500步 cosine学习率衰减。我实际的训练曲线大体是这个节奏训练步数验证loss生成效果20006.9左右输出的全是词碎片100005.6左右能拼出常见短词300004.7~5.0出现语法残缺的短句500004.4~4.6常见句式成形第一次生成出来一堆乱码是正常的这几乎是必经之路。评估模型别只盯着loss准备一组固定prompt每次训练到检查点时用同样的prompt采样几段对比看有没有改善。千万不要每次换prompt否则你没法判断哪个改动真正有效。能说人话的标准在从零项目里要放低它不需要达到GPT-4水平只要明显学会语言的基本形状比如空格、词频、标点位置开始合理就可以进入下一阶段。4. 从语言模型到推理模型CoT数据、SFT与GRPO的实操记录语言模型能生成句子离会推理还有一座山。这部分是build a reasoning model from scratch的核心也是我认为整个项目里信息密度最高的三个月。4.1 推理模型到底把推理放在了哪里观察推理模型的输出你会发现它在给出答案之前会先写一大段内部推理类似先看条件A再看条件B所以...。R1路线的基本配方可以概括为长思维链CoT 可验证的奖励 强化学习。这里有个很关键的理解预训练阶段模型学会的是接续上文生成而不是推理。推理能力很大程度是后训练阶段被奖励信号逼出来的。我做了对照实验直接把CoT数据喂给模型做SFT模型学会的是输出格式它知道应该先写思考再写答案但在训练集范围外并不会主动展开多步推理。只有用GRPO这类RL方法让模型在采样中尝试不同的中间步骤并用答案正确性做奖励它才会逐渐内化多写几步思考能提高答对概率这个因果。所以我的结论是推理从零复现的难点不在Transformer结构而在RL训练系统本身。4.2 在自己模型上复现CoT效果的三个步骤第一步构建小规模CoT数据集。我用的是数学选择题和简单逻辑题格式统一为「问题、思考、答案」三段。思考部分控制在3~6步总共约8000条一半是手写模板生成的另一半从公开数据集过滤得到。思考步骤不要写太长太长模型在RL阶段会学到凑字数。第二步SFT教格式。训练3个epoch学习率降到1e-4context拉长到1024以容纳思考过程。这一步不追求推理变强只追求一件事模型稳定输出思考再答案的格式。评估指标用一个格式合规率生成结果里有没有正确出现结尾的答案标记。合规率达到90%以上再进RL否则后面奖励信号会被格式问题搅浑。第三步GRPO强化学习。GRPO的核心是对于同一个问题采样一组回答组内归一化收益再用归一化后的优势更新策略同时用KL约束让它不要偏离SFT模型太远。示意代码如下def compute_advantages(rewards): # rewards: (group_size,) mean rewards.mean() std rewards.std() 1e-4 return (rewards - mean) / std # 对每个问题采样k条轨迹记录policy和ref的logprob # ref是冻结的SFT模型policy是从SFT模型复制出来继续训练的 def grpo_loss(logprobs, ref_logprobs, rewards, kl_beta0.02): norm_rewards compute_advantages(rewards) ratio torch.exp(logprobs - ref_logprobs) kl torch.exp(ref_logprobs - logprobs) - (ref_logprobs - logprobs) - 1 policy_loss -torch.mean(torch.clamp(ratio, 0.8, 1.2) * norm_rewards) return policy_loss kl_beta * kl.mean()奖励设计我调了很多版最后稳定到答案正确1.0格式错误没写思考段或没写答案标记-0.5超长截断-1.0。这里最想强调的一点是不要给思考过程质量打主观分一旦模型发现过程分可以被凑出来它就会用各种方式hack后面细说。4.3 用GRPO做推理微调时我踩到的坑这一小节全是真金白银的教训每一条都花了我至少一整天。坑一reward hacking。模型很快就学会在思考段里反复重复题目文字最后随便蒙一个答案。因为格式正确能拿到基础分重复文字能凑够长度模型成了复读机。对策就是我前面说的格式相关只给罚分不给正分唯一正分来源是答案正确。坑二奖励稀疏。开放题没法自动判分模型采样半天也不知道自己对没对学得很慢。我的解法是RL阶段只用选择题和简单计算题答案可以程序化校验开放题留在SFT阶段处理不放进RL。坑三KL系数失控。kl_beta设太小模型一两个小epoch后输出就开始退化、词频失衡。我最后把kl_beta固定在0.02到0.05之间并且给ratio加了clip(0.8, 1.2)能显著减缓退化。坑四采样长度失控。RL早期模型为了凑满分疯狂输出垃圾2000个token都不够。给生成设max_new_tokens512超长直接截断并按格式错误扣分。走完这趟我的体感是推理从零构建的日常工作80%在数据和奖励工程20%才是模型和训练技巧。外界看那些复现项目总觉得难点在模型有多聪明实际全在工程细节里。5. 关于算力、成本和值不值得从零开始的实在话最后聊点计划之外的实在话。很多读者卡在第一步就放弃了理由是我没有那么多GPU。这个心结值得打开。5.1 没有A100也能练起来我一直用本地RTX 4090辅助用一些云GPU按需实例配合Colab和Kaggle的免费配额做极小实验。核心诀窍是用bf16混合精度训练一个30M参数模型在512序列长度下显存占用能压到8GB以内。大致参考参数量序列长度精度显存含Adam状态参考训练时长10M512fp32约3GB1~3小时33M512bf16约6GB8~14小时60M1024bf16 梯度检查点约10GB1~2天工程上的小习惯很省事数据加载器记得开num_workers和pin_memoryloss不要每步都写盘100步写一次检查点至少留两个防止训练中途炸了直接回到解放前。别小看这些细节一次训练跑十几个小时因为没存检查点而从头再跑的感觉这辈子不想体验第二次。5.2 该停止较劲的时候就停止从零构建的边界诚实地说我最后训练出来的模型单纯从能力上讲没有任何生产价值它跑不过任何开源小模型。但这正是我推荐每个想做AI engineering的人都走一遍的原因它的价值是直觉。当你亲自调试过一次梯度爆炸、一次NaN loss、一次RL reward hacking之后你再排查线上模型的问题路径会清晰得多因为你脑内有完整的因果模型而不是一堆黑盒调用。做完基础版之后还有很多扩展方向手写RoPE旋转位置编码、RMSNorm、GQA分组查询注意力、KV cache、投机采样。每补一个模块对推理框架的理解就深一层。但注意从零构建是学习工具不是生产方案如果你的目标是做一个产品请立刻回到框架上。最后说点题外话。这个项目给我最大的回报不是那个能生成像样短句的小模型而是后来排查模型退化问题时我几乎不需要翻文档就能猜出是数据环节出了问题还是哪一层的数值溢出了。如果你也准备花三个月做这件事我的建议只有一个把每次实验的loss曲线和生成样例完整存下来三个月后回看你会感谢那个愿意从头搭起的自己。