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

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

一、现象长什么样

用 DAPO（Dynamic Sampling Policy Optimization）训练时，我们观察到：调大 gradient_accumulation_steps 后，训练行为明显变激进，loss 数值也跟着变大，但训练效果反而变差。具体日志：

gradient_accumulation_steps=1 -> loss ≈ 0.42, 收敛正常
gradient_accumulation_steps=4 -> loss ≈ 1.68 (=0.42×4), 训练发散/不稳

现象特征：

• loss 随 gradient_accumulation_steps 近似线性放大（×N）；
• 但理论上，梯度累积只是"把 N 个小批的梯度加起来再更新一步"，等效 batch = per_device_batch × grad_accum × world_size，loss 的标量数值应当和累积步数无关（或被归一化掉），只是更新频率变了；
• DAPO 的 loss 没被 gradient_accumulation_steps 除，于是梯度被放大约 N 倍，等价于学习率被偷偷放大 N 倍，训练不稳。

这是典型的"自定义 loss 没接入 Trainer 的梯度累积归一化"问题。

二、背景

HuggingFace Trainer 里，当开启梯度累积时，每个 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 倍，训练随累积步数变不稳。

具体：

1. 归一化缺失 ：loss.backward() 前没有 loss = loss / grad_accum_steps；
2. 梯度放大：N 个 micro-batch 各贡献未除的 loss 梯度，累加后 = 正确 × N；
3. 学习率等效放大：优化器 step 时用的梯度是 N 倍，等价于 LR × N；
4. 只在 grad_accum>1 暴露 ：grad_accum=1 时除以 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_outputs=False):
        # ... 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_accum=4 与 grad_accum=1 的训练应表现一致（仅更新频率不同），不再发散。

六、解决方案（第二层）：复用基类归一化逻辑，避免手写遗漏

第一层是补丁，但"手写除法"容易在重构时又被漏掉。第二层从结构上保证：DAPO 的 loss 走和基类一样的归一化路径，或者显式调用基类的归一化辅助函数：

class DAPOTrainer:
    def compute_loss(self, model, inputs, return_outputs=False):
        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_steps=4)
    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(f"OK: grad_accum=1 有效 loss={eff_1:.4f}, grad_accum=3 有效 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 训练变激进"，建议：

1. 确认 compute_loss 是否除以 grad_accum ：没有就加 loss / gradient_accumulation_steps。
2. 复用统一辅助函数 ：把归一化收成 _normalize_for_grad_accum，避免手写遗漏。
3. 验证数值一致 ：grad_accum=1 与 grad_accum=N 下，有效平均 loss 应相同。
4. 加不变量测试：锁住"loss 与 grad_accum 无关、梯度不被放大"。
5. 对照基类 ：HuggingFace Trainer 的 training_step 内有同样的除法，照此对齐。
6. 监控 grad_norm：调大 grad_accum 后 grad_norm 应稳定，不应线性放大。

九、排查清单

如果调大 gradient_accumulation_steps 后训练变激进/发散，按顺序查：

1. 看 loss 是否随 grad_accum 线性放大：是则归一化缺失。
2. 搜 compute_loss ：返回的 loss 是否有 / gradient_accumulation_steps。
3. 确认是自定义 Trainer：DAPO 等重写 loss 的 Trainer 最易漏这步。
4. 复用辅助函数：把归一化收成唯一函数，防重构遗漏。
5. 看 grad_norm：未归一化时 grad_norm 会随 grad_accum 放大。
6. 加断言/测试：锁住"有效 loss 与 grad_accum 无关"。
7. 对照基类 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，学习率会被静默放大，训练稳定性悄悄崩掉。