1.5 从零训一个 Transformer:将组件组织为可训练的模型#

前面几课介绍的 Token 与嵌入注意力与 Transformer 是构成语言模型的独立组件。

本节将它们组织为一个可稳定训练、长时间运行而不发散的语言模型,并给出一套业界通行的工程配置,这也是 Stanford CS336 作业 1从零实现语言模型的完整内容[^cs336]。

核心问题

  • 能解释训练循环中前向、交叉熵、反向传播、优化器更新与 zero_grad 的作用。
  • 能推导 AdamW 的更新公式,并说明解耦权重衰减与 L2 正则化的差异。
  • 能选择预热与余弦衰减参数,说明峰值学习率和训练后期步长的依据。
  • 能区分梯度裁剪、bf16、fp16 loss scaling 与参数初始化所抑制的失稳类型。
  • 能根据 batch、序列长度与训练 token 预算组织数据,并排查 loss 尖峰及检查点恢复问题。

一、训练循环:四步结构#

语言模型训练反复执行预测、度量误差、求梯度与更新参数,使随机初始化的权重逐步刻画语言统计规律。每一步的对象与作用,以及更新后的参数如何回到下一批前向计算,集中表示在图 1中;后续工程问题都围绕这一闭环在十亿参数和万亿 token 规模下能否稳定持续展开。

交叉熵损失的定义如下,对一个长为 $T$ 的序列,取模型赋予真实下一个 token 的对数概率之负值并求平均:

$$\mathcal{L}(\theta)=-\frac{1}{T}\sum_{t=1}^{T}\log p_\theta\big(x_t \mid x_{\lt t}\big)$$

该损失度量模型赋予正确 token 的概率与确定性 1 之间的差距:模型对正确 token 越确信(概率越接近 1),$-\log p$ 越接近 0,损失越小;反之,若模型只给正确 token 分配 0.001 的概率,该位置的损失即为 $-\log 0.001\approx 6.9$,惩罚显著。

训练的整体效果,就是在大量 token 上持续降低这一平均负对数概率。

将 $\mathcal{L}$ 换底为 2 后取指数,即得 1.4 · 资源核算 中的困惑度(perplexity),其数值等于模型在每个位置上等效面对的候选数量,困惑度 20 表示模型在每个位置上的不确定性相当于在 20 个等概率候选之间选择。

CS336 作业 1 的 OpenWebText 榜单排序依据的正是该指标1

概率分布的来源需要明确。模型最后一层输出的并非概率,而是词表中每一项对应一个实数分值 $z\in\mathbb{R}^{V}$(称为 logits)。要将其转换为和为 1 的概率分布,需经过一层 softmax:

$$p_\theta(x_t\mid x_{\lt t})=\frac{e^{z_{x_t}}}{\sum_{j=1}^{V} e^{z_j}}$$

将 softmax 代回交叉熵,单个位置的损失化简为 $\mathcal{L}_t=-z_{x_t}+\log\sum_j e^{z_j}$。其中第一项提升正确 token 的分值,第二项(称为 log-sum-exp)抑制所有 token 分值的整体水平,二者共同作用使正确 token 的分值相对突出。

该表达式在实现上存在一处必须处理的数值问题:$e^{z_j}$ 中只要某个 $z_j$ 较大(例如 90),$e^{90}$ 即溢出为 inf。因此成熟实现均先减去当前行的最大值 $z_{\max}$ 再取指数($e^{z_j-z_{\max}}$),结果在数学上完全等价且不会溢出。

PyTorch 的 F.cross_entropy 已将 softmax 与该稳定化过程融合为单一算子,实践中应直接将 logits 传入,而不手写 log(softmax(...)),后者既慢又易产生 NaN。

选用交叉熵而非其他损失有其理论依据。预测下一个 token本质上是分类问题,而交叉熵正是最大似然估计在分类任务上的形式:最小化 $-\sum\log p_\theta(x_t)$ 等价于最大化模型生成该训练语料的概率。

它在梯度上还具有良好性质,对 logits 求导后梯度恰为 $p_\theta-\mathbf{1}_{x_t}$(预测分布减去 one-hot 真值),形式简洁,且不像平方误差那样在预测严重偏离时梯度反而减小,因而优化过程既快又稳定。

训练闭环将预测误差反向传播为梯度,再由优化器持续修正模型参数。
图 1 Transformer 训练闭环。教程综合图;前向计算形成预测误差,反向传播产生梯度,优化器更新后的参数进入下一次前向计算。

以下是将上述四步落为代码的最小实现,与 CS336 作业 1 要求实现的主干一致1(省略模型定义与分布式部分以突出主线):

Python Transformer 最小训练循环
model = Transformer(...)                      # RoPE + SwiGLU + RMSNorm(pre-norm)
opt   = AdamW(model.parameters(), lr=peak_lr,
              betas=(0.9, 0.95), eps=1e-8, weight_decay=0.1)

