【Bug已解决】DAPO loss is not normalized by gradient_accumulation_steps 解决方案

【Bug已解决】DAPO loss is not normalized by gradient_accumulation_steps 解决方案一、现象长什么样用DAPODynamic Sampling Policy Optimization训练时我们观察到调大gradient_accumulation_steps后训练行为明显变激进loss 数值也跟着变大但训练效果反而变差。具体日志gradient_accumulation_steps1 - loss ≈ 0.42, 收敛正常 gradient_accumulation_steps4 - loss ≈ 1.68 (0.42×4), 训练发散/不稳现象特征loss 随gradient_accumulation_steps近似线性放大×N但理论上梯度累积只是把 N 个小批的梯度加起来再更新一步等效 batch per_device_batch × grad_accum × world_sizeloss 的标量数值应当和累积步数无关或被归一化掉只是更新频率变了DAPO 的 loss 没被gradient_accumulation_steps除于是梯度被放大约 N 倍等价于学习率被偷偷放大 N 倍训练不稳。这是典型的自定义 loss 没接入 Trainer 的梯度累积归一化问题。二、背景HuggingFaceTrainer里当开启梯度累积时每个 micro-batch 算出的 loss 会先除以gradient_accumulation_steps再backward这样累加 N 个 micro-batch 的梯度后总梯度恰好等于把这 N 个样本当一个大 batch 算一次的梯度。即loss loss / gradient_accumulation_steps loss.backward() # 每个 micro-batch 都除过累加后归一 if (step 1) % gradient_accumulation_steps 0: optimizer.step()这一步归一化是保证累积 N 步 一个大 batch在数值上等价的关键。如果某个 Trainer如 DAPO自己重写了compute_loss但忘了做这个除法那么 N 个 micro-batch 的梯度原样相加总梯度就变成正确值 × N等价于学习率被放大 N 倍。DAPO 因为要做动态采样过滤掉整组 reward 标准差为 0 的样本重写了 loss 计算却漏掉了这个归一化因子于是踩雷。三、根因根因一句话DAPO 的compute_loss重写了损失计算含动态采样过滤但没有像基类 Trainer 那样把每个 micro-batch 的 loss 除以gradient_accumulation_steps导致梯度累积时 N 个 micro-batch 的梯度被原样相加总梯度被放大 N 倍等价于学习率被放大 N 倍训练随累积步数变不稳。具体归一化缺失loss.backward()前没有loss loss / grad_accum_steps梯度放大N 个 micro-batch 各贡献未除的 loss 梯度累加后 正确 × N学习率等效放大优化器 step 时用的梯度是 N 倍等价于 LR × N只在 grad_accum1 暴露grad_accum1时除以 1 无影响一切正常一旦 1 就偏差且偏差随步数线性增长训练发散。本质是自定义优化器/损失路径绕过了 Trainer 的梯度累积归一化约定。四、最小可运行复现下面用纯 Python 模拟归一化缺失导致梯度被放大约 N 倍的机制def accumulate_no_norm(micro_losses, grad_accum): 旧实现直接累加每个 micro-batch 的 loss 梯度没除 grad_accum。 total 0.0 for l in micro_losses: total l # 漏了 / grad_accum return total def accumulate_with_norm(micro_losses, grad_accum): 正确实现每个 micro-batch loss 先除 grad_accum 再累加。 total 0.0 for l in micro_losses: total l / grad_accum return total def demo(): micro [0.42, 0.40, 0.44, 0.41] # 4 个 micro-batch g 4 no_norm accumulate_no_norm(micro, g) with_norm accumulate_with_norm(micro, g) print(f未归一化总梯度 {no_norm:.3f}) print(f归一化后总梯度 {with_norm:.3f}) print(f比值 {no_norm / with_norm:.1f}x ( grad_accum 倍等价于 LR 被放大)) if __name__ __main__: demo()输出未归一化总梯度 1.670 归一化后总梯度 0.418 比值 4.0x ( grad_accum 倍等价于 LR 被放大)第一行1.670 ≈ 0.42×4正是线上现象比值 4.0x 说明梯度被放大了gradient_accumulation_steps倍等价于学习率翻 4 倍。复现了核心 bug。五、解决方案第一层compute_loss 里除以 gradient_accumulation_steps第一层最直接在 DAPO 的compute_loss返回 loss 前除以累积步数与基类行为对齐class DAPOTrainer: def __init__(self, gradient_accumulation_steps: int 1): self.gradient_accumulation_steps gradient_accumulation_steps def compute_loss(self, model, inputs, return_outputsFalse): # ... DAPO 的动态采样 策略损失计算 ... per_token_loss self._dapo_loss(model, inputs) loss per_token_loss.mean() # 关键与 Trainer 基类一致按累积步数归一化 loss loss / self.gradient_accumulation_steps return (loss, outputs) if return_outputs else loss核心是loss loss / self.gradient_accumulation_steps。这样每个 micro-batch 的梯度被等比例缩小累加 N 个后总梯度回到大 batch 等价的正确值loss 标量也和grad_accum无关。修复后grad_accum4与grad_accum1的训练应表现一致仅更新频率不同不再发散。六、解决方案第二层复用基类归一化逻辑避免手写遗漏第一层是补丁但手写除法容易在重构时又被漏掉。第二层从结构上保证DAPO 的 loss 走和基类一样的归一化路径或者显式调用基类的归一化辅助函数class DAPOTrainer: def compute_loss(self, model, inputs, return_outputsFalse): per_token_loss self._dapo_loss(model, inputs) loss per_token_loss.mean() # 用统一辅助函数做归一化避免各处手写 loss self._normalize_for_grad_accum(loss) return (loss, outputs) if return_outputs else loss def _normalize_for_grad_accum(self, loss): 唯一真源梯度累积归一化。 g getattr(self, gradient_accumulation_steps, 1) or 1 return loss / g def demo(): t DAPOTrainer(gradient_accumulation_steps4) base 0.42 print(归一化后单 micro-batch loss , t._normalize_for_grad_accum(base).item() if hasattr(t._normalize_for_grad_accum(base), item) else t._normalize_for_grad_accum(base)) if __name__ __main__: demo()把归一化收成_normalize_for_grad_accum唯一函数任何重算 loss 的路径都调它避免有的路径除、有的路径忘除的漂移。七、解决方案第三层断言归一化生效 不变量测试第三层加护栏确保loss 数值与 grad_accum 无关这一不变量被锁定def effective_loss_after_accum(micro_losses, grad_accum): return sum(l / grad_accum for l in micro_losses) def test_loss_independent_of_grad_accum(): micro [0.42, 0.40, 0.44] # 不论累积步数多少累加后的有效平均 loss应一致 eff_1 effective_loss_after_accum(micro, 1) eff_3 effective_loss_after_accum(micro, 3) assert abs(eff_1 - eff_3) 1e-9, loss 应不随 grad_accum 变化 print(fOK: grad_accum1 有效 loss{eff_1:.4f}, grad_accum3 有效 loss{eff_3:.4f}) def test_grad_scale_correct(): # 模拟未归一化会被放大 N 倍归一化后不变 micro [0.42] * 4 bad sum(micro) # 未归一 good sum(x / 4 for x in micro) # 归一 assert abs(good - 0.42) 1e-9 and abs(bad - 1.68) 1e-9 print(OK: 归一化后梯度不被 grad_accum 放大) if __name__ __main__: test_loss_independent_of_grad_accum() test_grad_scale_correct()两个测试分别锁住有效 loss 与 grad_accum 无关和归一化后梯度不被放大 N 倍。任何把除法漏掉的改动都会让断言失败CI 直接拦下。八、落地建议如果你在 DAPO或任何自定义 Trainer上发现调大 grad_accum 训练变激进建议确认 compute_loss 是否除以 grad_accum没有就加loss / gradient_accumulation_steps。复用统一辅助函数把归一化收成_normalize_for_grad_accum避免手写遗漏。验证数值一致grad_accum1与grad_accumN下有效平均 loss 应相同。加不变量测试锁住loss 与 grad_accum 无关、梯度不被放大。对照基类HuggingFaceTrainer的training_step内有同样的除法照此对齐。监控 grad_norm调大 grad_accum 后 grad_norm 应稳定不应线性放大。九、排查清单如果调大 gradient_accumulation_steps 后训练变激进/发散按顺序查看 loss 是否随 grad_accum 线性放大是则归一化缺失。搜 compute_loss返回的 loss 是否有/ gradient_accumulation_steps。确认是自定义 TrainerDAPO 等重写 loss 的 Trainer 最易漏这步。复用辅助函数把归一化收成唯一函数防重构遗漏。看 grad_norm未归一化时 grad_norm 会随 grad_accum 放大。加断言/测试锁住有效 loss 与 grad_accum 无关。对照基类 Trainer其training_step里有同样的除法逻辑。十、小结DAPO 训练调大gradient_accumulation_steps后变激进/发散根因是DAPO 重写了compute_loss含动态采样过滤却漏掉了 Trainer 基类默认做的每个 micro-batch loss 除以gradient_accumulation_steps归一化。于是梯度累积时 N 个 micro-batch 的未除 loss 梯度被原样相加总梯度被放大 N 倍等价于学习率被偷偷放大 N 倍训练随累积步数线性变不稳。它只在grad_accum 1时暴露1时除以 1 无影响所以容易在调参时才发现。修复分三层第一层在compute_loss返回前加loss / gradient_accumulation_steps与基类对齐梯度回到大 batch 等价正确值第二层把归一化收成_normalize_for_grad_accum唯一辅助函数任何重算 loss 的路径都调它避免重构遗漏第三层加有效 loss 与 grad_accum 无关、梯度不被放大 N 倍不变量测试把回归在 CI 拦下。核心心法是任何自定义 Trainer 重写 loss 时都必须复刻基类的梯度累积归一化——否则梯度累积不再等价于大 batch学习率会被静默放大训练稳定性悄悄崩掉。