第 8 章　训练主循环、梯度累积、混合精度与裁剪

8.1 一步到底包含几次前向

在 nanoGPT 里，iter_num 主要计数优化器更新；一个更新内部可以包含多个微批次。令本进程累积次数为 $a$ ，第 $r$ 个微批次平均损失为 $ℓ_{r}$ ，目标是优化它们的平均值

L = \frac{1}{a} \sum_{r = 1}^{a} ℓ_{r}, \nabla_{θ} L = \frac{1}{a} \sum_{r = 1}^{a} \nabla_{θ} ℓ_{r} .

因此循环每次先把 loss 除以 $a$ ，再 backward，直到全部微批次完成才 step。梯度默认累加进 .grad；若在每个微批次后清零，就只剩最后一个微批次的梯度。若忘记除以 $a$ ，梯度尺度增大 $a$ 倍，裁剪和优化器行为都会改变。训练循环

这里的等价性需要条件：微批次包含相同数量的有效 token，模型对样本没有跨批次统计依赖，随机层处理方式也需可比。若一个微批次只有一个未忽略目标，另一个有 100 个，简单平均两项“各自均值”会给前者过大权重。应按有效 token 数加权。dropout 开启时，整批和分批的随机掩码消费也可能不同，不宜要求逐位一致。

本书在 dropout=0、double、相同样本下比较整批四条序列与两次各两条序列，参数梯度最大差为 $5.55 \times 10^{- 17}$ 。它验证了计算逻辑，不是显存或吞吐测试。

交互实验：有效 token 加权梯度累积

训练实验室 · 有效 token 与梯度累积

一批只有一个目标，另一批有三个，该各占一半吗？

把同一组有效目标拆成两个 microbatch（微批次），比较按 token 加权与按批次等权的损失、梯度和一次更新。默认例中，两种算法会把参数推向相反方向。

输入 · 改变有效目标数

微批次 A 有效 token 数微批次 A 目标值0–8 个有效目标；每个目标都使用这个 y 值。

微批次 B 有效 token 数微批次 B 目标值0–8 个有效目标；不计补齐或已被 mask 忽略的位置。

初始参数 w−5–5。两批都在这个固定参数点求导。学习率 η0–1。只影响最后一次 SGD 更新。

模型很小：所有输入 x=1，所以预测值是 w。每个有效目标使用半平方误差 ℓ=½(w−y)²，梯度是 w−y。这是可手算的标量例子，不是语言模型的交叉熵。

过程 · 先累加，最后更新

步骤 1 / 4 · 按有效 token 加权

这一轮共有 4 个有效目标、2 个非空微批次。有效 token 数决定每个微批次在整体目标中的权重。

数字显示 6 位小数，计算使用未舍入值。各表与公式可左右滑动；键盘聚焦后可用方向键查看。

微批次 A · 1 个有效目标

A1y=2.000000预测=0.000000

微批次 B · 3 个有效目标

B1y=-1.000000预测=0.000000
B2y=-1.000000预测=0.000000
B3y=-1.000000预测=0.000000

每批内部的 loss sum、count、mean 与梯度
微批次	有效数 n	损失和 S	均值 S/n	梯度和	均值梯度
A	1	2.000000	2.000000	-2.000000	-2.000000
B	3	1.500000	0.500000	3.000000	1.000000

token 加权：L = Σ Sᵣ / N = Σ (nᵣ / N) × Lᵣ

非空批等权：L_equal = Σ Lᵣ / K_nonempty

逐 token 梯度：∂ℓ/∂w = w − y

N 是这一轮所有有效目标的总数；K_nonempty 只数非空微批次。求和中的 Lᵣ 只对非空批定义。空批既不计算 0/0，也不作为“零均值”加入等权平均。

固定 w 不变：每批贡献先按权重缩放，再加进累积梯度
微批次	token 权重 n/N	加权梯度贡献	加权累计	非空批等权	等权梯度贡献	等权累计
A	0.250000	-0.500000	-0.500000	0.500000	-1.000000	-1.000000
B	0.750000	0.750000	0.250000	0.500000	0.500000	-0.500000

指标 · 与同一个大 batch 对照

