AIGC标识 构建大语言模型

dsh 生成学习文档, 一份面向纯 CPU 环境的大语言模型实战指南:从数据、分词、注意力、GPT 实现,
到预训练、推理加速、分类微调、指令微调、LoRA、评估与部署,全部代码均可直接运行。
但所有示例不依赖 GPU,不依赖网络下载,采用确定性生成的语料与数据集,因此任何人都能在普通电脑上复现的每一个结果。

序:如何用纯 CPU 构建大语言模型

本章导读

如果你手上只有一台普通笔记本、一块没有独立显卡的服务器,甚至是一台云上按小时计费的四核虚拟机,你依然可以从零把一个大语言模型"建"出来——把语料喂进去,把注意力写出来,把损失压下去,最后得到一个能续写、能分类、能听懂指令的模型。本书要回答的就是这件事:在完全没有 GPU 的环境里,一个大语言模型是怎么从零长出来的,每个环节的真实代价是多少,哪些地方可以省、哪些地方省不了。

本章先建立全局图景:三条构建路线、CPU 的能力边界、硬件预算,以及必须一次性配好的环境。

1.1 为什么值得在 CPU 上做这件事

训练大模型的主流叙事是"堆卡"。但把这件事拆开看,会发现大部分认知负担其实与硬件无关:为什么要有因果掩码、困惑度为什么等于交叉熵的指数、为什么要给 LayerNorm 和 bias 单独分组、指令微调该在哪里掩码、LoRA 为什么能只训百分之一的参数——这些问题的答案在 CPU 和 GPU 上完全一样。

在 CPU 上做这件事有三个实打实的好处。

第一,成本为零且随时可复现。 本书全部实验可以在四核 CPU、7 GB 内存的机器上跑完,总耗时以分钟计。没有排队、没有抢占、没有"这个实验要等下周有卡再说"。你的每一次改动都能立刻看到损失曲线怎么动。

第二,强制你把注意力放在"为什么"上。 GPU 太快了,快到你可以用一个错误的实现跑出一个看起来还不错的结果。CPU 慢,慢到每个数量级都要算清楚:为什么批大小取 8 而不是 64,为什么上下文取 128 而不是 1024,为什么梯度累积在内存受限时是救命稻草。这些算术不是负担,而是理解模型成本结构的唯一途径。

第三,训练与推理的路径都完整。 本书不只讲"怎么训",也讲 int8 动态量化、KV 缓存、TorchScript 导出、最小推理服务——这些恰好是 CPU 部署的主战场。很多团队的真实业务负载根本不在 GPU 上,而在 CPU 服务器上做小模型的批量推理。

1.2 三条构建路线与全书地图

一个大语言模型的诞生分三个阶段,本书的章节严格对应这三阶段:

阶段一:架构与数据。 先把文本变成词元,再把词元变成向量,然后实现注意力、层归一化、前馈网络、残差连接,最后拼成一个完整的 GPT。这一步产出的模型参数是随机初始化的,它只会输出乱码,但结构已经完全正确。对应第 2 章(文本数据)、第 3 章(注意力)、第 4 章(GPT 模型)。

阶段二:预训练。 在大量无标注文本上最小化"预测下一个词元"的交叉熵损失,得到一个基础模型。基础模型学会了语言的统计规律——词怎么搭配、句子怎么转折、代码长什么样——但它还不会"听话",你问它问题,它可能继续接话而不是回答。对应第 5 章(预训练)、第 6 章(生成与推理优化)、第 7 章(加载预训练权重)。

阶段三:微调。 用少量带标注的数据把基础模型对齐到具体任务上:分类微调把它变成判别器,指令微调把它变成助手,LoRA 让这件事在极小的参数预算内完成。对应第 8 章(分类微调)、第 9 章(指令微调)、第 10 章(LoRA)。

最后三章是工程闭环:第 11 章讲评估与调试,第 12 章讲部署,第 13 章用一个脚本把全流程串起来,附录 A、B 提供 PyTorch 速成与速查手册。

1.3 CPU 的真实边界

必须诚实地说清楚 CPU 能做什么、不能做什么。下面这张表是全书最重要的"预期管理",数字都是可以自己复算的估算。

任务 CPU(4 核 / 8 核) 结论
从零预训练 1M~20M 参数模型 分钟到小时级 完全可行,教学与验证首选
从零预训练 124M 参数模型(GPT-2 small) 需要数周至数月 不现实,请改用加载预训练权重再微调
微调 124M(分类头或 LoRA) 数小时 可行但慢,建议先冻结主干
微调 124M(全量) 数十小时 不推荐,除非只跑极少量步数
推理 124M(fp32,单序列) 每秒几个到几十个词元 完全可行
推理 124M(int8 量化) 比 fp32 快 1.5~3 倍 推荐做法
推理 1B 以上模型 内存放不下/极慢 建议换硬件或换策略

估算的依据很简单:训练一个词元大约需要 $6P$ 次浮点运算(前向 $2P$ 用于矩阵乘法,反向约 $4P$),其中 $P$ 是参数量。一台四核 CPU 在以 fp32 运行、充分利用 AVX2/AVX-512 时,实测有效算力大致在 10~40 GFLOPS 量级;而一块中端 GPU 的算力是这个数字的 200~500 倍。于是:

$$\text{训练时间} \approx \frac{6 \times P \times N_{\text{tokens}}}{\text{有效算力}}$$

以一个 $P=1$M 的教学模型、$N=2\times10^7$ 个词元、有效算力 20 GFLOPS 为例:

$$T \approx \frac{6 \times 10^6 \times 2\times 10^7}{20 \times 10^9} \approx 6000\ \text{秒}$$

也就是一个多小时。若把参数量放大到 124M,同样词元数就要 124 倍时间——约 8 天,而这还只是 2000 万词元,远低于 GPT-2 训练所用的量级。这就是为什么第 7 章要专门讲"加载预训练权重",也是为什么本书的教学配置选在 1M 参数量级。

1.4 硬件预算表

下表给出不同机器上"推荐做什么"。表中的内存占用按 fp32、AdamW、约 16 字节/参数的经验规则估算(参数 4 + 梯度 4 + 优化器一阶矩 4 + 二阶矩 4),激活值另计。

机器配置 可舒适训练的模型 训练内存(估算) 推荐任务
2 核 / 4 GB ≤ 0.5M 参数 约 8 MB + 激活 分词器、注意力单元测试、1 轮冒烟测试
4 核 / 8 GB 1M~3M 参数 16~50 MB + 激活 本书 TINY 配置的完整流程(推荐)
8 核 / 16 GB 3M~10M 参数 50~160 MB + 激活 SMALL 配置预训练、124M 的 LoRA 微调
16 核 / 32 GB 10M~50M 参数 0.2~0.8 GB + 激活 更长上下文与更大批大小;124M 分类微调

注意一个反直觉的事实:在 CPU 上,模型参数本身往往不是内存瓶颈,激活值和批大小才是。 一个 1M 参数的模型只要 4 MB 存 fp32 权重、16 MB 存完整训练状态,但一个 batch=32、序列长度 512、维度 128、4 层的激活张量轻松超过 100 MB。因此 CPU 调优的第一手段通常是减小批大小与序列长度,而不是换更小的模型。

需要特别提醒:CPU 训练对内存带宽比对算力更敏感。矩阵乘法在 CPU 上通常是"内存/缓存受限"而不是"运算受限",所以 int8 量化、更小的 dtype、更紧凑的批处理往往比调线程数收益更大。第 6 章会用实测数据说明这一点。

1.5 环境安装(一次配好,全书通用)

本书需要 Python 3.10 以上与 CPU 版 PyTorch。强烈建议使用虚拟环境,因为系统 Python 通常被发行版锁定(PEP 668 会直接拒绝安装)。

# 1) 在项目目录创建虚拟环境
cd /home/x01/dsh
python3 -m venv .venv

# 2) 安装 CPU 版 PyTorch(注意 index-url 指向 CPU wheel,体积小、无 CUDA 依赖)
.venv/bin/pip install --index-url https://download.pytorch.org/whl/cpu torch

# 3) 安装本书用到的轻量依赖
.venv/bin/pip install numpy

# 4) 可选:画损失曲线
.venv/bin/pip install matplotlib

安装完成后做一次自检:

import torch, os
print("PyTorch:", torch.__version__)
print("CUDA 可用:", torch.cuda.is_available())      # 纯 CPU 环境应为 False
print("逻辑核心:", os.cpu_count())
print("计算线程:", torch.get_num_threads())

自检输出里 CUDA 可用: False 是正常且预期的。书里所有代码都用同一个设备选择语句,它会在 CPU 机器上自动退化,不需要你改任何一行:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

如果你的机器有 GPU,这段代码会自动使用 GPU,但本书的所有示例都不依赖它。

1.6 CPU 性能调优总览

下面列出本书用到的全部 CPU 优化手段,详细原理与实测数据分散在后续章节,这里先给出全局清单,方便你随时回查。

线程与并行。 用 torch.set_num_threads(n) 把线程数设为物理核心数(不是逻辑核心数)。超线程带来的额外线程会因为争抢缓存而变慢。同时把 set_num_interop_threads(1),因为单进程训练不需要算子间并行,多开会额外占用核心。这两行必须在任何张量运算之前调用。

批大小与上下文长度。 计算量近似正比于 batch × context。CPU 上建议从 batch=8, context=128 起步;内存紧张时先减批大小,用梯度累积补回等效批大小(见第 5 章)。

梯度累积。 把 $K$ 个微批的梯度相加后再更新一次参数,等效批大小变成 $K$ 倍,而激活内存只需一个微批。这是 CPU 上"用小内存模拟大批量"的标准手段。

数据类型。 fp32 是稳妥默认。CPU 上 bfloat16 只有在支持 AVX512-BF16 或 AMX 指令集的较新处理器上才真正加速,否则可能更慢——务必实测而不是想当然。

int8 动态量化。 把线性层权重压成 8 位整数,模型体积降到约四分之一,内存带宽压力大幅下降,CPU 上通常有 1.5~3 倍加速。这是 CPU 推理最划算的一步。

KV 缓存。 自回归生成时缓存历史的键和值,把每步计算量从平方级降到线性级,实测常见 2~10 倍加速(第 6 章给出本项目的实测数值)。

推理模式。 用 torch.inference_mode() 代替 torch.no_grad(),再省几个百分点的开销。

数据加载。 CPU 训练中把 num_workers 保持为 0。数据只是取切片,极轻;开多进程反而会因为复制内存、IPC 通信、以及与计算线程争抢核心而变慢。

缓存中间结果。 分词是最慢的纯 Python 环节,把编码结果缓存成 .npy,重复实验时直接从磁盘加载,能省下几十秒到几分钟。

1.7 本书的代码组织与运行方式

全部代码放在 code/ 目录,每个文件都能单独运行,且都带 if __name__ == "__main__": 演示块:

文件 作用 单独运行命令
code/config.py 配置、设备、参数量与内存估算 python code/config.py
code/corpus.py 离线生成语料、分类数据、指令数据 python code/corpus.py
code/tokenizer.py 教学分词器与字节级 BPE python code/tokenizer.py
code/dataset.py 滑动窗口数据集与 DataLoader python code/dataset.py
code/attention.py 从朴素自注意力到多头 + KV Cache python code/attention.py
code/model.py 完整 GPT 模型与采样生成 python code/model.py
code/train.py 预训练循环(warmup/余弦/裁剪/累积) python code/train.py
code/finetune_classify.py 分类微调 python code/finetune_classify.py
code/finetune_instruct.py 指令微调 python code/finetune_instruct.py
code/lora.py LoRA 从零实现与合并 python code/lora.py
code/quantize.py int8 动态量化与基准测试 python code/quantize.py
code/deploy.py TorchScript 导出与最小服务 python code/deploy.py
code/run_all.py 端到端冒烟测试与实测数据采集 python code/run_all.py --quick
code/bench.py 专项基准测试(KV 缓存/量化/线程) python code/bench.py
code/verify_doc.py 本文档字数与代码校验 python code/verify_doc.py --check-code
code/build_doc.py 组装本文档(把 code/*.py 嵌入正文) python code/build_doc.py

运行方式有一个小细节值得说明:用 python code/xxx.py 这种方式运行脚本时,Python 会自动把脚本所在目录(也就是 code/)加到 sys.path[0],因此文件之间可以直接写 from model import GPTModel 这样简洁的导入,而不需要把项目打包成模块。这也是为什么本书坚持"从项目根目录运行 python code/xxx.py",而不是先 cd code。如果你在 code/ 目录里执行 python model.py,效果是一样的;但如果你把脚本复制到别处运行,导入就会失败。

第一条命令建议这样跑:

python code/config.py            # 看配置、参数量分解、内存估算
python code/corpus.py            # 生成语料与数据集(完全离线)
python code/run_all.py --quick   # 端到端跑一遍(空闲 4 核机器约 3 分钟)
python code/bench.py             # 专项基准:KV 缓存 / 量化 / 线程扩展性

关于规模的预期管理:本书 TINY 配置的预训练一次只要 1~2 分钟,
但它的生成质量并不会很好(1M 参数的诚实上限)。理解了"成本结构"与"诊断方法",
换到更大的机器或加载预训练权重时,你可以直接复用全书的所有代码与流程。

1.8 阅读路线图

完全新手:按顺序读第 2 章 → 第 3 章 → 第 4 章,边读边跑对应脚本,把形状搞熟;再读第 5 章把训练跑起来;最后挑第 8 章或第 9 章做一次微调。

有 Transformer 基础、只想跑通:直接从第 4 章开始看实现,跳到第 5 章跑预训练,然后第 6 章做推理优化。

想解决具体业务问题:分类任务看第 8 章 + 第 12 章;问答/助手看第 9 章 + 第 6 章;参数预算紧看第 10 章;延迟敏感看第 6 章 + 第 12 章。

只想快速体检代码是否正常:跑 python code/run_all.py --quick,它会把 13 个阶段全部走一遍并输出真实耗时表。

1.9 常见错误与排查

错误一:error: externally-managed-environment。 这是 Debian/Ubuntu 的 PEP 668 保护,不是 PyTorch 的问题。解决方式是用虚拟环境(本章 1.5 节的做法)。实在要用系统环境,可加 --break-system-packages,但那会污染系统包,不推荐。

错误二:装成了 CUDA 版 torch。 症状是安装体积好几 GB、torch.cuda.is_available() 为 False 但仍加载了 CUDA 库。原因是没有指定 --index-url https://download.pytorch.org/whl/cpu。检查方式:python -c "import torch; print(torch.__version__)",CPU 版通常显示 2.x.x+cpu。

错误三:ModuleNotFoundError: No module named 'model'。 说明你没有从项目根目录运行脚本,或者把代码文件挪了位置。请确保当前目录是 dsh/,命令是 python code/model.py。

错误四:装了 numpy 却仍报 No module named 'numpy'。 通常是把包装进了系统 Python,而运行时用的是虚拟环境的 Python。解决办法是统一用 .venv/bin/pip install 和 .venv/bin/python(或先 source .venv/bin/activate)。

错误五:跑完发现比预期慢十倍。 先检查三件事:是否忘了设置线程数、批大小是不是被设得过大导致内存换页(用 free -g 看 swap 是否在涨)、以及是否在循环里反复做分词等 Python 层操作。第 11 章的性能剖析方法可以精确定位。

1.10 本章小结

  • 大语言模型的构建分三阶段:架构与数据 → 预训练 → 微调;CPU 环境完全可以走完这三阶段,只是规模需要控制在百万参数级。
  • 训练成本可用 $6PN$ 估算,内存可用约 16 字节/参数估算;这两个公式是全书做取舍的基础。
  • CPU 上的主要瓶颈是内存带宽而不是算力,因此量化、KV 缓存、批大小控制的收益往往高于调线程。
  • 环境必须用虚拟环境安装 CPU 版 PyTorch;设备选择统一写成 torch.device("cuda" if ... else "cpu"),从而在任何机器上都能跑。
  • 全部代码都在 code/ 下,且每个文件都能单独运行,从项目根目录执行 python code/xxx.py。

练习

  1. 用 1.3 节的 $6PN$ 公式估算:在有效算力 25 GFLOPS 的机器上,训练一个 4M 参数模型、消耗 5 亿词元需要多少小时?(思路:代入公式得到 $6\times4\times106\times5\times108 / (25\times10^9)$ 秒,再换算成小时。)
  2. 运行 python code/config.py,找到 124M 配置的参数量分解,回答:嵌入层(词元嵌入+位置嵌入)占总参数的比例是多少?如果把上下文长度从 1024 提到 4096,参数量增加多少?(思路:位置嵌入项为 $L\times d$,注意输出头与词元嵌入是否共享权重会影响结论。)
  3. 你的机器有多少物理核心?分别用 torch.set_num_threads(1)、物理核心数、逻辑核心数 各跑一次 python code/attention.py,记录耗时,验证"线程数超过物理核心反而变慢"的判断。
  4. 假设内存只有 4 GB,需要训练一个 8M 参数模型,按 16 字节/参数估算训练状态需要多少内存?还能留下多少给激活值?你会优先调整哪些超参数?(思路:先算优化器与梯度占用,再从批大小和上下文长度上省。)

第 1 章 理解大语言模型

本章导读

本章不写模型代码,而是把"大语言模型到底是什么"这件事讲透:它的唯一训练目标、它的架构选择、它的参数都花在哪里、以及为什么这些数字决定了你在 CPU 上能做什么。

读完本章你应该能回答四个问题:为什么 GPT 只用了 Transformer 的解码器?为什么"预测下一个词"就足以学会语言?一个 124M 模型的 1.24 亿参数具体分布在哪些张量上?为什么同样的模型在你的笔记本上跑得动推理却训不动?

1.1 大语言模型的定义:一个被极度简化的目标

大语言模型(Large Language Model,LLM)本质上是一个参数化的概率分布 $p_\theta(x_t \mid x_1, \dots, x_{t-1})$:给定前面的词元序列,输出下一个词元上的概率分布。参数量 $\theta$ 从几千万到几千亿,结构大多是"仅解码器的 Transformer"。

整个训练目标只有一条:

$$\mathcal{L}(\theta) = -\frac{1}{N}\sum_{t=1}^{N} \log p_\theta(x_t \mid x_{<t})$$

这就是自回归语言建模的交叉熵损失,等价于最大化训练语料的似然。它简单到令人不安:没有人工标注、没有任务定义、没有特征工程,只有"预测下一个词"。但恰恰是这个目标迫使模型学会语法(否则预测不了下一个功能词)、事实(否则预测不了专有名词)、推理(否则预测不了"因此"后面的结论)。

有三个概念必须在这里区分清楚:

  • 词元(token):文本被切分后的最小单位,可以是字符、子词或字节。模型真正处理的对象是词元,不是字符也不是词。
  • 上下文长度(context length):模型一次能看到的最大词元数,记为 $L$。GPT-2 small 是 1024。
  • 参数(parameter):模型里所有可训练的张量元素,记为 $P$。

1.2 大语言模型能做什么

同一套参数,通过不同的"提示"或不同的"微调方式",可以完成差异极大的任务:

任务类型 输入 → 输出 典型做法
文本续写 提示 → 续写 直接生成
翻译 英文 → 中文 指令微调或提示
摘要 长文 → 短句 指令微调
问答 问题 → 答案 指令微调
代码生成 注释 → 代码 代码语料预训练
文本分类 文本 → 类别 分类微调(第 8 章)
检索/嵌入 文本 → 向量 取隐藏状态(本书不展开)

这种"一个模型、多种任务"的能力不是设计出来的,而是预训练目标的副作用:为了让下一个词元预测得足够准,模型必须隐式地学会做这些子任务。

1.3 构建 LLM 的三个阶段

阶段一:实现架构与准备数据。 分词器、嵌入层、注意力、前馈、残差、归一化,加上把文本切成训练样本的数据管道。产出的模型参数随机,输出乱码,但结构已经与 GPT-2 完全一致。这一步是理解全书的基础。

阶段二:预训练。 在海量无标注文本上最小化交叉熵。产出"基础模型"(base model)。它会续写、会补全代码,但不会遵循指令——你问"英国首都是哪里?",它可能接着编出"法国首都是哪里?"这样的下一个问题,因为在训练语料里,问题后面经常跟着另一个问题。

阶段三:微调。 用少量标注数据做对齐。两条主要路线:

  • 分类微调:把输出头换成"隐藏状态 → 类别数",用于判别任务;
  • 指令微调:保持生成目标,但数据换成"指令 + 输入 + 回答",让模型学会"回答"而不是"续写"。

本书在第 10 章还会讲第三条路:参数高效微调(LoRA),用 1% 的参数达到接近全量微调的效果。

1.4 为什么是"仅解码器"的 Transformer

2017 年的原始 Transformer 是编码器-解码器结构,用于机器翻译:编码器读入源语言(双向可见),解码器生成目标语言(只能看已生成部分)。此后架构出现两条分岔:

架构 代表 注意力 适合任务
仅编码器 BERT 双向 理解类:分类、抽取、匹配
编码器-解码器 T5、原始 Transformer 编码双向 + 解码因果 序列到序列:翻译、摘要
仅解码器 GPT 系列、LLaMA 因果(单向) 生成、通用任务、对话

GPT 选择仅解码器有三个理由。

第一,训练目标统一且无需标注。 因果语言建模只需要原始文本,任何网页、书籍、代码都是训练数据;而编码器常用的掩码语言建模(MLM)需要人工构造掩码,且无法直接用于生成。

第二,架构更简单,同样参数下更"通用"。 没有编码器-解码器之间的交叉注意力,参数全部用于单一堆叠;一旦规模足够大,它在理解和生成任务上都能达到甚至超过专用架构。

第三,推理时天然适合缓存。 因果注意力保证位置 $t$ 只依赖 $\le t$ 的信息,因此历史的键值可以缓存复用(第 6 章的 KV Cache),而双向注意力每一步都要重算整个序列。

代价是:仅解码器模型每生成一个词元都要前向一次,生成 $T$ 个词元的成本是 $T$ 次前向;这在 CPU 上尤其昂贵,也是第 6 章要重点优化的地方。

1.5 GPT 架构逐部件解剖

从输入到输出,数据流是这样的(设批大小 $B$、序列长度 $T$、嵌入维度 $d$、词表大小 $V$、层数 $n$):

步骤 操作 输出形状 参数量
1 输入词元 ID $(B, T)$ 0
2 词元嵌入查表 $(B, T, d)$ $V \cdot d$
3 位置嵌入查表 $(T, d)$ → 广播相加 $L \cdot d$
4 Dropout $(B, T, d)$ 0
5 $n$ × Transformer 块 $(B, T, d)$ $n \cdot P_{\text{layer}}$
6 末尾 LayerNorm $(B, T, d)$ $2d$
7 输出头线性层 $(B, T, V)$ $d \cdot V$

每个 Transformer 块的内部:

子步骤 操作 形状 参数量
5.1 LayerNorm $(B,T,d)$ $2d$
5.2 Q/K/V 投影 $(B,T,d)\times3$ $3d^2$($+3d$ 若带偏置)
5.3 拆头 → 打分 → softmax → 加权 $(B,h,T,d_h)$ 0
5.4 输出投影 $(B,T,d)$ $d^2 + d$
5.5 残差相加 $(B,T,d)$ 0
5.6 LayerNorm $(B,T,d)$ $2d$
5.7 前馈升维 $(B,T,4d)$ $4d^2 + 4d$
5.8 GELU $(B,T,4d)$ 0
5.9 前馈降维 $(B,T,d)$ $4d^2 + d$
5.10 残差相加 $(B,T,d)$ 0

其中 $h$ 是注意力头数,$d_h = d/h$ 是每个头的维度。这个表是全书最重要的一张表——第 4 章的代码就是把每一行翻译成 PyTorch。

1.6 参数量核算:124M 到底花在哪里

用上面的公式,我们逐项计算 GPT-2 small($V=50257$,$L=1024$,$d=768$,$n=12$,带 QKV 偏置):

项目 公式 数值
词元嵌入 $V \cdot d$ $50257 \times 768 = 38{,}597{,}376$
位置嵌入 $L \cdot d$ $1024 \times 768 = 786{,}432$
单层 QKV 投影 $3d^2 + 3d$ $3\times768^2 + 3\times768 = 1{,}771{,}776$
单层输出投影 $d^2 + d$ $590{,}592$
单层前馈 $8d^2 + 5d$ $4{,}723{,}200$
单层两个 LayerNorm $4d$ $3{,}072$
单层合计 — $7{,}088{,}640$
12 层合计 $n \cdot P_{\text{layer}}$ $85{,}063{,}680$
末尾 LayerNorm $2d$ $1{,}536$
输出头(不共享权重) $d \cdot V$ $38{,}597{,}376$
总计 — $\approx 163{,}046{,}400$

注意这里的总数是 163M,而不是常说的 124M。差在哪里?GPT-2 让输出头与词元嵌入共享同一张权重矩阵(weight tying),因此要减去一份 $V\cdot d$:

$$163{,}046{,}400 - 38{,}597{,}376 = 124{,}449{,}024 \approx 124\text{M}$$

这就是"124M"的来历。本书的 config.count_parameters_breakdown() 把这两条路都实现了:tie_weights=False 时给出 163M,True 时给出 124M。这个细节值得记住,因为它是"为什么我的参数量和论文对不上"的最常见原因。

再看参数分布的直觉:前馈层占了单层的约 67%,注意力占约 33%,而嵌入层单独就占了总参数的近四分之一(共享时)。这意味着两件事:

  1. 想缩小模型时,减小 $d$ 会同时缩小注意力和前馈(平方关系),是最有效的手段;
  2. 词表大小 $V$ 对参数量的影响是线性的,把词表从 50257 降到 2000(本书的做法)能省下大量参数,这也是 CPU 教学模型的关键取舍。

下面是本书 config.py 中实现这些估算的代码,包括参数量分解、KV 缓存与训练内存估算。这些公式在后续每一章都会用到。

"""
config.py —— 全局配置、设备选择、CPU 性能与参数量统计工具
=========================================================

本文件是整个项目的“公共基础设施”,它回答四个问题:

1. 模型有多大?(配置字典 + 参数量统计)
2. 用哪块设备跑?(永远能在纯 CPU 机器上安全退化)
3. CPU 上怎么用满核心?(线程设置)
4. 结果能不能复现?(随机种子)

所有其它模块都从这里导入配置,保证全文超参数一致。
"""

from __future__ import annotations

import os
import random
import time
from typing import Any, Dict

import numpy as np
import torch

# -------------------------------------------------------------------
# 1. 模型配置字典
# -------------------------------------------------------------------
# 说明:配置字典只描述“形状”,不描述数据。字段含义:
#   vocab_size     —— 词表大小(由分词器决定)
#   context_length —— 上下文窗口长度,即模型一次最多看多少个词元
#   emb_dim        —— 嵌入维度 d_model
#   n_heads        —— 注意力头数(必须能整除 emb_dim)
#   n_layers       —— Transformer 块的堆叠层数
#   drop_rate      —— Dropout 概率
#   qkv_bias       —— Q/K/V 线性层是否使用偏置(GPT-2 用 True)
# ------------------------------------------------------------------

# 教学用“微型”配置:4 核 CPU 上几分钟就能完成一次完整预训练。
GPT_CONFIG_TINY: Dict[str, Any] = {
    "vocab_size": 1000,
    "context_length": 128,
    "emb_dim": 128,
    "n_heads": 4,
    "n_layers": 4,
    "drop_rate": 0.1,
    "qkv_bias": False,
}

# 稍大一点:仍然适合 CPU,但已经能看出“规模带来的收益”。
GPT_CONFIG_SMALL: Dict[str, Any] = {
    "vocab_size": 2000,
    "context_length": 256,
    "emb_dim": 256,
    "n_heads": 8,
    "n_layers": 6,
    "drop_rate": 0.1,
    "qkv_bias": True,
}

# GPT-2 small(124M)官方配置:CPU 上可推理、可微调(很慢),不建议从零预训练。
GPT_CONFIG_124M: Dict[str, Any] = {
    "vocab_size": 50257,
    "context_length": 1024,
    "emb_dim": 768,
    "n_heads": 12,
    "n_layers": 12,
    "drop_rate": 0.1,
    "qkv_bias": True,
}


def config_from_tokenizer(tokenizer, base: Dict[str, Any] | None = None) -> Dict[str, Any]:
    """把分词器的真实词表大小写进配置,避免“词表对不上”的经典错误。

    分词器训练出来的 ``vocab_size`` 取决于语料和合并次数,写死配置很容易越界。
    这里统一用分词器的真实值覆盖 ``vocab_size``。
    """
    cfg = dict(base or GPT_CONFIG_TINY)
    cfg["vocab_size"] = int(tokenizer.vocab_size)
    return cfg


# ----------------------------------------------------------------
# 2. 设备选择与线程设置
# ----------------------------------------------------------------
def pick_device() -> torch.device:
    """返回可用设备:有 GPU 用 GPU,否则退回 CPU。

    在纯 CPU 机器上 ``torch.cuda.is_available()`` 恒为 ``False``,
    因此这一行永远解析为 ``cpu``;全文其余代码不需要任何改动。
    """
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def set_cpu_threads(n: int | None = None) -> int:
    """设置 PyTorch 在 CPU 上使用的计算线程数,返回实际生效值。

    - ``n=None``:使用 ``os.cpu_count()``(通常等于逻辑核心数)。
    - 经验法则:设为**物理核心数**左右最优;超过物理核心数会因超线程与
      调度开销导致轻微变慢。
    - ``set_num_interop_threads`` 必须在并行工作开始前调用,否则会抛
      ``RuntimeError``,因此这里用 try/except 包住。
    """
    if n is None:
        n = os.cpu_count() or 1
    n = max(1, int(n))
    torch.set_num_threads(n)
    try:
        torch.set_num_interop_threads(1)  # 单机推理/训练一般不需要算子间并行
    except RuntimeError:
        # 已经启动过并行工作,忽略即可。
        pass
    return torch.get_num_threads()


def set_seed(seed: int = 123) -> None:
    """固定随机种子:让“同一次实验”在 CPU 上可复现。"""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def describe_device() -> str:
    """返回一行人类可读的环境描述,方便写进实验日志。"""
    threads = torch.get_num_threads()
    interop = torch.get_num_interop_threads()
    gpu = "不可用(纯 CPU 环境)" if not torch.cuda.is_available() else "可用"
    return (
        f"PyTorch {torch.__version__} | CUDA {gpu} | "
        f"计算线程 {threads} | 算子间线程 {interop} | "
        f"逻辑核心 {os.cpu_count()}"
    )


class Timer:
    """极简计时上下文管理器:``with Timer("训练") as t: ...; print(t.elapsed)``"""

    def __init__(self, name: str = "", verbose: bool = True) -> None:
        self.name = name
        self.verbose = verbose
        self.elapsed = 0.0

    def __enter__(self) -> "Timer":
        self._t0 = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc, tb) -> bool:
        self.elapsed = time.perf_counter() - self._t0
        if self.verbose:
            print(f"[计时] {self.name}: {self.elapsed:.3f} 秒")
        return False  # 不吞异常


# --------------------------------------------------------------
# 3. 参数量统计(实测 + 公式推导)
# --------------------------------------------------------------
def count_parameters(model: torch.nn.Module, trainable_only: bool = False) -> int:
    """统计模型参数总量(实测值,用于和公式推导互相验证)。"""
    params = model.parameters()
    if trainable_only:
        return sum(p.numel() for p in params if p.requires_grad)
    return sum(p.numel() for p in params)


def count_parameters_breakdown(cfg: Dict[str, Any]) -> Dict[str, int]:
    """按“部件”推导参数量,用于解释“124M 到底花在哪里”。

    推导依据本书 ``GPTModel`` 的结构(与 GPT-2 一致):

    * 词元嵌入:      V * d
    * 位置嵌入:      L * d
    * 每层注意力:    qkv 投影  d * 3d  (+3d 偏置, 若开启)
                    输出投影  d * d   (+d  偏置)
    * 每层前馈:      d * 4d  (+4d) 与 4d * d (+d)
    * 每层两个 LayerNorm: 2 * 2d
    * 末尾 LayerNorm: 2d
    * 输出头:        d * V (bias=False)

    其中 V=vocab_size, L=context_length, d=emb_dim。

    若 ``cfg["tie_weights"]`` 为 ``True``,输出头与词元嵌入共享同一张矩阵
    (GPT-2 的做法),实际参数会少 ``V*d`` 个;默认关闭,便于逐项核对。
    """
    v = int(cfg["vocab_size"])
    l = int(cfg["context_length"])
    d = int(cfg["emb_dim"])
    n_layers = int(cfg["n_layers"])
    bias = 3 * d if cfg.get("qkv_bias", False) else 0

    tok_emb = v * d
    pos_emb = l * d
    attn_qkv = d * 3 * d + bias
    attn_out = d * d + d
    ff = (d * 4 * d + 4 * d) + (4 * d * d + d)
    ln_per_layer = 2 * (2 * d)
    per_layer = attn_qkv + attn_out + ff + ln_per_layer
    final_norm = 2 * d
    tied = bool(cfg.get("tie_weights", False))
    out_head = 0 if tied else d * v  # 权重共享时不再新增参数
    total = tok_emb + pos_emb + n_layers * per_layer + final_norm + out_head
    return {
        "词元嵌入": tok_emb,
        "位置嵌入": pos_emb,
        "单层注意力": attn_qkv + attn_out,
        "单层前馈": ff,
        "单层LayerNorm": ln_per_layer,
        "单层合计": per_layer,
        "全部Transformer层": n_layers * per_layer,
        "末尾LayerNorm": final_norm,
        "输出头": out_head,
        "总计": total,
    }


def human_count(n: float) -> str:
    """把参数量格式化成 124.4M 这种可读形式。"""
    for unit in ["", "K", "M", "B", "T"]:
        if abs(n) < 1000:
            return f"{n:.2f}{unit}" if unit else f"{int(n)}"
        n /= 1000.0
    return f"{n:.2f}P"


def human_bytes(n: float) -> str:
    """字节数格式化。"""
    for unit in ["B", "KB", "MB", "GB", "TB"]:
        if abs(n) < 1024:
            return f"{n:.2f} {unit}"
        n /= 1024.0
    return f"{n:.2f} PB"


def estimate_model_bytes(n_params: int, dtype_bytes: int = 4) -> int:
    """参数占用的裸内存:参数量 × 每参数字节数。

    fp32 -> 4 字节,bf16/fp16 -> 2 字节,int8 -> 1 字节。
    """
    return int(n_params) * int(dtype_bytes)


def estimate_kv_cache_bytes(
    n_layers: int,
    n_heads: int,
    emb_dim: int,
    context_length: int,
    batch_size: int = 1,
    dtype_bytes: int = 2,
) -> int:
    """KV Cache 内存估算公式。

    $$ \\text{bytes} = 2 \\times n_{layers} \\times B \\times n_{heads}
       \\times T \\times d_{head} \\times \\text{sizeof(dtype)} $$

    其中系数 2 表示 Key 与 Value 各一份,$d_{head} = d_{model} / n_{heads}$。
    KV Cache 是自回归推理的核心加速手段(见第 6 章),但它的内存随
    上下文长度**线性增长**,是长上下文推理的主要瓶颈。
    """
    d_head = emb_dim // n_heads
    return int(2 * n_layers * batch_size * n_heads * context_length * d_head * dtype_bytes)


def training_memory_estimate(cfg: Dict[str, Any], batch_size: int, dtype_bytes: int = 4) -> Dict[str, int]:
    """训练期内存粗估(CPU 同样吃内存,因此也需要估算)。

    规则(AdamW + fp32):
      * 参数                 : P
      * 梯度                 : P
      * AdamW 一阶/二阶动量  : 2P
      * 激活(近似)         : B * L * d * n_layers * K,经验系数 K≈10
    合计约 **16 字节/参数**,这与社区常说的“全量微调需要约 16 倍参数量内存”
    一致(混合精度可降到约 10 字节/参数)。
    """
    p = count_parameters_breakdown(cfg)["总计"]
    acts = int(batch_size * cfg["context_length"] * cfg["emb_dim"] * cfg["n_layers"] * 10 * dtype_bytes)
    return {
        "参数": estimate_model_bytes(p, dtype_bytes),
        "梯度": estimate_model_bytes(p, dtype_bytes),
        "优化器状态": estimate_model_bytes(2 * p, dtype_bytes),
        "激活(近似)": acts,
        "合计": estimate_model_bytes(4 * p, dtype_bytes) + acts,
    }


if __name__ == "__main__":  # pragma: no cover - 手动演示
    set_seed(123)
    print(describe_device())
    for name, cfg in [
        ("TINY", GPT_CONFIG_TINY),
        ("SMALL", GPT_CONFIG_SMALL),
        ("124M", GPT_CONFIG_124M),
    ]:
        bd = count_parameters_breakdown(cfg)
        print(f"\n=== {name} ===")
        for k, val in bd.items():
            print(f"  {k:<18}{val:>14,}  ({human_count(val)})")
        print("  训练内存粗估:", {k: human_bytes(v) for k, v in training_memory_estimate(cfg, 8).items()})
        print(
            "  KV Cache(bs=1, fp16):",
            human_bytes(
                estimate_kv_cache_bytes(
                    cfg["n_layers"], cfg["n_heads"], cfg["emb_dim"], cfg["context_length"]
                )
            ),
        )

1.7 内存换算:为什么 7 GB 内存只能训百万参数

推理内存只需要三部分:参数 + 激活(随批大小和序列长度线性增长)+ KV 缓存。参数字节数就是:

$$\text{参数内存} = P \times \text{每参数字节数}$$

fp32 是 4 字节,bf16/fp16 是 2 字节,int8 是 1 字节。124M 模型 fp32 需要约 500 MB,int8 只要约 125 MB——这就是量化在 CPU 上如此有吸引力的原因。

训练内存则要大得多,因为除了参数还要保存梯度和优化器状态。AdamW 的账目是:

项目 字节/参数 说明
参数 4 fp32 权重
梯度 4 与参数同形
一阶动量 $m$ 4 AdamW 状态
二阶动量 $v$ 4 AdamW 状态
合计 16 这就是"训练约为推理 4 倍"的来源

于是 124M 模型全量训练的静态内存约 $124\times10^6 \times 16 \approx 2$ GB,再加上激活值(与 batch × context × d × n 成正比),在 7 GB 内存的机器上会非常紧张。而那些 1M 参数的教学模型,静态部分只要 16 MB,激活值也不过几十 MB——这才是 CPU 训练的舒适区。

KV 缓存是推理独有的开销,公式为:

$$\text{KV bytes} = 2 \times n \times B \times h \times T \times d_h \times \text{sizeof(dtype)}$$

系数 2 表示 Key 和 Value 各一份。以 GPT-2 small($n=12, h=12, d_h=64$)在 $T=1024$、$B=1$、fp16 下计算:

$$2 \times 12 \times 1 \times 12 \times 1024 \times 64 \times 2 \approx 37.7\ \text{MB}$$

看起来不大,但它随上下文长度和批大小线性增长,当 $B=32$ 时就是 1.2 GB。这正是长上下文推理在 CPU 上最容易撞墙的地方。

1.8 CPU 可行性分析:算一笔实数

现在把前面所有数字串起来,回答"我的机器能训多大的模型"。

训练一个词元的浮点运算量约为 $6P$(前向 $2P$ + 反向 $4P$)。假设四核 CPU 的有效算力为 $F$ GFLOPS,则:

$$\text{每秒可处理词元数} \approx \frac{F \times 10^9}{6P}$$

对 $P=1$M、$F=20$:

$$\frac{20\times109}{6\times106} \approx 3300\ \text{词元/秒}$$

再乘上实际利用率(CPU 上的矩阵乘法很难达到峰值,经验折扣 30%~50%),得到约 1000~1600 词元/秒。本书教学配置下实测的数量级与此吻合(具体数字见第 13 章与 _build/measurements.md)。

由此推出两条实用结论:

  1. 想在一小时内看到收敛,训练词元预算应控制在千万量级以内,模型控制在百万参数量级;
  2. 批大小对总耗时影响很小,对内存影响很大。因为计算量正比于"批大小 × 序列长度",总词元数是固定的,批大小只决定并行度;但激活内存正比于批大小。因此在 CPU 上正确的策略是"小批 + 梯度累积",而不是"大批硬扛"。

1.9 常见错误

错误一:把"124M"当成参数量的精确值。 124M 是权重共享后的结果,若你的实现没有共享输出头与词元嵌入,算出来会是 163M。两者都对,只是设定不同。本书的配置里通过 tie_weights 显式控制。

错误二:混淆字符数、词数与词元数。 中文一个字通常是 1~3 个字节、0.5~1 个 BPE 词元;英文一个单词常被拆成 1~2 个词元。所有成本估算都必须基于词元数,用字符数估会差好几倍。

错误三:认为"上下文长度翻倍,成本只是翻倍"。 注意力打分的复杂度是 $O(T^2 d)$,上下文翻倍时注意力部分的计算量变成四倍。只有在很小的 $T$ 下,整体才近似线性——这正是 CPU 上把上下文限制在 128~256 的现实理由。

错误四:直接把训练内存估成参数量。 训练需要约 16 字节/参数,是 fp32 参数内存的 4 倍。用 4 字节估算会严重低估,然后在训练中途被 OOM 或疯狂换页打脸。

1.10 本章小结

  • LLM 的唯一训练目标是自回归交叉熵:给定前文预测下一个词元。所有能力都是这个目标的副产品。
  • GPT 属于仅解码器 Transformer,因果注意力使其天然支持 KV 缓存,但生成必须逐词元串行。
  • 参数量可以逐项精确核算:124M = 嵌入 38.6M + 位置 0.8M + 12 层 85.1M + 归一化 0.002M,输出头与词元嵌入共享。
  • 训练内存约 16 字节/参数,推理约 4 字节/参数(fp32),int8 可降到 1 字节/参数。
  • CPU 上的可行规模是百万参数量级;成本用 $6PN$ 与有效算力估算,实测折扣取 30%~50%。

练习

  1. 用 1.6 节的公式计算本书 TINY 配置($V=800$,$L=128$,$d=128$,$n=4$,无 QKV 偏置、不共享权重)的总参数量,并与 python code/config.py 的输出对照。(思路:逐项代入表格公式,注意 LayerNorm 每层两个。)
  2. 若把 TINY 的 $d$ 从 128 提到 256、层数不变,参数量大约变成几倍?计算量变成几倍?(思路:除嵌入层外的项都是 $d^2$ 量级,因此约 4 倍;注意力打分部分随 $d$ 线性。)
  3. 用 KV 缓存公式算:$n=4$、$h=4$、$d=128$、$T=512$、$B=16$、fp32 时的缓存大小是多少?如果改成 int8 缓存能省多少?(思路:代入公式,注意 $d_h = d/h$。)
  4. 假设你要在 8 GB 内存的机器上做全量微调,按 16 字节/参数估算最多能训多大模型?留 50% 内存给激活值后又是多少?(思路:先算静态内存上限,再对半。)

第 2 章 处理文本数据

本章导读

神经网络只会做浮点运算,它看不懂"猫"这个字,只认识数字。因此构建大语言模型的第一道工序是把文本变成整数序列,再把整数序列变成向量序列。这一步看起来琐碎,却直接决定了三件事:词表有多大、序列有多长、模型能处理哪些语言。

本章从零实现三种分词器(教学版、特殊词元版、字节级 BPE),解释为什么最终必须用字节级 BPE 才能优雅地处理中文;然后构造滑动窗口数据集,把长文本切成"输入-目标"对;最后把词元变成可训练的嵌入向量,并注入位置信息。

读完本章你应该能回答:为什么字节级 BPE 不会出现"未知词"?为什么词表大小是 256 加合并次数?为什么位置嵌入是必需的而不是可选的?

2.1 词嵌入:从离散符号到连续向量

假设词表里有 6 个词:["猫", "狗", "汽车", "香蕉", "跑", "吃"]。最朴素的编码是 one-hot:[1,0,0,0,0,0]、[0,1,0,0,0,0]……它有两个致命问题。

维度灾难。 词表有 5 万词时,每个词就要用 5 万维向量表示,而其中只有一位是 1。存储与计算都浪费了 99.998%。

语义信息为零。 任意两个不同词的 one-hot 向量内积都是 0,余弦相似度也是 0。"猫"和"狗"的距离与"猫"和"汽车"完全一样,模型无法从编码本身知道它们相关。

词嵌入(word embedding) 解决这两个问题:把每个词映射成一个低维稠密向量(比如 256 维),让语义相近的词在向量空间中彼此靠近。关键点在于——这些向量不是人设计的,而是训练出来的。模型在"预测下一个词"的过程中,会发现"猫"和"狗"经常出现在相似的上下文里,于是自动把它们推向相近的向量位置。

在 PyTorch 里,词嵌入就是一张形状为 $(V, d)$ 的矩阵 $E$,第 $i$ 行就是第 $i$ 个词的嵌入向量。查表操作等价于 one-hot 向量与 $E$ 相乘(但实现上直接索引,快得多):

# 查表:等价于 one_hot(i) @ E,但只花一次内存读取
embedding = E[token_id]     # 形状 (d,)

nn.Embedding(vocab_size, emb_dim) 干的就是这件事,它是一个可训练层,梯度会回传到对应的行上:只有被用到的词,其嵌入向量才会更新。

2.2 分词:三种粒度的取舍

分词(tokenization)是把字符串切成词元序列的过程。粒度有三种选择:

粒度 词表大小 序列长度 未知词 中文适配 语义单位
字符级 小(几千) 长 无 天然适配 太细
词级 巨大(几十万) 短 严重 需要分词工具 刚好
子词级(BPE 等) 可控(1 万~10 万) 中 无(字节级) 好 折中

词级分词的麻烦在于中文:中文没有空格,"南京市长江大桥"到底分成"南京市/长江大桥"还是"南京/市长/江大桥",需要额外的分词器,而分词器的错误会直接变成模型的错误。此外新词(网络用语、专有名词、代码标识符)永远在出现,词级词表会不断遇到未知词。

字节级 BPE(byte-level BPE) 是目前主流大模型的事实标准(GPT-2、GPT-3、LLaMA 都用它)。它有三条性质值得记住:

  1. 永远不会遇到未知词:任何字符都能表示成 UTF-8 字节,字节一共 256 个,全都在词表里;
  2. 词表大小可控:由合并次数决定,想多大就多大;
  3. 高频片段被合并:" the"、"的"、"模型" 这类高频序列会被合并成一个词元,从而缩短序列、降低计算量。

2.3 教学版分词器:看清"分词 → ID → 还原"的闭环

我们先用一个能完全看懂的简单分词器把闭环跑通。它的思路是:用正则按标点和空格切分,把所有片段收集起来构造词表,然后编码就是查表、解码就是反查表。

它有两个必须正视的缺陷,正好也是下面字节级 BPE 要解决的问题:

  • 词表外的词直接抛异常:原始实现遇到没见过的词会 KeyError;
  • 没有特殊词元:模型需要知道"文本在哪里结束"、"哪里是填充",这些都必须用专门的词元表示。

引入 <|unk|>(未知词)和 <|endoftext|>(文本结束)之后,程序不再崩溃,但代价是:所有未知词被折叠成同一个符号,模型丢失了信息。这正是词级方案的死结。

code/tokenizer.py 里的 SimpleTokenizerV1 与 SimpleTokenizerV2 就是这两个阶段,值得逐行读一遍——它们不到 40 行,却是理解分词的最短路径。

2.4 字节级 BPE:原理与实现

2.4.1 训练过程

字节级 BPE 的训练是一个反复"找最高频相邻对并合并"的循环:

  1. 把训练文本编码成 UTF-8 字节序列,此时每个词元都是 0~255 之间的整数;
  2. 统计所有相邻词元对的出现频次;
  3. 找出频次最高的一对 $(a, b)$,把它合并成一个新词元,新词元的 ID 从 256 开始递增;
  4. 回到第 2 步,直到达到目标词表大小。

举例:假设文本是"大语言模型大语言模型"。它的 UTF-8 字节序列里,"大"占 3 个字节,因此"大"+"语"的相邻对会反复出现。假如"大语言"这个词组出现很多次,BPE 就会依次合并出"大语"和"大语言"这样的词元。最终词表 = 256 个字节 + 每次合并产生的新词元 + 特殊词元:

$$\text{vocab_size} = 256 + n_{\text{merges}} + n_{\text{special}}$$

2.4.2 编码过程

编码时,把每个片段转成字节列表,然后按合并产生的先后顺序依次应用每条规则。先产生的合并优先级更高——这与训练时的贪心策略一致,也保证了训练与推理的一致性(这是一个非常重要的接口契约:合并顺序必须被保存下来)。

2.4.3 解码过程

解码时把每个 ID 还原成字节:小于 256 的直接就是一个字节;大于等于 256 的,要知道它是由哪一对合并而来的,于是递归展开,最后把字节序列用 UTF-8 解码回字符串。因为可能截断在多字节字符中间,解码时要用 errors="replace" 兜底,避免抛异常。

2.4.4 特殊词元

  • <|endoftext|>:文档边界。生成时它也是停止条件;指令微调时它标记回答结束。
  • <|pad|>:填充。批处理时把短序列补齐到同一长度,但填充位置不应参与损失计算,因此标签要置为 -100(第 9 章会详细讲)。

这两个词元不参与 BPE 合并,直接分配在词表末尾,并且编码时要优先识别(否则它们会被当成普通字符串拆成字节)。

2.4.5 压缩率:分词器质量的量化指标

对同一个语料,好的分词器应该用更少的词元表示它:

$$\text{压缩率} = \frac{\text{字节数}}{\text{词元数}}$$

压缩率越高,序列越短,训练与推理越省算力。英文文本通常约 3~4 字节/词元;中文因为一个汉字占 3 字节,压缩率会更低一些(本项目的实测值见 _build/measurements.md)。当你想调整词表大小时,压缩率是最直接的依据:把词表从 800 提到 4000,压缩率会上升,但嵌入层参数量也线性上升,需要权衡。

2.4.6 训练成本与工程处理

BPE 训练是纯 Python 实现的 $O(n_{\text{merges}} \times \text{len})$ 过程:每一次合并都要完整扫描一遍序列统计词对。序列有 4 万词元、要合并 500 次时,就是 2000 万次循环——在 CPU 上是几十秒的量级。

因此本书做了两个工程取舍:

  • 限制前 4 万个字符参与训练(max_chars)。前 4 万字符已足以统计出稳定的高频词对;
  • 把编码结果缓存到磁盘(.npy)。分词是训练前的一次性开销,缓存后重复实验只需几百毫秒,这对"反复调超参数"的流程极其重要。

下面是本章涉及的全部代码:数据与语料生成、三种分词器实现、滑动窗口数据集。

"""
corpus.py —— 完全离线、可复现的小语料与数据集构造
==================================================

原书为了演示,需要联网下载《伊索寓言》/《绿野仙踪》语料、GPT-2 权重以及
垃圾邮件数据集。本项目要求“纯 CPU、可离线、可复现”,因此这里用**确定性
生成**的方式造出三类数据:

1. ``build_corpus``            —— 预训练语料(中英混排的说明性文本)
2. ``build_spam_csv``          —— 分类微调数据集(垃圾/正常短信)
3. ``build_instruction_json``  —— 指令微调数据集(指令/输入/回答三字段)

这三个生成器都只用标准库,不依赖网络、不依赖 pandas,随机性由固定种子
控制,因此**任何人运行都会得到完全相同的文件**,便于对照文档中的输出。
"""

from __future__ import annotations

import csv
import json
import os
import random
from typing import Dict, List

DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")

# --------------------------------------------------------------
# 1. 预训练语料
# ---------------------------------------------------------------
# 语料主题围绕“大语言模型”,既包含中文也包含英文与代码片段。
# 这样做有两个好处:
#   (a) 字节级 BPE 分词器能展示它对中文(UTF-8 多字节)的处理能力;
#   (b) 语言模型学到的是有结构的文本,损失下降曲线更平滑、更有说服力。
# ---------------------------------------------------------------
_PARAGRAPHS: List[str] = [
    "大语言模型是一种以自回归方式预测下一个词元的神经网络,它把自然语言处理问题统一成了序列建模问题。",
    "Transformer 架构由编码器和解码器组成,而 GPT 系列只保留了解码器部分,因此被称为“仅解码器”模型。",
    "自注意力机制让序列中的每个位置都能直接看到其它位置的信息,从而绕开了循环神经网络的长距离依赖难题。",
    "因果掩码保证模型在预测第 t 个词元时只能使用前 t-1 个词元的信息,这正是自回归生成的核心约束。",
    "多头注意力把嵌入维度切成若干份,每一份独立做注意力计算,使模型能够同时关注语法、语义和位置等不同模式。",
    "层归一化把每一层的激活值重新标准化,使深层网络的训练更加稳定,也让学习率的选择不再那么敏感。",
    "残差连接把输入直接加到子层输出上,为梯度提供了一条“高速公路”,是能训练上百层网络的关键技巧。",
    "前馈网络对每个位置独立地做两次线性变换,中间维度通常放大到 4 倍,它承担了模型中大部分参数。",
    "字节对编码把常见字符序列合并成更长的子词,从而在词表规模和序列长度之间取得平衡。",
    "词元嵌入把离散的词元映射成连续向量,位置嵌入则告诉模型每个词元在序列中的顺序。",
    "语言模型的训练目标是最小化交叉熵损失,等价于最大化训练语料的似然,其指数形式就是困惑度。",
    "困惑度可以直观理解为模型在每一步平均需要在多少个候选词中犹豫,数值越低说明模型越确定。",
    "AdamW 优化器把权重衰减从梯度更新中解耦出来,配合余弦学习率衰减,是预训练事实上的标准配置。",
    "学习率预热在小批量训练初期逐步提高学习率,避免参数在优化器二阶矩估计还不稳定时被“撞飞”。",
    "梯度裁剪把所有梯度的整体范数限制在一个阈值以内,避免个别异常批次破坏已经学到的东西。",
    "梯度累积把多个小批次的梯度相加后再更新一次参数,使显存或内存有限的设备也能模拟大批量训练。",
    "在纯 CPU 环境中,瓶颈通常不是浮点算力而是内存带宽,因此批大小和上下文长度要谨慎选择。",
    "推理阶段最大的浪费是重复计算:KV 缓存把历史键值保存下来,使每步计算量从平方级降到线性级。",
    "动态量化把权重压缩成 8 位整数,在 CPU 上通常能同时获得更小的模型和更快的矩阵乘法。",
    "参数高效微调只训练一小部分新增参数,LoRA 用两个低秩矩阵的乘积去近似权重更新,效果常常接近全量微调。",
    "分类微调在语言模型顶部接一个线性层,用带标签的数据把生成模型改造成判别模型。",
    "指令微调让基础模型学会“听懂人话”,它使用指令、输入、回答三段式的数据格式。",
    "一个训练良好的模型应该同时具备低训练损失和低验证损失,两者差距扩大就是过拟合的信号。",
    "随机种子、数据顺序和初始化都会影响最终结果,因此任何实验结论都应报告多次运行的平均值。",
    "把大问题拆成小实验,是工程上最可靠的做法:先让 1000 万参数的模型跑通,再考虑放大。",
    "The decoder-only Transformer predicts the next token given all previous tokens in the sequence.",
    "Attention is a weighted average over values, where the weights come from query-key similarity.",
    "A language model is trained with a simple objective: maximize the log-likelihood of the corpus.",
    "Scaling laws suggest that loss decreases as a power law in model size, data size and compute.",
    "Quantization trades a small amount of accuracy for a large reduction in memory traffic.",
    "On a CPU, memory bandwidth is often the limiting factor, so smaller dtypes help more than expected.",
    "def train_step(model, batch):  # 一个最小训练步的伪代码",
    "    logits = model(batch)      # 前向传播:得到每个位置对全词表的打分",
    "    loss = cross_entropy(logits, targets)  # 与真实的下一个词元比较",
    "    loss.backward()            # 反向传播:计算每个参数的梯度",
    "    optimizer.step()           # 用梯度更新参数",
    "    optimizer.zero_grad()      # 清空梯度,准备下一个批次",
    "对话系统把历史消息拼接成一个长序列,并用特殊词元标出每条消息的角色与边界。",
    "评测一个生成模型不能只看损失,还要看它是否遵循指令、是否重复、是否顺畅。",
    "把模型存成检查点,除了参数还要保存优化器状态,否则断点续训会丢掉动量信息。",
    "部署时最关心的三件事是:延迟、吞吐和内存占用,它们三者往往互相牵制。",
]

_ENGLISH_WORDS = (
    "the model learns to predict the next token in a sequence of text, and it does so by "
    "adjusting millions of parameters with gradient descent on a large corpus of documents"
).split()

_DIALOGS: List[str] = [
    "用户:什么是注意力?\n助手:注意力是一种加权求和,权重由查询与键的相似度决定。",
    "用户:为什么要用位置嵌入?\n助手:因为自注意力本身对顺序不敏感,必须显式注入位置信息。",
    "用户:CPU 上能训练多大的模型?\n助手:用 4 核 CPU 可以顺利训练千万参数级的教学模型,更大的模型建议只做推理或微调。",
    "用户:怎样判断过拟合?\n助手:训练损失继续下降而验证损失开始上升,就说明模型在记忆训练集。",
]


def _random_english_sentence(rng: random.Random, length: int = 12) -> str:
    words = [rng.choice(_ENGLISH_WORDS) for _ in range(length)]
    words[0] = words[0].capitalize()
    return " ".join(words) + "."


def build_corpus(path: str | None = None, repeats: int = 60, seed: int = 123) -> str:
    """生成预训练语料并写入磁盘,返回语料全文。

    参数
    ----
    path     : 输出路径,默认为 ``<项目根>/data/corpus.txt``
    repeats  : 把基础段落重复多少轮。重复会让语料变大、模型更快“记住”,
               适合教学演示;真实预训练需要海量且不重复的语料。
    """
    path = path or os.path.join(DATA_DIR, "corpus.txt")
    rng = random.Random(seed)
    lines: List[str] = ["大语言模型预训练语料(教学用,完全离线可复现)", ""]
    for r in range(repeats):
        # 每轮都重新打乱顺序,避免模型只学到固定顺序
        block = list(_PARAGRAPHS)
        rng.shuffle(block)
        lines.append(f"=== 第 {r + 1} 轮 ===")
        for p in block:
            lines.append(p)
        for d in _DIALOGS:
            lines.append(d)
        for _ in range(3):
            lines.append(_random_english_sentence(rng))
        lines.append("")
    text = "\n".join(lines)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        f.write(text)
    return text


def load_corpus(path: str | None = None, build_if_missing: bool = True, repeats: int = 60) -> str:
    """读取语料;若文件不存在则自动生成(保证示例一条命令即可运行)。"""
    path = path or os.path.join(DATA_DIR, "corpus.txt")
    if not os.path.exists(path):
        if not build_if_missing:
            raise FileNotFoundError(f"语料不存在: {path}")
        return build_corpus(path, repeats=repeats)
    with open(path, "r", encoding="utf-8") as f:
        return f.read()


# -----------------------------------------------------------------
# 2. 垃圾短信分类数据集
# -----------------------------------------------------------------
_SPAM_PREFIX = ["", "【紧急】", "【通知】", "尊敬的用户,", "亲爱的客户,", "限时活动!"]
_SPAM_CORE = [
    "恭喜您中奖 10000 元,请点击链接领取",
    "免费领取最新款手机,仅剩少量名额",
    "您有一笔低息贷款额度待激活,额度高至 50 万",
    "点击链接领取优惠券,全场一折起",
    "您的账户异常,请立即登录验证身份",
    "免费试用会员 30 天,无需信用卡",
    "独家投资机会,日收益 10% 稳赚不赔",
    "加客服微信领取红包,先到先得",
    "urgent: claim your free prize now, click the link",
    "congratulations, you have won a lottery ticket",
    "lowest price guaranteed, buy now and save 90%",
]
_SPAM_SUFFIX = ["", " 退订回复 T", " 详情咨询客服", " 点击 www.example.com", " 回复 Y 立即办理", " 名额有限,速来"]

_HAM_PREFIX = ["", "你好,", "请问,", "提醒一下,"]
_HAM_CORE = [
    "明天上午十点开会,会议室在三楼",
    "我把文件的最终版本发到你邮箱了",
    "晚饭想吃什么?我六点下班回家",
    "机器学习作业下周一是截止日期",
    "这个周末一起去图书馆复习功课吧",
    "昨天的实验结果我已经整理成表格",
    "记得带上身份证,办理手续需要",
    "模型训练完成以后请把日志发我一份",
    "the meeting is moved to three o'clock tomorrow",
    "please review my pull request when you have time",
    "let us meet at the library after class",
    "i finished the data cleaning script yesterday",
]
_HAM_SUFFIX = ["", "。", ",谢谢!", ",麻烦你了", ",收到请回复", "。"]


def build_spam_csv(
    path: str | None = None,
    n_per_class: int = 600,
    seed: int = 123,
    balance: bool = True,
) -> str:
    """生成二分类数据集(标签为 ``ham`` / ``spam``),写入 CSV 并返回路径。

    真实数据(如 SMS Spam Collection)需要联网下载;这里用模板+随机组合造出
    一个**可分但并非平凡**的任务:类别关键词需要被模型学到,而模板本身带有
    噪声(前缀/后缀随机),因此模型必须依赖内容而不是某个固定模式。
    """
    path = path or os.path.join(DATA_DIR, "spam.csv")
    rng = random.Random(seed)
    rows: List[Dict[str, str]] = []
    for _ in range(n_per_class):
        text = rng.choice(_SPAM_PREFIX) + rng.choice(_SPAM_CORE) + rng.choice(_SPAM_SUFFIX)
        rows.append({"Label": "spam", "Text": text})
    for _ in range(n_per_class):
        text = rng.choice(_HAM_PREFIX) + rng.choice(_HAM_CORE) + rng.choice(_HAM_SUFFIX)
        rows.append({"Label": "ham", "Text": text})
    # 打乱顺序:避免“前半是 spam、后半是 ham”这种会污染训练的数据顺序
    rng.shuffle(rows)
    if balance:  # 显式检查类别平衡,方便事后复现实验
        n_spam = sum(1 for r in rows if r["Label"] == "spam")
        assert n_spam == len(rows) - n_spam, "数据集不平衡"
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=["Label", "Text"])
        writer.writeheader()
        writer.writerows(rows)
    return path


def load_spam_csv(path: str | None = None) -> List[Dict[str, str]]:
    """用标准库读取 CSV,返回 ``[{"Label":..., "Text":...}, ...]``。

    不依赖 pandas,因此在任何 CPU 环境都能跑;如果你装了 pandas,
    也可以 ``pd.read_csv(path)`` 得到完全等价的结果。
    """
    path = path or os.path.join(DATA_DIR, "spam.csv")
    if not os.path.exists(path):
        build_spam_csv(path)
    with open(path, "r", encoding="utf-8", newline="") as f:
        return list(csv.DictReader(f))


# ----------------------------------------------------------------
# 3. 指令微调数据集
# ----------------------------------------------------------------
_INSTRUCT_TEMPLATES = [
    ("把下面的英文翻译成中文。", "en", "zh"),
    ("将下面的中文翻译成英文。", "zh", "en"),
    ("把下面的句子改写成被动语态。", "active", "passive"),
    ("纠正下面句子中的语法错误。", "bad", "good"),
    ("用一句话概括下面这段话。", "long", "short"),
    ("回答下面的问题。", "q", "a"),
]

_PAIRS: List[Dict[str, str]] = [
    {
        "en": "The chef prepared the meal.",
        "zh": "厨师准备了这顿饭。",
        "active": "The chef prepared the meal.",
        "passive": "The meal was prepared by the chef.",
        "bad": "She go to school every day.",
        "good": "She goes to school every day.",
        "long": "注意力机制通过计算查询与键的相似度得到权重,再对值做加权求和,从而让每个位置都能直接访问序列中的其它位置。",
        "short": "注意力是一种基于相似度的加权求和。",
        "q": "什么是困惑度?",
        "a": "困惑度是交叉熵损失的指数形式,表示模型平均在多少个候选词之间犹豫。",
    },
    {
        "en": "The model predicts the next token.",
        "zh": "模型预测下一个词元。",
        "active": "The programmer fixed the bug.",
        "passive": "The bug was fixed by the programmer.",
        "bad": "He don't like coffee.",
        "good": "He doesn't like coffee.",
        "long": "字节对编码从单字节开始,反复把出现频率最高的相邻符号对合并成新符号,最终得到的词表既不会太大,也能覆盖常见词。",
        "short": "BPE 通过不断合并高频相邻符号来构造子词词表。",
        "q": "为什么需要位置嵌入?",
        "a": "因为自注意力对位置不敏感,必须显式加入顺序信息。",
    },
    {
        "en": "Training on a CPU requires patience.",
        "zh": "在 CPU 上训练需要耐心。",
        "active": "The teacher graded the exams.",
        "passive": "The exams were graded by the teacher.",
        "bad": "I has two dogs.",
        "good": "I have two dogs.",
        "long": "残差连接把子层的输入直接加到输出上,使梯度可以绕过非线性变换直接回传,这显著缓解了深层网络的梯度消失问题。",
        "short": "残差连接为梯度提供了一条直接回传的通路。",
        "q": "什么是 KV 缓存?",
        "a": "KV 缓存保存历史位置的键和值,使自回归生成每一步只需计算当前词元。",
    },
    {
        "en": "Quantization reduces memory traffic.",
        "zh": "量化减少了内存流量。",
        "active": "The team shipped the feature.",
        "passive": "The feature was shipped by the team.",
        "bad": "They was late yesterday.",
        "good": "They were late yesterday.",
        "long": "LoRA 冻结原始权重,只训练两个低秩矩阵 A 和 B,用它们的乘积近似权重更新量,从而把可训练参数减少几个数量级。",
        "short": "LoRA 用低秩矩阵近似权重更新,只训练极少量参数。",
        "q": "梯度累积有什么用?",
        "a": "它把多个小批次的梯度累加后再更新,用同样的内存模拟更大的批大小。",
    },
    {
        "en": "Layer normalization stabilizes training.",
        "zh": "层归一化使训练更稳定。",
        "active": "The student answered the question.",
        "passive": "The question was answered by the student.",
        "bad": "This are my book.",
        "good": "This is my book.",
        "long": "预训练让模型从海量无标注文本中学到语言的统计规律,微调则把这些通用能力对齐到具体任务上。",
        "short": "预训练学通用规律,微调对齐具体任务。",
        "q": "为什么要做学习率预热?",
        "a": "因为训练初期优化器的二阶矩估计还不稳定,过大的学习率会让参数发散。",
    },
]


def build_instruction_json(
    path: str | None = None,
    repeats: int = 220,
    seed: int = 123,
) -> str:
    """生成指令微调数据集(JSON 数组,每条含 instruction/input/output)。"""
    path = path or os.path.join(DATA_DIR, "instruction.json")
    rng = random.Random(seed)
    records: List[Dict[str, str]] = []
    for _ in range(repeats):
        for tmpl, inp_key, out_key in _INSTRUCT_TEMPLATES:
            pair = rng.choice(_PAIRS)
            records.append(
                {
                    "instruction": tmpl,
                    "input": pair[inp_key],
                    "output": pair[out_key],
                }
            )
    rng.shuffle(records)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(records, f, ensure_ascii=False, indent=1)
    return path


def load_instruction_json(path: str | None = None) -> List[Dict[str, str]]:
    path = path or os.path.join(DATA_DIR, "instruction.json")
    if not os.path.exists(path):
        build_instruction_json(path)
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def write_jsonl(records: List[Dict], path: str) -> str:
    """把记录写成 JSONL(每行一个 JSON),这是指令数据最常见的存储格式。"""
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        for r in records:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    return path


if __name__ == "__main__":  # pragma: no cover - 手动演示
    p1 = build_corpus()
    text = load_corpus()
    p2 = build_spam_csv()
    rows = load_spam_csv()
    p3 = build_instruction_json()
    recs = load_instruction_json()
    print("语料:", p1, "字符数:", len(text))
    print("分类数据:", p2, "条数:", len(rows), "示例:", rows[0])
    print("指令数据:", p3, "条数:", len(recs), "示例:", recs[0])
    print("指令数据样例格式化:\n" + "\n".join(f"  {k}: {v}" for k, v in recs[0].items()))
"""
tokenizer.py —— 三种分词器:教学版、特殊词元版、字节级 BPE(从零实现)
=====================================================================

分词(tokenization)是语言模型的第一道工序:把原始文本切成模型能处理的
离散符号。本文件按由浅入深的顺序给出三种实现:

1. ``SimpleTokenizerV1``  —— 按空格/标点切分,用固定词表映射(会丢未知词)
2. ``SimpleTokenizerV2``  —— 加入 ``<|unk|>`` / ``<|endoftext|>`` 特殊词元
3. ``BPETokenizer``      —— 字节级 BPE,从零训练,能处理中文与任意字符

为什么最终一定要用字节级 BPE?
    * 词表固定 → 永远不会遇到“未知词”(任何字符都能拆成 UTF-8 字节);
    * 词表可控 → ``vocab_size`` 由合并次数决定,可按 CPU 能力缩放;
    * 序列较短 → 高频子词被合并成一个词元,缩短序列、降低计算量。
"""

from __future__ import annotations

import json
import os
import re
from collections import Counter
from typing import Dict, Iterable, List, Sequence, Tuple

import numpy as np

DATA_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")

# ---------------------------------------------------------------
# 1. 教学版分词器:让读者看清“分词 → ID → 反分词”的闭环
# ---------------------------------------------------------------
class SimpleTokenizerV1:
    """最朴素的分词器:正则切分 + 词表查表。

    缺点:遇到词表外的词直接 ``KeyError``,因此在真实语料上不可用。
    """

    def __init__(self, vocab: Dict[str, int]) -> None:
        self.str_to_int = vocab
        self.int_to_str = {i: s for s, i in vocab.items()}

    def encode(self, text: str) -> List[int]:
        # 用非字母数字的字符作为分隔符,并保留标点本身
        preprocessed = re.split(r'([,.:;?_!"()\']|--|\s)', text)
        preprocessed = [item.strip() for item in preprocessed if item.strip()]
        return [self.str_to_int[s] for s in preprocessed]

    def decode(self, ids: Sequence[int]) -> str:
        text = " ".join(self.int_to_str[i] for i in ids)
        # 去掉标点前的多余空格,让输出更自然
        return re.sub(r'\s+([,.?!"()\'])', r"\1", text)


class SimpleTokenizerV2:
    """加入 ``<|unk|>`` 与 ``<|endoftext|>`` 的改进版:不再抛异常。"""

    def __init__(self, vocab: Dict[str, int]) -> None:
        self.str_to_int = vocab
        self.int_to_str = {i: s for s, i in vocab.items()}

    def encode(self, text: str) -> List[int]:
        preprocessed = re.split(r'([,.:;?_!"()\']|--|\s)', text)
        preprocessed = [item.strip() for item in preprocessed if item.strip()]
        preprocessed = [
            item if item in self.str_to_int else "<|unk|>" for item in preprocessed
        ]
        return [self.str_to_int[s] for s in preprocessed]

    def decode(self, ids: Sequence[int]) -> str:
        text = " ".join(self.int_to_str[i] for i in ids)
        return re.sub(r'\s+([,.?!"()\'])', r"\1", text)


def build_simple_vocab(text: str, extra_specials: Sequence[str] = ("<|endoftext|>",)) -> Dict[str, int]:
    """从一段文本构造 ``SimpleTokenizerV2`` 需要的词表。"""
    preprocessed = re.split(r'([,.:;?_!"()\']|--|\s)', text)
    words = sorted({item.strip() for item in preprocessed if item.strip()})
    vocab = {tok: i for i, tok in enumerate(["<|unk|>", *extra_specials, *words])}
    return vocab


# ----------------------------------------------------------------
# 2. 字节级 BPE(核心实现,约 120 行)
# -----------------------------------------------------------------
def _merge_pair(ids: List[int], pair: Tuple[int, int], new_id: int) -> List[int]:
    """把序列中所有相邻的 ``pair`` 替换成 ``new_id``(单趟线性扫描)。"""
    out: List[int] = []
    a, b = pair
    i = 0
    n = len(ids)
    while i < n:
        if i < n - 1 and ids[i] == a and ids[i + 1] == b:
            out.append(new_id)
            i += 2
        else:
            out.append(ids[i])
            i += 1
    return out


class BPETokenizer:
    """字节级字节对编码(Byte-level BPE)。

    训练过程(对应 ``train``):

    1. 把文本编码成 UTF-8 字节序列,得到 0~255 的初始词元;
    2. 统计所有相邻词元对的频次;
    3. 把频次最高的一对合并成一个新词元(ID = 256, 257, ...);
    4. 重复 2~3 直到达到目标词表大小。

    编码过程(对应 ``encode``):按**合并产生的先后顺序**依次应用每条合并规则。
    先产生的合并优先级更高,这与训练时的贪心策略一致。

    解码过程(对应 ``decode``):把每个 ID 递归展开成字节,再用 UTF-8 解码。
    """

    DEFAULT_SPECIALS = ["<|endoftext|>", "<|pad|>"]

    def __init__(
        self,
        merges: Dict[Tuple[int, int], int] | None = None,
        special_tokens: Sequence[str] | None = None,
    ) -> None:
        self.merges: Dict[Tuple[int, int], int] = dict(merges or {})
        self.special_tokens: List[str] = list(
            self.DEFAULT_SPECIALS if special_tokens is None else special_tokens
        )
        # 合并规则的“优先级”就是它在字典中的插入序号
        self._merge_list: List[Tuple[Tuple[int, int], int]] = list(self.merges.items())
        self._expand_cache: Dict[int, bytes] = {}
        self._special_ids: Dict[str, int] = {
            tok: 256 + len(self._merge_list) + i for i, tok in enumerate(self.special_tokens)
        }
        self._special_pattern = (
            re.compile("(" + "|".join(re.escape(t) for t in self.special_tokens) + ")")
            if self.special_tokens
            else None
        )

    # -- 基本属性 ---------------------------------------------------------
    @property
    def vocab_size(self) -> int:
        """词表大小 = 256 个字节 + 合并出的新词元 + 特殊词元。"""
        return 256 + len(self._merge_list) + len(self.special_tokens)

    @property
    def eot_id(self) -> int:
        """文本结束词元 ``<|endoftext|>`` 的 ID。"""
        return self._special_ids["<|endoftext|>"]

    @property
    def pad_id(self) -> int:
        return self._special_ids.get("<|pad|>", 0)

    def special_id(self, token: str) -> int:
        return self._special_ids[token]

    # -- 训练 -------------------------------------------------------------
    def train(self, text: str, vocab_size: int, verbose: bool = False) -> "BPETokenizer":
        """在 ``text`` 上训练 BPE,直到词表达到 ``vocab_size``。

        ``vocab_size`` 必须大于 ``256 + len(special_tokens)``。
        """
        n_merges = vocab_size - 256 - len(self.special_tokens)
        if n_merges <= 0:
            raise ValueError("vocab_size 太小:至少要大于 256 + 特殊词元个数")
        ids = list(text.encode("utf-8"))
        merges: Dict[Tuple[int, int], int] = {}
        for i in range(n_merges):
            counts = Counter(zip(ids, ids[1:]))
            if not counts:
                break
            pair, freq = counts.most_common(1)[0]
            if freq < 2:  # 再合并也没有收益,提前停止
                if verbose:
                    print(f"[BPE] 提前停止:最高频词对频次 {freq} < 2")
                break
            new_id = 256 + i
            ids = _merge_pair(ids, pair, new_id)
            merges[pair] = new_id
            if verbose and (i + 1) % 100 == 0:
                print(f"[BPE] 已合并 {i + 1} 次,当前最高频对 {pair} 频次 {freq}")
        self.merges = merges
        self._merge_list = list(merges.items())
        self._expand_cache = {}
        self._special_ids = {
            tok: 256 + len(self._merge_list) + i for i, tok in enumerate(self.special_tokens)
        }
        return self

    # -- 编码 / 解码 ------------------------------------------------------
    def _encode_chunk(self, chunk: str) -> List[int]:
        ids = list(chunk.encode("utf-8"))
        for pair, new_id in self._merge_list:
            ids = _merge_pair(ids, pair, new_id)
        return ids

    def encode(self, text: str, allow_special: bool = True) -> List[int]:
        """把字符串编码成 ID 列表。特殊词元(如 ``<|endoftext|>``)会被识别为单个 ID。"""
        if not text:
            return []
        if allow_special and self._special_pattern is not None:
            parts = self._special_pattern.split(text)
        else:
            parts = [text]
        ids: List[int] = []
        for part in parts:
            if part == "":
                continue
            if part in self._special_ids:
                ids.append(self._special_ids[part])
            else:
                ids.extend(self._encode_chunk(part))
        return ids

    def _expand(self, token_id: int) -> bytes:
        """把单个 ID 展开成字节序列(带缓存,避免重复递归)。"""
        cached = self._expand_cache.get(token_id)
        if cached is not None:
            return cached
        if token_id < 256:
            result = bytes([token_id])
        else:
            # 反向查表:新 ID -> 它由哪一对合并而来
            pair = None
            for p, nid in self._merge_list:
                if nid == token_id:
                    pair = p
                    break
            if pair is None:
                raise ValueError(f"非法词元 ID: {token_id}")
            result = self._expand(pair[0]) + self._expand(pair[1])
        self._expand_cache[token_id] = result
        return result

    def decode(self, ids: Iterable[int], skip_special: bool = True) -> str:
        """把 ID 列表还原成字符串,遇到不完整字节用替换字符兜底。"""
        id_to_special = {v: k for k, v in self._special_ids.items()}
        out: List[bytes] = []
        for i in ids:
            i = int(i)
            if i in id_to_special:
                if not skip_special:
                    out.append(id_to_special[i].encode("utf-8"))
                continue
            out.append(self._expand(i))
        return b"".join(out).decode("utf-8", errors="replace")

    # -- 持久化 -----------------------------------------------------------
    def save(self, path: str) -> str:
        os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
        payload = {
            "merges": [[int(a), int(b), int(nid)] for (a, b), nid in self._merge_list],
            "special_tokens": self.special_tokens,
        }
        with open(path, "w", encoding="utf-8") as f:
            json.dump(payload, f, ensure_ascii=False)
        return path

    @classmethod
    def load(cls, path: str) -> "BPETokenizer":
        with open(path, "r", encoding="utf-8") as f:
            payload = json.load(f)
        merges = {(int(a), int(b)): int(nid) for a, b, nid in payload["merges"]}
        return cls(merges=merges, special_tokens=payload["special_tokens"])

    def __repr__(self) -> str:  # pragma: no cover - 调试友好
        return f"BPETokenizer(vocab_size={self.vocab_size}, merges={len(self._merge_list)})"


# ---------------------------------------------------------------
# 3. 便捷函数:训练 / 缓存
# ---------------------------------------------------------------
def train_bpe_on_corpus(
    corpus_path: str | None = None,
    vocab_size: int = 800,
    out_path: str | None = None,
    max_chars: int = 40000,
    verbose: bool = True,
) -> BPETokenizer:
    """在语料的前 ``max_chars`` 个字符上训练 BPE 并保存。

    为什么要限制 ``max_chars``?BPE 训练是纯 Python 实现的 $O(\\text{merges}
    \\times \\text{len})$ 过程,语料越大越慢。前 4 万字符已经足以统计出稳定的
    高频词对,对教学与中小规模训练完全够用。
    """
    from corpus import load_corpus  # 同目录导入(运行方式:python code/xxx.py)

    out_path = out_path or os.path.join(DATA_DIR, "bpe.json")
    text = load_corpus(corpus_path)
    subset = text[:max_chars]
    tok = BPETokenizer()
    tok.train(subset, vocab_size=vocab_size, verbose=verbose)
    tok.save(out_path)
    if verbose:
        print(f"[BPE] 训练完成:词表 {tok.vocab_size},已保存到 {out_path}")
    return tok


def load_or_train_bpe(path: str | None = None, vocab_size: int = 800) -> BPETokenizer:
    """加载已保存的分词器;不存在则训练一个。"""
    path = path or os.path.join(DATA_DIR, "bpe.json")
    if os.path.exists(path):
        return BPETokenizer.load(path)
    return train_bpe_on_corpus(vocab_size=vocab_size, out_path=path)


def tokenize_corpus(
    tokenizer: BPETokenizer,
    text: str | None = None,
    cache_path: str | None = None,
    force: bool = False,
) -> np.ndarray:
    """把整段语料编码成 ``int32`` 数组并缓存到 ``.npy``。

    编码是训练前的一次性开销。缓存后,重复实验只需几百毫秒即可加载,
    这对“反复调超参数”的 CPU 实验流程非常重要。
    """
    from corpus import load_corpus

    cache_path = cache_path or os.path.join(DATA_DIR, "corpus_ids.npy")
    if os.path.exists(cache_path) and not force:
        return np.load(cache_path)
    text = text if text is not None else load_corpus()
    ids = np.array(tokenizer.encode(text), dtype=np.int32)
    np.save(cache_path, ids)
    return ids


def compute_compression_stats(tokenizer: BPETokenizer, text: str) -> Dict[str, float]:
    """统计压缩率:平均每个词元对应多少字节 / 多少个字符。

    这是评估分词器质量最直观的指标:压缩率越高,同样文本的序列越短,
    训练和推理越省算力。
    """
    ids = tokenizer.encode(text)
    n_bytes = len(text.encode("utf-8"))
    return {
        "字符数": len(text),
        "字节数": n_bytes,
        "词元数": len(ids),
        "字节/词元": n_bytes / max(1, len(ids)),
        "字符/词元": len(text) / max(1, len(ids)),
    }


if __name__ == "__main__":  # pragma: no cover - 手动演示
    from corpus import load_corpus

    text = load_corpus()

    # (1) 教学版分词器:只在英文片段上演示,词表来自语料本身
    sample = "The chef prepared the meal. The meal was good."
    vocab = build_simple_vocab(text[:20000])
    tok_v2 = SimpleTokenizerV2(vocab)
    ids = tok_v2.encode(sample)
    print("SimpleTokenizerV2 词表:", len(vocab), "编码:", ids)
    print("解码:", tok_v2.decode(ids))

    # (2) 字节级 BPE
    bpe = train_bpe_on_corpus(vocab_size=800, verbose=True)
    for s in [
        "大语言模型",
        "attention is all you need",
        "The chef prepared the meal.",
    ]:
        enc = bpe.encode(s)
        print(f"\n原文: {s}\n编码: {enc}\n解码: {bpe.decode(enc)!r}")
    print("\n压缩统计:", compute_compression_stats(bpe, text[:20000]))
    ids_arr = tokenize_corpus(bpe, text)
    print("语料词元总数:", ids_arr.shape, "前 20 个 ID:", ids_arr[:20].tolist())
"""
dataset.py —— 滑动窗口数据集与 DataLoader
=========================================

语言模型的监督信号是“自己”:把词元序列切成固定长度的窗口,
输入是窗口内的前 n 个词元,目标是**右移一位**的同长度序列。

    输入 : [大, 语言, 模型, 是]
    目标 : [语言, 模型, 是, 一种]

``stride`` 控制窗口之间的重叠程度:
    * stride == max_length:窗口互不重叠,数据利用率最低;
    * stride == 1        :几乎每个位置都作为一个新窗口,数据量最大但重复计算多;
    * 折中值(如 max_length // 2):CPU 训练最常用的选择。
"""

from __future__ import annotations

from typing import List, Tuple

import torch
from torch.utils.data import DataLoader, Dataset


class GPTDatasetV1(Dataset):
    """把长序列切成 ``(input, target)`` 对。"""

    def __init__(self, txt: str, tokenizer, max_length: int, stride: int) -> None:
        self.input_ids: List[torch.Tensor] = []
        self.target_ids: List[torch.Tensor] = []

        token_ids = tokenizer.encode(txt)
        if len(token_ids) <= max_length:
            raise ValueError(
                f"语料太短:只有 {len(token_ids)} 个词元,窗口需要 {max_length + 1} 个"
            )

        # 单趟滑动窗口切分,步长为 stride
        for i in range(0, len(token_ids) - max_length, stride):
            input_chunk = token_ids[i : i + max_length]
            target_chunk = token_ids[i + 1 : i + max_length + 1]
            self.input_ids.append(torch.tensor(input_chunk, dtype=torch.long))
            self.target_ids.append(torch.tensor(target_chunk, dtype=torch.long))

    def __len__(self) -> int:
        return len(self.input_ids)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
        return self.input_ids[idx], self.target_ids[idx]


def create_dataloader_v1(
    txt: str,
    batch_size: int = 4,
    max_length: int = 256,
    stride: int = 128,
    shuffle: bool = True,
    drop_last: bool = True,
    num_workers: int = 0,
    tokenizer=None,
) -> DataLoader:
    """构造 ``DataLoader``。

    关于 ``num_workers`` 的 CPU 实践建议
    ------------------------------------
    * **默认 0**:在纯 CPU 训练中,数据准备本身很轻(只是取切片),
      而多进程会复制内存、增加调度与 IPC 开销,通常**不会更快**。
    * 只有当数据需要昂贵的在线预处理(如实时分词、图像解码)时,
      才把 ``num_workers`` 设为 2~4。
    * 另外,多进程 DataLoader 与 ``torch.set_num_threads`` 会争抢核心,
      若必须使用,请把工作进程数 × 线程数控制在物理核心数以内。
    """
    if tokenizer is None:
        raise ValueError("必须传入 tokenizer")
    dataset = GPTDatasetV1(txt, tokenizer, max_length, stride)
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        drop_last=drop_last,
        num_workers=num_workers,
    )


def inspect_batch(dataloader: DataLoader, tokenizer=None) -> None:
    """打印一个批次的形状与内容,用于确认“输入/目标确实错开一位”。"""
    data_iter = iter(dataloader)
    x, y = next(data_iter)
    print("输入形状:", tuple(x.shape), " dtype:", x.dtype)
    print("目标形状:", tuple(y.shape), " dtype:", y.dtype)
    if tokenizer is not None:
        print("输入文本:", tokenizer.decode(x[0].tolist())[:120].replace("\n", "\\n"))
        print("目标文本:", tokenizer.decode(y[0].tolist())[:120].replace("\n", "\\n"))
    else:
        print("输入 ID:", x[0][:12].tolist())
        print("目标 ID:", y[0][:12].tolist())
    assert x.shape == y.shape, "输入与目标形状必须一致"
    assert torch.equal(x[0][1:], y[0][:-1]), "目标应当是输入右移一位"


def build_lm_dataloaders(
    text: str,
    tokenizer,
    cfg: dict,
    batch_size: int = 8,
    stride: int | None = None,
    train_ratio: float = 0.9,
    num_workers: int = 0,
) -> Tuple[DataLoader, DataLoader]:
    """按比例切分训练集/验证集,并返回两个 DataLoader。

    注意:切分必须在**文本层**完成而不是在窗口层完成,否则同一个窗口的
    重叠部分会同时出现在训练集和验证集里,验证损失会被严重低估。
    """
    split_idx = int(train_ratio * len(text))
    train_text = text[:split_idx]
    val_text = text[split_idx:]
    stride = stride or cfg["context_length"] // 2
    train_loader = create_dataloader_v1(
        train_text,
        batch_size=batch_size,
        max_length=cfg["context_length"],
        stride=stride,
        shuffle=True,
        drop_last=True,
        num_workers=num_workers,
        tokenizer=tokenizer,
    )
    val_loader = create_dataloader_v1(
        val_text,
        batch_size=batch_size,
        max_length=cfg["context_length"],
        stride=stride,
        shuffle=False,
        drop_last=False,
        num_workers=num_workers,
        tokenizer=tokenizer,
    )
    return train_loader, val_loader


if __name__ == "__main__":  # pragma: no cover - 手动演示
    from corpus import load_corpus
    from tokenizer import load_or_train_bpe

    tok = load_or_train_bpe()
    text = load_corpus()
    loader = create_dataloader_v1(
        text, batch_size=4, max_length=32, stride=16, shuffle=False, tokenizer=tok
    )
    print("窗口数量:", len(loader.dataset))
    inspect_batch(loader, tokenizer=tok)

2.5 滑动窗口采样:语言模型的监督信号

语言模型不需要人工标注,因为答案就在原文里。把词元序列切成固定长度的窗口,输入是窗口内容,目标是右移一位的窗口:

原始序列: [大, 语言, 模型, 是, 一, 种, 自, 回, 归, 模, 型]
输入 (T=4):[大, 语言, 模型, 是]
目标 (T=4):[语言, 模型, 是, 一]

这样一次前向就同时产生 4 个训练信号(每个位置都在预测下一个词元),这是 Transformer 相对 RNN 的巨大效率优势:训练时可以并行处理整条序列。

stride 决定窗口之间的重叠程度:

stride 窗口数 数据利用率 计算重复
max_length 最少 低(边界信息被切掉) 无重复
max_length // 2 中等 高 有一半重叠
1 最多 最高 大量重复

本书默认用 context_length // 2,这是 CPU 训练中最常用的折中:既让模型看到足够多的相邻上下文,又不至于把数据集撑到爆炸。

2.5.1 划分训练集与验证集:一个容易搞错的细节

必须在文本层切分,而不是在窗口层切分:

split_idx = int(0.9 * len(text))
train_text, val_text = text[:split_idx], text[split_idx:]

原因很直白:相邻窗口重叠了一半词元。如果在窗口层随机划分,同一个句子会同时出现在训练集和验证集里,验证损失会被严重低估,你会以为模型泛化得很好,而实际上它在背答案。这是一个在真实项目里反复出现的隐蔽错误。

2.5.2 DataLoader 的 CPU 参数选择

num_workers 的取值在 CPU 场景下与"直觉"相反:默认应该是 0。原因有三条:

  1. 数据准备极轻——__getitem__ 只是从预先生成的张量列表里取一个切片;
  2. 多进程要复制数据集(每个 worker 一份),内存开销成倍增长;
  3. worker 进程与 PyTorch 计算线程争抢同一批核心,在 4 核机器上会明显拖慢训练。

只有当预处理很重(在线分词、图像解码)时,才值得把 num_workers 设为 2~4,并且要把"进程数 × 每进程线程数"控制在物理核心数以内。

2.6 词元嵌入与位置嵌入

2.6.1 词元嵌入

nn.Embedding(V, d) 给出形状 $(V, d)$ 的矩阵。输入 $(B, T)$ 的 ID,输出 $(B, T, d)$ 的向量。注意参数量是 $V\times d$:$V=50257$、$d=768$ 时是 3860 万参数,这是 124M 模型里最大的一块。

2.6.2 位置嵌入:为什么必需

自注意力有一个容易被忽略的性质:它对输入顺序是不敏感的。如果把输入序列打乱,注意力算出的每个位置的输出只是"跟着打乱",而不是"变了"。因为注意力权重只取决于查询和键的内容相似度,完全不知道谁在前谁在后。

但语言是有顺序的:"猫追狗"和"狗追猫"是两个意思。因此必须显式地把位置信息注入进去。两种做法:

方案 做法 优点 缺点
学习式位置嵌入 nn.Embedding(L, d) 查表相加 简单、可训练、GPT 系列采用 超过训练长度就无法外推
正弦式位置编码 用固定的 sin/cos 公式生成 理论可外推到更长序列 表达力稍弱,实际不如学习式
旋转位置编码 RoPE 对 Q/K 做旋转变换 外推性好,现代模型主流 实现复杂,本书不展开

本书采用学习式位置嵌入:

tok_embeds = tok_emb(x)                                    # (B, T, d)
pos_embeds = pos_emb(torch.arange(T))                      # (T, d) → 广播
x = drop(tok_embeds + pos_embeds)                          # (B, T, d)

关键约束:T 不能超过 context_length。一旦越界,nn.Embedding 会直接抛 IndexError: index out of range。这也是第 6 章 KV 缓存实现里必须检查 pos_offset 的原因。

2.7 运行与验证

在项目根目录执行:

python code/corpus.py      # 生成语料与数据集,全部离线
python code/tokenizer.py   # 训练 BPE、演示编码解码与压缩率
python code/dataset.py     # 演示滑动窗口与 DataLoader

tokenizer.py 的预期输出(节选,具体词表大小取决于语料):

SimpleTokenizerV2 词表: 1234 编码: [1, 2, 3, ...]
[BPE] 训练完成:词表 800,已保存到 data/bpe.json

原文: 大语言模型
编码: [140, 251, 88, 190, ...]
解码: '大语言模型'

压缩统计: {'字符数': 20000, '字节数': 54321, '词元数': 18765, ...}
语料词元总数: (123456,) 前 20 个 ID: [...]

dataset.py 的预期输出:

窗口数量: 1234
输入形状: (4, 32)  dtype: torch.int64
目标形状: (4, 32)  dtype: torch.int64
输入文本: ...
目标文本: ...

注意脚本里的一条断言:torch.equal(x[0][1:], y[0][:-1])——它验证"目标确实是输入右移一位",是防止数据管道写错的第一道保险。

2.8 常见错误

错误一:IndexError: index out of range in self(在 Embedding 层)。 两种可能:位置索引超过了 context_length;或词表不匹配——用词表 800 的分词器编码出 ID=850,却把模型建成了 vocab_size=800。记住:模型配置里的 vocab_size 必须来自分词器的真实值,本书用 config_from_tokenizer() 强制这一点。

错误二:编码后再解码,文本对不上。 常见原因是解码时跳过了某些特殊词元(skip_special=True),或在多字节字符中间截断。检查方法:decode(encode(s)) == s 对纯文本应当成立(不含特殊词元时)。

错误三:BPE 训练完发现词表和预期差很多。 如果目标 vocab_size 很小(比如 300),合并次数只有 44 次,模型会退化成接近字节级;如果语料太短,最高频词对的频次很快降到 2 以下,训练会提前停止(代码里会打印"提前停止")。合理区间是 800~8000。

错误四:训练损失一直在 6.9 附近不动(TINY 配置下约 6.7)。 这通常是数据管道的问题,而不是模型的问题:检查 x 和 y 是不是同一个张量(没右移)、词表是否与模型不匹配、shuffle 是否误设为 False 导致每轮都是同一批数据。

错误五:验证损失比训练损失还低。 在语言模型上这不一定是好事,通常是验证集划分太小或发生了数据泄漏(窗口重叠跨过切分点)。本书在文本层切分,并且验证集用 shuffle=False,就是为了避免这类假象。

2.9 本章小结

  • 词嵌入把离散符号映射成低维稠密向量,语义相近的词在训练中自动靠近;nn.Embedding 只是一张可训练的查表矩阵。
  • 三种分词粒度各有取舍,中文场景下字节级 BPE 是事实标准:无未知词、词表可控、高频片段被合并。
  • 字节级 BPE 的核心是"统计最高频相邻对并合并",词表大小 = 256 + 合并次数 + 特殊词元;训练与编码必须共享同一套合并顺序。
  • 滑动窗口把长文本切成输入-目标对,目标右移一位;训练/验证必须在文本层切分,避免重叠泄漏。
  • 位置嵌入不是可选项——自注意力对顺序不敏感,必须显式注入位置信息,且位置索引不能超过 context_length。
  • CPU 上 num_workers=0 是更快的选择,分词结果应缓存到磁盘。

练习

  1. 把 train_bpe_on_corpus 的 vocab_size 分别设为 500、800、2000、4000,各跑一次 python code/tokenizer.py,记录压缩率与语料词元总数,画出"词表大小 vs 压缩率"的趋势,并解释边际收益递减的原因。(思路:高频词对被优先合并,剩下的都是低频组合。)
  2. 把 max_chars 从 40000 降到 5000,观察 BPE 训练时间和压缩率如何变化,并解释为什么用更少数据训练出的分词器在完整语料上压缩率会下降。
  3. 修改 dataset.py 中的 stride(分别取 max_length、max_length//2、1),记录窗口数量,并说明为什么窗口数增加并不等于"训练更充分"。
  4. 故意把 create_dataloader_v1 的 tokenizer 换成词表更大的分词器,观察模型前向时报什么错,并写出两种修复方案。(思路:一是重建模型匹配新词表,二是在配置里同步更新 vocab_size。)

完整生成文档可登录: http://www.x01wq.cn 浏览文档部分。

posted on 2026-09-30 13:39  x01  阅读(6)  评论(0)    收藏  举报

导航