CUDA运行正常，CPU下PowBackward0反向传播返回NaN问题求助

一、为何仅CPU环境触发NaN问题？ #

• 数值精度差异：CUDA默认用FP32单精度浮点数，CPU常默认FP64双精度，两者浮点数舍入、精度处理逻辑有细微差别。计算grad_x**2 + grad_y**2时，CPU的高精度运算会暴露被CUDA精度掩盖的极小数值误差（比如接近0的负数），反向传播时触发NaN。
• Autograd底层实现不同：CPU和CUDA的PowBackward0（平方操作反向传播）实现逻辑有差异，CUDA版本对边缘情况（如输入接近0）做了宽松处理，CPU版本计算更严格，容易在极小值的导数计算中产生NaN。
• LBFGS迭代特性放大偏差：LBFGS是二阶优化器，会积累历史梯度信息，CPU环境的数值精度差异导致的微小偏差，会在迭代中被放大，最终触发反向传播的NaN问题。

二、正向无NaN但反向出现的原因 #

• 反向导数计算的特性：正向计算sqrt(a)（a=grad_x²+grad_y²）时，只要a≥0就不会有NaN，但反向传播时sqrt的导数是1/(2*sqrt(a))。当a趋近于0时，CPU高精度运算下这个导数会趋近于无穷大，与其他数值相乘后溢出产生NaN；而CUDA单精度会将无穷大截断为极大有限值，不会触发NaN。
• 正向打印的局限性：你打印的是正向传播的最终结果，但反向传播依赖计算图中保存的中间张量。可能存在中间张量被in-place修改，或者打印时的精度截断掩盖了极小值，导致正向看起来正常，但反向传播时这些极小值被放大引发异常。
• LBFGS线搜索的影响：LBFGS的线搜索会尝试不同步长，过程中可能生成极端中间参数值，这些值正向传播时表现正常，但反向传播时会触发导数计算异常。

解决建议 #

• 添加数值稳定项：在sqrt内加入极小常数，避免a趋近于0：

  self.grad_mag = torch.sqrt(self.grad_x**2 + self.grad_y**2 + 1e-8)
  

• 统一数值精度：强制CPU环境使用FP32，和CUDA保持一致：

  grad_x = grad_x.float()
  grad_y = grad_y.float()
  self.grad_mag = torch.sqrt(grad_x**2 + grad_y**2)
  

• 检查in-place操作：确保未对grad_x、grad_y等张量做in-place修改，这类操作会破坏Autograd的中间缓存，引发反向传播数值异常。
• 调整LBFGS参数：减小max_iter、调整tol避免迭代过度发散，或者限制history_size减少历史梯度积累量。

内容的提问来源于stack exchange，提问作者Cedric Martens