我们与 Apple 的 Metal 工程团队合作，很高兴地宣布支持在 Mac 上进行 GPU 加速的 PyTorch 训练。到目前为止，Mac 上的 PyTorch 训练仅利用 CPU，但随着即将发布的 PyTorch v1.12 版本，开发人员和研究人员可以利用 Apple 芯片 GPU 实现显着更快的模型训练。这解锁了在 Mac 上本地执行机器学习工作流（如原型设计和微调）的能力。
Metal 加速
通过将 Apple 的 Metal Performance Shaders (MPS) 作为 PyTorch 的后端，可以实现 GPU 加速训练。MPS 后端扩展了 PyTorch 框架，提供了在 Mac 上设置和运行操作的脚本和功能。MPS 通过针对每个 Metal GPU 系列的独特特性进行微调的内核来优化计算性能。新设备将机器学习计算图和原语映射到 MPS Graph 框架和 MPS 提供的微调内核。
Apple 芯片上的训练优势
每台 Apple 芯片 Mac 都具有统一内存架构，使 GPU 可以直接访问完整的内存存储。这使得 Mac 成为一个出色的机器学习平台，使用户能够在本地训练更大的网络或批次大小。这降低了与基于云的开发相关的成本或对额外本地 GPU 的需求。统一内存架构还减少了数据检索延迟，提高了端到端性能。
在下面的图中，您可以看到加速 GPU 训练和评估与 CPU 基线相比的性能提升
加速 GPU 训练和评估速度比仅 CPU 快（倍数）
入门
要开始使用，只需在您的 Apple 芯片 Mac（运行 macOS 12.3 或更高版本，并使用原生版本 (arm64) 的 Python）上安装最新的 预览版（Nightly）构建 即可。
您还可以在 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 的大致性能。