Installing and running pytorch on M1 GPUs (Apple metal/MPS)

Image sourced from MacRumors

Hey everyone! In this article I’ll help you install pytorch for GPU acceleration on Apple’s M1 chips. Let’s crunch some tensors on Apple metal!

We’re in exciting times for the future of computing and AI. In 2020, when Apple released their new macbooks running on the M1 chips, one question on my mind with everything learned from Apple’s big reveal was whether pytorch could crunch tensors on those kickass GPUs and neural engines. Fast forward to 2022 and the pytorch team announced a new version of pytorch that takes advantage of M1 GPUs. Soumith and his team in collab with the Apple metal team have done an incredible job to make this possible. Kudos to them! While it’s early stage, it’s very exciting so far! Ideally, the instructions on the website should have been enough but I ran in to some roadblocks whilst setting up my mac to run pytorch on m1 GPUs. I am writing this article to provide guidance for anyone who might face similar issues. I hope you find this helpful.

On installing python for m1…

Tl-dr: Use conda

Installing python packages on the m1 macs itself is quite an extreme sport. I was at this for 2 days and a lot of tears when I set up my mac. I ended up resorting to installing and managing python packages on the the Rosseta’s x86 emulator for the terminal to run the intel version of homebrew and python versions and packages via pyenv. However, you can also build package wheels yourself if you’re up for it and have the time. My word of advice is to read lots of user experiences online before you start modifying your system🥲.

I use poetry to manage packages when building apps. When I am working on data science projects through, I prefer to use Anaconda. Conda’s great for data science and ML. You should be fine with either the GUI or installation script — but I opted to go with the GUI installation. As at the time of this writing, you will need to run pytorch on m1 GPUs via the preview (nightly) version of pytorch with the recommendation from pytorch’s website:

Installing PyTorch

Tl-dr: If you’re in a hurry and don’t want to read the rest of the article, then use the recommended pip wheel to install pytorch. It works!

pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

Verify your installation

Python 3.10.4 (main, Mar 31 2022, 03:37:37) [Clang 12.0.0 ] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch; torch.backends.mps.is_available()
True
>>>

Alright, now let’s see the details on the options available to us.

Option 1: Install with conda

At the moment, the pytorch recommends that you install pytorch, torchaudio and torchvision with conda.

conda install pytorch torchvision torchaudio -c pytorch-nightly

However, it’s ridden with errors (for now). For me in particular, I couldn’t find torchvision and torchaudio in the nightly channel for pytorch.

Collecting package metadata (current_repodata.json): done
Solving environment: failed with initial frozen solve. Retrying with flexible solve.
Collecting package metadata (repodata.json): done
Solving environment: failed with initial frozen solve. Retrying with flexible solve.
PackagesNotFoundError: The following packages are not available from current channels:
- torchvision
- torchaudio
Current channels:
- https://conda.anaconda.org/pytorch-nightly/osx-arm64
- https://conda.anaconda.org/pytorch-nightly/noarch
- https://repo.anaconda.com/pkgs/main/osx-arm64
- https://repo.anaconda.com/pkgs/main/noarch
- https://repo.anaconda.com/pkgs/r/osx-arm64
- https://repo.anaconda.com/pkgs/r/noarch
To search for alternate channels that may provide the conda package you're
looking for, navigate to
https://anaconda.org
and use the search bar at the top of the page.

I found though while there were no suitable packages found for torchaudio and torchvision, when I install pytorch only with conda install pytorch -c pytorch-nightly it works.I was also able to use the MPS backend. Right after that, you can install torchvision and torchaudio via their pip wheels

pip3 install --pre torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

And then you’re good to go! Verify your installation by checking availability of the MPS backend in your python shell or jupyter notebook.

Python 3.10.4 (main, Mar 31 2022, 03:37:37) [Clang 12.0.0 ] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch; torch.backends.mps.is_available()
True
>>>

Option 2: Install from pre-built wheels (Easiest and recommended)

You can also install from pre-built wheels from the pytorch wheel archive. You’ll notice when you’re trying to install the stable version of PyTorch, the team actually recommends using the wheels to install the packages for M1 Macs

