第 18 讲（拓展）：全流程实战 —— 从零训练一个会写故事的小模型

对应：CS336 作业 1 的完整流程（原课第 18、19 讲是嘉宾讲座，这里用一讲”全流程实战”代替）。 配套代码：code/lecture18/ 预计阅读时间：1.5 小时（加上训练时间：CPU 上 10~20 分钟） 硬件说明：CPU、Apple 芯片（MPS）、NVIDIA GPU 都可以。

0. 本讲导航

0.1 学完本讲，你应该能够

1. 把前面各讲的零件串成一条完整的流水线：数据 → 分词器 → 二进制 token 文件 → 规划 → 训练 → 评估与生成。
2. 在训练之前估算参数量、FLOPs、显存、用时，判断数据够不够。
3. 读懂（并能自己写出）一个”能用”的训练脚本：优化器分组、学习率调度、梯度裁剪、评估、检查点、断点续训、日志。
4. 用正确的指标评估一个语言模型，并和简单的基线比较。
5. 在训练出问题时，知道先检查哪些地方。

0.2 先修知识

几乎所有前面的讲次：第 1 讲（BPE）、第 2 讲（资源核算、训练循环、检查点）、第 3 讲（Transformer、稳定性）、第 9 讲（数据量与重复数据）、第 12 讲（困惑度与每字节比特数）、第 15 讲（批量生成）。

0.3 本讲地图

第1节 全流程地图：每个零件来自哪一讲
第2节 数据：语料、划分、分词器、二进制文件          ← 01_prepare_data.py
第3节 训练之前先算账                                ← 02_plan.py
第4节 训练脚本逐段讲解                              ← 03_train.py
第5节 评估与生成                                    ← 04_generate_and_eval.py
第6节 出了问题怎么办：调试清单
第7节 从这里出发：怎样放大
第8节 全课回顾

1. 全流程地图

CS336 作业 1 的最后一步，是把自己写的所有部件装在一起，在 TinyStories 上训练一个能写故事的小模型。这一讲就是这个流程的完整实现，而且每一个零件都来自前面的讲次：

这一讲的代码大部分是”胶水”：把这些零件按正确的顺序、正确的参数连起来。真实的工程工作很多就是这样的胶水——而 bug 也最常出现在这些连接处（第 6 节）。

2. 数据

2.1 语料

作业 1 用 TinyStories（Eldan & Li 2023）：几百万个用 GPT-3.5/4 写的、只用简单英文词汇的儿童故事。它的妙处在于：只有几百万参数的小模型，就能学会写出通顺、前后连贯的故事——在笔记本电脑上就能看到”语言模型学会语言”的全过程。

我们不下载它（如果你已经有了，可以用 --input），而是用 story_corpus.py 生成一个风格类似的语料：8 种情节（丢了玩具、交新朋友、下雨天、学本领、分享、克服害怕、种花、过生日），每种都有多种写法、对话、心情和结尾的道理，名字和物品随机组合。默认生成 5 万个故事（约 21 MB）。它比第 3 讲的合成故事丰富得多，但仍比真实的 TinyStories 规整——句子来自模板——所以小模型学得比较快。

2.2 按”故事”划分训练集和验证集

最后 2% 的故事作为验证集。按故事划分，而不是在整个文本的某个字节处切一刀：否则同一个故事的前半段在训练集、后半段在验证集，验证损失会被低估（一种轻微的”泄漏”，第 12 讲）。

2.3 分词器

用第 1 讲的 BPE，在训练集开头的 2 MB 上训练，<|endoftext|> 作为特殊 token。大语料上只用一部分文本训练分词器就够了（常见的合并在一小部分数据里就能学到）；作业 1 在 TinyStories 上用 1 万的词表，这里语料的用词少得多，1024 就够了（01 号文件会提示合并”用完了”的情况）。

在我的运行中：词表 1024，平均每个 token 约 4.1 个字节（第 1 讲：GPT-2 分词器在英文上也约 4 字节 / token）；最常见的 token 是 ”.”、”,”、” a”、” and”、” was”、” the”……；训练集约 500 万个 token。

2.4 uint16 二进制文件

