PyTorch 宣布，通过与 Apple 的 Metal 工程团队合作，目前已实现在搭载 Apple M1 芯片的 Mac 上使用 GPU 加速训练。在这之前，在 Mac 上进行 PyTorch 训练只能使用 CPU，但随着 PyTorch v1.12 即将发布，开发者和研究者可以利用 Apple M1 GPU 的强大性能，从而显著提升模型训练速度。
对于 Mac 来说，此功能极大丰富了在 Mac 上执行机器学习工作流程的能力，例如在本地进行原型设计和微调。
![]()
根据 PyTorch 官方博客的介绍，PyTorch 通过使用 Apple 的 Metal Performance Shaders (MPS) 作为 PyTorch 的后端来启用 GPU 加速训练。MPS 后端扩展了 PyTorch 框架，提供了在 Mac 上设置和运行操作的脚本和功能。MPS 使用针对每个 Metal GPU 系列的独特特性进行微调的内核来优化计算性能。新设备将机器学习计算图和原语映射到 MPS Graph 框架和 MPS 提供的调整内核上。
此外，由于搭载 Apple M1 芯片的 Mac 具有统一的内存结构，因此 GPU 可以直接访问完整的内存存储。这使得 Mac 成为机器学习的绝佳平台，使用户能够在本地训练更大的网络或批大小，并降低了与基于云算力的开发相关的成本或对额外本地 GPU 的需求。统一内存架构还减少了数据检索延迟，提高了端到端的性能。
下图展示了与 CPU 基线相比，GPU 加速训练带来了显著的性能提升：
![]()
可以看到，使用 GPU 可将模型训练速度提升约 7 倍，评估 (evaluation) 速度则最高提升约 20 倍。
Apple 于 2022 年 4 月使用配备 Apple M1 Ultra、20 核 CPU、64 核 GPU、128GB 内存 和 2TB SSD 的 Mac Studio 进行了测试。系统为 macOS Monterey 12.3，安装了预发布的 PyTorch 1.12、测试模型为 ResNet50 (batch size=128), HuggingFace BERT (batch size=64) 和 VGG16 (batch size=64)
虽然 PyTorch 已官宣适配 Apple M1， 不过毕竟尚未正式推出，所以并非所有基于 PyTorch 的模型都能使用 M1 的 GPU 加速。大家可以自行测试体验，环境要求：需要在搭载 M1 系列芯片的 Mac 上安装原生版本 (arm64) 的 Python，以及安装最新预览版 PyTorch(Preview build)，并将系统升级至 macOS 12.3 预览版或更新的版本。
![]()
