20_scaling_law_viz.html

Scaling Law:用小模型预测 124M 的 loss

第 6 章花了 4090 一整天,把 124M 训了 10B token。能不能先训一批小模型(每个几分钟到两小时),算出大模型会训到多少? Scaling law 的做法:参数量 N 从 0.6M 到 25M,每个配几档训练 token 数 D,拟合 L(N, D) = E + A/N^α + B/D^β,再把第 6 章的 N 和 D 代进去。这次的预测是 3.36,实测 3.03,偏高 0.33;这一关把误差从哪来也一起查清楚。配套代码 phase5-fullstack/18_scaling.py,全部在 RTX 4090 上训练,和第 6 章同一套 GPT-2 结构、分词器与 FineWeb-Edu 数据。

STEP 1

一批小模型

STEP 2

N 和 D 两个方向

STEP 3

拟合一个公式

STEP 4

算力怎么分

STEP 5

外推到 124M

这一关 N 0.6M … 25M 非嵌入参数 · 5 档 D 50M … 1.6B 训练 token · 4 档 + 16 亿 检验 49M · 85M 不参加主拟合 目标 124M · 10B 第 6 章,N=85.1M batch 65,536 token / 步 评测 327,680 val token,所有点同一份 N 只数非嵌入参数(Kaplan 口径),124M 的非嵌入部分是 85.1M

① 一批小模型,各训几档 token

5 种宽度 / 深度,每种训 4 次:5000 万、1 亿、2 亿、4 亿 token;较大的三种再各训一次 16 亿。每次学习率都按自己的总步数从头做 warmup + 余弦退火。选一个模型,看它各次训练的 val loss 曲线(红色虚线是发散、不参加拟合的那次):

–

层数 × 宽度

–

N(非嵌入参数)

–

算上嵌入的总参数

–

训练合计用时

↳ 代码:18_scaling.py train(模型与第 6 章 04_gpt2_124m.py 相同,层数 / 宽度可调)· 网格 runs/queue_scaling.sh

为什么每档 D 都要从头训,不能训一次 4 亿、中途记 loss?

学习率的余弦退火要在训练结束时降到底,loss 才降到这个 D 能达到的水平。训 4 亿 token 的那次,在第 1 亿 token 时学习率还很高,loss 比专门训 1 亿、退火到底的那次要高。 Chinchilla 论文(Hoffmann et al., arXiv 2203.15556)指出 Kaplan 等人对所有模型用了同样的训练 token 数和学习率计划,认为这是两篇结论不同的原因之一;Chinchilla 自己让退火长度和训练 token 数匹配。图里 4 亿那条曲线在 1 亿处的值,可以和单独训 1 亿那条的终点比一比。

为什么 N 不算嵌入参数?

这批小模型的词表都是 50,304,嵌入表 50,304 × 宽度:宽度 128 的模型,嵌入占 –,Transformer 层只有 0.6M。算上嵌入,小模型的"大小"几乎全由词表决定。Kaplan et al.(arXiv 2001.08361)图 6 显示,只数非嵌入参数时,不同深宽比的模型落在同一条曲线上,所以用这个口径。 Chinchilla 数的是全部参数,两篇的系数因此不能直接互相套用。

有两次训练的 loss 后半程往上走,是怎么回事?

L4_d192 训 4 亿 token、L6_d256 训 16 亿 token 这两次,val loss 先降后升。原因是注意力分数失控,可以用 qk-norm 修好。下图是 L4_d192 · 4 亿 token 这次的原结构与加 qk-norm 对照(同配置、同随机种子):

原结构加 qk-norm

↳ 下一步:把各次训练的终点摆进一张 N × D 的表,看 loss 往哪个方向降得快。

② N 和 D 两个方向

每格是一次训练的最终 val loss,颜色越深越低。点一格,下面画出同一行(同样 N、D 变)和同一列(同样 D、N 变)的变化,横轴都是对数:

固定 N,D 变固定 D,N 变

为什么 D = 50M 那一列的 loss 都挤在一起?

–

↳ 下一步:两个方向各自按幂律下降,还有一个降不下去的底。把它写成一个公式来拟合。