for step in range(max_steps):
    # 1. 数据:从 np.memmap 的 token 大数组里随机切窗口
    x, y = get_batch(train_tokens, batch_size, seq_len, device)

    # 2. 学习率:warmup + 余弦,每一步先设好本步 lr
    for g in opt.param_groups:
        g["lr"] = lr_schedule(step, peak_lr, warmup_steps, total_steps)

    # 3. 前向 + 交叉熵
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        logits = model(x)                     # (B, T, vocab)
        loss = cross_entropy(logits.view(-1, V), y.view(-1))

    # 4. 反向 → 裁剪 → 更新 → 清零
    opt.zero_grad(set_to_none=True)
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    opt.step()

    if step % ckpt_every == 0:                # 5. 可复现 checkpoint
        save_checkpoint(model, opt, step, path)

常见错误:遗漏 zero_grad。PyTorch 的 .backward() 对梯度执行累加而非覆盖。若不清零,本步梯度会叠加于上一步之上,等效于一个不断放大的学习率,表现为 loss 先异常上升、随后变为 NaN。

梯度累积(gradient accumulation)则有意利用这一累加行为:连续对若干小 batch 执行 backward 而不清零,累积到目标规模后再 step,从而在有限显存下获得较大的有效 batch。

二、优化器 AdamW:为什么不是 SGD#

动机:为每个参数分配自适应步长#

最基本的更新规则是 SGD:$\theta \leftarrow \theta-\eta\,g$,全部参数共用同一个学习率 $\eta$。问题在于 LLM 的损失面高度病态:不同方向的曲率差异极大,有的方向陡峭、单步即越过最优点,有的方向平坦、多步几乎不移动,同时还存在大量稀疏且尺度悬殊的梯度分量(如词表中罕见 token 的嵌入、深层网络中的微弱信号)。

单一学习率无法同时适应这些方向,取值偏大则陡峭方向发散,取值偏小则平坦方向停滞。

Adam 的解法是为每个参数维护各自的自适应步长:一方面记录梯度的方向惯性(一阶动量,用于平抑随机波动),另一方面估计梯度的历史幅度(二阶动量,据以缩小历史幅度大的方向的步长、放大历史幅度小的方向的步长)。由此不同尺度的参数均可采用合适的步长推进,收敛快且稳定,这也是几乎所有 LLM 弃用 SGD 的原因。

具体而言,SGD 对所有方向施加同等大小的步长,在陡峭方向易于越过最优点、在平坦方向进展缓慢。一阶动量令连续同向的更新逐步累积、增强推进效率,同时使单步的噪声梯度不足以立即改变整体方向。二阶动量则依据各方向梯度平方的历史累积自动调节步长:某方向历史梯度平方大(长期陡峭),则缩小该方向步长;某方向历史梯度平方小(长期平坦),则放大步长。

Adam 同时实现这两项调节,因而在参数量巨大、各方向曲率差异悬殊的 LLM 损失面上,成为兼顾速度与稳定性的近乎唯一可行的选择。其代价是需为每个参数额外存储 $m$ 与 $v$ 两份状态,显存开销约为参数量的两倍,这正是 1.4 · 资源核算优化器状态占据大部分显存的成因。

核心公式#

记第 $t$ 步梯度为 $g_t$,AdamW 维护一阶矩 $m_t$(动量)与二阶矩 $v_t$(梯度平方的滑动平均):

$$m_t=\beta_1 m_{t-1}+(1-\beta_1)\,g_t,\qquad v_t=\beta_2 v_{t-1}+(1-\beta_2)\,g_t^2$$

因为 $m,v$ 初始化为 0,早期会偏向 0,需做偏差校正:

$$\hat m_t=\frac{m_t}{1-\beta_1^{\,t}},\qquad \hat v_t=\frac{v_t}{1-\beta_2^{\,t}}$$

最后更新参数,注意末尾那一项 $\lambda\theta_{t-1}$:

$$\theta_t=\theta_{t-1}-\eta\left(\frac{\hat m_t}{\sqrt{\hat v_t}+\epsilon}+\lambda\,\theta_{t-1}\right)$$

前一项 $\hat m_t/(\sqrt{\hat v_t}+\epsilon)$ 即方向惯性除以幅度的自适应步长;后一项 $\lambda\theta$ 是权重衰减,在每步中令参数向 0 收缩少许,以抑制过拟合。

解耦权重衰减到底解耦了什么#

这是 AdamW 相对于早期 Adam+L2 的唯一而关键的改动,出自 Loshchilov & Hutter《Decoupled Weight Decay Regularization》(arXiv:1711.05101, ICLR 2019)2。论文的核心观察在于两种正则方式在不同优化器下并不一致。

SGD 而言,L2 正则(在损失中加入 $\frac{\lambda}{2}\lVert\theta\rVert^2$)与权重衰减(更新时直接执行 $\theta\leftarrow(1-\eta\lambda)\theta$)完全等价。

但对 Adam 二者并不等价:若将 L2 惩罚并入梯度 $g_t$,它会随梯度一同进入分母 $\sqrt{\hat v_t}$ 而被自适应缩放,梯度大的参数分母大、其 L2 惩罚被除得偏小、衰减反而偏弱,导致正则强度在参数之间被系统性扭曲。

