从零实现一个小型语言模型

这一章，我们终于让模型自己写一句话。不过，先别请它写小说。我们给它一个很小、也很容易核对的世界：

小林选了红色杯子，放到桌上。最后小林拿走了红色杯子。

句子不长，但里面有一个跨越多字的关系：后面拿走的颜色，应该与前面选中的颜色一致。我们会训练上一章的 Transformer，比较模型大小、可用上下文和训练文本质量的影响，再把训练好的参数保存、加载，用来生成文本。

1. 先约定什么算成功

如果只问“生成的句子像不像中文”，很多错误很容易混过去：

小林选了红色杯子，放到桌上。最后小林拿走了绿色杯子。

它读起来通顺，却破坏了本实验的前后对应规则。所以我们同时观察两类指标：

• 平均交叉熵与困惑度：模型预测所有目标词元的表现。
• 末次颜色准确率：在提供真实前文时，模型对第二次颜色的最高分预测是否正确。

后者只检查这一个位置，不等于整句生成准确率，也不是对现实世界事实的综合评价。我们还会直接查看生成文本，找出重复、错色和无法停止等现象。

2. 做一份可以完整检查的小语料

我们组合四个人名、四种颜色、四类物品、四个地点，得到 $4^4=256$ 条自编文本。每条都是 26 个字符，结构相同，但四个槽位的取值不同。

例如，物品可以是杯子、盒子、本子或袋子；地点可以是桌上、架上、柜里或门边。后半句重复前面的人名、颜色和物品，形成明确的依赖关系。

固定随机种子后，划分为 192 条训练、32 条验证、32 条测试。三个集合没有完全相同的句子，但共享词汇与句式。这是在检查一个有限世界中的组合规律，不能把结果当作开放文本写作能力。

为便于观察，我们用字符级编码，不再训练 BPE。词表由训练文本中的字符，加 PAD、BOS、EOS 三项组成，共 35 项。新提示出现词表外字符时，代码会明确报错；这个模型暂时不能处理任意中文输入。

每条文本编码后，在开头加 BOS、末尾加 EOS。输入与目标错开一位，因此一条文本提供 27 个目标：

输入：BOS 小 林 选 了 红 色 …… 子 。
目标： 小 林 选 了 红 色 杯 …… 。 EOS

一轮完整训练共有 $192\times27=5184$ 个目标。PAD 只用于补齐批次，不计入损失；BOS 是输入，EOS 则是需要预测的目标。

3. 切块以后，模型到底能看到哪里

本实验比较块长 32 和块长 8。块长 32 足以容纳一条文本的全部输入；块长 8 则把它拆成四段。

图 1：横向位置按预测目标编号。第一次颜色在目标索引 4，第二次颜色在索引 21；短块模型预测后者时，访问不到前者所在的块。

每个块单独输入模型，没有跨块循环状态。块内仍有因果掩码，当前位置只能查看当前及更早的输入。块长 8 是一个上限，靠近块开头的位置实际上只能看到更短的前缀。

划分发生在每条文本内部，不会把一条训练文本的结尾与另一条开头错误地连成自然续写。验证和测试使用同样的分块规则。

每块在右侧补 PAD，目标中的 PAD 通过 ignore_index=PAD 排除。这里有效位置都位于填充之前，因果掩码已经使它们无法访问右侧 PAD；我们也不会使用 PAD 查询位置的预测。若改成左填充或非因果网络，就不能直接照搬这套处理。

用于比较的平均损失按所有有效目标累计，再除以总目标数。不能先算每块平均损失后简单求平均，否则最后那些只有三个目标的短块，可能与包含八个目标的块获得相同权重。

4. 四组设置，各自改变什么

所有网络都使用两层 Pre-LN Transformer 块、四个注意力头和可学习位置表，沿用上一章的实现。最大位置表长度统一为 32；短块方案只使用前八个位置。

颜色错配方案把训练集中所有第二次颜色收集起来，打乱后放回去。这样，文本数量、长度乃至整个训练集的字符计数都不变，但 192 条中有 142 条的前后颜色不再一致。验证和测试仍然使用原始的一致文本。

