模型规模、数据规模与计算预算

假设你有一笔固定的计算预算，可以训练一个小模型很久，也可以训练一个大模型一会儿。应该选哪一个？

第二十一章把模型宽度从 32 增加到 64，两者都训练了 40 轮。这个比较告诉我们“多花计算后发生了什么”，却没有回答“相同计算预算该怎样花”。要回答后一个问题，得先把账算清楚。

这一章会碰到几个很大的数字。不用被零的个数吓住：我们仍从那个只有两层的小 Transformer 出发，一项项数参数、数矩阵乘法，再看看研究者怎样从大量实验中寻找规模规律。

1. 四个容易被混在一起的数量

数量	它在问什么	本书小实验里的例子
参数量 $N$	有多少个可学习的数	宽度 32 的模型有 28,480 个
训练词元量 $D$	训练过程共读了多少词元	同一句重复读十次，也会累计十次
计算量 $C$	执行了多少计算	本章统计训练中的密集矩阵乘法
内存占用	同时要把多少内容放在内存里	权重之外，还有梯度、优化器状态和中间结果

它们彼此有关，但不能互相替代。参数少的模型可能读了很多数据；同样多的训练词元，长上下文可能更贵；模型文件能放进显存，也不代表训练时能放得下。

还有一个常被省略的词：唯一。192 条不同文本重复抽样 6,400 次，累计见过很多样本，并没有变成 6,400 条独立的新材料。报告里最好同时给出可用数据量、实际不同样本数和累计处理量。

2. 数参数：宽度翻倍，为什么不只贵一倍

沿用第二十章的网络。词表大小为 $V$ ，宽度为 $d$ ，最大位置数为 $P$ ，有 $L$ 个 Transformer 块。输入嵌入和输出投影不共享参数，前馈隐藏层宽度为 $4d$ 。

从几个最熟悉的矩阵开始：

部分	参数量
输入词元嵌入	$Vd$
位置嵌入	$Pd$
每块的 Q、K、V 与输出投影	$4d^2$
每块的两层前馈权重	$8d^2$
每块的前馈偏置和两个 LayerNorm	$9d$
最后的 LayerNorm	$2d$
输出词表投影	$dV$

这里的注意力投影没有偏置，前馈层有偏置；每个 LayerNorm 有一组缩放和一组平移。把它们加起来：

$\begin{aligned} N={}&2Vd+Pd\\ &+L(12d^2+9d)+2d \end{aligned}$

代入 $V=35$ 、 $P=32$ 、 $L=2$ ：

宽度	参数量	相对宽度 32
32	28,480	1 倍
64	106,112	约 3.73 倍

网络里有很多 $d\times d$ 大小的矩阵。宽和高都翻倍，元素数就接近四倍；嵌入、偏置等项只随 $d$ 线性增长，所以总量没有恰好乘四。

这个公式属于当前实现。换成共享嵌入、门控前馈层、不同归一化或不同注意力结构，就需要重新列账，不能拿参数量公式当成 Transformer 的固定定义。

3. 数计算：一次矩阵乘法要忙多久

先说清楚计数约定

一个形状为 $(m,k)$ 的矩阵乘以 $(k,n)$ 的矩阵，会得到 $mn$ 个输出，每项需要累加 $k$ 次乘积。常见的估算把一次乘法和一次加法计为 2 次浮点运算：

$C_{\mathrm{matmul}}\approx 2mkn$

FLOPs 在这里表示浮点运算的次数；FLOP/s 才是每秒运算速度。硬件的峰值速度不等于训练过程的持续速度：数据搬运、算子调度和等待都会花时间。

下面采用上述 $2mkn$ 约定，只统计密集矩阵乘法。它适合比较这个小网络的主要计算，但不是完整的硬件工作量账单。

给小模型算一遍

一个批次有 $B$ 个块，每块长度为 $T$ 。在一层中：

Q、K、V 和输出投影合计约 $8BTd^2$ 。
两次前馈投影合计约 $16BTd^2$ 。
注意力的 QK 与权重乘 V 合计约 $4BT^2d$ 。

固定头数后，各头维度加起来仍是 $d$ 。再加上最后的词表投影，前向矩阵乘法量为：

$\begin{aligned} C_{\mathrm{forward}}={}&L(24BTd^2+4BT^2d)\\ &+2BTdV \end{aligned}$

