Accelerate 文档

在 Mac 上加速 PyTorch 训练

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

在 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 机器上,此功能默认启用。要禁用它,请在 accelerate launch 命令中传递 --cpu 标志,或在回答 accelerate config 问卷时回答相应问题。

您可以直接运行以下脚本,在支持 MPS 的 Apple 芯片机器上进行测试

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

需要注意的几个问题

  1. 分布式设置 gloonccl 无法与 mps 设备一起工作。这意味着目前只能使用 mps 设备类型的单个 GPU。

最后,请记住,Accelerate 仅集成了 MPS 后端,因此如果您对 MPS 后端的使用有任何问题或疑问,请向 PyTorch GitHub 提交问题。

< > 在 GitHub 上更新