Mac端用Lightning Trainer运行MPNN遇MPS算子未实现错误求助

核心问题原因 #

你设置的PYTORCH_ENABLE_MPS_FALLBACK=1环境变量未生效，大概率是设置时机错误——这个变量必须在导入PyTorch、PyTorch Lightning等相关库之前配置，否则不会起作用。

可行解决方案 #

1. 提前设置环境变量（最有效） #

将环境变量配置代码放在所有PyTorch相关导入的最前面：

import os
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'

# 之后再导入相关库
import torch
import pytorch_lightning as pl
# 导入你的MPNN模型、数据加载器等模块

2. 强制特定算子回退到CPU #

如果上述方法仍有部分算子不兼容，可以手动将涉及scatter_reduce的操作切换到CPU执行。比如在MPNN模型的前向传播中做设备切换：

class MPNNModel(pl.LightningModule):
    def forward(self, x, edge_index):
        # 先执行MPS支持的操作
        # 遇到scatter_reduce相关操作时，临时切换到CPU
        x_cpu = x.to('cpu')
        edge_index_cpu = edge_index.to('cpu')
        result_cpu = self.scatter_based_operation(x_cpu, edge_index_cpu)
        # 计算完成后切回MPS设备
        return result_cpu.to('mps')

3. 暂时使用CPU训练（兜底方案） #

如果MPS兼容问题暂时无法解决，可直接切换到CPU训练，修改Trainer配置：

trainer = pl.Trainer(
    logger=False,
    enable_checkpointing=True,
    enable_progress_bar=True,
    accelerator="cpu",  # 改为CPU加速器
    devices=1,
    max_epochs=20,
)

4. 更新PyTorch版本 #

部分算子的MPS支持会在新版本中修复，尝试升级到最新稳定版：

pip install --upgrade torch torchvision torchaudio

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