当前查看：按有效 token 加权。所有损失与梯度均在更新前的同一个 w 处计算。

有效目标总数 N42 个非空微批次

所有目标的损失和3.500000在两条路线中都使用同一组目标

当前路线的平均损失0.875000先核对分母与权重，再比较数字

当前路线的累积梯度 g0.250000梯度下降沿 −g 方向更新

一次更新后的 w-0.025000w_new = w − η × g

两条路线的梯度差0.750000绝对差；不能用 loss 较小代替检查

w_new = 0.000000 − 0.100000 × (0.250000) = -0.025000

相同目标、相同参数点、相同学习率，只比较一轮一次更新
路线	平均损失	梯度	更新后的 w
合并为一个大 batch	0.875000	0.250000	-0.025000
按有效 token 加权	0.875000	0.250000	-0.025000
微批次均值等权	1.250000	-0.500000	0.050000

按有效 token 加权与合并整批一致；微批次均值等权改变了目标权重。短批中的单个 token 可能因此占更大比重。

解释 · 累积的应是同一个目标

若目标是“所有有效 token 等权”，就应把各批损失和相加，再除以全轮有效目标数。等价做法是：每个非空批的均值乘以 n/N，再反向传播。只除以微批次数，通常得到另一个加权目标。

本实验在开始时清零累积梯度；两批都使用同一个 w，没有在中途更新参数。只有所有贡献收齐后，才用同一学习率做一次普通 SGD。逐批先更新再计算下一批，已经改变了求导点，不能拿来验证这里的等价关系。

什么时候可以与一个大 batch 对得上？

这里固定有效目标与逐目标损失，没有 dropout、跨样本统计、权重衰减或随机抽样。线性求和让按 token 加权的梯度与合并整批一致，实际浮点加法顺序仍可能带来微小差异。页面用解析梯度，CPU 契约测试另用 PyTorch autograd 核对。

同样的加权原则也适用于逐 token 交叉熵，但这不意味着本页运行了语言模型训练。这里不模拟分布式通信、混合精度、梯度裁剪或调度器，也不处理完整训练循环的尾批状态。

不操作也能读懂：w=0 时，A 的 1 个目标 y=2，损失和为 2、梯度和为 −2；B 的 3 个目标 y=−1，损失和为 1.5、梯度和为 3。整体梯度为 (−2+3)/4=0.25，批均值等权却为 (−2+1)/2=−0.5。学习率 0.1 时，一次更新分别得到 −0.025 与 0.05。

依据：Accelerate 1.10.1：变长 token 的梯度累积；PyTorch 2.6.0：backward 的梯度累加。核对于 2026-10-14。

8.2 优化器为什么分成两组

configure_optimizers 先保留 requires_grad=True 的去重参数，再按维数分类。维数不少于 2 的矩阵使用配置的 weight decay；一维及标量参数不衰减。于是 token 和位置 embedding、各线性矩阵被衰减，LayerNorm 缩放及各种偏置不衰减。它不是按参数名字写死一张排除表，也不是“所有 embedding 都不衰减”。优化器构造

AdamW 维护梯度的一阶矩 $m_{t}$ 和二阶矩 $v_{t}$ ，用 $β_{1}, β_{2}$ 控制历史平滑。省略实现细节后，更新可理解为

θ_{t + 1} = (1 - η_{t} λ) θ_{t} - η_{t} \frac{{\hat{m}}_{t}}{\sqrt{{\hat{v}}_{t}} + ϵ} .

这里 $θ_{t}$ 是第 $t$ 次更新前的参数， ${\hat{m}}_{t}$ 、 ${\hat{v}}_{t}$ 是经过初始偏差校正的矩估计， $ϵ > 0$ 是避免分母过小的常数。若把当前梯度记为 $g_{t}$ 、从全零矩状态开始并令更新次数 $t$ 从 1 计数，则 $m_{t} = β_{1} m_{t - 1} + (1 - β_{1}) g_{t}$ 、 $v_{t} = β_{2} v_{t - 1} + (1 - β_{2}) g_{t}^{2}$ ，校正为 ${\hat{m}}_{t} = m_{t} / (1 - β_{1}^{t})$ 、 ${\hat{v}}_{t} = v_{t} / (1 - β_{2}^{t})$ ；平方、平方根和除法均逐参数元素进行。 $η_{t}$ 为学习率， $λ$ 为该参数组的权重衰减。衰减与梯度统计解耦；它不等同于先向 loss 加一个 L2 项再任由 Adam 的自适应缩放处理。不同组的 $λ$ 可以为 0，但本脚本每步把同一学习率写入两组。

