从零训练小型语言模型：数据、模型、训练与踩坑全流程

最近刚把一门深度学习课程的作业一完整跑通主题是Training也就是从头训练一个小型语言模型。先说结论这件事没有想象中那么玄但也没有网上教程写得那么轻松。你会在数据预处理、模型初始化、训练循环、评测生成这几个环节里反复踩坑然后把“为什么模型不收敛”这个问题从玄学变成工程。这篇内容适合两类人一类是正在做类似课程作业、需要一份完整流程参照的同学另一类是自己想从零训练一个几十M级别语言模型、但不知道从哪里下手的开发者。我会把整条链路拆开讲清楚包括数据管线、模型结构、训练配置以及我实际跑的时候遇到的故障和排查思路照着走基本能复现。1. 作业一到底在训练什么目标拆解与整体流程1.1 为什么“从头训练”值得认真做一遍很多人接触语言模型是从微调开始的下载一个预训练权重加载 tokenizer然后在自己的小数据集上跑几个 epoch。这么做很快但你对模型内部发生的事情几乎没有感知。微调时你以为自己在“教模型新知识”实际只是在调整最后几层分布的偏向而真正的前向计算、自注意力掩码、损失函数怎么作用到每一个 token 上这些细节都藏在别人写好的代码里。作业一恰好把这一切重新拉回到自己能控制的范围内。它不给预训练权重要求你自己定义词表、自己构造训练样本、自己写模型前向、自己写训练循环最后用生成结果来验收。这个过程的本质是在一个极小规模上复现一个现代生成式语言模型的完整生命周期跑完一遍你再看任何开源代码感觉会完全不同。1.2 任务边界我们不是在复现一个超大模型这份作业并不是要你训练一个 GPT-3 级别的东西它的核心约束是“在单卡、可接受时间内跑完”。我当时给自己定的目标是参数量控制在 50M 以内大概是一个 6 层、8 头、768 维的 decoder-only transformer训练语料控制在 1 亿 token 左右文本以英文为主单卡训练时间控制在 8 小时以内方便反复调参。这个规模放在今天看非常小但它足够支撑起完整的自回归语言模型行为。你会看到 loss 曲线平稳下降看到 generate 函数真的能根据前缀写出语法上像样的句子。这种正反馈对理解模型非常重要比拿一个 7B 模型跑推理要有教育意义得多。1.3 我把整条流程拆成了四个阶段实际操作中我按下面的顺序推进每一步都有明确的输入输出数据阶段原始文本 - 清洗 - 分词 - 切块 - 二进制训练集模型阶段定义 token embedding、位置编码、transformer block、输出头训练阶段实现训练循环、优化器、调度器、loss 计算、日志记录评测阶段计算验证集困惑度、编写生成函数、人工检查生成质量。这四个阶段顺序不能乱。我见过有人先把模型写得特别复杂然后再回头处理数据结果上下文长度、padding 方式都对不上返工成本非常高。先把数据和模型接口定死训练循环只是一个 for 循环的事。2. 数据管线从原始文本到自回归训练样本2.1 用 BPE 分词而不是直接按字符拆对英文语料来说最偷懒的方式是按空格切词或者干脆按字符切。按空格切会导致词表爆炸任何一个标点组合都会产生新 token按字符切会让序列变得特别长模型要学的东西过于底层。作业一明显希望你去实现一个更合理的子词分词其中BPEByte Pair Encoding是最标准的选择。BPE 的核心逻辑不复杂先把所有文本拆成单字符然后不断统计相邻字符对的频次把频次最高的对合并成一个新符号重复直到词表达到目标大小。它的优点是既保留了完整单词作为高频 token又能用子词单元覆盖生僻词。我当时把词表大小设成了 8K理由是训练语料只有 1 亿 token8K 已经足够覆盖绝大多数常见词又不会让 embedding 矩阵占用过多显存。分词阶段最容易忽略的是统一预处理规则。比如文本里有 HTML 标签、多余空格、特殊 Unicode 符号必须在训练和生成时使用完全相同的清洗流程。否则训练时看到的是dont生成时你输入don’t分词结果可能完全错开模型就变成了“哑巴”。2.2 如何把语料切成 context 块训练样本的格式是一个固定长度的 token 序列比如 256 或 512。原始语料长度一般远大于这个值所以要做滑窗切块。最简单的做法是每个 block 都由连续 token 组成block 之间不重叠def make_blocks(tokens, block_size): n len(tokens) // block_size blocks [] for i in range(n): start i * block_size block tokens[start:start block_size] blocks.append(block) return blocks但这里有个细节如果语料内部有章节边界直接硬切会把不相关的上下文拼在一起增加训练难度。我试过两种策略最终效果差异不小无重叠切块训练样本数少但每个样本内部语义相对完整带 stride 的滑动切块样本数多但会有大量重叠相当于部分数据被重复学习。对 1 亿 token 的语料来说我建议采用无重叠切块然后加上一层“在句子边界附近优先切分”的逻辑。具体做法是先定位所有句号、问号、感叹号选择离目标边界最近的自然断点作为实际切点。这样每个 block 基本是完整句子序列模型学起来要顺很多。2.3 批次构造里的两个隐蔽细节第一个隐蔽细节是标签和输入的错位关系。自回归语言模型的训练目标是给定当前位置之前的 token预测当前位置的 token。所以在构造 batch 时输入是tokens[:, :-1]标签是tokens[:, 1:]两条序列长度一致都等于block_size - 1。如果直接拿完整 block 当输入再拿同一个 block 当标签模型会去学“复制当前位置”训练 loss 会极低但生成完全乱套。第二个隐蔽细节是padding 后的 loss 掩码。当 batch 里存在不同长度文本时短文本会 pad 到统一长度。计算交叉熵时必须把 padding 位置的 loss 忽略掉否则模型会花大量精力去学习“预测 padding 符号”而 padding 符号本身毫无语义。实现时可以直接用cross_entropy(ignore_indexpad_id)或者自己构造 mask 后做 masked mean。我建议在进入训练之前先用一个随机权重的小模型跑一遍单 batch 的前向和反向确认没有维度错误再开始全量训练这一步能帮你省下至少半小时的调试时间。3. 模型搭建那些容易写错但影响极大的组件3.1 从零写 Decoder-Only Transformer 的模块划分我在搭建模型时没有用任何现成 transformer 库全部用基础张量操作实现。整个模型分成四个模块Embeddingtoken 映射到向量PositionalEncoding注入位置信息TransformerBlock自注意力 前馈网络 残差 LayerNormLMHead最后一层映射回词表 logits。需要特别留意的连接顺序是token embedding 加上位置编码后进入一堆 transformer block最后过一个 LayerNorm再接输出头。很多人会把最后一个 LayerNorm 漏掉导致训练的稳定性明显下降。3.2 权重初始化与残差连接前期不收敛的元凶模型能收敛一部分功劳来自正确的初始化。我当时沿用了一套被广泛验证过的方案embedding 层和输出头共享权重数值在[-0.1, 0.1]之间均匀初始化线性层权重标准差设为0.02偏置初始化为 0attention 里的q/k/v输出层以及 MLP 第二个全连接层额外除以sqrt(2 * num_layers)作用是抑制深层残差分支的累积方差。如果你不做这最后一步6 层以上的模型非常容易出现早段 loss 不降、甚至训练刚开始就发散的情况。原因在于残差连接是逐层累加的每个 block 的贡献都在叠加最后一层输出分布方差会过大softmax 之后过于尖锐梯度传回去就变成了异常值。3.3 位置编码和 attention mask 怎么配合我用的位置编码是最常见的sinusoidal直接加到 token embedding 上不需要额外学习。实现时可以预计算一个(max_len, d_model)的位置矩阵每次 forward 时取前seq_len行加进去。更关键的是 attention mask。decoder-only 模型要求每个位置只能看到自己和之前的位置不能看到未来。标准写法是用torch.triu生成一个上三角为True的布尔矩阵然后在注意力分数上加-infmask torch.triu(torch.ones(seq_len, seq_len, devicex.device), diagonal1).bool() scores scores.masked_fill(mask.unsqueeze(0).unsqueeze(0), float(-inf)) scores torch.softmax(scores, dim-1)这个 mask 如果写反了后果是模型会偷看未来 token训练 loss 会异常低但生成时完全不能用。我在这个点上栽过一次后面会详说。4. 训练循环优化器、调度器、梯度裁剪的细节4.1 AdamW 参数与 weight decay 的正确设置优化器我选了 AdamW因为它分离了权重衰减和自适应学习率在大模型训练里比普通 Adam 稳定得多。参数配置是参数我使用的值说明lr3e-4峰值学习率不能太大betas(0.9, 0.95)一阶和二阶动量系数eps1e-8防止除零weight_decay0.1只作用于权重矩阵不作用于 bias 和 normweight decay 不要对所有参数一视同仁。对 bias、LayerNorm 的 scale 和 shift 参数做 weight decay 没有意义反而可能损害性能。实现时可以传一个param_groups按参数名过滤。学习率为什么不能太大小型模型本身 batch 不大梯度噪声相对明显如果峰值学习率超过1e-3很容易在某一步跨过 loss 的“高墙”直接进入发散区域。我看到过很多调参失败的案例最后都把学习率降下来就好了。4.2 学习率预热和余弦退火语言模型训练对学习率非常敏感我使用了如下的调度策略第一阶段预热。前 1000 步学习率从0线性升到3e-4。预热的作用是让模型在前期避免大步长带来的震荡尤其是随机初始化导致的 loss 极高区域。第二阶段余弦退火。从第 1000 步到总步数学习率按余弦曲线从3e-4降到1e-5。退火让模型后期走得稳能进一步压 loss。总步数根据语料计算总token数 / (batch_size * block_size)。比如 1 亿 token、batch size 64、block size 256每个 step 处理 16384 token总共约 6100 步。那预热可以设成 500 步占比 8% 左右。4.3 梯度裁剪与梯度累积梯度裁剪是必须做的。语言模型的 loss 曲线偶尔会出现一个尖刺如果不做裁剪这一大步会直接把权重推到坏点后面很难恢复。我通常把 global grad norm 裁剪到1.0grad_norm torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)另一个技巧是梯度累积。如果你的显存只能放下 batch size 32但实验目标想对比 batch size 64 的效果可以累积 2 个 step 再更新一次参数。实现时注意只有累积到指定步数时才调用optimizer.step()和optimizer.zero_grad()同时 loss 除以累积次数避免 loss 值被放大。5. 踩坑复盘loss 不降、生成乱码、显存爆掉的排查链路5.1 我遇到的三个典型故障信号第一次完整跑训练时日志里出现了三个非常典型的问题loss 在第 2000 步后几乎不动稳定在某个值附近既不下降也不上升perplexity 评估结果异常高甚至比随机猜测还差生成函数输出的内容是乱码和重复片段完全没有语义结构。这三个信号有时是同一个根因造成的有时是叠加问题。如果只看 loss 曲线很容易误判成“模型容量不够”实际上问题可能出在数据预处理或注意力 mask 上。5.2 用一组“最小诊断实验”定位问题我先写了一个非常小的诊断脚本只取 4 个样本block size 设为 32训练 10 步观察 loss 能不能降到极低。这个实验的目的是检验“模型是否有能力记住这条数据”如果连这个都做不到说明代码有结构性错误而不是调参问题。我当时跑出来的结果是 loss 能降但降到一定值后开始上下波动。进一步检查后发现两个问题第一标签构造时我没做错位输入和标签完全一致模型在训练时学的是复制。解决后 loss 明显下降但验证集仍然不太对。第二注意力 mask 把未来位置-inf加错了维度masked_fill时把整个序列全遮住了导致 attention 权重偏向平均分布。这意味着每个 token 其实只看得到有限的上下文模型等于在一个“近视”状态下训练。修正 mask 之后验证集困惑度立刻掉了十几个点。这个诊断思路可以复制先小、先慢、先看代码逻辑再看超参数。不要一上来就调学习率那样只会让 bug 藏得更深。5.3 生成乱码的根因排查生成乱码这件事一开始我以为是模型欠拟合后来发现是两个小细节在作怪。第一个是采样方式。生成时如果总是取argmax小模型很容易陷入重复循环如果用温度采样但温度设得太低比如0.2概率分布尖锐输出也会非常刻板。我后来把温度调到0.8并加入 top-k50 采样生成结果一下子自然了很多。第二个是特殊 token 处理。我的分词器里没有单独定义 BOS/EOS训练样本只是普通的连续 token 序列。生成时直接从一个固定前缀开始模型会在句尾持续输出自然语言。但如果训练时语料中有大量文本边界被直接切掉模型就会学到“话说到一半强行续写”细看会觉得每句话都不完整。解决方案还是回到数据的句子边界切分。5.4 显存爆掉的应对顺序OOM 是我的老朋友。我总结了一套按性价比排序的应对清单先缩小 block size从 512 降到 256显存占用几乎减半再缩小 batch size但配合梯度累积保证有效 batch 不变开启混合精度训练fp16 能让显存减少约一半同时加快训练最后才考虑 checkpoint 或更复杂的显存优化手段因为代码复杂度会上升。有一句话值得反复提能放进显存的模型才是好模型。不要为了追求更大的 block size 而牺牲模型层数对一个小型语言模型来说深度带来的收益通常大于单条序列长度。6. 本地复现与后续扩展建议6.1 一份可以直接照抄的配置清单如果你在本地或单卡服务器上复现下面这一份配置是我在作业中反复调出来的稳定组合可以作为起始点# 数据 vocab_size: 8192 block_size: 256 train_tokens: 100000000 # 模型 num_layers: 6 num_heads: 8 d_model: 768 d_ff: 2048 dropout: 0.1 tie_embeddings: true # 训练 optimizer: adamw lr: 0.0003 min_lr: 0.00001 warmup_steps: 600 weight_decay: 0.1 grad_clip_norm: 1.0 batch_size: 64 grad_accumulation: 2 total_epochs: 1 # 生成 temperature: 0.8 top_k: 50这份配置的总参数量大约 40M 到 50M在一个 24G 显存的卡上跑起来非常宽松。如果是 16G 显存把block_size降到 128 或batch_size减半即可。6.2 如果继续往下做我会先碰这几个方向作业验收之后我顺手做了几个小扩展每一个都只改一个变量但能学到完全不同的东西换数据集把英文语料换成代码或中文语料观察分词器和生成风格的变化增大词表从 8K 升到 32K看模型困惑度和训练速度的权衡加评估任务除了困惑度接一组下游任务做 zero-shot 评估尝试参数高效微调在训练好的小模型上用 LoRA 微调一个特定风格对比与全量微调的差距。跑完这个作业最值得记住的不是 loss 曲线多漂亮而是你终于知道每一个旋钮在改变什么。下次再看到别人的训练日志你能从学习率、batch size、梯度范数这些数字里读出问题再看到开源框架的实现你也不会把它当成黑盒。这就是作业一的价值所在。