四组都训练 40 轮，每轮 12 次参数更新。长块的批量为 16，短块的批量为 64；由于每条文本正好拆成四个短块，各组每轮遍历的真实目标总数相同。

这控制了轮数、目标数量和更新次数，没有控制相同计算量。宽模型参数更多，短块的注意力矩阵更小，运行成本当然不同。第五篇讨论计算预算时，我们还会换一种更严格的比较方式。

优化器使用 AdamW，学习率为 0.003，权重衰减为 0.01，并把梯度总范数裁剪到至多 1。每轮观察验证交叉熵，保存截至当时最好的参数；最后按验证损失选择四个方案中的一个，再运行测试集。

5. 生成时，怎样从分数中选择

模型给出词表上每一项的分数后，可以选择最大值，也可以采样。为了控制采样的分布，我们引入三个参数。

温度记作 $\tau>0$：

$$p_i=\frac{\exp(z_i/\tau)} {\sum_j\exp(z_j/\tau)}$$

$z_i$ 是原始分数。温度低于 1 时，较高分项的优势通常被放大；高于 1 时，分布变得更平缓。它不会改变词元分数的排序，更不会自动提高事实准确率。

代码约定 temperature=0 直接采用贪心选择。这是一个分支约定，不是把零代入分母。

Top-k 只保留概率最高的 $k$ 个候选。Top-p 则把候选按概率排序，取累计概率首次达到 $p$ 的最短前缀。比如概率为 0.5、0.3、0.15、0.05，取 $p=0.75$ 时，要保留前两项，再重新归一化。

两种筛选方法都可能把正确但低概率的词元删掉，也可能保留高概率的错误。本文实现先做温度变换，再做 top-k，最后在剩余分布上做 top-p；同时使用时，要注意这个顺序。

图 2：用固定分数 3、2、1、0 演示温度的作用，不代表训练模型某一步的实际输出。

生成从 BOS 或给定提示开始，每次加入一个词元，遇到 EOS 停止，并额外设置 60 个新词元的上限。PAD 和 BOS 不允许作为普通生成结果再次出现。

当序列超过块长时，这份生成代码只取最近一个窗口，并从位置 0 重新编号。它会重新计算整个窗口，没有使用缓存。训练与评估则按固定块边界切分，所以长文本生成时的窗口边界并不完全相同。尤其是短块模型，生成表现可能同时受可见历史不足、位置重置和这种边界差异影响。

6. 完整训练与生成代码

环境沿用 Python 3.12 和固定版本依赖。在仓库根目录准备好环境后运行：

python -m pip install -r requirements.txt
python src/examples/small_language_model.py

实验在 CPU 上完成，不需要下载模型或数据。所有网络参数都从随机初始化开始，PyTorch 负责基本张量运算与自动微分。

如果单独下载代码，请把模型定义 tiny_transformer.py与训练脚本 small_language_model.py放在同一目录。前者已在上一章完整展示；本章直接导入同一实现。

"""第 21 章：训练、比较、保存和生成短文本；与 tiny_transformer.py 放在同一目录。"""

import argparse
from collections import Counter
from copy import deepcopy
from hashlib import sha256
from itertools import product
from math import exp
from pathlib import Path
from random import Random
from tempfile import TemporaryDirectory
import torch
from torch import nn
from tiny_transformer import TinyTransformer

PAD, BOS, EOS = 0, 1, 2
COLOR_INDEX = 21  # 原句中第二次颜色的位置，输入加 BOS 后仍是相应目标的索引。


def corpus():
    records = [
        f"{name}选了{color}色{thing}，放到{place}。最后{name}拿走了{color}色{thing}。"
        for name, color, thing, place in product(
            ["小林", "小周", "小陈", "小李"], ["红", "蓝", "绿", "白"],
            ["杯子", "盒子", "本子", "袋子"], ["桌上", "架上", "柜里", "门边"])
    ]
    assert all(len(text) == 26 for text in records)
    Random(2026).shuffle(records)
    return records[:192], records[192:224], records[224:]


