PyNative Trainer 在完成梯度累积后会无条件调用 optimizer。当本步 loss 或梯度中出现 NaN/Inf 时,非有限梯度仍会进入 optimizer,导致参数、master weight、动量等优化器状态被污染;部分优化器的内部 global step 也会继续推进。
期望行为:
当前 Trainer._optimizer_update() 在 _calculate_global_grad_norm() 返回后直接调用 self.optimizer(grads)。当 global_norm 为 NaN 时,原有 clip_coef < 1 判断不会拦截非有限梯度,因此 optimizer 仍会更新。
Trainer._optimizer_update()
_calculate_global_grad_norm()
self.optimizer(grads)
global_norm
clip_coef < 1
此外,仅检查梯度范数不足以覆盖“loss 为 NaN/Inf、但梯度暂时有限”的情况;分布式场景也必须同步跳过决定,避免不同 rank 在 optimizer collective 处分叉。
r2.1.0-beta1
git diff --check
说明:当前验证为 CPU 单元测试与静态分布式控制流模拟,尚未执行真实多卡 NPU 回归。
upstream/r2.1.0-beta1
57468751333e52a5d03fd2a03b13777bc869dcd7
Checklist
问题详细描述
PyNative Trainer 在完成梯度累积后会无条件调用 optimizer。当本步 loss 或梯度中出现 NaN/Inf 时,非有限梯度仍会进入 optimizer,导致参数、master weight、动量等优化器状态被污染;部分优化器的内部 global step 也会继续推进。
期望行为:
复现与原因
当前
Trainer._optimizer_update()在_calculate_global_grad_norm()返回后直接调用self.optimizer(grads)。当global_norm为 NaN 时,原有clip_coef < 1判断不会拦截非有限梯度,因此 optimizer 仍会更新。此外,仅检查梯度范数不足以覆盖“loss 为 NaN/Inf、但梯度暂时有限”的情况;分布式场景也必须同步跳过决定,避免不同 rank 在 optimizer collective 处分叉。
影响范围
r2.1.0-beta1验证
git diff --check:通过。说明:当前验证为 CPU 单元测试与静态分布式控制流模拟,尚未执行真实多卡 NPU 回归。
版本信息
upstream/r2.1.0-beta1:57468751333e52a5d03fd2a03b13777bc869dcd7