Using PyTorch with M1 Chips on GPU
On 18th May 2022, PyTorch announced support for GPU-accelerated PyTorch training on Mac. It will be made available with PyTorch v1.12 release, but is available in the Preview(Nightly) build right now.
I tried it out on my Macbook Air M1, and decided to share the steps to set up the Preview(Nightly) build of PyTorch and give it a spin.
I already have miniconda set up on my Macbook. You can find the project and steps to set it up easily here.
In order to install PyTorch Preview(Nightly) build using conda, I used the suggested command by PyTorch here:
Now, to make use of MPS acceleration, we need to move our Tensors and Models to M1 GPU, like we do for training on NVIDIA GPU.
For moving models to NVIDIA GPU, we use CUDA:
Similarly for M1, we move it to device “mps” using:
To move tensors to NVIDIA GPU, we use CUDA:
Similarly, to move tensors to M1 GPU, we use MPS:
There are a few PyTorch operations, which have still not been migrated to support MPS. For such operations we need to enable Fallback on CPU option in PyTorch using an environment variable. This will slow down overall training speed.
In Conda Manager, we can add this environment variable:
conda env config vars set PYTORCH_ENABLE_MPS_FALLBACK=1
conda activate torch-nightly
The operations (list available here) that are yet to be migrated to supporting MPS:
I will try to benchmark a few trainings on Google Colab and my Macbook M1 to see how they perform. If you have access to GPU and M1, please do add your benchmarks in comments.
This is a very amateur article, and may be full of mistakes and holes, please do suggest changes (please BNBR), and I will be quick to implement those.
References:
1. https://pytorch.org/blog/introducing-accelerated-pytorch-training-on-mac/
2. https://pytorch.org/docs/master/notes/mps.html
3. https://sebastianraschka.com/blog/2022/pytorch-m1-gpu.html
4. https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html#setting-environment-variables
