Step Law：联合预测预训练的 batch size 与学习率#

TL;DR#

本页内容7 个部分

Step Law：从多规模联合搜索拟合规则，给定模型大小与训练量后预测全局 token batch 和峰值学习率

1. 模型规模与训练量不能确定完整训练配置#

Chinchilla讨论怎样分配模型参数量与训练 token 数。确定这两个量后，训练配置仍未完整：同样训练 100B tokens，一次更新使用多少 tokens、每次使用多大的学习率，都会改变优化过程。B 表示十亿。

例如，全局 batch 为 1M tokens 时，100B tokens 大约对应 100,000 次更新；batch 为 2M 时，则只有约 50,000 次。这里 M 表示百万，假设没有跳步并忽略末尾不足额批次。不能只用相同 token 预算，认定两个训练过程等价。

Step Law 将这个选择写成一个预测任务：输入模型大小与训练量，联合预测全局 batch 和峰值学习率，减少在各个规模上分别进行完整超参数搜索的成本。它使用训练观测拟合经验关系。[1]

2. 由参数量与训练量预测超参数#

公式使用以下变量和单位：

量	含义	代入公式时的单位
$N$	按方法定义统计的非词表参数量，non-vocabulary parameters	实际参数个数
$D$	计划训练的 token 总数	实际 token 个数
$\hat\eta_{\max}$	预测的峰值学习率	学习率系数
$\hat B$	预测的一次更新使用的全局 batch	tokens / update

$\hat\eta_{\max}$ 和 $\hat B$ 分别是峰值学习率和全局 batch 的拟合预测值，不保证在新的训练设置下达到最低损失。 $N$ 不应直接用包含词表参数的模型总参数量替代；模型名称中的“1B”也不足以确定这里的输入。 $D$ 是完整训练预算，不是当前已经训练到的 token 数。[1]

在其他训练条件固定时，用 $\ell_{N,D}(\eta_{\max},B)$ 表示训练结束后的损失， $\eta_{\max}$ 是候选峰值学习率， $B$ 是候选全局 token batch。联合优化目标为：

(\eta_{\max}^{*},B^{*})=\underset{\eta_{\max},B}{\arg\min}\;\ell_{N,D}(\eta_{\max},B).

实际搜索只覆盖有限配置，最小值应理解为给定搜索范围和评价规则下的较优结果。Step Law 用以下拟合函数直接预测这对配置：

\begin{aligned} \hat\eta_{\max}(N,D)&=1.79N^{-0.713}D^{0.307},\\ \hat B(D)&=0.58D^{0.571}. \end{aligned}

系数与指数使用公开的式（1）和系数表； $N,D$ 都代入实际计数，不能把 1B、100B 直接写成数值 1、100。[1]

2.1 用同一配置计算峰值学习率#

假设非词表参数量为 $N=10^9$ ，训练预算为 $D=10^{11}$ ，即 1B 参数、100B tokens。以下是对公开公式的代入，不是新的训练实验：

\hat\eta_{\max}=1.79(10^9)^{-0.713}(10^{11})^{0.307}\approx1.6325\times10^{-3}.

这里得到的是整个日程的峰值，不是训练全过程恒定使用的学习率。给定日程后，还要计算各步的实际学习率；阶段安排与选择见学习率配方。

2.2 计算全局 token batch 与更新次数#

将同一训练预算 $D=10^{11}$ 代入 batch 预测公式：

\hat B=0.58(10^{11})^{0.571}\approx1{,}107{,}715\ \text{tokens / update}.

预测结果约为每次更新 1.108M tokens。忽略序列数取整和末尾不足额批次时，更新次数估计为：

\text{预计更新次数}=\frac{D}{\hat B}\approx\frac{10^{11}}{1{,}107{,}715}\approx90{,}276.

学习率与 batch 应作为一对候选检查。只采用预测学习率，却将 batch 改成显著不同的值，就改变了公式对应的配置。

2.3 模型或训练量改变时，预测怎样变化#

由同一组幂律直接计算：

非词表参数量 $N$	训练量 $D$	预测峰值学习率	预测全局 batch，tokens
1B	100B	$1.6325\times10^{-3}$	1,107,715
2B	100B	$9.9591\times10^{-4}$	1,107,715
1B	200B	$2.0196\times10^{-3}$	1,645,569

固定 $D$ ，参数量翻倍时，预测学习率的变化倍率为 $\text{学习率倍率}=2^{-0.713}\approx0.6101$ ；本拟合式中的 batch 不变。固定 $N$ ，训练量翻倍时，预测结果分别为 $\text{学习率倍率}=2^{0.307}\approx1.2371$ 与 $\text{batch 倍率}=2^{0.571}\approx1.4856$ 。

固定模型大小、增大计划训练量时，公式同时预测更高的峰值学习率和更大的 batch。 不能将它解释为固定 batch、改变训练长度时也必然应提高学习率，更不能据此在训练中随已消费 token 自动增大学习率。

3. 从联合搜索得到预测公式#

3.1 联合搜索的固定条件#

建立训练观测时，先固定一组 $N,D$ ，在不同的 $(\eta_{\max},B)$ 组合下从头训练，再比较最终损失。更换另一组 $N,D$ 后重复这一过程，才得到跨规模拟合需要的观测。