AdamW 的修正是将 $\lambda\theta$ 从梯度与动量中分离出来,使其不经过 $\sqrt{\hat v_t}$,而直接乘以学习率作用于参数(即上式中独立的一项)。由此,权重衰减 $\lambda$ 与学习率 $\eta$ 的最优取值可以相互独立地调整、互不干扰,这正是解耦的含义,也是 AdamW 成为 LLM 默认优化器的直接原因。

这一扭曲可进一步说明。在早期 Adam+L2 中,一个梯度长期较大的参数(如高频词的嵌入)其分母 $\sqrt{\hat v_t}$ 也大,L2 项 $\lambda\theta$ 被该大分母除后正则强度趋近于零;而一个梯度较小的参数(罕见词)分母小,L2 惩罚反被放大。其结果是更新越频繁的参数衰减越弱,与对所有参数施加均匀收缩的预期恰好相反。

AdamW 使权重衰减绕过分母后,每个参数每步都乘以同一个 $(1-\eta\lambda)$ 因子,正则化由此恢复到均匀且可预测的状态。由此可得一条实践推论:在 AdamW 中通常不对 bias、LayerNorm/RMSNorm 的增益(gain)等一维参数施加权重衰减,将其向 0 收缩缺乏正则意义,反而会削弱归一化层的表达能力;

实现上一般按参数维度划分为两个 param group,仅对二维以上的权重矩阵设定 $\lambda$。

LLM 上的常见取值及其依据:GPT-3 采用 $\beta_1=0.9$、$\beta_2=0.95$、$\epsilon=10^{-8}$、weight decay $=0.1$[^gpt3]。其中 $\beta_2=0.95$ 低于 Adam 默认的 0.999。$\beta_2$ 越大,二阶矩 $v_t$ 的历史窗口越长,一次异常大的梯度会被长期保留、持续压低后续步长。

取 $\beta_2=0.95$ 使 $v_t$ 更快遗忘历史,对偶发大梯度(loss 尖峰的诱因)响应更灵敏、恢复更快,是大规模训练的经验取值。

一步 AdamW 的数值算例#

将公式代入具体数字以说明自适应步长的作用机制。设某参数在第一步($t=1$)取得梯度 $g_1=0.2$,超参取 $\beta_1=0.9$、$\beta_2=0.95$、$\eta=10^{-3}$、$\epsilon=10^{-8}$,暂不计权重衰减:

  • 一阶矩:$m_1=0.9\times 0+0.1\times 0.2=0.02$;二阶矩:$v_1=0.95\times 0+0.05\times 0.2^2=0.002$;
  • 偏差校正:$\hat m_1=\dfrac{0.02}{1-0.9}=0.2$,$\hat v_1=\dfrac{0.002}{1-0.95}=0.04$;
  • 更新量:$\eta\cdot\dfrac{\hat m_1}{\sqrt{\hat v_1}+\epsilon}=10^{-3}\times\dfrac{0.2}{0.2+10^{-8}}\approx 10^{-3}$。

最后一步中,$\sqrt{\hat v_1}=\sqrt{0.04}=0.2$ 恰等于 $\hat m_1$,比值约为 1,因此无论该方向梯度的绝对值为 0.2 还是 200,第一步的更新量都约等于学习率 $\eta$ 本身。这即 Adam 自适应机制的核心:它将各方向的步长归一化到与学习率同量级,梯度的绝对尺度被 $\sqrt{\hat v}$ 约去。

由此 Adam 的学习率具有一项良好性质,其数值近似等于每步参数变化量的上界,使超参在不同模型与任务间的迁移更为容易,这也是 Adam 相较 SGD 更节省调参成本的深层原因。

三、学习率调度:warmup + 余弦衰减#

学习率 $\eta$ 并非常数,而是随训练步数变化的一条曲线。业界标准配置为线性 warmup + 余弦衰减,其完整形状见图 2,CS336 作业 1 要求实现的也是该方案1

学习率先预热、再达到峰值、最后平滑衰减,兼顾训练初期稳定性与后期收敛。
图 2 预热与余弦衰减学习率。教程综合图,依据 Brown 等 2020 年论文第 2.3 节的训练配置组织;训练初期以小步长控制更新尺度,达到峰值后逐步衰减以改善后期收敛。

warmup:早期采用较小学习率的原因#

训练初始阶段权重为随机初始化,模型对数据尚无任何拟合,前几步的梯度幅度大、方向近似噪声。若此时直接使用满额峰值学习率,极易在单步内将参数推入损失发散的区域,导致 loss 变为 NaN 且无法恢复3

warmup 的作用是限制早期更新的幅度:令学习率从接近 0 线性上升至峰值,以最初数百至数千步的小步长使网络从随机初始区域进入良态区域,随后再逐步提速。GPT-3 的 warmup 覆盖训练的前 3.75 亿 token,此后才进入余弦衰减4