def corrupt_colors(texts):
    # 打乱第二次颜色的对应关系，保持文本长度和整个训练集的字符计数不变。
    colors = [text[COLOR_INDEX] for text in texts]
    Random(2028).shuffle(colors)
    return [text[:COLOR_INDEX]+color+text[COLOR_INDEX+1:]
            for text, color in zip(texts, colors)]


def vocabulary(texts):
    return ["<PAD>", "<BOS>", "<EOS>"]+sorted(set("".join(texts)))


def make_windows(texts, vocab, block_size):
    ids = {token: index for index, token in enumerate(vocab)}
    rows = []
    for text in texts:
        sequence = [BOS]+[ids[char] for char in text]+[EOS]
        for start in range(0, len(sequence)-1, block_size):
            x = sequence[start:start+block_size]
            y = sequence[start+1:start+block_size+1]
            # 最后一块的输入长度必须与可预测目标数量一致。
            x = x[:len(y)]
            marks = [start+i == COLOR_INDEX for i in range(len(y))]
            padding = block_size-len(y)
            rows.append((x+[PAD]*padding, y+[PAD]*padding, marks+[False]*padding))
    return tuple(torch.tensor([row[i] for row in rows], dtype=dtype)
                 for i, dtype in enumerate([torch.long, torch.long, torch.bool]))


@torch.no_grad()
def evaluate(model, data):
    model.eval()
    x, y, marks = data
    loss_sum, target_count, color_correct, color_count = 0.0, 0, 0, 0
    for start in range(0, len(x), 64):
        batch = slice(start, start+64)
        logits = model(x[batch])
        loss_sum += nn.functional.cross_entropy(
            logits.flatten(0, 1), y[batch].flatten(), ignore_index=PAD, reduction="sum").item()
        target_count += (y[batch] != PAD).sum().item()
        color_correct += ((logits.argmax(-1) == y[batch]) & marks[batch]).sum().item()
        color_count += marks[batch].sum().item()
    return loss_sum/target_count, color_correct/color_count


def train(texts, validation, vocab, width, block_size, epochs):
    torch.manual_seed(2026)
    model = TinyTransformer(len(vocab), max_context=32, width=width, heads=4, layers=2).double()
    data = make_windows(texts, vocab, block_size)
    val_data = make_windows(validation, vocab, block_size)
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.003, weight_decay=0.01)
    generator = torch.Generator().manual_seed(2027)
    batch_size = 16 if block_size == 32 else 64
    best_loss, best_epoch, best_state = float("inf"), 0, None
    history = []
    x, y, _ = data
    for epoch in range(1, epochs+1):
        model.train()
        for batch in torch.randperm(len(x), generator=generator).split(batch_size):
            optimizer.zero_grad(set_to_none=True)
            logits = model(x[batch])
            loss = nn.functional.cross_entropy(logits.flatten(0, 1), y[batch].flatten(),
                                               ignore_index=PAD)
            if not torch.isfinite(loss):
                raise RuntimeError("损失不是有限数，请检查数据、掩码与学习率")
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
        train_loss, _ = evaluate(model, data)
        val_loss, accuracy = evaluate(model, val_data)
        history.append((train_loss, val_loss, accuracy))
        if val_loss < best_loss:
            best_loss, best_epoch = val_loss, epoch
            best_state = deepcopy(model.state_dict())
    model.load_state_dict(best_state)
    return model, best_epoch, history


def run(epochs=40):
    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    train_texts, validation, test = corpus()
    vocab = vocabulary(train_texts)
    configs = [
        ("基础模型", 32, 32, train_texts),
        ("加宽模型", 64, 32, train_texts),
        ("短上下文", 32, 8, train_texts),
        ("颜色错配", 32, 32, corrupt_colors(train_texts)),
    ]
    results = []
    for name, width, block_size, texts in configs:
        model, epoch, history = train(texts, validation, vocab, width, block_size, epochs)
        val_loss, accuracy = evaluate(model, make_windows(validation, vocab, block_size))
        results.append((name, block_size, model, epoch, history, val_loss, accuracy))
    selected = min(results, key=lambda row: row[5])
    return (train_texts, validation, test), vocab, results, selected


