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