本实现先计算完整的注意力分数矩阵，再做因果遮挡，因此仍按完整的 $T^2$ 计数。掩掉上三角，并没有让已经发生的矩阵乘法消失。

对这里参与求导的矩阵乘法，反向还要分别求两个矩阵输入的梯度，大致是前向的两倍。按当前网络的运算路径，前向加反向的矩阵乘法账本为：

$C_{\mathrm{step}}=3C_{\mathrm{forward}}$

本书也用 PyTorch 的运算记录核对了 mm、bmm 与 addmm 三类操作，得到同样的计数。这个等式限定在当前实现与计数范围内：Softmax、归一化、GELU、损失、偏置加法、优化器更新、数据处理、验证以及内存访问都没有计入。

取 $B=16$ 、 $T=32$ ，得到：

宽度	每次更新的矩阵 FLOPs
32	91,521,024
64	334,036,992

固定预算为小模型更新 400 次的计算量：

$C_{\max}=36{,}608{,}409{,}600$

约为 36.608G 矩阵 FLOPs，其中 G 表示 $10^9$ 。加宽模型最多完整更新 109 次，使用约 99.46% 的上限；剩余预算不足以再做一步。

同一计算预算下，两种宽度模型每步花费与可执行更新次数的对比。

图 1：宽度增大，每一步更贵，留给更新次数的空间就更少。本图统计矩阵乘法，未把验证等成本包含在内。

PAD 虽然不算损失，仍然可能花计算

本例每条文档有 26 个字符，加上 EOS，一共预测 27 个真实目标，但网络每次处理 32 个位置。小模型更新 400 次，因此：

$\begin{aligned} D_{\mathrm{positions}}&=400\times16\times32\\ &=204{,}800\\ D_{\mathrm{targets}}&=400\times16\times27\\ &=172{,}800 \end{aligned}$

这两个数用途不同：前者帮助计算实际矩阵尺寸，后者是损失的有效目标量。忽略 PAD 的损失，并不会自动省去 PAD 位置上的网络运算。这也是上一章讨论 packing 的一个理由。

4. 常见近似： $C\approx6ND$

讨论较大的密集语言模型时，经常见到：

$C\approx6ND$

它把训练中的主要矩阵计算粗略看成：每个训练词元经过约 $N$ 个参数，前向约 $2N$ 次运算，前向加反向约 $6N$ 次。

这个近似很方便，但 $N$ 的具体口径——例如是否排除某些嵌入参数——要跟所用模型和文献说明一致。注意力的序列长度相关项、输出层相对大小、填充以及其他运算，也可能让估算偏离实际。

拿一个纯假设的例子：7B 参数、140B 训练词元，B 表示十亿。按此近似：

$\begin{aligned} C&\approx6\times(7\times10^9)\\ &\quad\times(140\times10^9)\\ &=5.88\times10^{21}\ \text{FLOPs} \end{aligned}$

这个数字不是训练耗时。若要估计时间，还得知道具体设备和软件栈能持续达到怎样的吞吐，以及通信、检查点、验证等额外时间。也不要把 140B 自动理解为 140B 个从未重复的词元。

下一章的实验不用这个粗估来分配预算，而采用上一节与实现对应的矩阵计算公式。

5. 内存：权重文件只是其中一项

模型推理时，如果权重使用 BF16，每个参数约 2 字节。7B 参数的权重本身约为 14GB，按十进制 GB 计算。实际推理还需要激活、缓存等空间，不能据此断言 14GB 的设备一定够用。

训练还要存更多状态。一个便于理解的 FP32 Adam 账本是：

内容	每个参数对应的字节数
FP32 权重	4
FP32 梯度	4
FP32 一阶估计	4
FP32 二阶估计	4
合计	16

于是同一个 7B 参数模型，仅这四项就约 112GB，仍未包括激活和临时工作区。

有些混合精度实现使用 2 字节的参数、2 字节的梯度，再保留 4 字节主权重和两个 4 字节优化器状态，也会得到每参数 16 字节。但实现可以不同，不能只看“用了 BF16”就认定全部状态都缩成一半。

我们的小实验为了数值核对使用双精度，它的实际状态内存也不能直接套用 FP32 表格。脚本打印的 FP32 内存是指定假设下的账本。

