來源:機器之心
機器之心報道
編輯:澤南、蛋醬
對于 Mac 用戶來說,這是令人激動的一天。
今年 3 月,蘋果發布了其自研 M1 芯片的最終型號 M1 Ultra,它由 1140 億個晶體管組成,是有史以來個人計算機中最大的數字。蘋果宣稱只需 1/3 的功耗,M1 Ultra 就可以實現比桌面級 GPU RTX 3090 更高的性能。
隨著用戶數量的增長,人們已經逐漸接受使用 M1 芯片的計算機,但作為一款 Arm 架構芯片,還有人在擔心部分任務的兼容性問題。
昨天,通過與蘋果 Metal 團隊工程師合作,PyTorch 官方宣布已正式支持在 M1 版本的 Mac 上進行 GPU 加速的 PyTorch 機器學習模型訓練。
此前,Mac 上的 PyTorch 訓練僅能利用 CPU,但隨著即將發布的 PyTorch v1.12 版本,開發和研究人員可以利用蘋果 GPU 大幅度加快模型訓練。現在,人們可以在 Mac 上相對高效地執行機器學習工作,例如在本地進行原型設計和微調。
蘋果芯片的 AI 訓練優勢
PyTorch GPU 訓練加速是使用蘋果 Metal Performance Shaders ( MPS ) 作為后端來實現的。MPS 后端擴展了 PyTorch 框架,提供了在 Mac 上設置和運行操作的腳本和功能。MPS 使用針對每個 Metal GPU 系列的獨特特性進行微調的內核能力來優化計算性能。新設備將機器學習計算圖和原語映射到 MPS Graph 框架和 MPS 提供的調整內核上。
每臺搭載蘋果自研芯片的 Mac 都有著統一的內存架構,讓 GPU 可以直接訪問完整的內存存儲。PyTorch 官方表示,這使得 Mac 成為機器學習的絕佳平臺,讓用戶能夠在本地訓練更大的網絡或批大小。
這降低了與基于云算力的開發相關的成本或對額外的本地 GPU 算力需求。統一內存架構還減少了數據檢索延遲,提高了端到端性能。
可以看到,與 CPU 基線相比,GPU 加速實現了成倍的訓練性能提升:
上圖是蘋果于 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)。性能測試是使用特定的計算機系統進行的,反映了 Mac Studio 的大致性能。
有開發者推測,鑒于谷歌云服務中使用的英偉達 T4 在 FP32 任務上的浮點性能為 8 TFLOPS,而 M1 Ultra 的圖形計算能力大概在 20 TFLOPS 左右。在最有利情況下,可以期望的 M1 Ultra 速度提升或可達到 2.5 倍。
若想使用最新的加速能力,你需要在使用 M1 系列芯片的 Mac 電腦上安裝原生版本(arm64)的 Python,并將系統升級至 macOS 12.3 預覽版或更新的版本。
開發者親測:加速效果顯著
雖然官方已宣布提供支持,但目前還不是所有在 PyTorch 上的模型都能用 M1 芯片集成的 GPU 加速,你也可以花幾分鐘進行一下測試。
機器學習研究者,捷克理工大學博士 Dmytro Mishkin 對多個模型的推理進行了測試,結果顯示,大多數圖像分類架構都提供了很好的加速。對于一些自定義代碼(比如 kornia),可能無法正常工作。
各個測試結果如下:
首先是經典的卷積神經網絡 VGG16,從 2.23 秒提升到 0.5 秒:
接下來是大部分芯片發布會上都會跑的 Resnet50,它在 M1 GPU 上的速度較慢,不升反降,從 0.549 秒到 0.592 秒:
但 ResNet18 的提速驚人,從 0.243 秒到 0.024 秒:
AlexNet 的速度對比為 0.126 秒 vs0.005 秒,速度提升了幾十倍:
嘗試一下視覺 transformer 模型,在 M1 CPU 上的速度是 1.855 秒,在 M1 GPU 上則運行崩潰了 ……
EfficientNetB0 實現了 2.5 倍的加速:
EfficientNetB4 實現了 3.5 倍加速:
ConvMixer 運行良好,從 11 秒提速到 2.8 秒:
Dmytro Mishkin 也表示,使用 M1 芯片集成的 GPU 加速只需要預熱一下模型,沒有同步命令。和 CUDA 不同,無需異步執行。
威斯康星大學麥迪遜分校助理教授 Sebastian Raschka 也對 M1 芯片的 GPU 機器學習能力進行了一番測試,他使用的芯片是 M1 和 M1 Pro。
看上去,M1 CPU 似乎比 M1 GPU 更快。但 LeNet-5 是一個非常小的網絡,而 MNIST 是一個非常小的數據集。如果用 rescaled CIFAR-10 圖像再試一次,結果如下:
與 M1 Pro CPU(正數第二行)和 M1 Pro GPU(倒數第二行)相比,M1 Pro GPU 訓練網絡的速度提高了一倍。
可見,M1 系列芯片的 GPU 加速結果非常可觀,在部分情況下已能滿足開發者的需求。不過我們知道在 M1 Ultra 這樣的芯片中也有 32 核的神經網絡引擎,目前卻只有蘋果自己的 Core ML 框架支持使用該部分獲得加速。
不知啟用了 Neural Engine 之后,M1 芯片的 AI 推理速度還能提升多少?