还有一层原因与优化器直接相关。AdamW 早期的二阶矩 $\hat v_t$ 从 0 起步、可用样本很少,估计误差大,分母 $\sqrt{\hat v_t}$ 的波动会使自适应步长剧烈起伏;此时若叠加满额学习率,极易偏离良态方向。warmup 期间学习率较小,恰好为 $\hat v_t$ 积累足够梯度样本、把二阶矩估计准确争取了时间。

因此 warmup 步数通常与二阶矩的历史窗口长度 $\tfrac{1}{1-\beta_2}$ 处于同一量级,$\beta_2=0.95$ 对应约 20 步的窗口,实践中 warmup 常设为数百至数千步以留足缓冲。一条常见的工程经验是:在改用更大 batch 或更长序列后若 loss 一开训即发散,首选延长 warmup,其效果往往优于降低峰值学习率。

余弦衰减:后期降低学习率的原因#

warmup 结束后,学习率沿余弦曲线平滑下降。前期学习率高,用于快速逼近损失面上的低损失区域;后期学习率低,用于在该区域内精细收敛。相较阶梯式突降,余弦的平滑过渡避免了每次调降学习率时的震荡。设峰值 $\eta_{\max}$、终值 $\eta_{\min}$、warmup 步数 $t_w$、总步数 $T$,当 $t\gt t_w$ 时:

$$\eta(t)=\eta_{\min}+\tfrac12\big(\eta_{\max}-\eta_{\min}\big)\left(1+\cos\!\frac{(t-t_w)\pi}{T-t_w}\right)$$

工程上常将终值设为峰值的约 10%($\eta_{\min}\approx0.1\,\eta_{\max}$)而不衰减至 0:末期保留少量学习率以继续吸收数据,其收敛效果优于完全停止更新,图 2的曲线尾部停于 0.1 即出于此。

余弦调度在实践中存在一处局限:其曲线形状由总步数 $T$ 决定,须在训练开始前确定总步数才能计算每一步的学习率。而实际中常需先训练一段、再依据表现决定是否延长,一旦中途改变 $T$,整条余弦曲线即需重算,已有进度难以衔接。

为此近年出现 WSD(warmup–stable–decay)调度:warmup 后令学习率长期保持于恒定平台,仅在最后一小段(如末尾 10%–20% 的步数)快速衰减。该方案的平台期可随时延长或截断,只需在收尾前接入短衰减段即可,因而适用于算力预算不确定、或需从同一中间 checkpoint 分叉出多个不同长度实验的场景。

两种调度机制一致:均依靠开头的 warmup 与结尾的 decay 维持稳定,区别仅在中段分别采用余弦下降与恒定平台。

峰值学习率的确定:一条经验规律是模型规模越大,峰值学习率越小,大模型的损失面更陡峭,同等步长更易越过最优点。GPT-3 175B 的峰值仅 $6\times10^{-5}$[^gpt3],而百万参数级的小模型常可用到 $10^{-3}$ 量级。峰值学习率还与 batch size 联动(大 batch 梯度更稳定,可配以稍大的学习率)。

实践中先在小规模上扫描确定,再依上述规律外推到目标规模。

四、稳定性技巧:各自抑制一类发散#

大规模训练发散的模式为数不多,每一项稳定性技术对应抑制其中一类。

梯度裁剪:抑制梯度爆炸#

训练中偶尔会出现某批数据或某种参数组合,使某一步的梯度范数骤增数个数量级,单步即令模型发散。梯度裁剪step 之前先计算所有参数梯度拼接后的全局范数 $\lVert g\rVert$,若其超过阈值 $c$(业界常取 1.0)则对整体等比缩放:

$$g \leftarrow g\cdot\min\!\left(1,\ \frac{c}{\lVert g\rVert}\right)$$

关键在于按全局范数缩放,仅限制幅度而不改变方向:所有参数的梯度同乘一个 $\le 1$ 的系数,更新方向保持不变,仅整体长度被限于阈值以内。GPT-3、PaLM 等均采用该技术,阈值取 1.045。它对梯度施加上界,与前述的 zero_grad 是两个独立的操作。

此处有两点需注意。其一,须按全局范数裁剪,而非逐参数、逐元素裁剪:逐元素地将每个梯度分量截断至 $[-c,c]$ 会改变梯度方向、破坏其作为最速下降方向的性质,而全局范数缩放为等比例操作,方向严格不变,仅将整个梯度向量的长度压回阈值内。其二,全局范数须跨所有参数计算,且在分布式训练中跨所有 GPU 求和后再开方;若各卡各自裁剪,阈值的语义即被破坏。

实践中通常将梯度范数作为监控指标记录:正常训练时它应在一个平稳区间内小幅波动,一旦某步范数骤增至平时的数十乃至上百倍,往往预示 loss 尖峰。此时裁剪既起到限幅作用,也起到预警作用。

混合精度:bf16 与 fp16 及 loss scaling 的取舍#

半精度计算速度更快、显存占用减半(1.4 · 资源核算 已述),但半精度的表示范围较窄,处理不当会使梯度下溢为 0