另一个大头是中间结果。朴素注意力会存放形状约为 $(B,h,T,T)$ 的分数或权重，其中 $h$ 是头数。长度翻倍，这部分元素数会乘四。后面系统篇会讨论更节省存储的注意力实现、重计算和分布式方法；现在先记住：参数量相同，训练内存也可以因批次和序列长度差很多。

6. 缩放规律：用小实验猜大实验，但先讲条件

研究者会训练许多不同规模的模型，观察验证损失怎样随参数、数据和计算量改变。在一定范围内，变化可以用较平滑的函数描述。这类经验关系通常称为缩放规律。

Kaplan 等人在 2020 年的研究系统考察了语言模型损失与模型规模、数据规模和计算量的关系。Hoffmann 等人在 2022 年的 Chinchilla 研究则重新估计固定训练计算量下的模型与数据分配，提出在他们考察的条件下，两者应大致同步扩大。

一个常用于帮助理解的拟合形式是：

$\mathcal L(N,D) =\mathcal L_\infty+\frac{A}{N^\alpha}+\frac{B_0}{D^\beta}$

$A$ 、 $B_0$ 、 $\alpha$ 、 $\beta$ 都是需要从实验估计的正数； $\mathcal L_\infty$ 是拟合中的极限项，不是已经被证明的语言本质常数。用 $B_0$ 是为了与前面的批次大小 $B$ 区分。

这条式子表达一种直觉：增加参数可以降低一部分损失，增加数据可以降低另一部分损失，但收益逐渐变小。

如果固定 $C\approx6ND$ ，那么：

$D\approx\frac{C}{6N}$

模型越大，预算能支持的训练词元就越少。代回拟合式，模型不足与数据不足可能在中间达到一个折中：

在人为设定的示意损失函数中，固定预算下过小或过大的模型都不占优，中间出现折中点。

图 2：使用无量纲的假设函数绘制，取 $ND=1$ 、极限项为 1、两个系数均为 0.5、两个指数均为 0.5。它只展示折中关系，不是论文拟合曲线，也不是本书小模型测得的缩放规律。

这不等于存在一个随时适用的“每个参数必须配多少词元”。最佳分配依赖数据、架构、训练计划、损失定义和预算范围；训练计算最省的方案，也不一定让长期部署成本最低。若模型要服务大量请求，较小模型多训练一些，可能换来更低的推理成本。这需要把预期调用量也列入目标，而不是只优化一次训练。

7. 外推：曲线很顺，不代表永远顺

把小实验拟合出的曲线延伸到远处，叫外推。它有吸引力：大实验昂贵，最好能提前判断方向。风险也在这里：延伸后的区域还没有被观察过。

几种变化尤其值得警惕：

小实验还有许多新数据可读，大实验却只能反复读同一批内容。
小模型使用短上下文，大模型切换到长上下文或不同注意力结构。
数据过滤、语言比例或分词器变了，损失的比较口径也随之改变。
下游能力存在明显门槛，平均预测损失的平滑改善没有直接变成任务成功。
评测材料与训练文本重叠，使损失看起来比实际泛化更好。

因此，缩放关系适合帮助选择下一组值得尝试的实验，也需要在新规模上检查。我们下一章的三种训练种子、几万参数和模板短句，只能支持这个小任务中的比较；即使画出很平滑的曲线，也不能宣称复现了大模型缩放定律。

8. 动手算一份预算

compute_budget.py 只依赖 Python 标准库：

python3 src/examples/compute_budget.py

每次训练预算上限：36,608,409,600 矩阵 FLOPs
宽度 32：参数=28,480；每步=91,521,024；更新=400
  预算使用=100.00%；处理位置=204,800；真实目标=172,800
  假设 FP32 参数/梯度/Adam 状态：0.435 MiB
宽度 64：参数=106,112；每步=334,036,992；更新=109
  预算使用=99.46%；处理位置=55,808；真实目标=47,088
  假设 FP32 参数/梯度/Adam 状态：1.619 MiB
另一个假设例子：7B 参数、140B 训练词元，6ND≈5.88e+21 FLOPs
同一个 7B 模型：仅 BF16 权重≈14.0 GB；按每参数 16 字节的训练状态≈112.0 GB

MiB 用 $2^{20}$ 字节计算，GB 用 $10^9$ 字节计算。比较设备容量和程序报告时，要留意单位是否相同。