脚本通过检查 AdamW 签名是否含 fused，并要求 device_type=='cuda' 才启用 fused 实现。签名检测解决的是 API 是否提供此参数，不证明特定硬件、dtype 或版本一定能成功运行。CPU 本书实测输出为 using fused AdamW: False。

8.3 AMP 改计算精度，不等于把所有状态都减半

CUDA 上的 torch.amp.autocast 为部分运算选择指定精度。模型默认参数仍是 FP32，并没有调用 model.half()。因此不能简单把参数、梯度、Adam 状态、激活全部按两字节估算。CPU 分支用空上下文，运行普通 FP32；本书梯度核对额外把模型转成 double。

GradScaler 仅在 dtype 为 float16 时启用。它把 loss 放大后反向，使某些很小的梯度较不容易下溢；更新前再恢复真实尺度。若出现非有限梯度，scaler 可跳过优化器更新并调整缩放因子。BF16 与 FP16 的数值范围不同，本脚本 BF16 不启用这个 scaler。不要把“低精度”视为一种统一行为。

固定代码调用 torch.cuda.amp.GradScaler，在测试版本产生 API 弃用警告，但 CPU、dtype=float32 下它是禁用的空操作。本书保留原版执行并记录警告，没有借机替换源码再声称原版通过。

8.4 先还原梯度，再裁剪，再更新

全局梯度范数为

g = \sqrt{\sum_{p} \sum_{j} (\partial L / \partial p_{j})^{2}} .

当阈值为 $c$ 时，梯度整体乘以不大于 1 的系数，约为 $min (1, c / g)$ 。这样保留方向、限制范数。源码在裁剪前调用 scaler.unscale_(optimizer)，否则阈值会作用到人为放大的梯度上，实际裁剪会随 loss scale 改变。然后 clip_grad_norm_，再 scaler.step、scaler.update，最后 zero_grad(set_to_none=True)。

裁剪不会恢复已经出现 NaN 的前向，也不会保证每步都收敛。它限制梯度，不直接限制 AdamW 的最终参数更新，因为后者还依赖历史矩和权重衰减。把阈值调小到训练几乎不动，不是解决根因的可靠方法。

8.5 日志中 loss 与时间的细节

循环打印 loss.item() * gradient_accumulation_steps，撤销最后一个微批次的缩放。它只得到最后一个微批次损失，不是这次更新所有微批次损失的精确平均。要记录精确平均，应在循环中累计每个未缩放 loss，但这可能增加同步成本或状态管理。

loss.item() 会要求把标量带回 CPU，CUDA 中常形成同步点。dt 的测量又与前后循环、验证和保存交错，所以不能把普通训练日志中的每个时间都当成纯模型前后向耗时。稳定性能测量用第 13 章专门说明的预热、同步和口径。

补充练习。 不改变全局 token 数，把 batch_size 减半、累积次数加倍，是否保证训练完全不变？

答案。 理想、无随机层、等长有效样本情况下，平均梯度可以相等；实际数据抽样次序、dropout、浮点加法顺序和系统开销会改变。更小微批次也可能降低硬件利用率。应分别验证梯度近似、训练指标和速度，不用“等效 batch”一词包办三项结论。

第 8 章 训练主循环、梯度累积、混合精度与裁剪 ​

8.1 一步到底包含几次前向 ​

交互实验：有效 token 加权梯度累积 ​

输入 · 改变有效目标数

过程 · 先累加，最后更新

指标 · 与同一个大 batch 对照

解释 · 累积的应是同一个目标

8.2 优化器为什么分成两组 ​

8.3 AMP 改计算精度，不等于把所有状态都减半 ​

8.4 先还原梯度，再裁剪，再更新 ​

8.5 日志中 loss 与时间的细节 ​

本章来源 ​