在 Mac 上加速 PyTorch 訓練

隨著 PyTorch v1.12 的釋出，開發者和研究人員可以利用 Apple 晶片的 GPU 來顯著加快模型訓練速度。這使得在 Mac 上本地執行機器學習工作流，如原型設計和微調，成為可能。Apple 的 Metal Performance Shaders (MPS) 作為 PyTorch 的後端實現了這一點，可以透過新的 "mps" 裝置來使用。它會將計算圖和原語對映到 MPS Graph 框架和 MPS 提供的調優核心上。更多資訊請參考官方文件 在 Mac 上引入加速的 PyTorch 訓練 和 MPS 後端。

使用 Apple 晶片進行訓練和推理的好處

1. 使使用者能夠在本地訓練更大的網路或更大的批次大小。
2. 由於統一記憶體架構，減少了資料檢索延遲，併為 GPU 提供了對全部記憶體儲存的直接訪問。因此，提高了端到端的效能。
3. 降低了與基於雲的開發相關的成本或對額外本地 GPU 的需求。

先決條件：要安裝支援 mps 的 torch，請遵循這篇不錯的 Medium 文章 GPU 加速來到 M1 Mac 上的 PyTorch。

開箱即用

在裝有支援 MPS 的 Apple 晶片 GPU 的 MacOs 機器上，此功能預設啟用。要停用它，請在 accelerate launch 命令中傳遞 --cpu 標誌，或在回答 accelerate config 問卷時回答相應問題。

您可以直接執行以下指令碼，在支援 MPS 的 Apple 晶片機器上進行測試

accelerate launch /examples/cv_example.py --data_dir images

需要注意的幾個問題

1. 分散式設定 gloo 和 nccl 無法與 mps 裝置一起工作。這意味著目前只能使用 mps 裝置型別的單個 GPU。

最後，請記住，Accelerate 僅集成了 MPS 後端，因此如果您對 MPS 後端的使用有任何問題或疑問，請向 PyTorch GitHub 提交問題。