多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

从零构建迷你LLM:Transformer预训练、SFT与GRPO推理实战

从零构建迷你LLM:Transformer预训练、SFT与GRPO推理实战 最近很多朋友在后台找我聊同一个话题天天调大模型接口prompt写得比小说还长可真要自己动手搞一个模型心里完全没底。这个ai-engineering-from-scratch项目就是把我从零构建一个迷你大语言模型LLM、再进一步训练出基础推理能力的全过程原原本本记录下来。项目参考了《Build a Large Language Model from Scratch》这类经典实践思路但没有停留在调库调参的层面而是把分词、模型架构、训练循环、对齐、强化学习一块块自己搭起来。如果你是想深入理解AI工程底层的算法工程师刚转行AI的后端或全栈工程师或者单纯想搞明白“推理模型到底是怎么训练出来的”的爱好者这篇文章值得你花20分钟看完。1. 项目拆解从零构建一个AI系统到底“零”在哪里1.1 三个层次的从零你选哪个“从零”这个词在AI圈已经被用滥了。有人把pip install transformers也叫从零有人把能跑开源模型叫从零。但真正动手做过的人才会意识到这里至少存在三个层次。第一层是API层。你调用GPT只需要写prompt关注输入输出这个层次下你理解的“模型”是一个黑盒。第二层是框架层。你用PyTorch和HuggingFace搭模型注意力用现成模块训练用Trainer这个层次你已经能看到模型骨架但很多细节还是被框架隐藏。第三层是纯手搓层。你从随机初始化权重开始自己实现embedding、注意力、层归一化、反向传播连数据加载和训练循环都是一行行写出来的。我这次做的项目定位在第二层到第三层之间网络结构完全手写训练流程自己掌控梯度自动求导借助PyTorch。为什么这么选因为纯手搓反向传播在这个阶段的学习收益边际已经很低了而且容易把大量时间耗在debug矩阵维度上反而忽略了真正关键的工程决策——数据怎么处理、损失怎么设计、训练策略怎么选。真正让你理解AI工程的从来不是记住某个矩阵怎么相乘而是你亲手做完一代代实验后积累的判断力。1.2 为什么这个时间点值得自己搭一个推理模型过去的惯性认知是训练大模型是巨头的事普通人玩不起。这话一半对一半不对。如果你非要复现千亿参数的模型确实需要几千万的成本和成百上千张卡。但如果你要复现的是“机制”——比如强化学习如何让模型产生推理能力那么几十亿token的数据和一张消费级显卡就完全够用了。现在的推理模型是热点业界那些能输出超长思考链的模型确实惊艳。但你如果只知道在prompt里加“请一步一步思考”那是被表象带偏了。真正让模型“学会推理”的是强化学习奖励信号驱动的行为塑造而不是提示词魔法。自己搭一个小模型从语言建模开始经过SFT再到GRPO强化学习你能亲眼看到那个“顿悟”时刻长什么样。这种体感是任何论文和教程都给不了你的。1.3 项目的形态、路线图和验收标准这个项目分三个阶段推进。阶段一预训练一个迷你语言模型参数量控制在30M到100M之间用公开的TinyStories或FineWeb-Edu子集做训练数据。阶段二做监督微调SFT让模型学会指令跟随的基本格式也就是能像对话式AI一样有问有答。阶段三用规则奖励加GRPO强化学习训练推理能力让模型面对数学题时学会输出思考链并提高正确率。验收标准我定得很具体模型能正确完成两位数以内的加法能在格式约束下先输出推理过程再给答案训练过程中能观察到“思考长度增长”的现象。这三个标准看起来不起眼但任何一个都涉及完整的数据管线、训练管线和评估管线。把这三条跑通你就已经掌握了从头构建AI系统的核心工程链路后续无论做多大参数量思路都是同一条。2. 第一步是喂数据分词器与数据管线2.1 数据质量比数据数量更决定下限很多第一次训练模型的人一上来就追求数据量疯狂堆GB级别的语料。但实际做过训练的人都会告诉你脏数据对模型能力的腐蚀往往比数据量不足更致命。我第一版用的数据集是直接爬的文本里面有大量HTML标签、乱码、重复段落。模型训了几千步loss能下降但生成的文本里偶尔会蹦出“”这种垃圾token这就是数据管线的锅。清洗逻辑其实不复杂核心是去掉三类东西非自然语言内容、超长和超短样本、高度重复的段落。超长样本会拖慢训练效率超短样本又缺乏语义上下文而重复样本会让模型背下来而不是学会。实际操作上我直接采用了FineWeb-Edu的处理思路按长度过滤去掉少于50字符的残渣按重复度过滤对连续重复的n-gram做剔除最后加一个困惑度过滤用现成的语言模型给每条样本打分把那些“读起来根本不像人话”的样本丢掉。数据量虽然只剩原始语料的三成左右但训练稳定性和生成质量都有明显改善。2.2 自己实现一个BPE分词器三步搞定分词器是通往模型的第一个关卡。词级别分词有一个致命问题词表爆炸任何新词都无法被表示。字符级别分词又让序列过长训练效率太低。BPEByte Pair Encoding是这两者的折中它把文本拆成可复用的子词单元。实现BPE核心就三步。第一步把所有训练语料按词频统计。第二步把每个词拆成字符序列同时记录词频。第三步迭代合并出现频率最高的相邻字符对直到词表达到目标大小。比如“low”和“lowest”这两个词经过合并后可能共享“low”这个子词新词“lower”也能用已有token组合出来这就是BPE的泛化能力来源。实操中不建议完全从零写pretokenizer直接参考GPT-2的regexp逻辑就行——先按空格和标点做预切分再在切分后的词上应用BPE。词表大小我试过1024、4096和8192三档。结果非常符合直觉1024的词表生成文本很僵硬模型被迫频繁拼接子词8192词表生成流畅但embedding矩阵和输出层参数显著变大在小模型上占用太多容量。最后选的是4096算是对参数量和表达能力的平衡。分词器训练好之后记得固定住special token[PAD]、[BOS]、[EOS]这些必须先占好位置否则后续SFT阶段再加token会让embedding维度发生变化整个训练要重跑。这个坑我踩过不止一次。2.3 数据配比、重复epoch与训练顺序的取舍数据集准备好之后还有两个容易被忽略的工程点。第一是配比如果既想保留故事类语料的流畅语言又想加入数学语料强化推理两类数据不能简单混在一起而应该按比例采样让每个batch里都有固定比例的不同类型样本。第二是重复epoch数小模型数据量不足时多跑几个epoch能明显提升指标但重复过多次会过拟合训练分布导致生成内容像复读机。我的经验是语言建模阶段一个epoch为主最多跑1.5个epochSFT阶段数据量小可以跑2到3个epoch。至于课程学习从小样本开始逐渐增大难度在小型实验里收益不明显我建议别在这个阶段浪费太多精力。到这里你可能会问数据、分词都准备好了接下来是不是直接写模型其实还差一个环节——把训练数据切成固定长度序列。我一开始按句子切结果一个batch里样本长度参差不齐padding浪费了大量算力。后来改成把整个语料拼接成一条长流再按固定长度比如512 token切成块训练效率提升非常明显。这个“拼接切块”的思路在miniGPT和nanoGPT项目里都在用看起来土但确实是语言模型预训练的标准姿势。3. 从零搭建Transformer把每一层吃透3.1 架构选型为什么是RoPE和SwiGLU模型的骨架是Transformer。但Transformer的细节经过这几年的迭代已经和原始论文差了很多。如果你照着《Attention Is All You Need》原版实现在小规模训练上会吃不少亏。第一个关键选择是位置编码。原始Transformer用sinusoidal绝对位置编码而现在主流是旋转位置编码RoPE。RoPE的核心思想是把位置信息编码成旋转角度在注意力计算时对Q和K向量做旋转操作这样位置编码天然具备相对位置的性质——模型能理解“token A在token B前面第三位”而不是“token A的位置是7”。实现也不复杂核心就是把向量按维度分成两半一半结合cos分量一半结合sin分量然后利用旋转变换的矩阵形式参与计算。第二个关键选择是激活函数。原版用ReLU我这次用SwiGLU。SwiGLU把门控机制引入MLP简单理解就是让网络学会“选择”哪些信息值得通过。在相同参数量下SwiGLU的困惑度比ReLU低不少这已经是社区验证过的结论。改动很小收益却很大没必要守旧。第三个关键选择是归一化。LayerNorm换成RMSNorm归一化只做缩放不做均值平移省掉了一半的统计计算效果几乎没有差别。另外要注意pre-norm和post-norm的差异pre-norm把归一化放在子层之前训练更稳定是现在所有主流模型的选择post-norm是原版论文的做法在小模型上容易训练不稳定。这些选型加起来模型看起来还是Transformer但实际已经是现代Transformer了。3.2 从batch到loss一个训练step的完整旅程理解了架构更重要的是把前向、反向的流程彻底打通。我来描述一个batch的实际旅程。输入是一批文本序列shape是[batch_size, seq_len]。先经过token embedding层每个token id变成一个d_model维的向量shape变成[batch_size, seq_len, d_model]。接着是N层Transformer block每一块内部依次做RMSNorm、多头注意力、残差连接、RMSNorm、SwiGLU MLP、残差连接。最后一层的输出经过final norm再和token embedding矩阵做转置相乘得到每个位置的logitsshape是[batch_size, seq_len, vocab_size]。然后计算交叉熵损失。这里有一个易错点预测的是“下一个token”所以input_ids要左移一位作为label。也就是输入[BOS] hello world模型在第一个位置要预测“hello”在第二个位置要预测“world”。如果这一位对不齐loss会大幅度偏离正确范围。训练时日志里loss稳定在8到10对应4096词表、随机初始化说明label对齐没问题如果一开始loss就有几十甚至上百先查这一步。反向传播阶段PyTorch自动求导帮我们算了梯度。但光有梯度还不够工程上还需要三件套梯度裁剪、学习率调度和优化器状态。梯度裁剪设max_norm1.0防止个别batch产生超大梯度把权重推出盆地学习率用warmup加cosine decaywarmup阶段线性升到峰值之后按余弦曲线衰减到峰值的十分之一优化器用AdamWweight decay定在0.1。这里要注意weight decay不能作用在embedding矩阵上否则词向量会被压缩到异常范围这个可以通过参数分组来控制。3.3 超参、算力预算与小模型训练的节奏感这一节可能是最容易劝退新手的。很多人动手之前先问“我一张4090能训多大模型”答案比想象中宽松。我的配置是30M参数模型、batch size约0.5M token、学习率峰值6e-4、上下文长度512。在单张4090上吞吐大概4000到6000 token/s训练10亿token大约需要两到三天。对于机制验证来说10亿token完全够用。配置项我的取值备注模型参数量30M左右6层Transformerd_model 512词表大小4096BPE子词上下文长度512预训练阶段固定训练token量约10亿小模型配大语料峰值学习率6e-4warmup 500步batch token数0.5M512×1024优化器AdamWweight decay 0.1梯度裁剪max_norm1.0必开这里有一个值得记住的经验法则训练算力预算约等于6乘以参数量乘以训练token数。也就是说10亿token训练30M模型已经算是“超配”模型会有余力去拟合更复杂的行为这恰恰是我们做机制验证想要的。小模型配大语料是复现推理能力涌现的一个关键前提。很多人在这一步抠抠搜搜只给模型几千万token那当然只能看到语言模型看不到任何涌现现象。3.4 训练监控别只盯loss曲线训练过程中我最常被问到的问题是“loss降得挺稳是不是就等着收了”我的答案是loss只看大盘细节还是得靠生成样例。训练语言模型时我习惯每500步打印一批固定prompt的生成结果比如让模型续写几点希望的句子。loss和生成质量并不总是同步有时loss平滑下降生成内容却全是重复的“的的的”这说明模型在走捷径。监控指标上除了loss还要看梯度范数。梯度范数如果长期在1.0以上说明训练不稳定如果长期接近0又说明模型可能已经饱和或者学习率太低。另外学习率调度的步数要跟总训练步数匹配warmup步数一般设成总步数的1%到3%。这些监控项看习惯了你对训练状态的判断会比盯任何工具面板都准。4. 从“复读机”到“思考者”SFT与强化学习训练推理能力4.1 SFT先把对话格式焊进模型预训练完成的模型本质上是一个“高级复读机”。你给它一句上文它大概率能续出合理内容但你让它“回答问题”它不会理你因为它根本没见过“问题-回答”这种对话结构。SFT阶段就是把这种结构通过监督学习焊进模型。数据格式我采用最朴素的模板指令部分加上固定的system提示回答作为目标输出。训练时有一个关键技巧loss mask。对于指令部分我们不计算损失只对回答部分计算交叉熵。如果指令部分也算loss模型会把“背问题”当成主要任务生成的回答质量会明显下降。实现上就是构造一个和label同样形状的mask矩阵指令部分置0回答部分置1。SFT看起来简单但有一个隐性风险模型会过拟合模板。训练数据里如果只有“请解释...”“请回答...”这类指令模型被微调后会变得只会回答固定句式的指令。缓解办法是在数据里混入多样化的指令模板和任务类型甚至加入少量纯文本续写样本防止模型把“对话能力”和“特定模板”绑定在一起。4.2 GRPO让模型在探索中学会推理SFT之后的模型能对话了但遇到数学题只会“一本正经地胡说八道”因为监督学习只能教会它模仿没法教会它尝试、犯错、修正。要让模型真正获得推理能力需要强化学习。首选方案是GRPOGroup Relative Policy Optimization相比传统PPO它的最大优势是不需要单独训练一个critic价值模型直接把同一组prompt采样出的多个回答的奖励做组内归一化作为优势估计。对资源有限的小团队来说省掉一个价值模型的训练成本意义极大。整体流程是这样的对一个数学prompt让当前策略模型采样G个回答我一般用8到16个。每个回答通过规则函数打分。规则分两部分格式分和正确性分。格式分要求回答必须包含思考标签和答案标签缺少就不得分正确性分看最终答案是否等于标准答案。然后对这一组回答的得分做标准化得到每个回答的优势值。得分高于组平均的回答其生成概率会被强化低于组平均的会被抑制。这就完成了策略的一次更新。为了防止模型为了刷奖励而彻底偏离语言模型基线还要在损失里加KL惩罚项约束新策略和参考策略之间的分布距离。实现的时候损失函数比SFT要绕一些核心是策略梯度强化信号乘以选择该token的对数概率的梯度再按batch累加。GRPO和PPO的差异在于优势值不做价值网络预测只用组内相对值因此实现代码少一大截也更好调。4.3 亲眼看到“推理爆发”训练现场记录这一段是整段实践最让人兴奋的部分。我第一阶段用的是小学算术题风格的数据集每个题目的标准答案都预先算好。训练初期模型输出的思考过程非常敷衍基本就是两三句话然后直接给个错误答案格式分勉强拿到正确率惨不忍睹。在GRPO跑了大约300个step之后我注意到一个现象平均思考长度开始稳步上升模型输出的推理过程从两三句慢慢变成五六句再从五六句变成十几句回答正确率也像爬楼梯一样往上跳。这就是论文里描述过的“aha moment”——模型在reward信号驱动下“发现”了多思考几步能提高正确率于是主动开始输出更长、更细致的解题过程。如果你在自己的项目里没观察到这个现象不要气馁。它有几个前提条件第一任务是可验证的奖励信号要密集且准确第二采样数量G不能太小组内对比才有区分度第三prompt的多样性要足够。我的经验是第一版GRPO每组只采4个回答几乎不涨改成16个之后信号明显稳定下来。这也解释了为什么论文里那些涌现现象动辄需要几千步小样本时策略梯度的噪声太大模型根本不知道哪个动作是好的。5. 踩坑实录训练过程中的典型问题与排查清单5.1 loss不降、loss爆炸、loss在降但生成是垃圾这三类问题是训练新手遇到最多的分别对应不同的病灶。loss完全不降先查三件事确认label对齐确认学习率范围在1e-4到1e-3之间确认数据加载没有把文本截断成“只有半个token”的情况。很多次我发现问题是tokenizer的padding没有对齐batch内长度导致输入尾部全是[PAD]。loss爆炸通常出现在学习率过高或者梯度裁剪缺失的场景先把max_norm1.0加上再把学习率降到当前的1/3试一次。最诡异的是loss缓慢下降但生成内容完全不像话。这种情况九成是数据问题。比如语料里有大量噪声内容或者重复文本过多模型学的是“高频噪声”而不是“语言结构”。把验证生成样本打印出来逐条看比盯loss曲线有效得多。我一直坚持每天训练日志里存几组生成样例这是最廉价的坏模型预警机制。5.2 显存不够与训练速度慢三个性价比最高的优化小模型训练最常见的工程瓶颈是显存。30M模型理论显存占用不到1GB但实际训练时AdamW优化器要额外保存一阶二阶动量梯度、激活值也会占大量显存。三项优化手段按性价比排序第一梯度累积把batch切小每N个小batch累积一次梯度再更新效果等价于大batch显存占用不变第二bf16混合精度消费级显卡的bfloat16能直接减半显存同时保持训练稳定注意要把loss计算保持在fp32防止下溢第三flash attention这项优化能把注意力计算的显存复杂度从平方级降成线性级PyTorch里一个参数就能开启。如果以上做了还慢问题很可能出在数据加载上。DataLoader的num_workers设为0会让CPU预处理成为GPU的瓶颈GPU利用率上不去。正确做法是num_workers设为4到8并开启prefetch让数据喂送和模型计算流水线化。5.3 推理阶段还在踩的坑采样和解码好不容易训完模型很多人卡在decode阶段。常见问题有两个。第一生成到EOS没完没了地续写这通常是因为SFT阶段没有把EOS纳入训练目标或者训练数据里EOS出现频率太低。第二temperature设得太高生成内容乱七八糟设得太低又变成复读机。对推理类任务采样时建议temperature在0.6到1.0之间top-p取0.9到1.0。再加一个min_p过滤可以有效杀掉那些极低概率的一次性token让生成质量明显上一个台阶。我后来在工程里做了一个简单但很实用的设计把采样参数写进配置对象所有解码路径测试、评估、日志样例共用同一套采样逻辑。这样你在实验记录里写“temperature0.8, top_p0.95”时关键细节不会被遗漏。看似微不足道但实验可复现性就是这样一点点磨出来的。5.4 训练中断与恢复checkpoint是你的命根训练跑到凌晨一张显卡崩了没有checkpoint全线重来。这种痛苦应该是每个训练过模型的人都懂的心头痛。我的建议从一开始就执行每500步存一次checkpoint至少保留最近3份。checkpoint里保存的不只是模型权重还有优化器状态、学习率调度器的步数、随机数种子。恢复训练时直接加载保证训练进度能接上。实在想省事也要保证每1000步保存一次。另外正式长训之前先拿2%的数据跑一遍10分钟的小规模冒烟测试确认数据加载、loss计算、checkpoint保存这些链路都通了再放开整夜训练。这个习惯帮我省了至少三个周末。5.5 排查清单速查表现象排查方向典型解法loss完全不降label对齐、学习率、数据加载检查input/label移位确认学习率在1e-4量级loss爆炸学习率过高、梯度爆炸开梯度裁剪把学习率除以3loss在降但生成垃圾数据脏、重复过多重新清洗数据打印生成样例逐条看显存OOM激活值占用过大梯度累积、bf16、flash attentionGPU利用率低数据加载慢调高num_workers开启prefetch生成不停EOS学习不到位SFT数据里强化EOS推理时设置max_len训练中断checkpoint缺失每500步存盘跑前做冒烟测试跑完整个项目我个人印象最深的不是模型最终在数学评测上拿到的分数而是那个深夜的某一次训练日志在一个没有见过的问题上模型的推理过程比大多数训练样本都要长它似乎在尝试一种训练集里从来没有出现过的解题路径。那一刻我突然理解了强化学习塑造出来的行为确实能超越蒸馏和模仿的天花板——它是模型在奖励信号下的自我探索。如果你也想走一遍这条路我的建议是不要一上来就想“我要做开源大模型”把目标缩得越小越好先做一个能正确算加法的小模型再做能分出step by step的小模型。每增加一个能力都是完整走一遍数据、训练、评估、部署的工程闭环。这个闭环的熟练程度才是“从零构建AI”真正留给你的财富。另外给一个小技巧写代码时把所有超参数和路径集中到一个config文件里实验中做好记录别指望自己的记忆力这比任何高级技巧都更能确保你的实验结果可复现。
返回列表