梯度累积与裁剪#

1. 累积的到底是什么梯度#

一次大 batch 更新需要同时保存较多激活。梯度累积让参数在一个窗口内保持不变，逐个 microbatch 做反向传播，将梯度加到同一份缓冲区，最后只更新一次。

设第 $k$ 个 microbatch 有 $n_k$ 个有效元素，其平均损失为 $\bar{\mathcal L}_k$。这里的元素可以是样本，也可以是排除 padding 后参与监督的 token；一个窗口中的定义必须相同。令 $N=\sum_{k=1}^K n_k$，则：

$$\mathcal L_{\text{window}}=\frac{1}{N}\sum_{k=1}^K n_k\bar{\mathcal L}_k,\qquad g_{\text{window}}=\sum_{k=1}^K\frac{n_k}{N}\nabla_\theta\bar{\mathcal L}_k.$$

因此每次反向传播的损失应乘 $n_k/N$。只有各 microbatch 的有效元素数相同，才简化成除以 $K$。尾部不足 $K$ 批时，用实际窗口分母并执行更新；不能丢掉尾部或继续除以固定 $K$。

单卡每批 8 个样本、累积 4 批、4 卡数据并行，且各卡样本数相同，有效 batch 为 $8\times4\times4=128$。若各卡有效 token 数不同，DDP 的 rank 平均不自动等于全局 token 平均。设数据并行卡数为 $R$，全局有效 token 数为 $N_{\text{global}}$，rank $r$ 的第 $k$ 批有 $n_{r,k}$ 个有效 token；若 DDP 最终对各 rank 梯度取 $1/R$ 平均，每个本地 mean loss 应乘 $R n_{r,k}/N_{\text{global}}$，这样最终得到全局 token 平均梯度。该系数依赖通信归约约定；已有框架可能已做这层处理，不能重复乘。

1.1 什么时候等价于大 batch#

数学上的梯度等价要求参数在窗口内固定、损失聚合权重一致、只更新一次，并且前向计算不依赖被拆开的 batch 统计。BatchNorm 会改变这种关系；LayerNorm 不消除 Dropout、数据增强和浮点求和顺序造成的差异。少除分母会先放大梯度，对 Adam 或含裁剪的更新也不能直接称为“学习率放大 $K$ 倍”。

分布式 no_sync() 可减少前几次反向的同步，但必须同时覆盖相应前向；通信实现见分布式训练篇。可核对的窗口代码统一放在训练循环。

2. 裁剪在累积之后做什么#

令 $g$ 为待更新参数的梯度拼接向量，$c>0$ 为全局 L2 范数阈值。按范数裁剪为：

$$g_{\text{clip}}=g\min\!\left(1,\frac{c}{\|g\|_2}\right),$$

其中 $g=0$ 时保持零向量。它限制输入优化器的梯度范数；Adam 的逐坐标归一化和权重衰减仍会改变最终参数更新，不能由此推出参数变化也被限制为 $\eta c$。

手算例子：$g=(3,4)$，范数为 $5$。取 $c=1$，得到 $g_{\text{clip}}=(0.6,0.8)$，范数为 $1$。若先分别裁剪每个 microbatch，再累积，一般不等于裁剪窗口的总梯度。

2.1 AMP 与监控顺序#

累积完成 → 解除 loss scaling → 记录并裁剪真实尺度梯度 → 检查并更新 → 清零。一个窗口中的缩放因子应保持不变；不能中途 unscale 后继续加入已缩放梯度。对应约定见 PyTorch AMP 梯度累积。

clip_grad_norm_ 返回裁剪前范数。每步只调用一次，同时记录该返回值与触发裁剪的频率。阈值需结合模型、损失分母和历史基线判断，没有通用的“正常梯度范数区间”。

3. 累积、重算与分片各解决什么#

累积通过减小单次 microbatch 降低激活峰值；激活重算通过保存边界、反向重算减少激活存储；分布式训练的显存核算解释参数、梯度和优化器状态的分片。累积本身不会减少模型状态，也不保证某个规模能在指定显卡上训练。

参考资料#

• PyTorch 2.9 AMP examples：累积窗口、unscale 与裁剪顺序。
• 梯度裁剪接口：全局范数与返回值约定。