混合精度训练的三项经典技术出自 Micikevicius 等《Mixed Precision Training》(arXiv:1710.03740, ICLR 2018)6:其一,保留一份 fp32 主权重作为累加基准;其二,前向与反向计算采用半精度;

其三,loss scaling,反向前将 loss 乘以一个大因子(如 $2^{14}$),把过小的梯度整体平移进半精度可表示的范围,step 前再除回。

保留 fp32 主权重的必要性在于:训练后期单步更新量常小至参数值的百万分之一量级,若以半精度存储参数,这一微小增量会因尾数位数不足而被直接舍入,导致更新失效、模型停止改进。解决方式是将参数的权威副本始终以 fp32 保存、所有微小更新累加于其上(单次虽小,累积后仍有效),仅在每步前向与反向时临时转为半精度以执行对精度不敏感的大矩阵乘。

loss scaling 则专门应对 fp16 的下溢:反向前将 loss 乘以一个大因子,梯度随之被放大同样倍数、整体平移进 fp16 可表示的区间,step 更新前再除回该因子,二者在数学上完全等价。现代框架多采用动态 loss scaling,因子过大导致溢出时自动减半、长期无溢出时试探性加倍,从而免去人工调参。

两种半精度格式的关键区别如下:

格式指数位尾数位动态范围是否需要 loss scaling
fp16510窄(易下溢)需要
bf1687与 fp32 相同通常不需要

bf16 以精度(尾数仅 7 位)换取与 fp32 相同宽度的动态范围,梯度基本不会下溢,因而无需 loss scaling。自 NVIDIA Ampere(A100)起,bf16 成为现代混合精度训练的事实标准7。fp16 尾数更多、数值更精确,但范围较窄,必须配合 loss scaling 才能稳定。

上文代码中的 torch.autocast(dtype=torch.bfloat16) 即采用 bf16 方案。

范围之所以比精度更为关键,在于训练中最需避免的并非计算精度略有损失,而是梯度小至被表示为 0、从更新中彻底消失。fp16 可表示的最小正规数约为 $6\times 10^{-8}$,深层网络中许多梯度分量本就落在该量级以下,一旦下溢为 0,对应参数即不再接收更新信号;

而 bf16 的最小正规数与 fp32 同量级(约 $10^{-38}$),这些小梯度均可安全表示,代价仅为每个数值损失若干位有效数字。对训练而言,保留全部梯度但精度略低,优于精度高但丢失部分梯度。这也解释了现代训练普遍选用 bf16 的原因:它免去了 loss scaling 中动态调整缩放因子、且因子选取不当会反复触发溢出或下溢的调参负担。

需注意 bf16 尾数仅 7 位、有效数字约 2–3 位,某些对精度敏感的算子(softmax 归一化、LayerNorm 的求和、loss 累加、优化器状态更新)仍须在 fp32 中执行torch.autocast 的内部白名单正据此将这些算子自动提升回 fp32,而将矩阵乘、卷积等对精度不敏感的算子保留在 bf16。

合理初始化:从源头抑制尖峰#

权重初始化的尺度直接决定早期训练能否稳定。其核心要求为方差守恒:信号在数十乃至上百层中前向传播、梯度再反向传播,每经过一层,数值的方差都不应被系统性地放大或缩小,放大会导致逐层爆炸(梯度 inf),缩小会导致逐层消失(深层无法学习)。

要使一层线性变换 $y=Wx$ 的输出方差与输入持平,权重方差须与输入维度 $d$ 成反比,因此常见初始化将权重标准差设为 $1/\sqrt{d}$ 量级,这即 Xavier 与 Kaiming 初始化的推导依据。

Transformer 还引入一个额外问题:残差连接将每一层的输出逐层累加,经 $L$ 层后残差主干上的方差累积至约 $L$ 倍,深层激活因此被持续放大。修正方式是对残差分支的输出额外施加 $1/\sqrt{2L}$ 等因子的缩小以抵消这一累积。自 GPT-2 起即采用这类按深度缩放残差的方法,使上百层的网络也能从第一步起稳定训练。

这不仅是初始化阶段的技巧。2024 年的 arXiv:2410.05052《Initialization of LLMs via Reparameterization to Mitigate Loss Spikes》提出以重参数化初始化源头抑制 loss 尖峰8

其观察是,尖峰往往源于某些参数在训练中被推至极端尺度,而通过重参数化使参数尺度天然受约束,即可在梯度裁剪等事后补救之外,于初始化阶段就抑制隐患。这表明稳定性本身至今仍是活跃的研究前沿:从初始化、归一化、优化器到数值精度,各环节均在被重新审视。

五、关键超参与数据:几个数值的确定依据#

将训练循环运行起来,还需确定若干数值。这些数值并非随意试探而来,各有其确定依据。

训练 token 数:由扩展律确定#

这是最易出错的一项。训练时间并非越长越好,应喂入的 token 数 $D$ 由 3.2 · Chinchilla 计算最优 给出。