def choose(logits, generator, temperature=0.8, top_k=None, top_p=1.0):
    if temperature < 0 or not 0 < top_p <= 1:
        raise ValueError("temperature 不小于零，top_p 在 (0, 1] 内")
    logits = logits.clone()
    logits[[PAD, BOS]] = float("-inf")
    if temperature == 0:
        return int(logits.argmax())
    logits /= temperature
    if top_k is not None:
        if not 1 <= top_k <= len(logits):
            raise ValueError("top_k 必须在 1 与词表大小之间")
        indices = logits.topk(top_k).indices
        keep = torch.zeros_like(logits, dtype=torch.bool)
        keep[indices] = True
        logits[~keep] = float("-inf")
    sorted_logits, indices = logits.sort(descending=True)
    probabilities = sorted_logits.softmax(-1)
    # 保留使累计概率首次达到阈值的那个词元，至少保留第一项。
    remove = probabilities.cumsum(-1)-probabilities >= top_p
    sorted_logits[remove] = float("-inf")
    chosen = torch.multinomial(sorted_logits.softmax(-1), 1, generator=generator)
    return int(indices[chosen].item())


@torch.no_grad()
def generate(model, vocab, block_size, prompt="", temperature=0.8, top_k=8, top_p=0.9):
    model.eval()
    ids = {token: index for index, token in enumerate(vocab)}
    if any(char not in ids for char in prompt):
        raise ValueError("提示中有词表未收录的字符")
    sequence = [BOS]+[ids[char] for char in prompt]
    generator = torch.Generator().manual_seed(2029)
    for _ in range(60):
        inputs = torch.tensor([sequence[-block_size:]])
        token = choose(model(inputs)[0, -1], generator, temperature, top_k, top_p)
        if token == EOS:
            return "".join(vocab[index] for index in sequence[1:]), "EOS"
        sequence.append(token)
    return "".join(vocab[index] for index in sequence[1:]), "长度上限"


def save_model(path, model, vocab, block_size, name, epochs, best_epoch, train_texts):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({"config": model.config, "state": model.state_dict(), "vocab": vocab,
                "block_size": block_size, "name": name, "epochs": epochs,
                "best_epoch": best_epoch,
                "torch_version": str(torch.__version__),
                "training_sha256": sha256("\n".join(train_texts).encode()).hexdigest()}, path)


def load_model(path):
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    model = TinyTransformer(**checkpoint["config"]).double()
    model.load_state_dict(checkpoint["state"])
    model.eval()
    return model, checkpoint["vocab"], checkpoint["block_size"]