标定设置包括 AdamW、序列长度 2,048、2,000 步线性 warmup，以及随后衰减到固定 $10^{-5}$ 的 cosine 日程。预测值的含义依赖这套设置；不能只保留峰值、却任意更换终值规则。[1]

这里研究的是超参数与最终训练损失的关系。搜索中出现较宽的低损失区域，只说明所测配置对一定范围的超参数变化不敏感，不能据此认定任意训练设置下的搜索目标都全局凸。

3.2 为什么取对数后可以回归#

先采用一般幂律形式，令 $c,d>0$ 为比例系数， $a,b,g$ 为待拟合指数：

\hat\eta_{\max}=cN^aD^b,\qquad \hat B=dD^g.

对正值取自然对数，乘积变成加和，幂次变成乘法：

\begin{aligned} \log\hat\eta_{\max}&=\log c+a\log N+b\log D,\\ \log\hat B&=\log d+g\log D. \end{aligned}

将每组实验的较优配置取对数后，可以用线性回归拟合截距和指数，再通过指数运算得到原尺度上的预测值。官方拟合实现会保留接近组内最低 smooth loss 的配置，进行 bootstrap 重采样，并汇总拟合系数；不只从每组抽取一个最小点。Smooth loss 指平滑后的损失记录。[2]

例如，拟合出 $a=-0.713$ ，就表示在其他输入不变时， $N$ 扩大 $k$ 倍，预测学习率扩大 $k^{-0.713}$ 倍。这解释了参数翻倍算例中的 0.6101；指数来自观测拟合，不是由算力公式 $C\approx6ND$ 推出的常数。

4. 将 token batch 转成训练配置#

预测 batch 约为 1,107,715 tokens，但训练通常要求整数条序列，以及可实现的数据并行和累积配置。假设每条输入都有 2,048 个计数 token，没有 padding、尾批或跳步，则：

\text{预测全局序列数}=\frac{\hat B}{2{,}048}\approx540.88.

本算例选择接近预测值的 544 条全局序列，这一选择用于演示配置换算，并非官方工具的输出。设数据并行副本数为 $R$ ，每个副本每次前向处理 $m$ 条序列，累积 $K$ 次，每条长度为 $S$ 。全局 token batch 为：

B_{\mathrm{impl}}=RmKS.

例如，8 个数据并行副本、每个副本 microbatch 为 2 条、累积 34 次：

\begin{aligned} \text{全局序列数}&=8\times2\times34=544,\\ B_{\mathrm{impl}}&=544\times2{,}048=1{,}114{,}112\ \text{tokens / update},\\ \text{预计更新次数}&=\frac{10^{11}}{1{,}114{,}112}\approx89{,}758. \end{aligned}

每个副本一次只处理 2 条，不是一次处理 544 条。这里乘的是数据并行副本数，不是包含张量并行设备在内的 GPU 总数；计数与聚合详见工程基础。

这时 batch 比连续预测值约大 0.58%。它是接近连续预测值的可执行配置，取整后对应的最优学习率仍可能变化。实际训练还需配置完整日程，并检查邻近候选的验证损失；末尾不足额批次按实际消费的 token 记账。

5. 验证对象与方法局限性#

评价预测配置时，应比较训练结束后的损失。 文献在保留的 1B 参数、100B tokens 测试配置上，报告预测配置对应的 loss 与搜索得到的最优结果相差约 0.094%；这不是“学习率或 batch 的误差只有 0.094%”，也不是所有规模的统一保证。[1]

这个方法有三项与拟合设计直接相关的局限性：

完整日程会改变标定关系。 改成 WSD、不同 warmup 或不同终值后，原系数是否适用需要验证；预测峰值不能替代整个日程。
省略 $N$ 的 batch 模型是一项经验简化。 $\hat B=dD^g$ 不显式包含模型大小、硬件或通信成本。它不能直接回答哪个 batch 的端到端耗时最低，也不等于临界 batch 的定义。
预测依赖经验拟合的适用范围。 输入的参数计数、数据配方、优化器或规模超出已验证条件时，需要检查外推。实现取整后的配置也应作为新的候选评测。

使用 Step Law 为 scaling ladder 选择超参数时，先用同一计数口径确定每层的 $N,D$ ，再生成候选学习率与 batch，完成训练后记录实际配置和损失。这样才能区分规模变化与超参数选择造成的差异。

参考资料#

Predictable Scale: Part I, Step Law — Optimal Hyperparameter Scaling Law in Large Language Model Pre-training，Li 等，2025，arXiv:2503.04715v7。§3.1–3.4：目标、标定设置、式（1）、表 2 与验证；附录 A.1：单位与峰值学习率口径。本文使用式（1）及表 2 的负 $N$ 指数计算；§3.4.1 有一句对 $N$ 趋势的文字描述与该符号不一致。
Step Law 官方拟合实现，commit 32a8f4c。code/fit_tool.py：近优配置筛选、序列数到 token batch 的换算，以及 log 空间回归和 bootstrap。正文用公开的舍入系数手算，不承诺与工具中的完整系数逐位一致。