我們與 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 的大致效能。