Hoffmann 等 2022(arXiv:2203.15556)训练了 400 余个参数量从 70M 到 16B 的模型,拟合出计算最优比例:每个参数约配 20 个训练 token9。据此,70B 模型的计算最优数据量约为 $1.4$ 万亿 token;

在同等算力预算下,规模较小而训练更充分的 Chinchilla 70B 超越了规模更大而训练不足的 Gopher 280B,这一结果修正了当时模型越大越好的普遍认识。

该比例源于一个简单的权衡。训练的总算力预算大致为参数量 $N$ 与训练 token 数 $D$ 的乘积(业界常用估算为 $C\approx 6ND$ FLOPs,1.4 · 资源核算 已推导系数 6 的来历)。

预算 $C$ 固定时,$N$ 与 $D$ 构成一对此消彼长的变量:预算过度投入参数量则可用 token 数受限,过度投入 token 数则模型规模受限。Chinchilla 的贡献在于在这条等预算线上找到使最终 loss 最低的平衡点,其位置恰在 $D\approx 20N$ 附近。

此前的 GPT-3 时代普遍将预算过度倾斜于参数量、数据供给严重不足(GPT-3 175B 仅训练约 3000 亿 token,按 20:1 应为 3.5 万亿),因而那些巨型模型实际训练不充分、算力利用效率偏低。

需注意 20:1 是训练侧的计算最优,未必是部署侧的最优。若一个模型需服务海量请求,推理成本会长期累积,此时通常有意超训,将远超 20:1 的 token 喂给一个规模偏小的模型,以训练阶段的额外开销换取推理阶段的长期节省。

近年开源小模型普遍训练至每参数数十 token 即出于此逻辑,Llama 系列将 7B 级模型训练至万亿 token 量级是典型例证。因此 20:1 是理解权衡的参照点,而非必须严守的定式。

2024–2025 年这一比例被反复验证:arXiv:2404.10102 的复现指出原论文部分置信区间偏窄,但约 20 token/参数的核心结论稳固10;arXiv:2509.23963 进一步在不同算力预算下重新拟合,发现 20:1 甚至比原论文更稳健11

主流训练配方中的参数实例#

模型或课程优化与数据设置具体例子
GPT-3 175BAdamW,$\beta_1=0.9$、$\beta_2=0.95$,峰值学习率 $6\times10^{-5}$前 3.75 亿 token 完成 warmup,随后按余弦计划衰减学习率4
Chinchilla 70B计算最优点约为 $D\approx20N$$N=70\text{B}$ 对应 $D\approx1.4\text{T}$,以更充分训练超过 Gopher 280B9
PaLM 540B高频 checkpoint 与尖峰恢复检测尖峰后回到约 100 步前的检查点,并跳过约 200 至 500 个批次5
Stanford CS336自实现 AdamW 与可序列化检查点保存模型、$m$、$v$、step 与数据游标,重载后复现同一更新轨迹1

epoch 通常 $\lt 1$:数据不重复#

传统深度学习以 epoch(将数据集完整遍历 $N$ 遍)计量训练量。前沿 LLM 预训练则受 token 数约束而非受数据集约束,高质量语料本就稀缺,训练量以已处理的 token 数或 step 数计量,多数大模型仅单遍训练、不重复数据(即 epoch $\lt 1$)12

这一点与图像分类时代将数据集遍历数百遍的做法形成对比,原因有二。其一,互联网级语料规模极大,在给定算力预算下、按扩展律应喂入的 token 数往往在覆盖完整语料前即耗尽预算,因而连一遍也无法完成。其二,重复处理同一批数据会助长记忆而非泛化,模型开始逐字复现训练样本,既浪费算力又带来隐私与数据污染风险。

当高质量 token 确实耗尽、不得不重复时,实证给出了一条清晰界线:重复至约 4 个 epoch 仍近似等效于同等质量的新数据,超过此界后每一遍的边际收益急剧衰减、直至趋近于零[^dataconstrained]。因此与其将一份小语料重复 40 遍,不如重复 4 遍后将剩余算力投入更大的模型或更充分的数据清洗。

这也是数据配比、去重、质量过滤在预训练中与调参同等重要、甚至更为重要的原因,在同等算力下,喂入 token 的质量差一个档次,最终模型能力即差一个档次,这一点将在 第 5 章 · 从零预训练 展开。

batch、序列长度与数据加载#

batch size 在 LLM 中以每步 token 数而非样本条数计量。GPT-3 最小的模型约用 0.5M token/batch,最大的 175B 逐步提升至约 3.2M token/batch(上下文窗口 2048)4

较大的 batch 梯度更稳定、更利于并行,但存在临界批量(critical batch size):batch 过小时梯度噪声大、每步方向偏差大,增大 batch 可通过平均消除噪声、提升每步方向的准确性;而一旦噪声已被抑制得足够低,继续增大 batch 只是重复确认同一方向、对收敛几乎无贡献,仅徒增算力消耗。

因此存在一个收益拐点,超过它后应将算力投向他处而非继续增大 batch。显存不足而又需较大有效 batch 时,采用梯度累积(见第一节旁注),连续对若干小 batch 执行 backward 而不清零,累积到目标规模后再 step,以时间换取显存。

