Mac设备PyTorch中MPS GPU无法启用及pickle错误求助

一、MPS设备未被识别（accelerator=None）的处理 #

• 先确认MPS支持是否正常：PyTorch 1.12.1对MPS的支持还处于早期阶段，需确保你的MacOS版本在12.3以上，且安装的PyTorch确实支持MPS。运行以下代码验证：

  import torch
  print(torch.backends.mps.is_available())
  print(torch.backends.mps.is_built())
  

• 手动指定设备：不要依赖自动检测，直接在训练代码里显式设置设备为MPS，示例代码如下：

  device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
  model.to(device)
  

• 若使用PyTorch Lightning这类框架，需在Trainer中明确指定accelerator='mps'，早期版本自动检测容易出现兼容性问题。

二、Pickle本地对象错误的处理 #

该错误是因为get_cosine_schedule_with_warmup生成的lr_lambda是本地函数，MPS环境下多进程数据加载无法序列化（pickle）本地对象，而Colab使用CUDA，进程通信机制不同所以没有这个问题。

• 将lr_lambda移到模块全局作用域：不要在函数内部定义这个lambda，放到代码顶层，示例如下：

  import math
  
  def lr_lambda(current_step):
      warmup_steps = 100  # 替换为你实际的warmup步数
      total_steps = 1000  # 替换为你实际的总训练步数
      if current_step < warmup_steps:
          return float(current_step) / float(max(1, warmup_steps))
      progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
      return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))
  
  scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lr_lambda)
  

• 替换调度器：改用torch.optim.lr_scheduler.CosineAnnealingWarmRestarts，它能实现类似的余弦退火+warmup效果，且无需自定义lambda函数，避开pickle问题。
• 临时关闭多进程数据加载：将DataLoader的num_workers参数设为0，虽然会减慢数据加载速度，但能快速验证是否是多进程pickle导致的错误。

三、验证流程 #

先修改学习率调度器代码解决pickle错误，再手动指定MPS设备，重新运行训练命令python train_torch.py --train --max_epochs 2，观察设备识别状态和训练是否正常。

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