在 Mac 上加速 PyTorch 训练

随着 PyTorch v1.12 版本的发布，开发者和研究人员可以利用 Apple 芯片 GPU 来显著加速模型训练。这解锁了在本地 Mac 上执行机器学习工作流程（如原型设计和微调）的能力。Apple 的 Metal Performance Shaders (MPS) 作为 PyTorch 的后端实现了这一点，并且可以通过新的 "mps" 设备使用。这将计算图和原语映射到 MPS Graph 框架和 MPS 提供的调整内核上。更多信息请参考官方文档 介绍在 Mac 上加速 PyTorch 训练 和 MPS 后端。

使用 Apple 芯片进行训练和推理的优势

1. 使用户能够在本地训练更大的网络或批量大小
2. 由于统一内存架构，减少了数据检索延迟，并为 GPU 提供了对完整内存存储的直接访问。因此，提高了端到端性能。
3. 降低了与基于云的开发或对额外本地 GPU 的需求相关的成本。

先决条件：要安装支持 mps 的 torch，请参考这篇 Medium 文章 GPU 加速来到 M1 Mac 上的 PyTorch。

开箱即用

在配备 MPS 功能 Apple 芯片 GPU 的 MacOS 机器上，默认启用此功能。要禁用它，请将 --cpu 标志传递给 accelerate launch 命令，或在回答 accelerate config 问卷时回答相应的问题。

您可以直接运行以下脚本，在启用 MPS 的 Apple 芯片机器上进行测试

accelerate launch /examples/cv_example.py --data_dir images

需要注意的一些注意事项

1. 分布式设置 gloo 和 nccl 不适用于 mps 设备。这意味着目前只能使用 mps 设备类型的单个 GPU。

最后，请记住，Accelerate 仅集成了 MPS 后端，因此如果您在使用 MPS 后端方面有任何问题或疑问，请在 PyTorch GitHub 上提交 issue。