Apple M1的MPS后端运行PyTorch模型出现报错求助

出现的警告与错误信息 #

UserWarning: The operator ‘aten::sgn.out’ is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at /Users/runner/work/_temp/anaconda/conda-bld/pytorch_1670525498485/work/aten/src/ATen/mps/MPSFallback.mm:11.)
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass

UserWarning: Error detected in ConvolutionBackward0. Traceback of forward call that caused the error

(Triggered internally at /Users/runner/work/_temp/anaconda/conda-bld/pytorch_1670525498485/work/torch/csrc/autograd/python_anomaly_mode.cpp:119.)
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass

File “/Users/user/miniconda3/envs/torch/lib/python3.10/site-packages/torch/autograd/init.py”, line 197, in backward
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
RuntimeError: Function ‘ConvolutionBackward0’ returned nan values in its 0th output.

系统配置 #

• MacOS: Ventura 13.2
• 芯片：Apple M1
• Python 3.10.9
• PyTorch 1.13.1

解决步骤 #

针对aten::sgn.out不支持MPS的警告 #

1. 升级PyTorch版本：PyTorch 1.13对MPS的支持尚不完善，升级到2.x版本（如2.0及以上），新版本补充了更多MPS兼容算子，大概率修复sgn算子的支持问题。
2. 手动指定设备执行：若暂时无法升级，可将使用sgn的代码段切换到CPU执行，完成后转回MPS设备，示例：

   # 假设x是MPS设备上的张量
   x_cpu = x.cpu()
   x_sgn = torch.sign(x_cpu)  # sgn等价于sign算子
   x_sgn_mps = x_sgn.to('mps')
   

针对ConvolutionBackward0返回NaN的错误 #

1. 降低学习率：过高的学习率易引发梯度爆炸，导致NaN出现。尝试将学习率下调一个数量级（如从1e-3改为1e-4），观察问题是否消失。
2. 检查输入数据：确认训练/验证数据中无NaN、Inf值，同时检查数据预处理流程（如归一化、标准化）是否正确，避免输入分布异常。
3. 启用梯度裁剪：在反向传播后添加梯度裁剪，限制梯度最大范数，防止梯度爆炸：

   loss.backward()
   torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)  # 可根据模型调整max_norm值
   optimizer.step()
   

4. 检查卷积层配置：确保卷积层权重初始化合理（如默认的Kaiming初始化），排查是否使用了易产生NaN的操作（如未加稳定项的对数、除法等）。
5. 切换到CPU验证：若上述方法无效，将模型和数据切换到CPU运行，若CPU运行正常，说明是PyTorch MPS后端的卷积反向传播bug，可升级PyTorch或等待官方修复。

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