ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

从零训练大语言模型:Transformer、思维链与推理增强全攻略

从零训练大语言模型:Transformer、思维链与推理增强全攻略 如果三年前有人告诉我我会为了搞懂AI工程而亲手从零训练一个语言模型再让它学会推理我一定觉得他是在开玩笑。但事实是当我想搞清楚注意力机制为什么有效Loss降到多少才算正常加一层残差到底影响什么这些问题时发现所有现成框架都在替我遮住答案。于是我决定走一条笨路从数据清洗、分词器实现、Transformer骨架、训练调度到评测全链路自己动手搭一遍。这篇文章就是这条路的完整复盘适合那些不想只做调包侠、想真正理解大语言模型构建与推理模型训练细节的人。这轮折腾大概持续了两个月中间翻过车、断电丢过检查点、也看着Loss毫无悬念地发散过。但把这些过程写下来之后我觉得最值得分享的反而不是什么惊艳的结果而是那些原来这一步是这么回事的瞬间。接下来按实际推进顺序把整套从零构建AI工程管线的思路、参数和踩坑经验一次说清楚。1. 为什么我坚持从零开始而不是直接调包1.1 跑通Demo带来的错觉市面上的大语言模型教程绝大多数教的是加载预训练权重写三行代码跑一个推理。我一开始也觉得这就算入门了。但有一次我想调整模型对某个特定领域的长文本理解能力试了几种办法都没效果回头一看我连模型前向传播里数据究竟怎么流动的都没搞明白。这时候才意识到跑通Demo只是使用AI离做AI工程还差着十万八千里。从零开始在这个语境里不是为了证明自己有多硬核而是为了把黑箱一个一个拆开。当你亲手写过一次自注意力、亲手做过一次分词表的合并、亲手调过一次学习率后面再遇到问题你能猜到一个大致的排查方向而不是只能上网搜为什么我的模型输出一堆乱码。1.2 从零到底指的是哪三个层面第一个层面是数据。真正去构建一个像样的训练语料而不是拿现成的jsonl跑一遍了事。你需要理解数据规模、清洗规则和分布都会直接影响模型能力。第二个层面是模型骨架。我选择自己实现一个极简的Transformer包括embedding、位置编码、注意力掩码、LayerNorm这些组件。哪怕你只是照着论文改写也会发现那些在框架里一键完成的事情背后有大量细节。第三个层面是训练循环。AdamW的weight decay和Adam的weight decay不是一回事warmup和余弦退火为什么要配合使用梯度裁剪的阈值怎么定这些问题只有自己写训练循环时才会逼着你去面对。1.3 这条路真正教会我的事显存和训练速度之间的权衡不是所有地方都值得用更大batch也不是所有层都应该用float32。验证集应该从第一天就搭好而不是等模型能跑通了再临时补。模型越大越好这句话在小规模实验里完全不成立很多时候有三个亿token的数据一个120M参数的模型已经能给你非常明确的反馈信号。这些体会看起来平淡但都是我用实打实的训练时间和debug时间换来的。2. 第一版训练管线怎么搭数据、分词器与一个手搓的Transformer骨架2.1 语料选择、清洗与规模估算我先做了一个很朴素的决定用约2GB的混合文本作为初始语料。中文部分我选了一些公开领域的长文本英文部分用了百科类子集。清洗规则并不复杂但每条都来自实际训练中的教训去掉包含乱码的段落、压缩连续空白符、过滤掉超过2000字的异常长文、把全角符号统一成半角。至于数据量到底要多少有个粗略的经验值可以参考。一个120M参数左右的模型训练大概需要0.5B到1B个token如果你的模型规模只有20M到30M那0.1B到0.3B token就够了。模型规模参数量参考建议训练token量单卡训练耗时参考RTX 4090微型20M约100M2~3小时小型60M约300M5~8小时常规120M约600M12~20小时进阶350M约1.5B3~5天这个表是我实测下来的节奏前提是序列长度512、batch size适中、用了混合精度。如果你没有这么多时间把数据量砍一半也能看到足够明显的训练趋势只是最终效果会差一些。2.2 分词器里藏着三个容易忽略的坑分词器是很多人习惯直接调库跳过的一步但它的影响比想象中大得多。第一个坑是词表大小。太小的词表会把词拆得过碎模型需要学更多组合规则太大的词表又会让embedding矩阵占用大量显存。我用的是BPE词表定在32000对小模型来说是个比较平衡的值。第二个坑是训练语料和真实语料分布不一致。如果分词器训在英文语料上却拿中文文本去做训练会发现同样的token串根本无法覆盖你的输入。我的做法是直接在混合语料上训练分词器让中文和英文按一定比例都出现在BPE的合并过程里。第三个坑是特殊token。padding token、eos token、unk token这几个看似不起眼但漏掉任何一个都会在训练或推理阶段冒出来奇怪的现象比如模型生成到一半突然结束或者把所有短句补成一样长导致注意力都被padding位置吸走。2.3 Transformer骨架自注意力、前馈与归一化的顺序这个阶段我写了一个非常精简的Transformer类核心就是多头注意力加前馈网络。手写注意力部分是整个项目里最值得做的事因为你会真正理解为什么要有掩码、为什么要做缩放。class SelfAttention(nn.Module): def __init__(self, hidden_dim, num_heads): super().__init__() self.num_heads num_heads self.head_dim hidden_dim // num_heads self.qkv nn.Linear(hidden_dim, 3 * hidden_dim) self.out_proj nn.Linear(hidden_dim, hidden_dim) def forward(self, x, maskNone): batch_size, seq_len, _ x.shape qkv self.qkv(x).reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim) q, k, v qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] q, k, v q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) scores q k.transpose(-2, -1) / (self.head_dim ** 0.5) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn_weights torch.softmax(scores, dim-1) out attn_weights v out out.transpose(1, 2).reshape(batch_size, seq_len, -1) return self.out_proj(out)LayerNorm的位置我一直坚持放在残差连接之前也就是Pre-LN结构。相比Post-LNPre-LN在深层的训练中更稳定尤其我用的是小学习率配合较大warmup的时候几乎不会出现某些层梯度过大导致整体崩掉的情况。关于这个顺序几个技术社区一直吵个没完但我自己的对比实验里Pre-LN在小模型上的收敛速度确实更友好。3. 训练中的工程细节学习率、Loss曲线与评测集设计3.1 学习率、批量大小与显存预算怎么联动训练启动阶段我用的参数是120M参数模型batch size 16序列长度512AdamW优化器初始学习率3e-4warmup steps 1000然后余弦退火到1e-5以下。有个很关键的经验学习率不是单独调出来的它和batch size、序列长度以及数据规模是绑在一起的。如果你把batch size翻倍梯度估计更稳学习率通常可以适当调大一点如果把序列长度从512涨到1024每个step的信息量变大此时如果学习率不变很容易出现loss震荡。我试过从3e-4直接拉到6e-4结果在大概两千步之后loss就开始拉高并且再也回不来。所以对于中小规模模型3e-4是一个很稳妥的起步值如果训练曲线太平了再考虑往上调。3.2 怎么看Loss曲线三张图判断训练状态很多新人面对loss曲线只会问一句这正常吗。我的建议是把训练过程拆成三个信号来看。正常状态loss平滑下降前20%的step降得最快后面慢慢进入平台期。这种情况下不要频繁干扰训练让它跑完就好。震荡状态loss在下降大趋势中来回抖动抖动幅度超过0.1的时候要怀疑学习率偏高或数据里有大量污染样本。先检查数据清洗再考虑降学习率。发散状态loss直接往上冲或者出现NaN。这种情况多半是梯度爆炸或数据里有NaN值。我用的是两个手段梯度裁剪设到1.0同时把batch里的异常样本打日志打出来定位。3.3 评测集模型不是背语料而是能干活我刚开始训练的时候只看loss结果发现loss降得挺漂亮但问几个问题模型回答得完全不像话。原因是单纯的语言模型loss低只能说明它对词序列的预测能力还行不代表它理解语义。于是我从第一天开始就配了一套很小的评测集大概100条人工构造的任务包括完形填空、简单指令、事实性问答。实操上的做法是每个checkpoint跑一遍这100条任务把回答记录下来和参考答案并排对比。这个做法很土但非常有效。你会清楚看到某个阶段模型学会了把问题复述一遍但不会回答再过一段时间才学会了提取关键信息给出答案。评测集的价值不只是给你一个分数而是给你一个随时可以回看的进度记录。有一次我改了一下数据清洗规则总体loss只降了一点点但评测集分数从47涨到了61这种信号如果只看loss是完全发现不了的。4. 从语言模型到推理模型思维链数据与轻量RL的实战记录4.1 推理模型和普通语言模型差在哪里推理模型这个词最近很火但说白了就是让模型在给出最终答案之前先输出一段可见的思考过程并且在这段思考中完成步骤拆分、验算和纠错。普通语言模型看到三个苹果吃了一个半还剩多少可能直接输出1.5个或2个但推理模型会先写总数3吃掉1.5剩余3-1.51.5验算1.51.53所以答案是1.5。这个能力不是靠把模型做大便能自动出现的而是训练数据的形态发生了根本变化。你需要让模型见过足够多的、带有中间步骤的样本甚至要让它看到中途发现错了再改正的过程。4.2 构造思维链数据的几个关键细节我为了让这个120M的小模型具备初步推理能力专门构造了一套思维链数据数量不大只有1000条精修样本但效果非常明显。构造过程中有三件事是最重要的。第一是数据格式要固定。我用的是指令问答中间穿插思考过程的模板每一条都明确切分为问题、思考、答案三部分。第二是难度要分层。如果全部是鸡兔同笼那种难题模型可能完全摸不到门路如果全部是三岁小孩都能答的问题模型也学不到验证这一环。我的做法是简单、中等、困难各占三分之一让模型先学会模仿思考格式再学会复杂步骤。第三是必须做答案校验。思维链数据里如果混入了推理错误但答案正确的样本模型会学到错误的中间逻辑。我每条数据都人工跑了一遍逻辑链确保思考过程的每一个推导步骤都能推出最终答案。4.3 用LoRA做SFT再用DPO强化推理偏好直接全量微调120M模型不是不行但我想保留一部分基座模型原本的能力就采用了LoRA做轻量SFT。关键参数是rank8alpha16dropout0.05只对注意力层的q和v投影做适配。SFT之后模型已经能输出带思考过程的回答但偶尔会思考过长、绕来绕去还跑偏。接着我用了DPODirect Preference Optimization进一步优化beta取值0.1构造了约500对偏好数据偏好对里正确且简洁的推理过程作为chosen冗长但不严谨的过程作为rejected。训练后的结果让我很意外小模型在50道数学推理题上的正确率从SFT前的28%到SFT后的53%DPO之后到了64%。虽然这个数字放在大模型圈子里不值一提但证明了从零构建推理模型这条路对中小规模模型完全可行。阶段推理题正确率回答平均长度特点基座模型28%12字直接给答案经常出错SFT后53%87字有思考过程有时过长DPO后64%63字过程精简正确率提升5. 放大实验后我踩过的三个坑显存、并行与检查点5.1 显存不够时先减什么而不是急着换卡很多人一上来就想上更大模型结果显存直接爆掉。我的经验是先按这个顺序排查序列长度、batch size、是否开了梯度检查点、是否用了混合精度。序列长度和注意力计算是平方关系我从512降到384显存立刻省出接近四分之一loss曲线并没有明显变差。batch size降到8之后可以继续训练只是收敛会稍微慢一点。如果你用的是单卡梯度检查点也建议打开虽然每个step会慢一点点但显存占用可能直接减半。混合精度基本是标配用bf16比fp16在小模型上更稳尤其是我这种从头训练的场景bf16的指数范围更大不容易出现梯度溢出。5.2 多卡并行最容易出问题的三个环节从单卡切到多卡时最大的问题反而不是模型代码而是数据加载和梯度同步。第一个坑是数据采样重复。多卡环境中如果每个进程都用自己的随机种子会出现同一份数据被不同卡重复采样模型等于在同样的样本上翻倍训练。正确做法是在DataLoader层面用DistributedSampler并且保证每张卡拿到不重叠的数据切片。第二个坑是梯度平均没对齐。DDP默认会对梯度做平均但你如果手动改过loss的scale很容易把梯度扩大或缩小N倍导致loss曲线在切换多卡后突然变得不平滑。第三个坑是通信瓶颈。小模型在8卡上训练很多时候速度提升不是8倍而是只有3到4倍原因就是频繁的梯度同步把时间都花在了通信上。解决思路是适当增大batch size、减少step数量让通信次数降下来。5.3 日志、检查点与一次让我后悔许久的断电我必须坦白一件事有一回训练跑到了3500步loss已经降到很不错的状态然后机房断电了。我发现自己根本没有配置自动保存检查点那一刻整个人是崩溃的。3500步的训练时间、电费、等待全部重来。那次之后我把检查点逻辑彻底改了每500步存一次完整权重每100步存一次optimizer状态和loss日志同时把关键指标推到可视化面板。磁盘占用并不大但对长训练来说这份保险是必需品。还有一个小建议训练日志不要只记录loss要把学习率、梯度范数、显存占用、每个卡的吞吐量全部记下来。后面排查问题的时候这些数据比loss本身有用得多。6. 一份可以照抄的从零路线图以及我最后悔的三件事6.1 一台GPU就能跑通的完整项目参数如果你也想从零训练一个带基础推理能力的模型我给你一套可以直接参考的配置这是我跑完一个完整项目后觉得性价比最合适的组合模型130M参数12层Transformerhidden size 768数据约300M token混合中文和英文分词器BPE词表32000训练单卡RTX 4090或类似24GB显存显卡bf16混合精度batch size 16序列长度512总时长约12到15小时推理增强1000条思维链样本做LoRA微调500对偏好样本做DPO整套流程加起来大约一周的空闲时间就能从零跑到最终评测。我强烈建议你把第一版模型的规模压到这个量级别一上来就追求7B或13B因为你在小模型上踩过的每一个坑到了大模型阶段都会以更昂贵的方式重现。6.2 按周拆解的时间路径与成本预估阶段时间核心任务第1-2周数据收集、清洗、分词器训练建好训练语料和评测集第3-4周Transformer实现、训练循环跑通基座模型的完整训练第5-6周思维链数据构造、LoRA微调让模型学会输出推理过程第7周DPO优化与评测迭代提高推理正确率、控制输出长度算力成本方面如果你用自己的显卡主要成本是电费和时间如果租云GPU按单卡4090的市场行情算十几个小时的训练成本相对可控完全在个人项目可承受的范围内。你真正要付出的其实是调试数据和跑实验的精力这部分省不掉。6.3 我最后悔没早点知道的三件事第一件验证集真的要从第一天就搭好。我一开始偷懒没有建评测集等到训练跑了一大半才开始想怎么定义模型好不好结果前面所有的checkpoint都没有可比性。第二件检查点和日志比模型结构更重要。我再也不想体验断电后一切归零的感觉。你永远不知道训练到什么时候会遇到意外所以请在最开始就把自动保存做好。第三件一次只改一个变量。我有一段时间同时调了学习率、批次大小、数据清洗规则和模型层数结果loss变差了我完全不知道是哪一步造成的。后来强制自己每次只动一个参数训练过程才变成了可控的对照实验。如果你正在考虑做一个类似的项目我的体会是不要急着追求复杂架构也不要迷信更大的模型。先把一条最朴素的从数据到评测的管线跑通再慢慢往里加入推理、偏好优化这些东西。那条看似绕远的从零开始的路实际上是最快把AI工程变成自己能力的过程。
返回列表