# MacOS Conda binaries are for x86_64 only, for M1 please use wheels

What’s a wheel?

A wheel is a type of built distribution. In this case, built means that the wheel comes in a ready-to-install format and allows you to skip the build stage required with source distributions. — Brad Solomon, Real Python

You can either install the wheels by downloading it from pytorch/pip or building and installing a wheel from the source code itself. In this option, we we will download and install the pytorch wheel already built from the source code. You can do this by:

pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

The --pre flag tells pip to install “Include pre-release and development versions. By default, pip only finds stable versions.”

Note: Be sure that you are executing pip for the virtual environment. I use pyenv to manage my python versions and found that pip3 was referring to my system pip. So I ran my command with pip instead. It depends on your setup

You’ll find that the MPS backend is also available when you check for it.

You learn more about python wheels through Brad’s excellent primer from RealPython.

Option 3: Build pytorch (and torchaudio, torchvision) from source

If you’re up for it, you can also build the packages from the source code. We will create our own wheels for pytorch, torchaudio and torchvision to install the nightly version of pytorch on M1 gpus.

Let’s build our binaries from source for installation. Here’s your guide curated from pytorch, torchaudio and torchvision repos. The documentation is yet to be updated for installation on MPS devices — so I had to make some modifications as you’ll see below:

Step 1: Create a conda environment

conda env create --name pytorchm1
conda activate pytorchm1 

Step 2: Configure installation dependencies

### if you do not have a c++ compiler installed. You'll need it...
brew install gcc
# install common dependencies
conda install astunparse numpy ninja pyyaml setuptools cmake cffi typing_extensions future six requests dataclasses
# Add these packages if torch.distributed is needed
conda install pkg-config libuv

Optional: Install mkl and mkl-include

I couldn’t install mkl and mkl-include but was still able to build my binaries for pytorch, torchvision and torchaudio. I ignored this error though since mkl and mkl-include are packages for optimizing math on intel CPUs and I am not running an intel CPU 🤷‍♂️.

But if you do need to check off the boxes, then:

conda install mkl mkl-include

…Or you could just build it from source :-)

If you’d like to learn more about mkl, check out the API docs.

Step 3: Download sources for pytorch, torchaudio and torchvision

Once your conda env is activated, create a folder for your packages and then clone the repos into that folder

mkdir packages && cd packages
git clone --recursive https://github.com/pytorch/vision.git torchvision
git clone --recursive https://github.com/pytorch/audio.git torchaudio
# get pytorch from the source 
git clone --recursive https://github.com/pytorch/pytorch pytorch

You need to include submodules, which is why you should clone recursively.

Step 4: Install pytorch nightly

Installation instructions were sourced the pytorch readme. However I made some edits to enable installation with Apple’s MPS. That’s just because I noticed that when you use the installation scripts from the pytorch readme, you’d get Apple’s MPS (Metal Performance Shaders) deactivated by default. It can be changed with USE_MPS=1 USE_PYTORCH_METAL=1 (AlbanD mentions that

Let’s build PyTorch!

cd pytorch
# if you are updating an existing checkout
git submodule sync
git submodule update --init --recursive --jobs 0
# checkout the nightly branch of pytorch
git checkout nightly # only if you want to use the preview build
# install pytorch nightly
export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}
MACOSX_DEPLOYMENT_TARGET=10.9 CC=clang CXX=clang++ USE_MPS=1 USE_PYTORCH_METAL=1 python setup.py install
# cd back into your packages folder
cd ..

You should have something like this at the tail end of your terminal outputs when it’s done:

writing manifest file 'torch.egg-info/SOURCES.txt'
Copying torch.egg-info to /Users/chrisdare/opt/anaconda3/envs/pytorchm1/lib/python3.9/site-packages/torch-1.13.0a0+git00c4d7f-py3.9.egg-info
running install_scripts
Installing convert-caffe2-to-onnx script to /Users/chrisdare/opt/anaconda3/envs/pytorchm1/bin
Installing convert-onnx-to-caffe2 script to /Users/chrisdare/opt/anaconda3/envs/pytorchm1/bin
Installing torchrun script to /Users/chrisdare/opt/anaconda3/envs/pytorchm1/bin

