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

最近刚把一门深度学习课程的作业一完整跑通，主题是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 我把整条流程拆成了四个阶段

实际操作中，我按下面的顺序推进，每一步都有明确的输入输出：

1. 数据阶段：原始文本 -> 清洗 -> 分词 -> 切块 -> 二进制训练集；
2. 模型阶段：定义 token embedding、位置编码、transformer block、输出头；
3. 训练阶段：实现训练循环、优化器、调度器、loss 计算、日志记录；
4. 评测阶段：计算验证集困惑度、编写生成函数、人工检查生成质量。

这四个阶段顺序不能乱。我见过有人先把模型写得特别复杂，然后再回头处理数据，结果上下文长度、padding 方式都对不上，返工成本非常高。先把数据和模型接口定死，训练循环只是一个 for 循环的事。

2. 数据管线：从原始文本到自回归训练样本

2.1 用 BPE 分词而不是直接按字符拆

对英文语料来说，最偷懒的方式是按空格切词，或者干脆按字符切。按空格切会导致词表爆炸，任何一个标点组合都会产生新 token；按字符切会让序列变得特别长，模型要学的东西过于底层。作业一明显希望你去实现一个更合理的子词分词，其中BPE（Byte Pair Encoding）是最标准的选择。

BPE 的核心逻辑不复杂：先把所有文本拆成单字符，然后不断统计相邻字符对的频次，把频次最高的对合并成一个新符号，重复直到词表达到目标大小。它的优点是既保留了完整单词作为高频 token，又能用子词单元覆盖生僻词。我当时把词表大小设成了 8K，理由是训练语料只有 1 亿 token，8K 已经足够覆盖绝大多数常见词，又不会让 embedding 矩阵占用过多显存。

分词阶段最容易忽略的是统一预处理规则。比如文本里有 HTML 标签、多余空格、特殊 Unicode 符号，必须在训练和生成时使用完全相同的清洗流程。否则训练时看到的是don't，生成时你输入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_index=pad_id)，或者自己构造 mask 后做 masked mean。

我建议在进入训练之前，先用一个随机权重的小模型跑一遍单 batch 的前向和反向，确认没有维度错误再开始全量训练，这一步能帮你省下至少半小时的调试时间。

3. 模型搭建：那些容易写错但影响极大的组件

3.1 从零写 Decoder-Only Transformer 的模块划分

我在搭建模型时没有用任何现成 transformer 库，全部用基础张量操作实现。整个模型分成四个模块：

• Embedding：token 映射到向量；
• PositionalEncoding：注入位置信息；
• TransformerBlock：自注意力 + 前馈网络 + 残差 + LayerNorm；
• LMHead：最后一层映射回词表 logits。

需要特别留意的连接顺序是：token embedding 加上位置编码后，进入一堆 transformer block，最后过一个 LayerNorm，再接输出头。很多人会把最后一个 LayerNorm 漏掉，导致训练的稳定性明显下降。

3.2 权重初始化与残差连接：前期不收敛的元凶

模型能收敛，一部分功劳来自正确的初始化。我当时沿用了一套被广泛验证过的方案：

• embedding 层和输出头共享权重，数值在[-0.1, 0.1]之间均匀初始化；
• 线性层权重标准差设为0.02，偏置初始化为 0；
• attention 里的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的布尔矩阵，然后在注意力分数上加-inf：

mask = torch.triu(torch.ones(seq_len, seq_len, device=x.device), diagonal=1).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 稳定得多。参数配置是：

weight 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.0：

grad_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 我遇到的三个典型故障信号

第一次完整跑训练时，日志里出现了三个非常典型的问题：

1. loss 在第 2000 步后几乎不动，稳定在某个值附近，既不下降也不上升；
2. perplexity 评估结果异常高，甚至比随机猜测还差；
3. 生成函数输出的内容是乱码和重复片段，完全没有语义结构。

这三个信号有时是同一个根因造成的，有时是叠加问题。如果只看 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-k=50 采样，生成结果一下子自然了很多。

第二个是特殊 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、梯度范数这些数字里读出问题；再看到开源框架的实现，你也不会把它当成黑盒。这就是作业一的价值所在。