③ 拟合 L(N, D) = E + A/N^α + B/D^β

E 是数据本身的不确定性,模型再大、数据再多也降不下去;A/N^α 是模型太小欠的账,B/D^β 是数据太少欠的账。按 Chinchilla §3.3,对 log L 做 Huber 回归、L-BFGS 求解,拟合用 21 个点(发散的两个不算)。拖两个滑块,看预测的 loss 由三部分怎么组成:

N(非嵌入参数)–

D(训练 token)–

E 数据的底 A/N^α 模型太小 B/D^β 数据太少

–

预测 loss

–

模型项

–

数据项

–

④ 算力怎么分:固定 C,N 和 D 此消彼长

训练算力约 C ≈ 6·N·D(每个 token 前向 2N、反向 4N 次浮点运算)。C 固定时,模型越大能训的 token 越少。拖 C,曲线是公式预测的 loss 随 N 的变化,最低点就是这个预算下的算力最优模型大小:

算力 C(FLOPs)–

–

最优 N

–

最优 D

–

D / N(每个参数几个 token)

–

这个预算下的最低 loss

↳ 代码:18_scaling.py fit 末尾:N_opt = G · (C/6)^(β/(α+β)),G = (αA / βB)^1/(α+β)

这里算出的 D/N 为什么远大于 Chinchilla 的 20?

–

第 6 章的 124M 按这个标准训多了还是少了?

–

↳ 下一步:把第 6 章的 N = 85.1M、D = 10B 代进公式,和那次真实训练对照。

⑤ 外推到 124M,和真实训练对照

下面几个点都没参加拟合,按离拟合范围的远近排:往 D 方向外推的 16 亿 token 训练、往 N 方向外推的 49M 与 124M 结构(4 亿 token),以及第 6 章训了 10B token 的 124M。灰色区间是 bootstrap 的 90% 范围(拟合点有放回地重抽、重新拟合 200 次)。先选用哪批点拟合,再点一行:

实测 − 预测(负数 = 实测比预测好)

为什么外推总是偏高,主拟合的 bootstrap 区间也没盖住实测?

–

第 6 章的 124M 和这批小模型的训练配方一样吗?

结构、分词器、数据、上下文长度、AdamW 设置、学习率峰值(宽度 768 按本章规则也是 6e-4)与余弦退火到 10% 都相同。不同的是两处:第 6 章每步 524,288 个 token(本章 65,536,8 倍), warmup 715 步(本章 100 步)。本章按同一配方训了 124M 结构 4 亿 token 的点,公式对它的预测也偏高(见表中「12 层 × 768」一行), 说明误差不全来自配方差异。要把配方的影响单独拆出来,得按本章配方把 124M 训满 10B token,在这块 4090 上约一天。

重评的 124M 为什么是 3.0301,不是第 6 章写的 3.02?

评测口径不同。第 6 章训练脚本每次只在 val shard 开头取 20 批 × 8 × 1024 = 163,840 个 token; 本章所有点(包括这个 124M)统一用开头 40 批 = 327,680 个 token,bf16 前向。拟合点和检验点必须用同一个口径,所以 124M 也按本章口径重评。

达标本关实测

↳ 跑法:bash runs/queue_scaling.sh(CUDA,4090 上约 7.1 小时)· python 18_scaling.py fit runs/s_*.json --exclude L4_d192_D400M L6_d256_D1600M --out runs/ch20_fit.json(CPU;加 --max-d 400e6 / --holdout 1e9 得到另两组拟合)

🎉 Scaling Law · 通关

你用一批小模型拟合出 L(N, D),外推到第 6 章训了一整天的 124M,看到外推越远误差越大、bootstrap 盖不住形状上的偏差,还排查了一次注意力分数失控的训练。下一关回到后训练:只给判分器,用 RL 教 124M 做加法。

① 一批小模型,各训几档 token

② N 和 D 两个方向

③ 拟合 L(N, D) = E + A/Nα + B/Dβ

④ 算力怎么分:固定 C,N 和 D 此消彼长

⑤ 外推到 124M,和真实训练对照

③ 拟合 L(N, D) = E + A/N^α + B/D^β