把整个语料编码一次，token 编号存成 uint16（词表 < 65536，每个 token 2 字节）。训练时用 np.memmap 打开，不把整个文件读进内存——即使文件有几百 GB 也没关系（第 2 讲）。一个常见的坑：写和读的 dtype 必须一致，二进制文件本身不记录类型，读错了不会报错，只会得到一堆错误的数字。

3. 训练之前先算账

02_plan.py 用第 2、3、9 讲的方法，在训练前把账算清楚。默认的 cpu 预设（d_model 192、4 层、6 个头、SwiGLU 隐藏维度 512、上下文 256、批次 32、3000 步）在我的电脑上：

几点观察：

• 小模型里嵌入层占的比例很大（第 3 讲）；
• 激活值远大于参数和优化器状态：批次 32 × 长度 256 时约 950 MB，和批次大小成正比。显存不够时，先减小批次，或者用梯度累积（第 7 讲）；
• 实测吞吐量因机器差别很大（你的数字很可能不同）；实测比估算更可靠；
• 数据：约 4.9 轮，略超过”4 轮以内接近新数据”的经验（第 9 讲）；每个参数只看 11 个 token，按 Chinchilla 的标准还没训练够。想练得更好，就同时加步数和数据（--num_stories）。

先算账、再训练的习惯，在大规模训练中价值千金：一次训练动辄上百万美元，没有人会在不知道”需要多少显存、要跑多久”的情况下按下开始键。

4. 训练脚本逐段讲解

03_train.py 的主循环只有十几行，但每一行都有讲究：

lr = lr_cosine_schedule(step, lr_max, lr_min, warmup_steps, max_steps)   # 1. 学习率
for g in opt.param_groups: g["lr"] = lr
x, y = get_batch(train_data, batch_size, context_length, device, rng)    # 2. 取批次
logits = model(x)                                                         # 3. 前向
loss = F.cross_entropy(logits.reshape(-1, V), y.reshape(-1))
opt.zero_grad(set_to_none=True)                                           # 4. 清空梯度
loss.backward()                                                           # 5. 反向
grad_norm = clip_grad_norm_(model.parameters(), grad_clip)                # 6. 梯度裁剪
opt.step()                                                                # 7. 更新

优化器：AdamW，β = (0.9, 0.95)（比默认的 0.999 更”健忘”，大模型训练中更稳定）。权重衰减分组：矩阵（线性层、嵌入）做权重衰减 0.1，RMSNorm 的增益向量不做——把增益往 0 拉没有道理。

学习率：线性预热 100 步到 3e-3，再余弦衰减到 3e-4。小模型能承受比大模型大得多的学习率（第 3、11 讲：最优学习率随宽度变小）；预热让 Adam 的二阶矩估计先稳定下来，避免一开始就迈大步。

梯度裁剪：总范数超过 1.0 就等比例缩小。日志里记录了裁剪前的梯度范数——它是判断训练是否健康的重要信号（第 6 节）。

评估：每 250 步在验证集上用固定的 20 个批次估计损失。固定批次，曲线上的起伏就只来自模型的变化，而不是”这次抽到的题比较难”。

检查点：保存模型、优化器状态（Adam 的一阶、二阶矩！只存模型的话，恢复后 Adam 要重新”热身”）、步数（学习率调度要用）、随机数状态（恢复后抽到的批次和不中断时完全一样）。--resume 从最新的检查点接着训练；另外单独保存验证损失最好的检查点。

日志：每一步写一行 CSV：步数、学习率、训练损失、验证损失、梯度范数、token 数、用时。记录一切——训练出问题时，日志是唯一的线索。

混合精度与编译（可选）：--amp 在 CUDA 上用 BF16 做矩阵乘法（第 2 讲，损失仍用 FP32 计算）；--compile 用 torch.compile 融合算子（第 6 讲）。

训练时应该看到什么：损失在前几十步迅速下降（先学会最常见的 token 和标点），然后变慢；验证损失和训练损失接近（这个语料的模板很多，过拟合要到很多轮之后才明显）；梯度范数在预热期间可能较大，之后稳定在一个较小的范围。在我的冒烟测试里（1 万个故事的数据、只训练 60 步），验证损失已经从约 7 降到 1.04 nat/token。

5. 评估与生成

04_generate_and_eval.py 做四件事：

