如何让RLlib在搭载Apple M3 Pro的MacBook Pro上使用GPU训练

• 确保使用支持MPS的RLlib版本
  Apple Silicon的MPS加速支持在RLlib 2.9.0及以上版本中才得到完善支持，先升级到对应版本：

  pip install --upgrade ray[rllib]
  

• 明确指定框架并配置GPU参数
  RLlib不会自动启用MPS设备，需要在命令中指定使用的框架，并强制开启GPU支持：

  • 若使用PyTorch，执行以下命令：

    rllib train --algo DQN --env CartPole-v1 --stop '{"training_iteration": 30}' --framework torch --config '{"num_gpus": 1, "torch_device": "mps"}'
    

  • 若使用TensorFlow，执行以下命令：

    rllib train --algo DQN --env CartPole-v1 --stop '{"training_iteration": 30}' --framework tf --config '{"num_gpus": 1}'
    

• 注意小环境的GPU使用特性
  CartPole属于计算量极小的简单环境，GPU调度的开销可能超过加速收益，RLlib会默认选择CPU执行。如果要观察明显的GPU占用，建议换成更复杂的环境（如Breakout-v4）测试：

  rllib train --algo DQN --env Breakout-v4 --stop '{"training_iteration": 30}' --framework torch --config '{"num_gpus": 1, "torch_device": "mps"}'
  

• 验证Ray的GPU资源识别情况
  可以通过以下代码确认Ray是否正确识别MPS设备：

  import ray
  ray.init(num_gpus=1)
  print(ray.available_resources())
  

  输出中包含GPU: 1即表示资源识别正常。

内容的提问来源于stack exchange，提问作者W.P. McNeill