M1 MacBook原生运行PyTorch出现准确率低、loss异常问题求助

1. MPS加速后端适配问题 #

• 原因：原生PyTorch for M1默认调用Metal Performance Shaders(MPS)作为GPU加速后端，早期版本的MPS算子存在大量精度溢出、逻辑实现错误的问题，尤其是损失计算、反向传播相关的算子适配不全，会直接导致loss不下降、出现nan、准确率异常低的情况
• 解决办法：
  • 临时验证后端问题：在代码最开头添加import os; os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"，开启不支持的算子自动回退到CPU运行，如果训练恢复正常即可确认是MPS适配问题
  • 升级PyTorch到最新稳定版：1.13及之后版本的PyTorch已经修复了绝大多数MPS算子的适配问题，不要使用测试阶段的原生PyTorch编译版
  • 临时切换CPU后端保证稳定：把代码中指定设备的.to("mps")修改为.to("cpu")，小规模训练场景下CPU运行的速度和MPS差距不大，稳定性更高

2. Mini-forge环境依赖冲突 #

• 原因：Mini-forge默认安装的arm64版本依赖库（如numpy、PIL等）如果和PyTorch版本不匹配，会导致数据预处理阶段的数值范围异常，输入模型的张量不符合训练要求，最终训练失效
• 解决办法：
  • 卸载当前环境的PyTorch相关包：conda uninstall pytorch torchvision torchaudio
  • 使用PyTorch官方源重装适配版本：conda install pytorch::pytorch torchvision torchaudio -c pytorch，不要使用conda-forge源的非官方编译版本
  • 排查依赖冲突：运行pip check检查是否存在版本冲突，将冲突的依赖调整到兼容版本

3. 浮点精度适配问题 #

• 原因：M1芯片GPU默认优先使用FP16精度计算，很多通用示例代码没有针对FP16做溢出适配，容易出现梯度下溢/上溢的问题，导致loss为nan或者不收敛
• 解决办法：
  • 关闭自动混合精度：如果代码中开启了torch.cuda.amp相关的混合精度逻辑，直接关闭该配置，强制使用FP32训练
  • 给损失函数增加数值稳定配置：比如在交叉熵损失中添加小的epsilon参数，避免出现log0的异常计算

修复效果验证步骤 #

• 运行MNIST示例前3个epoch，观察测试集准确率是否超过97%
• 运行CIFAR10示例前100步训练，观察loss是否正常下降，无nan值出现

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