1. 训练曲线：从 log.csv 画出验证损失和学习率（字符画）；
2. 整个验证集上的损失：用互不重叠的窗口计算，报告 nat/token、比特/token、比特/字节（第 12 讲：不同的分词器之间只有”每字节”可以比较）和困惑度；
3. 基线：一元模型（只看每个 token 的频率）和二元模型（只看前一个 token）。Transformer 比二元模型好出来的部分，来自更长的上下文——记住主角的名字、情节的顺序、引号的配对。在冒烟测试中（60 步），二元模型是 1.85 nat/token，Transformer 已经是 1.04；
4. 生成：用第 15 讲的批量生成（KV 缓存），在温度 0.5、0.8、1.2 下各写几个故事。

损失低，不等于会写故事。冒烟测试里只训练了 60 步，验证损失已经远低于二元模型，可写出来的”故事”还是一团乱麻（名字前后不一、句子拼接错乱）——这个语料很规整，很多 token 很好猜（例如 “Once upon a” 后面几乎一定是 ” time”），平均损失被这些”送分题”拉低了。真正的难点（名字前后一致、情节连贯、正常结束）只占很小一部分 token，要训练得更久才能学会。所以一定要亲眼看生成的样本：

• 主角的名字是不是前后一致？
• 情节是不是连贯（丢了玩具 → 找 → 找到）？
• 对话的引号是不是配对？
• 是不是正常结束（生成了 <|endoftext|>）？
• 温度低时更保守、更模板化；温度高时更多样、更容易出错。

完整训练 3000 步之后你会看到什么，以你自己运行的结果为准。

6. 出了问题怎么办：调试清单

训练不收敛、损失是 NaN、生成的是乱码……按这个顺序检查：

1. 先在很小的数据上过拟合：用一个批次反复训练，损失应该能降到接近 0。降不下去，说明模型、损失或优化器有 bug（第 2 讲）；
2. 数据：解码几个训练批次看看（tokenizer.decode(x[0])）——是正常的文本吗？输入和目标是不是恰好错开一位？memmap 的 dtype 对吗？
3. 损失的初始值：随机初始化时，损失应该约等于 ln(词表大小)（1024 → 6.93）。远大于它，说明初始化有问题（第 2 讲）；
4. 学习率：损失忽上忽下或变成 NaN，先把学习率减小 3~10 倍；看梯度范数是不是突然暴涨（第 3 讲的稳定性技巧：更小的学习率、QK-Norm、z-loss、更长的预热）；
5. 评估：评估时有没有 model.eval() 和 torch.no_grad()？验证集和训练集是不是用同一个分词器编码的？
6. 生成：生成时用的分词器和训练时是同一个吗？特殊 token 的编号对吗？提示词的编码方式（开头有没有空格）和训练数据一致吗？
7. 断点续训：恢复后的第一个损失应该和中断前接近；如果突然变大，多半是优化器状态或学习率调度没恢复。

7. 从这里出发：怎样放大

• 更大的数据：真实的 TinyStories（--input），然后是 OpenWebText 这样的网页数据（作业 1 的第二个数据集）——难度会上一个台阶；
• 更大的模型、GPU：--preset gpu --amp；显存不够时用梯度累积；多张卡时用数据并行（第 7 讲的 DDP）或 FSDP；
• 更好的超参数：在小模型上扫学习率和批次大小，按第 11 讲的方法（μP、规模定律）迁移到大模型；
• 更严格的评估：除了验证损失，还可以看下游任务（第 12 讲）；
• 后训练：在训练好的模型上做 SFT，让它听指令写故事（第 15 讲）。

8. 全课回顾

回到第 1 讲提出的问题：给定固定的计算和数据预算，你能训练出的最好的模型是什么？ 这门课从头到尾都在回答它：

1. 基础（第 1~4 讲）：分词、资源核算、Transformer 的每一个部件，以及注意力的替代方案与 MoE；
2. 系统（第 5~8、10 讲）：GPU 的工作方式、核函数、并行、推理——让每一个 FLOP 都算数；
3. 规模定律（第 9、11 讲）：用小实验预测大模型，选择模型大小、数据量、超参数；
4. 数据（第 12~14 讲）：怎样评估，数据从哪里来，怎样过滤、去重、混合；
5. 对齐（第 15~17 讲）：SFT、RLHF、DPO、RLVR，以及多模态；
6. 全流程（本讲）：把它们装在一起。