展开完整代码：参数、矩阵计算与状态内存

"""第 24 章：TinyTransformer 的参数、矩阵乘法预算和状态内存账本。"""


def parameters(width, vocab=35, context=32, layers=2):
    # 输入嵌入与输出投影不共享；注意力投影无偏置，前馈层有偏置。
    return 2*vocab*width + context*width + layers*(12*width**2+9*width) + 2*width


def matrix_flops(width, batch=16, length=32, vocab=35, layers=2):
    # 一乘一加计 2 FLOPs；只统计密集矩阵乘法，不计逐元素运算。
    forward = layers*(24*batch*length*width**2 + 4*batch*length**2*width)
    forward += 2*batch*length*width*vocab
    return 3*forward  # 前向一次，加两个矩阵输入各自的反向。


BUDGET = 400*matrix_flops(32)


def main():
    print(f"每次训练预算上限：{BUDGET:,} 矩阵 FLOPs")
    for width in [32, 64]:
        count, cost = parameters(width), matrix_flops(width)
        steps = BUDGET//cost
        print(f"宽度 {width}：参数={count:,}；每步={cost:,}；更新={steps}")
        print(f"  预算使用={steps*cost/BUDGET:.2%}；"
              f"处理位置={steps*16*32:,}；真实目标={steps*16*27:,}")
        print(f"  假设 FP32 参数/梯度/Adam 状态：{16*count/1024**2:.3f} MiB")
    n, tokens = 7_000_000_000, 140_000_000_000
    print(f"另一个假设例子：7B 参数、140B 训练词元，6ND≈{6*n*tokens:.2e} FLOPs")
    print(f"同一个 7B 模型：仅 BF16 权重≈{2*n/10**9:.1f} GB；"
          f"按每参数 16 字节的训练状态≈{16*n/10**9:.1f} GB")


if __name__ == "__main__":
    main()

9. 停下来试一试

模型宽度翻倍，为什么参数量通常增长超过两倍？
如果预算允许 109.59 次更新，为什么脚本只执行 109 次？
同样的网络与序列长度，把批次大小翻倍，单步矩阵计算量如何变化？固定预算下更新次数呢？
把序列长度翻倍，哪些矩阵计算项乘二，哪些乘四？
每条文档有 27 个有效目标，却处理 32 个位置。为什么不能把公式中的 $T$ 直接改成 27？
一个模型的 BF16 权重有 14GB，为什么不能直接说它只需要 14GB 的训练显存？
同一批数据重复十次，与十倍不同数据，在 $D$ 的累计计数上可能相同。为什么学习效果未必相同？
为下一章预算表加入验证、检查点和墙钟时间。如果按实际时间比较，需要额外控制哪些条件？

参考答案与提示

多个参数矩阵的两边都与宽度有关，所以有 $d^2$ 项。
一次完整更新是本实验的最小分配单位，不能超过预算。要采用不同大小的最后一批，需要重新明确计算与优化协议。
单步计算量乘二；完整更新次数大约减半，整数取整可能留下一点预算。
含 $BTd^2$ 、 $BTdV$ 的项乘二；含 $BT^2d$ 的注意力项乘四。
实现仍在完整 32 个位置上做密集矩阵计算。排除损失与省掉计算不是同一操作。
训练还有梯度、优化器状态、激活、临时工作区等；14GB 只说明了指定精度下的权重大小。
新信息量、重复程度、覆盖范围与记忆风险不同，累计处理次数无法独自说明这些性质。
固定设备、线程数、精度和软件版本，避免混入其他负载，分别报告训练与评估时间，必要时重复计时。算术量相近也可能因硬件利用率不同而耗时不同。

10. 继续阅读

Kaplan 等，Scaling Laws for Neural Language Models：看研究者怎样组织规模实验并拟合损失趋势。
Hoffmann 等，Training Compute-Optimal Large Language Models：固定训练计算量下的模型与数据分配。
PyTorch Profiler 文档：了解性能记录与 FLOPs 估计的覆盖范围。记录到的部分操作计数不能自动代表整机工作量。

预算已经列好。下一章让模型真正花掉这笔预算，并检查：更大、更多、质量更好，分别带来了什么。

Keyboard shortcuts

从机器学习到大模型：原理、实践与前沿