def show_generation(model, vocab, block_size):
    prompt = "小林选了红色杯子，放到桌上。最后小林拿走了"
    for label, prefix, temperature in [("贪心续写", prompt, 0.0), ("采样生成", "", 0.8)]:
        text, stop = generate(model, vocab, block_size, prefix, temperature)
        print(f"{label}：{text}（{stop}）")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--epochs", type=int, default=40)
    actions = parser.add_mutually_exclusive_group()
    actions.add_argument("--save", type=Path)
    actions.add_argument("--load", type=Path)
    args = parser.parse_args()
    torch.set_num_threads(1)
    if args.load:
        show_generation(*load_model(args.load))
        return
    if args.epochs < 1:
        parser.error("--epochs 必须大于零")
    datasets, vocab, results, selected = run(args.epochs)
    counts = Counter(char for text in datasets[0] for char in text)
    counts["<EOS>"] = len(datasets[0])
    probabilities = torch.tensor([counts[token]+0.1 for token in vocab], dtype=torch.float64)
    probabilities /= probabilities.sum()
    indices = {token: i for i, token in enumerate(vocab)}
    targets = [indices[token] for text in datasets[1] for token in list(text)+["<EOS>"]]
    baseline_loss = -probabilities[targets].log().mean().item()
    print(f"训练/验证/测试文本：{len(datasets[0])} / {len(datasets[1])} / {len(datasets[2])}")
    print(f"字符词表大小：{len(vocab)}；每轮训练目标数={len(datasets[0])*27}")
    print(f"一元基线验证交叉熵={baseline_loss:.3f}，困惑度={exp(baseline_loss):.3f}")
    corrupt = corrupt_colors(datasets[0])
    wrong = sum(a[COLOR_INDEX] != b[COLOR_INDEX] for a, b in zip(datasets[0], corrupt))
    print(f"颜色错配训练文本：{wrong} / {len(corrupt)}")
    for name, block_size, model, epoch, _, loss, accuracy in results:
        count = sum(p.numel() for p in model.parameters())
        print(f"{name}：参数={count}，块长={block_size}，最佳轮次={epoch}，"
              f"验证交叉熵={loss:.3f}，困惑度={exp(loss):.3f}，末次颜色准确率={accuracy:.3f}")
    name, block_size, model, best_epoch, _, _, _ = selected
    loss, accuracy = evaluate(model, make_windows(datasets[2], vocab, block_size))
    print(f"固定方案：{name}；测试交叉熵={loss:.3f}，"
          f"困惑度={exp(loss):.3f}，末次颜色准确率={accuracy:.3f}")
    prompt = "小林选了红色杯子，放到桌上。最后小林拿走了"
    for other_name, other_block, other_model, _, _, _, _ in results[2:]:
        text, stop = generate(other_model, vocab, other_block, prompt, temperature=0)
        print(f"{other_name}续写：{text}（{stop}）")
    with TemporaryDirectory() as folder:
        path = args.save or Path(folder)/"tiny-lm.pt"
        selected_training = corrupt if name == "颜色错配" else datasets[0]
        save_model(path, model, vocab, block_size, name, args.epochs, best_epoch, selected_training)
        restored, restored_vocab, restored_block = load_model(path)
        probe = torch.tensor([[BOS, vocab.index("小"), vocab.index("林")]])
        with torch.no_grad():
            equal = torch.equal(model(probe), restored(probe))
        assert restored_vocab == vocab and restored_block == block_size
        print(f"保存再加载后输出一致：{equal}")
        if args.save:
            print(f"已保存：{args.save}")
        show_generation(restored, vocab, restored_block)


if __name__ == "__main__":
    main()

固定环境下的输出为：

训练/验证/测试文本：192 / 32 / 32
字符词表大小：35；每轮训练目标数=5184
一元基线验证交叉熵=3.284，困惑度=26.685
颜色错配训练文本：142 / 192
基础模型：参数=28480，块长=32，最佳轮次=40，验证交叉熵=0.214，困惑度=1.239，末次颜色准确率=1.000
加宽模型：参数=106112，块长=32，最佳轮次=40，验证交叉熵=0.211，困惑度=1.234，末次颜色准确率=1.000
短上下文：参数=28480，块长=8，最佳轮次=40，验证交叉熵=0.364，困惑度=1.439，末次颜色准确率=0.188
颜色错配：参数=28480，块长=32，最佳轮次=38，验证交叉熵=0.265，困惑度=1.303，末次颜色准确率=0.219
固定方案：加宽模型；测试交叉熵=0.212，困惑度=1.236，末次颜色准确率=1.000
短上下文续写：小林选了红色杯子，放到桌上。最后小林拿走了绿色盒子。最后盒子。最后盒子。最后盒子。最后盒子。最后盒子。最后盒子。最后盒子。最后盒子。最后盒子。最后盒子。最后盒子。（长度上限）
颜色错配续写：小林选了红色杯子，放到桌上。最后小林拿走了绿色杯子。（EOS）
保存再加载后输出一致：True
贪心续写：小林选了红色杯子，放到桌上。最后小林拿走了红色杯子。（EOS）
采样生成：小李选了蓝色本子，放到门边。最后小李拿走了蓝色本子。（EOS）

