編輯：澤南、蛋醬
對於 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 推理速度還能提升多少？
參考內容：
https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/
https://sebastianraschka.com/blog/2022/pytorch-m1-gpu.html
