已开启
[Bug]: PyNative 非有限 loss/梯度仍执行优化器更新 #2532
Java不加糖创建于  8月20日
Java不加糖成员
8月20日 创建

Checklist

问题详细描述

PyNative Trainer 在完成梯度累积后会无条件调用 optimizer。当本步 loss 或梯度中出现 NaN/Inf 时,非有限梯度仍会进入 optimizer,导致参数、master weight、动量等优化器状态被污染;部分优化器的内部 global step 也会继续推进。

期望行为:

  • 任意 rank 检测到本步 loss 或全局梯度范数非有限时,所有 rank 一致跳过本次 optimizer 调用;
  • 本步梯度仍需清零,避免污染下一步;
  • 正常有限值训练路径、梯度裁剪接口和已有返回值保持兼容。

复现与原因

当前 Trainer._optimizer_update() 在 _calculate_global_grad_norm() 返回后直接调用 self.optimizer(grads)。当 global_norm 为 NaN 时,原有 clip_coef < 1 判断不会拦截非有限梯度,因此 optimizer 仍会更新。

此外,仅检查梯度范数不足以覆盖“loss 为 NaN/Inf、但梯度暂时有限”的情况;分布式场景也必须同步跳过决定,避免不同 rank 在 optimizer collective 处分叉。

影响范围

  • MindFormers PyNative Trainer
  • 单卡及数据并行、流水线并行等多 rank 训练
  • 目标分支:r2.1.0-beta1

验证

  • 新增 finite、NaN/Inf loss、NaN/Inf gradient、梯度清零、DTensor-like loss、远端 rank 非有限标志、梯度累积和 PP 调用链单元测试;
  • 定向测试:21 passed;
  • 仓库固定 Pylint 3.2.6:通过;
  • git diff --check:通过。

说明:当前验证为 CPU 单元测试与静态分布式控制流模拟,尚未执行真实多卡 NPU 回归。

版本信息

  • upstream/r2.1.0-beta1: 57468751333e52a5d03fd2a03b13777bc869dcd7
  • MindSpore / 硬件环境:不依赖特定硬件的代码路径问题;本地回归使用 CPU 环境
likedislike
JJava不加糖成员
8月20日 关联了pull request:【r2.1.0-beta1】【pynative】【bugfix】检测到 NaN/Inf 时跳过优化器更新
JJava不加糖成员
8月20日 关联了pull request:【pynative】【bugfix】检测到 NaN/Inf 时跳过优化器更新