贯穿始终的几条经验：

• 先算账：FLOPs、显存、数据量，在动手之前心里有数；
• 小规模实验 → 外推：规模定律、μP、数据配比、“模拟重复轮数”，都是同一个思想；
• 数据决定一切：从 BPE 的训练语料到偏好数据的标注者，模型学到的就是数据里的东西；
• 代理指标会骗人：困惑度低不等于会写故事，奖励模型分数高不等于用户满意，训练指标正常不等于模型没坏——一定要看真实的样本。

9. 练习

练习 1（初始损失）：词表大小为 1024 的模型随机初始化时，损失大约是多少？如果你看到初始损失是 15，可能出了什么问题？

随机初始化时输出接近均匀分布，损失 ≈ ln(1024) ≈ 6.93。初始损失 15 说明输出的 logits 数值太大（分布很”尖”而且是错的）：可能是输出层初始化的标准差太大、缺了最后的归一化层、或者嵌入的尺度不对（第 2 讲的初始化）。

练习 2（数据轮数）：训练集有 500 万个 token，批次 32、长度 256，要训练 1 万步，一共是几轮？按第 9 讲的经验，这样训练合适吗？怎样改进？

10000 × 32 × 256 ≈ 8200 万个 token，约 16 轮。远超”4 轮以内接近新数据”的经验，后面的训练收益会迅速递减，还可能过拟合（验证损失开始上升）。改进：生成更多数据（--num_stories），或减少步数，或换更小的模型。

练习 3（显存）：把批次从 32 加到 128、上下文从 256 加到 512，激活值显存大约变成多少倍？（提示：激活值中既有 ∝ T 的项，也有 ∝ T² 的注意力分数。）

批次 ×4；∝ T 的项 ×2，∝ T² 的项 ×4。总的来说在 ×8（全是 ∝ T 的项）和 ×16（全是 ∝ T² 的项）之间，用 02_plan.py --batch_size 128 --context_length 512 可以直接算出来。所以加长上下文比加大批次更”费”显存；FlashAttention 可以避免存下 T×T 的注意力矩阵（第 5 讲）。

练习 4（每字节比特数）：模型 A 用词表 1024 的分词器，验证损失 1.2 nat/token，平均 4.1 字节/token；模型 B 用字节级分词（每个 token 1 字节），验证损失 0.35 nat/token。哪个更好？

换算成比特/字节：A = 1.2 / ln2 / 4.1 ≈ 0.42；B = 0.35 / ln2 / 1 ≈ 0.50。A 更好。每 token 的损失不能直接比较，因为不同分词器的”一个 token”包含的信息量不同（第 12 讲）。

练习 5（断点续训）：如果检查点里只保存了模型参数，没有保存优化器状态和步数，恢复训练时会发生什么？

没有步数：学习率调度会从头开始（重新预热、回到最大学习率），相当于一次”学习率重启”，损失可能先上升。没有优化器状态：Adam 的一阶、二阶矩从 0 开始，前几步的更新方向和大小都不准（偏差修正会让最初几步的步长接近 lr，而不是按历史梯度缩放），可能造成一次损失的跳升。没有随机数状态：抽到的批次和不中断时不同（结果不可复现，但通常不影响效果）。

练习 6（动手）： (a) 学习率扫描：--lr_max 取 1e-3、3e-3、1e-2，各训练 1000 步（用不同的 --run_name），比较验证损失； (b) 故意制造一个 bug：在 get_batch 里把目标写成和输入一样（不错开），训练损失会怎样？生成会怎样？ (c) 把生成换成 top-p 采样（p = 0.9），和纯温度采样比较故事的质量。

10. 延伸阅读

• Eldan & Li，TinyStories: How Small Can Language Models Be and Still Speak Coherent English?（2023）。
• CS336 作业 1 的说明（stanford-cs336/assignment1-basics）：从 BPE 到训练循环的完整规格。
• Andrej Karpathy 的 nanoGPT、llm.c：极简但完整的训练代码，很适合对照阅读。
• Kaplan et al.（2020）、Hoffmann et al.（2022）：模型大小与数据量怎样搭配（第 9 讲）。
• 本课程的其他讲义：每一个零件的”为什么”都在对应的讲次里。