Home Setting up PyTorch Development for Mac M1/M2 ARM
Post
Cancel

Setting up PyTorch Development for Mac M1/M2 ARM

Want to build pytorch on an M1 mac? Running into issues with the build process? This guide will help you get started.

There are issues with building PyTorch on Mac M1/M2 ARM devices due to conflicts with protobuf that comes with OSX 12 and 13.

The following instructions are based off the pytorch official guide:

1
2
3
4
5
6
7
8
9
10
11
# create a conda environment
conda create -n pytorch-build python=3.10
conda activate pytorch-build

conda install astunparse numpy ninja pyyaml setuptools cmake cffi \
 typing_extensions future six requests dataclasses

# These are useful for development
conda expecttest hypothesis mypy pytest ninja

git clone --recursive https://github.com/pytorch/pytorch

Params do the following:

BUILD_CAFFE2=0: skip unnecessary build

and USE_ONNX=0: USE_ONNX=0 prevents trying to build protobuf

REL_WITH_DEB_INFO=1: include debugging symbols.

MACOSX_DEPLOYMENT_TARGET=12.3: target a macOS that has MPS available

If your planning on developing, enabling caching will speed up the build process (after the first build):

1
2
3
4
5
6
7
8
9
brew install ccache
# increase max size of cache
ccache -M 25Gi  # -M 0 for unlimited
# unlimited number of files
ccache -F 0

export CMAKE_C_COMPILER_LAUNCHER=ccache
export CMAKE_CXX_COMPILER_LAUNCHER=ccache
export CMAKE_CUDA_COMPILER_LAUNCHER=ccache

Finally, build!

1
2
export CMAKE_PREFIX_PATH=${CONDA_PREFIX:-"$(dirname $(which conda))/../"}
MACOSX_DEPLOYMENT_TARGET=12.3 BUILD_CAFFE2=0 USE_ONNX=0 REL_WITH_DEB_INFO=1 python setup.py develop --cmake

and install

1
pip install -e .

and confirm that it works:

1
2
3
4
5
6
7
8
9
10
11
12
13
import torch

device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using {device} device")

x = torch.rand(5, 3)
print(x)

Many thanks to @damian0815 for the MPS build tip.

This post is licensed under CC BY 4.0 by the author.