序列长度直接影响显存与注意力开销:注意力开销随序列长度呈平方增长,序列长度翻倍则注意力的计算量与显存约增至四倍。常见做法是先短后长,前期用短序列快速训练、控制成本,先充分学习语言的基本规律,后期再增大上下文长度、专门训练跨长距离依赖的能力。

长上下文本身是一个独立主题,第 14 章 · 长上下文 将专门讨论如何将窗口从数千扩展至数十万乃至上百万。

数据加载与 packing 方面,CS336 作业 1 的做法是将整个语料 tokenize 为一个大型 token 数组、存为 np.memmap(内存映射,无需全部载入内存),get_batch 每次随机切取 (batch, seq_len) 的窗口1

为避免算力浪费,短文档会被打包(packing),多条文档首尾相接填满一个序列窗口,而非补充大量 padding 空转。

六、loss 尖峰:诊断与恢复#

即便已施加梯度裁剪,超大模型训练中 loss 仍会偶尔突然跳高再(可能)回落,此即 loss 尖峰(spike)。PaLM 540B 的训练日志显示,即便全程开启梯度裁剪,仍观察到约 20 次尖峰,且这些尖峰仅在大模型上出现、间隔无规律5

Google 团队通过一项关键实验揭示了其成因:从尖峰发生前约 100 步的 checkpoint 重启,并跳过尖峰前后所见的约 200–500 个数据 batch再续训,结果尖峰不再复现5

这表明尖峰既非单纯由坏数据引起,也非单纯由坏参数引起,而是特定数据批次与特定参数状态这一组合恰好触发的,以不同参数状态处理同一批数据、或以不同批数据处理同一状态,均不会引发尖峰。这一组合触发的判断具有实际意义:它意味着无需精细筛选并剔除个别有害样本(那样既治标不治本,又可能误删有效数据),只需在恢复时错开这一时序组合即可。

尖峰的规模依赖性值得强调:小模型上几乎不出现尖峰,因其损失面相对平缓、参数对更新的敏感度较低;模型规模越大、损失面越陡峭,越易在某个方向上被一批数据推出稳定区。因此小规模训练通过不代表大规模训练同样稳定,许多稳定性问题仅在大规模下暴露,这也是大模型训练须全程监控 loss 曲线与梯度范数、并将 checkpoint 存储得足够密集的原因。

除回退加跳批外,业界还有若干缓解手段:临时降低学习率后再回升、对个别层采用更保守的初始化、或采用前述重参数化初始化从源头抑制尖峰[^initspike]。但没有任何单一手段可一劳永逸,运维层面的密集存档与可回退能力始终是最后保障。

要在尖峰初现时即予以捕获,日常须监控几条关键曲线:训练 loss(正常状态为平滑下降,突然的台阶式跳高即尖峰)、梯度全局范数(如前所述,它往往比 loss 更早示警),以及参数与激活的数值范围(是否逼近 bf16/fp16 的表示上限)。

成熟的训练框架会将这些指标实时输出至监控面板,并配以 NaN 自动停机机制,一旦检测到 loss 或梯度变为 NaN/inf,立即停训并保留现场,避免脏数值污染后续 step。捕获之后的处置流程如图 3所示:

损失尖峰应沿数据与数值链路定位,并从确定健康的检查点重放验证。
图 3 损失尖峰的诊断与恢复路径。教程综合图;监控首先捕获异常,再沿数据和数值链定位原因,最后从健康检查点重放并验证修复。

上述流程的前提是可复现的 checkpoint。恢复流程只有在能精确回到过去某一步时才成立,因此 checkpoint 必须完整保存 model 权重、optimizer 状态(AdamW 的 $m,v$)、当前 step、学习率与数据游标,缺一不可。

若仅保存模型权重而丢失优化器的一阶、二阶矩,续训时 AdamW 相当于从零冷启动,收敛轨迹将与尖峰前无法对齐。CS336 作业 1 专门要求实现可序列化、能原样加载 model 与 optimizer的 checkpoint[^cs336],正是出于此。

低精度与稳定性的权衡仍在演化。一个易被忽视之处是:若训练与推理采用不同的数值精度或算子实现,同一份权重给出的概率分布会存在细微差异;在强化学习后训练(模型自行采样、再据采样结果更新自身)中,这一差异会被逐步放大为训练所依据的策略与推理时实际执行的策略之间的偏差。

2025 年的 arXiv:2510.04212 分析了低精度(含 Flash Attention 场景)训练失败的机理13,arXiv:2510.26788 则重新讨论 fp16 在训练-推理一致性中的作用,指出在某些场景下 fp16 反比 bf16 更利于对齐,即bf16 为事实标准之外仍存在值得权衡的空间14

这些工作表明,如何做到长时间稳定训练且训练与推理一致,尚无定论,仍是工程与研究交叉的活跃领域。