Step 5: Install torchaudio

Installation instructions were sourced the torchaudio readme

# install pending torchaudio dependencies
brew install pkg-config
# move to your torch audio repo and install
cd torchaudio
CC=clang CXX=clang++ python setup.py install
cd ..

You should have something like this at the tail end of your terminal outputs when it’s done:

Installed /Users/chrisdare/opt/anaconda3/envs/pytorchm1/lib/python3.9/site-packages/torchaudio-0.12.0a0+6776299-py3.9-macosx-11.1-arm64.egg
Processing dependencies for torchaudio==0.12.0a0+6776299
Searching for torch==1.13.0a0+git00c4d7f
Best match: torch 1.13.0a0+git00c4d7f
Adding torch 1.13.0a0+git00c4d7f to easy-install.pth file
Installing convert-caffe2-to-onnx script to /Users/chrisdare/opt/anaconda3/envs/pytorchm1/bin
Installing convert-onnx-to-caffe2 script to /Users/chrisdare/opt/anaconda3/envs/pytorchm1/bin
Installing torchrun script to /Users/chrisdare/opt/anaconda3/envs/pytorchm1/bin
Using /Users/chrisdare/opt/anaconda3/envs/pytorchm1/lib/python3.9/site-packages
Searching for typing-extensions==4.1.1
Best match: typing-extensions 4.1.1
Adding typing-extensions 4.1.1 to easy-install.pth file
Using /Users/chrisdare/opt/anaconda3/envs/pytorchm1/lib/python3.9/site-packages
Finished processing dependencies for torchaudio==0.12.0a0+6776299

Step 6: Install torchvision

The steps for installing torchvision are also similar

# install
conda install -c conda-forge ffmpeg
conda install pillow
# install vision
cd torchvision
MACOSX_DEPLOYMENT_TARGET=10.9 CC=clang CXX=clang++ python setup.py install

You should have something like this at the tail end of your terminal outputs when it’s done:

Processing torchvision-0.13.0a0+b969cca-py3.9-macosx-11.1-arm64.egg
creating /Users/chrisdare/opt/anaconda3/envs/pytorchm1/lib/python3.9/site-packages/torchvision-0.13.0a0+b969cca-py3.9-macosx-11.1-arm64.egg
Extracting torchvision-0.13.0a0+b969cca-py3.9-macosx-11.1-arm64.egg to /Users/chrisdare/opt/anaconda3/envs/pytorchm1/lib/python3.9/site-packages
Adding torchvision 0.13.0a0+b969cca to easy-install.pth file
Installed /Users/chrisdare/opt/anaconda3/envs/pytorchm1/lib/python3.9/site-packages/torchvision-0.13.0a0+b969cca-py3.9-macosx-11.1-arm64.egg
Processing dependencies for torchvision==0.13.0a0+b969cca
Searching for torchvision==0.13.0a0+b969cca
Reading https://pypi.org/simple/torchvision/
/Users/chrisdare/opt/anaconda3/envs/pytorchm1/lib/python3.9/site-packages/pkg_resources/__init__.py:122: PkgResourcesDeprecationWarning:  is an invalid version and will not be supported in a future release
warnings.warn(
No local packages or working download links found for torchvision==0.13.0a0+b969cca
error: Could not find suitable distribution for Requirement.parse('torchvision==0.13.0a0+b969cca')

The error at the bottom has to do with versioning and installing dependencies for torchvision. You may need to handle dependency installations manually. It should be easy with a conda or pip install.

Let’s crunch some tensors!

To check if the backend is available, use

torch.backends.mps.is_available()

Test your pytorch installation by running the example image classifier code from the pytorch website.

Be sure to use mps as the device for pytorch:

DEVICE = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
## and then move your model and data to the device before you train or eval. Have fun folks!

And that’s it! You’re done installing pytorch to run on m1 GPUs. Have a good time crunching tensors!

One more thing: Staying in the loop

The next stable version of python is expected to have support for the MPS backend.
To learn more, and stay updated, visit the pytorch discuss forum for MPS.

Shouts to AlbanD, you’re doing an awesome job!