7. 读懂这几组结果

先看损失。只统计字符频率的一元基线，验证困惑度约为 26.685；训练后的 Transformer 明显更低。这个比较表明，它学到了一些依赖上下文的规律，而不只是总体字符频率。

图 3：训练数据不同的两条曲线不能直接当作同一道题的分数；颜色错配模型在不一致的文本上训练，在一致文本上验证。

再看更具体的颜色指标。基础和加宽模型在这次验证中都预测对了全部 32 个末次颜色；短块方案只有 6 个正确，颜色错配方案只有 7 个正确。

图 4：每组只有 32 个验证样本，分数存在明显的离散性。这里报告一次固定划分和种子的结果，没有把它当作稳定排名。

颜色错配模型尤其值得停下来看看。它的困惑度只有 1.303，乍看并不高，却在很多颜色位置犯错。原因之一是：一条文本有 27 个目标，第二次颜色只占其中一个；大量固定句式都很容易预测，会拉低平均损失。任务中最关心的错误，不一定在平均指标里占很大比重。

实际续写也把问题展示出来了：颜色错配模型把红杯子写成绿杯子，却仍能顺利写完句子；短块模型既换了颜色和物品，又陷入“最后盒子”的循环。停止上限让程序能够退出，并不修复内容。

加宽模型的验证损失略低，因此按既定规则选中了它。它有约 3.7 倍参数，但这次差异很小，两个一致长块模型的颜色准确率也相同。一次实验不足以证明增加宽度普遍值得；还需要多次种子、不同数据和计算成本的比较。

基础和加宽模型的最优轮次都在第 40 轮，也提醒我们：实验恰好在这里结束，不能据此断言更长训练一定无益。

8. 保存模型，也保存它怎样解释输入

默认运行会在临时目录保存再加载一次，检查输出一致后清理文件。想保留可重复使用的结果，可以指定路径：

python src/examples/small_language_model.py --save runs/tiny-lm.pt
python src/examples/small_language_model.py --load runs/tiny-lm.pt

第二条命令只加载并生成，不再训练四组模型。

保存内容包括最佳权重、网络配置、字符词表、块长、方案名称、训练轮数、最佳轮次、PyTorch 版本和训练文本的校验值。只保存一堆权重，还不足以知道某个编号代表哪个字、模型需要多大的位置表。

这里保存的是可用于推理的模型快照。若要从中断处精确继续训练，还需要优化器状态、随机数状态以及数据进度等信息；第五篇会进一步讨论训练检查点。

这个小模型已经走完了数据、目标、结构、优化、评估、保存和生成的完整流程。不过，它仍然生活在我们编写的 256 条模板文本里。换成新闻、代码或多轮对话，会遇到远更多的数据和训练问题。这正是接下来大模型预训练要面对的内容。

9. 自测与扩展实验

1. 26 字的文本为什么产生 27 个训练目标，而不是 28 个？
2. 三个数据集没有相同完整句子，是否就代表测试了开放领域泛化？
3. 为什么短块模型和基础模型的参数量相同？
4. 142 条颜色被改错后，训练集字符计数为什么仍然不变？
5. 相同轮数、目标数和更新次数，是否代表相同计算预算？
6. 困惑度低但末次颜色准确率差，说明平均指标遗漏了什么？
7. 概率为 0.5、0.3、0.15、0.05，top-p 取 0.9，会保留哪些候选？
8. 只修改采样温度，能否修复训练文本中的错误对应关系？
9. 分别用 20、40、80 轮训练，并换几个数据划分与训练种子，报告波动，而不只保留最好的一次。
10. 增加第二种句式、不同长度的插入语，再重新划分数据。先检查代码中的固定颜色索引和长度假设；怎样调整指标，才能继续测到同一关系？

实现细节可继续对照 PyTorch 的交叉熵说明和模型保存与加载说明。这里的文本、四组对照和结果都是本书的教学实验，不是对某个公开大模型的复现。