小结#

  • 训练一个 Transformer 由预测下一个 token的四步循环(前向 → 交叉熵 → 反向 → AdamW step + zero_grad)与一整套稳定性工程共同构成
  • 优化器采用 AdamW($\beta_2\approx0.95$、解耦权重衰减,使权重衰减与学习率可各自独立调整),学习率采用 warmup + 余弦衰减(早期缓升以防发散,末期衰减至峰值约 10%)
  • 稳定性依赖全局范数梯度裁剪(阈值 1.0)、bf16 混合精度(动态范围同 fp32、无需 loss scaling,fp16 才需要)与合理初始化
  • 超参方面,训练 token 数由扩展律确定(约 20 token/参数),epoch 通常 $\lt 1$、数据不重复
  • 真正的困难在于超大规模下长时间稳定运行而不发散,loss 尖峰须依靠完整的 checkpoint 回退与跳批加以恢复
  • 训练本质上是一项工程任务,稳定性优先于结构上的复杂设计

上一节1.4 · 资源核算下一节1.6 · 混合专家 MoE本章总览第 1 章 · 大模型基础

参考文献#

作业概览 https://deepwiki.com/stanford-cs336/assignment1-basics/1-overview


  1. Stanford CS336《Language Modeling from Scratch》, Assignment 1(BPE、含 RoPE/SwiGLU/RMSNorm 的 Transformer、交叉熵、自写 AdamW、warmup+余弦 lr、全局范数梯度裁剪、np.memmap 数据加载、可序列化 checkpoint). 课程主页 https://cs336.stanford.edu/。 ↩︎ ↩︎ ↩︎ ↩︎ ↩︎

  2. Loshchilov & Hutter, Decoupled Weight Decay Regularization, arXiv:1711.05101, ICLR 2019. https://arxiv.org/abs/1711.05101 ↩︎

  3. APXML, Stabilization Techniques Revisited (Clipping, LR, Warmup)(随机初始化下立即满额 lr 会致混乱更新、loss 爆炸/NaN;warmup 作早期敏感度节流). https://apxml.com/courses/how-to-build-a-large-language-model/chapter-24-identifying-mitigating-training-instabilities/stabilization-techniques-revisited ↩︎

  4. Brown et al., Language Models are Few-Shot Learners (GPT-3), arXiv:2005.14165, 2020(AdamW $\beta_1{=}0.9,\beta_2{=}0.95$;全局范数裁剪 1.0;175B 峰值 lr $6\times10^{-5}$,warmup 覆盖前 3.75 亿 token 后余弦衰减). https://arxiv.org/abs/2005.14165 ↩︎ ↩︎ ↩︎ ↩︎

  5. Chowdhery et al., PaLM: Scaling Language Modeling with Pathways, arXiv:2204.02311, 2022(开梯度裁剪仍约 20 次尖峰;回退约 100 步 checkpoint、跳约 200–500 batch 后不再复现). https://arxiv.org/pdf/2204.02311 ↩︎ ↩︎ ↩︎ ↩︎

  6. Micikevicius et al., Mixed Precision Training, arXiv:1710.03740, ICLR 2018(fp32 master weights、loss scaling、fp16 累加进 fp32 三要素). https://arxiv.org/abs/1710.03740 ↩︎

  7. MemX Glossary, Mixed Precision Training: FP16 and BF16 Explained(bf16 8 指数/7 尾数、范围同 fp32、通常免 loss scaling;自 Ampere 起为事实标准). https://memx.app/glossary/mixed-precision-training/ ↩︎

  8. Initialization of LLMs via Reparameterization to Mitigate Loss Spikes, arXiv:2410.05052, 2024(从初始化/重参数化源头抑制 loss 尖峰). https://arxiv.org/pdf/2410.05052 ↩︎

  9. Hoffmann et al., Training Compute-Optimal Large Language Models (Chinchilla), arXiv:2203.15556, 2022(约 20 token/参数;70B Chinchilla 同算力胜 280B Gopher). https://arxiv.org/abs/2203.15556 ↩︎ ↩︎

  10. Chinchilla Scaling: A replication attempt, arXiv:2404.10102, 2024(复现第三种估计法,指出原置信区间偏窄,~20 token/参数结论成立). https://arxiv.org/html/2404.10102v1 ↩︎

  11. Evaluating the Robustness of Chinchilla Compute-Optimal Scaling, arXiv:2509.23963, 2025(不同算力预算下 20:1 比例甚至比原论文更稳健). https://arxiv.org/pdf/2509.23963 ↩︎

  12. Brenndoerfer, Data-Constrained Scaling: Training LLMs Beyond the Data Wall(前沿预训练 token 受限、单遍不重复;耗尽后重复至约 4 epoch 近似等效新数据). https://mbrenndoerfer.com/writing/data-constrained-scaling-llm-training-data-limits ↩︎

  13. Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention, arXiv:2510.04212, 2025(低精度含 Flash Attention 场景训练失败机理). https://arxiv.org/abs/2510.04212 ↩︎

  14. Defeating the Training-Inference Mismatch via FP16, arXiv:2510.26788, 2025(重议 fp16 在训练-推理一致性中的角色). https://arxiv.org/abs/2510.26788 ↩︎