7. 训练循环：一次参数更新

有了分词器、数据和模型，下一步是把它们连成一个训练循环。

这一页不逐行解释 base_train.py，而是回答一个更实际的问题：从按下命令到模型参数发生变化，中间到底发生了什么？

启动命令先做了什么

scripts/base_train.py 在进入训练循环前，做了很多一次性准备：

1. 初始化计算设备和随机种子。
2. 决定使用 cuda、mps 还是 cpu。
3. 读取分词器，得到词表大小。
4. 根据 --depth 构建模型配置。
5. 在元设备（meta device）上创建模型，只分配形状，不分配真实数据。
6. 把模型移动到真实设备。
7. 初始化权重。
8. 根据模型规模自动计算训练预算、批量大小和学习率。
9. 创建优化器和数据加载器。
10. 可选地用 torch.compile 编译模型。

读模型创建和初始化scripts/base_train.py第 126–164 行↗

为什么先在 meta 设备上创建

大模型直接在 GPU 上创建会同时分配参数、梯度、优化器状态等多份内存，初始化过程本身就可能超出显存。

nanochat 先在 meta 设备上创建模型。这个设备只记录张量的形状和类型，不存真实数值。然后用 to_empty() 在目标设备上分配内存，再调用 init_weights() 填入初始值。

这样可以把“创建结构”和“分配真实内存”分开，避免初始化阶段的显存峰值。

深度决定训练预算

nanochat 的目标是让不同深度的模型都处在合适的训练配置上，而不是让用户手动调几十个超参数。

读自动超参数计算scripts/base_train.py第 249–316 行↗

主要规则如下：

1. 训练词元数

目标词元数 ≈ target_param_data_ratio × 可缩放参数量

默认比例是 12。runs/speedrun.sh 为了在更短时间内超过 GPT-2，使用 --target-param-data-ratio=8，也就是稍微少训练一些数据。

“可缩放参数量”不是所有参数，而是代码选定的 Transformer 矩阵参数加输出层参数。具体定义在 num_scaling_params() 中。

2. 批量大小

以 d12 为参考模型，最佳批量大小约为 2^19 = 524288 个词元。模型变大、训练词元变多时，批量大小按大约 D^0.383 增长，再取最接近的 2 的幂。

这是一个经验规律，不是必须遵守的物理定律。你仍然可以用 --total-batch-size 手动覆盖。

3. 学习率

批量大小改变后，学习率按平方根缩放：

学习率缩放 = √(当前批量 / 参考批量)

更大的批量通常可以使用更大的学习率。nanochat 对 AdamW 和 Muon 都使用这个近似，虽然 Muon 的缩放并没有被完全严格证明。

4. 权重衰减

权重衰减根据批量和训练词元数共同调整，目的是让训练过程中的有效正则化强度保持稳定。

这些自动计算让 --depth 成为主要的复杂度旋钮。代价是，当你手动改变批量或训练步数时，需要理解这些联动，不能只看一个参数。

多卡训练与梯度累积

一次参数更新使用的总词元数，等于：

每张卡的批大小 × 序列长度 × 卡数 × 梯度累积步数

例如：

device-batch-size = 8
max-seq-len = 256
GPU 数量 = 1
total-batch-size = 4096

则：

每张卡每次前向 = 8 × 256 = 2048 个词元
需要的梯度累积 = 4096 ÷ 2048 = 2

训练循环会连续做两次前向和反向传播，把梯度累加，然后才更新一次参数。这样可以在显存有限时模拟更大的批量。

nanochat 没有使用 PyTorch 的 DistributedDataParallel 包装器。它在优化器内部同步梯度，并把优化器状态分片到各张卡上。对初学者来说，先把它理解成“每张卡处理不同数据，梯度最终合在一起更新”就够了。

一次参数更新的时序

图 1：一次参数更新的时序。多个 micro step 的梯度先累加，最后才执行一次优化器更新。

读一次训练步scripts/base_train.py第 505–544 行↗

训练循环每一轮还会做这些事：

• 定期在验证集上计算 BPB。
• 定期计算 CORE 评分。
• 定期用模型生成样本文本，肉眼观察效果。
• 定期保存模型、优化器和训练状态。
• 记录吞吐量、显存和 MFU。

训练日志里看什么

一行日志可能包含很多数字：

step 00100/00200 | loss: 4.12 | lrm: 0.95 | dt: 120ms | tok/sec: 34,000 | mfu: 12.5

字段	含义
step	当前参数更新次数
loss	训练损失的平滑值
lrm	当前学习率相对初始学习率的倍数
dt	一次更新耗时
tok/sec	每秒处理多少词元
mfu	模型浮点运算利用率，衡量硬件利用效率

mfu 在 CPU 上没有意义，会是 0 或不可用。不要拿 CPU 的 mfu 和 GPU 比较。

训练结束会发生什么

训练循环会多跑一轮用于最后的评估和保存。最终检查点包含：

• 模型参数。
• 优化器状态。
• 模型配置和用户命令。
• 验证集 BPB。
• 数据加载器位置。
• 训练循环状态。

其中优化器状态由各张卡分别保存，所以你会看到带 rank 的文件名。恢复训练时必须找齐对应 rank 的文件。

模型保存到：

~/.cache/nanochat/base_checkpoints/<model-tag>/

下一阶段的 SFT 会从这里读取。

官方脚本怎样组合参数

runs/speedrun.sh 的核心命令是：

torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- \
  --depth=24 \
  --target-param-data-ratio=8 \
  --device-batch-size=16 \
  --fp8 \
  --run=$WANDB_RUN

它使用 8 张 H100、d24 模型、FP8 训练，并把数据参数比降到 8，以缩短达到 GPT-2 能力的时间。

CPU 版本则是：

python -m scripts.base_train \
  --depth=6 \
  --head-dim=64 \
  --window-pattern=L \
  --max-seq-len=512 \
  --device-batch-size=32 \
  --total-batch-size=16384 \
  --num-iterations=5000

两者的参数差别很大，但训练循环是同一套。先在小模型上理解流程，再换大规模参数，是最稳妥的顺序。

本章要记住的三件事

1. 训练开始前会先完成计算设备、分词器、模型、优化器和数据加载器的准备。
2. 一次参数更新可能包含多次前向和反向传播，梯度累积用于模拟更大的批量。
3. --depth 会联动模型宽度、训练词元数、批量大小和学习率；改它之前先看自动计算规则。

下一章解释更新参数时到底发生了什么：优化器。