PyTorch 官宣適配 Apple M1,支援 GPU 加速,訓練速度提升 7 倍

2022-05-20 09:00:15

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(),並將系統升級至 macOS 12.3 預覽版或更新的版本。

展開閱讀全文