与苹果公司的金属工程团队合作,我们很高兴地宣布支持在 Mac 上进行 GPU 加速的 PyTorch 训练。到目前为止,Mac 上的 PyTorch 训练仅利用 CPU,但随着即将发布的 PyTorch v1.12 版本的推出,开发人员和研究人员可以利用苹果硅 GPU 进行显著更快的模型训练。这解锁了在 Mac 上本地执行机器学习工作流程(如原型设计和微调)的能力。
金属加速
使用苹果的 Metal Performance Shaders (MPS) 作为 PyTorch 的后端来启用加速 GPU 训练。MPS 后端扩展了 PyTorch 框架,提供脚本和能力来设置和运行在 Mac 上的操作。MPS 通过针对每个 Metal GPU 家族的独特特性进行优化的内核来优化计算性能。新的设备将机器学习计算图和原语映射到 MPS Graph 框架和 MPS 提供的调优内核上。
苹果硅芯片的培训优势
每款苹果硅芯片的 Mac 都拥有统一的内存架构,为 GPU 提供直接访问整个内存存储的能力。这使得 Mac 成为机器学习的优秀平台,使用户能够本地训练更大的网络或批量大小。这降低了与基于云的开发或需要额外本地 GPU 相关的成本。统一的内存架构还减少了数据检索延迟,提高了端到端性能。
在下面的图表中,您可以查看与 CPU 基准相比,加速的 GPU 训练和评估的性能提升:
与仅 CPU 相比,加速的 GPU 训练和评估速度提升(倍数)
开始使用
要开始使用,只需在运行 macOS 12.3 或更高版本且具有原生版本(arm64)Python 的 Apple 硅 Mac 上安装最新预览版(夜间版)。
您还可以在 Apple 的 Metal 页面上了解更多关于 Metal 和 MPS 的信息。
* 苹果公司在 2022 年 4 月使用生产型 Mac Studio 系统进行测试,该系统配备 Apple M1 Ultra、20 核心 CPU、64 核心 GPU、128GB RAM 和 2TB SSD。测试使用 macOS Monterey 12.3、预发布 PyTorch 1.12、ResNet50(批量大小=128)、HuggingFace BERT(批量大小=64)和 VGG16(批量大小=64)。性能测试是在特定的计算机系统上进行的,反映了 Mac Studio 的大致性能。