Developing Pytorch Geometric on M1 Apple Silicon

graphs
code
development
environment
Author

Ravi Kalia

Published

March 3, 2024

TLDR: My investigation indicates that the library pytorch geometric (also referred to as pyg, using package name torch_geometric) on Apple Silicon (AS) - particularly the M1 chip - has partial developer support.

A bit more: It is now feasible to perform a full developer installation of torch_geometric using PyTorch on AS with Apple’s native GPU compute framework, Metal Performance Shaders - mps - (which includes its own kind of embedded BLAS). The only remaining challenge is testing with a developer installation. The test suite does not currently build for AS.

I’ve been talking with the pyg-team and they are interested in supporting AS more fully. It’s a matter of time and resources. rusty1s (the lead developer) has been very helpful and responsive, adding a test runner for Apple Silicon.

Developing Pytorch Geometric on M1 Apple Silicon

Image of Macbook Pro M1

Photo by TheRegisti on Unsplash

I’ve recently begun contributing to an open-source Python project, specifically the intriguing geometric machine learning library, pytorch geometric (torch_geometric). After forking the repository, my goal is to develop and contribute code.

As we’ll explore below, the installation process for Apple Silicon (AS) in my case, the M1 Pro chip can be involved. This complexity arises from pre-compiled dependencies specific to chip architecture, including C/C++ extensions for base packages and the lack of support for AS in Intel’s Math Kernel Library, oneMKL.

The pyg-team does provide pre-built distributions for most chip architectures, excluding AS. (Notably project receives technical support from Nvidia and Intel.)

This project’s core dependency is PyTorch which runs on AS with Apple’s mps framework. However, as the project is not tested for this architecture - issues can arise from specific edge cases.

Apple released its own shared memory array based framework mlx and there is a nascent geometric deep learning project specifically for mlx.

Returning to torch_geometric, in this post, we’ll put some effort into the installation process for different use cases on AS devices. That is, for two kinds of roles:

  • User
  • Developer (including a full set of developer dependencies)

User Installation

The installation of the package as a user, utilizing a wheels distribution, is generally straightforward across most CPU/GPU architectures, including AS. Here are the steps:

  1. Check for conda, install if missing.
  2. Establish a clean Python virtual environment.
  3. Activate the newly created virtual environment.
  4. Install the latest version of PyTorch.
  5. Proceed with the installation of PyTorch Geometric.
  6. Check that it imports without errors.

There’s a gist of what worked for me. The cell below uses the %load magic command to load the bash script into the code cell, followed by running it with insertion of the %%bash magic command at the top of the code block.

%%bash

# %load https://gist.githubusercontent.com/project-delphi/38d1db47ed28dde3c8418d5f435c865c/raw/665c079a1f6eff72207309ed97dd6b49194df812/pyg_user_install.sh
# set variables here
DIR="$HOME/Code/throwaway/pytorch-geometric-user-install"
PYTHON_VERSION=3.11
RECENT_TORCH_VERSION=2.2.0

# install miniconda for apple silicon, if not already installed
if [ -d "$HOME/anaconda3" ] || [ -d "$HOME/miniconda3" ]
then
    echo "Conda is installed"
else
    curl -O https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh
    sh Miniconda3-latest-MacOSX-arm64.sh -b -u  > /dev/null 2>&1
fi

mkdir -p "$DIR"
cd "$DIR"
conda create --yes  -p $DIR/.venv python=$PYTHON_VERSION > /dev/null 2>&1
eval "$(conda shell.bash hook)"
conda activate $DIR/.venv
pip install -q --upgrade pip

##### TORCH BUILD AND INSTALL ON M1, to use GPUs #####
pip install -q numpy # to remove user warning with torch install
pip install -q mpmath==1.3.0 # bugfix
xcode-select --install  > /dev/null 2>&1 # if xcode not installed

###### install torch ######
pip install -q --pre torch==$RECENT_TORCH_VERSION torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu 

# install torch geometric
pip install -q torch-geometric

# check
python --version
python -c "import torch; print(f'torch version: {torch.__version__}')"
python -c "import torch_geometric as pyg; print(f'torch geometric version: {pyg.__version__}')"
Conda is installed
Python 3.11.8
torch version: 2.2.0
torch geometric version: 2.5.1

Which is great for using torch_geometric on my M1 Pro.

The problem is that I want to develop (and not just use) locally - that is to have an editable local install on my 2021 16 inch Macbook Pro M1 Pro. Sadly the M1 architecture does not have official support by pyg-team 😔 (I suspect the issues observed also apply to later Apple Silicon too.)

Let’s see how far we can get with the installation process.

Developer Installation

Non Apple Silicon Machines

From the project contributing guidelines, the instructions are clear.

  1. Install a recent version of PyTorch
  2. Optionally install some dependencies if changes make use of them
  3. Be sure to uninstall the pytorch_geometric package
  4. Clone the repo
  5. Run an editable install command for the repo
  6. Run pytest: local testing is kind of essential for a developer install.

Develop away - (as long as it’s a supported architecture).

M1, Possibly All Apple Silicon

This M1 hardware problem for developers has been noted.

For developers, if feature development on pytorch_geometric makes use of the listed package dependencies, several M1/AS issues have been raised.

(It’s likely that the package torch-spline-conv, another package dependency, is also an issue for M1/AS users - though no issues mentioning this are given.)

These dependencies are being subsumed into other packages (for example torch.scatter); at some point they won’t be a problem.

Why is Apple Silicon a Problem?

To develop locally, I need a an editable install version of pytorch geometric. This editable install needs additional dependencies (for fuller developer functionality such as testing, type checking, compiling, linting and documentation) some of which depend on C++ extensions which are not compiled for the M1/AS architecture. The project founder (@rusty1s) has noted that M1 was not supported from the onset - when it wasn’t available on github actions, and there are no plans to support it now. Later Apple Silicon is supported, but developer build are variable.

The Solution

pyg-team suggested earlier that M1 users wanting the fuller editable version of the package can use the cmake & ninja build systems to create libraries and dependencies that target M1 - this will give a working modifiable install of pytorch geometric. Some OS and compiler flags need to be set.

Let’s see if I can do this and get a development environment setup.

What I’ll do is as follows:

  1. check the OS & Hardware
  2. make sure to uninstall all versions of pytorch geometric for all locations.
  3. create and activate a clean Python virtual environment (seems that conda is the best way to go)
  4. install a specific version of PyTorch using conda, as recommended by Apple
  5. build dependencies for clang and macos on my M1 Pro
  6. build the editable install
OS and Hardware

Let’s start by listing my hardware:

%%bash
echo "Operating System: $(uname -s)"
echo "Hardware: $(uname -m)"
echo "macOS Version: $(sw_vers -productVersion)"
echo "Chipset: $(sysctl -n machdep.cpu.brand_string)"
Operating System: Darwin
Hardware: arm64
macOS Version: 14.4
Chipset: Apple M1 Pro
Editable Developer Install

I’ve put this gist together from responses in github issues raised relating to AS. It installs, but doesn’t pass all tests - it seems because the tests are composed from objects that don’t work well with Apple Silicon, rather than anything fundamentally broken in AS.

%%bash
# %load https://gist.githubusercontent.com/project-delphi/b3b5cc91386997ff882f0a3f04a4b89a/raw/6146ec6cf950c2eeb184e0da3e59ff6fdd69550a/pytorch_geometric_apple_silicon_developer_install.sh
# set variables here
DIR="$HOME/Code/throwaway/pytorch-geometric-developer-install"
PYTHON_VERSION=3.11
RECENT_TORCH_VERSION=2.2.0
GITHUB_USERNAME="project-delphi"
MIN_MACOSX_DEPLOYMENT_TARGET=$(sw_vers -productVersion)

# install miniconda for apple silicon, if not already installed
if [ ! -d "$HOME/anaconda3" ] && [ ! -d "$HOME/miniconda3" ]
then
    echo "installing conda..."
    curl -O https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh
    sh Miniconda3-latest-MacOSX-arm64.sh -b -u  > /dev/null 2>&1
fi

mkdir -p "$DIR"
cd "$DIR"
conda create --yes  -p $DIR/.venv python=$PYTHON_VERSION  > /dev/null 2>&1
eval "$(conda shell.bash hook)"
conda activate $DIR/.venv

pip install -q --upgrade pip

##### TORCH BUILD AND INSTALL ON M1, to use GPUs #####
pip install -q numpy # to remove user warning with torch install
pip install -q mpmath==1.3.0 # bugfix
xcode-select --install  > /dev/null 2>&1 # if xcode not installed

###### install pytorch ######
pip install -q --pre torch==$RECENT_TORCH_VERSION torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu

# install dev build dependencies
pip install -q cmake
pip install -q ninja wheel
pip install -q git+https://github.com/pyg-team/pyg-lib.git
MACOSX_DEPLOYMENT_TARGET=$MIN_MACOSX_DEPLOYMENT_TARGET CC=clang CXX=clang++ python -m pip -q --no-cache-dir  install  torch-scatter
MACOSX_DEPLOYMENT_TARGET=$MIN_MACOSX_DEPLOYMENT_TARGET CC=clang CXX=clang++ python -m pip -q --no-cache-dir  install  torch-sparse
MACOSX_DEPLOYMENT_TARGET=$MIN_MACOSX_DEPLOYMENT_TARGET CC=clang CXX=clang++ python -m pip -q --no-cache-dir  install  torch-cluster
MACOSX_DEPLOYMENT_TARGET=$MIN_MACOSX_DEPLOYMENT_TARGET CC=clang CXX=clang++ python -m pip -q --no-cache-dir  install  torch-spline-conv

# clone the forked repository and rebase to original
git clone "https://github.com/$GITHUB_USERNAME/pytorch_geometric.git"  2>/dev/null
cd pytorch_geometric
if ! git remote | grep -q 'upstream'; then
    git remote add upstream "https://github.com/pyg-team/pytorch_geometric"
fi
git fetch upstream  -q
git rebase upstream/master

# build dev install
MACOSX_DEPLOYMENT_TARGET=$MIN_MACOSX_DEPLOYMENT_TARGET CC=clang CXX=clang++ python -m pip install -q --no-cache-dir -e ".[dev,full]"  #> /dev/null 2>&1

# check
python --version
python -c "import torch; print(f'torch version: {torch.__version__}')"
python -c "import torch_geometric as pyg; print(f'torch geometric version: {pyg.__version__}')"

So a kind of success. I can install a full developer version pytorch_geometric on my M1 Pro.

However, we also need to check regarding testing (of which there are several undocumented flavours in the project).

For the standard set of tests, today I get this:

%%bash
# set variables here
DIR="$HOME/Code/throwaway/pytorch-geometric-developer-install"
cd "$DIR"
eval "$(conda shell.bash hook)"
conda activate $DIR/.venv
# install missing packages needed for testing
pip install -q matplotlib-inline ipython
pytest -q --tb=no | tail -n 1
Fatal Python error: Segmentation fault

Current thread 0x00000001f82fbac0 (most recent call first):
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/torch/_ops.py", line 755 in __call__
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/pyg_lib/partition/__init__.py", line 35 in metis
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/pytorch_geometric/torch_geometric/testing/decorators.py", line 224 in withMETIS
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/pytorch_geometric/test/distributed/test_dist_link_neighbor_loader.py", line 140 in <module>
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/assertion/rewrite.py", line 178 in exec_module
  File "<frozen importlib._bootstrap>", line 690 in _load_unlocked
  File "<frozen importlib._bootstrap>", line 1147 in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 1176 in _find_and_load
  File "<frozen importlib._bootstrap>", line 1204 in _gcd_import
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/importlib/__init__.py", line 126 in import_module
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/pathlib.py", line 584 in import_path
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/python.py", line 520 in importtestmodule
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/python.py", line 573 in _getobj
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/python.py", line 315 in obj
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/python.py", line 589 in _register_setup_module_fixture
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/python.py", line 576 in collect
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/runner.py", line 388 in collect
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/runner.py", line 340 in from_call
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/runner.py", line 390 in pytest_make_collect_report
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/pluggy/_callers.py", line 102 in _multicall
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/pluggy/_manager.py", line 119 in _hookexec
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/pluggy/_hooks.py", line 501 in __call__
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/runner.py", line 565 in collect_one_node
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/main.py", line 839 in _collect_one_node
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/main.py", line 976 in genitems
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/main.py", line 981 in genitems
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/main.py", line 981 in genitems
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/main.py", line 981 in genitems
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/main.py", line 981 in genitems
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/main.py", line 813 in perform_collect
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/main.py", line 349 in pytest_collection
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/pluggy/_callers.py", line 102 in _multicall
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/pluggy/_manager.py", line 119 in _hookexec
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/pluggy/_hooks.py", line 501 in __call__
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/main.py", line 338 in _main
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/main.py", line 285 in wrap_session
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/main.py", line 332 in pytest_cmdline_main
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/pluggy/_callers.py", line 102 in _multicall
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/pluggy/_manager.py", line 119 in _hookexec
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/pluggy/_hooks.py", line 501 in __call__
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/config/__init__.py", line 174 in main
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/lib/python3.11/site-packages/_pytest/config/__init__.py", line 197 in console_main
  File "/Users/ravikalia/Code/throwaway/pytorch-geometric-developer-install/.venv/bin/pytest", line 8 in <module>

Extension modules: numpy.core._multiarray_umath, numpy.core._multiarray_tests, numpy.linalg._umath_linalg, numpy.fft._pocketfft_internal, numpy.random._common, numpy.random.bit_generator, numpy.random._bounded_integers, numpy.random._mt19937, numpy.random.mtrand, numpy.random._philox, numpy.random._pcg64, numpy.random._sfc64, numpy.random._generator, torch._C, torch._C._fft, torch._C._linalg, torch._C._nested, torch._C._nn, torch._C._sparse, torch._C._special, scipy._lib._ccallback_c, scipy.sparse._sparsetools, _csparsetools, scipy.sparse._csparsetools, scipy.linalg._fblas, scipy.linalg._flapack, scipy.linalg.cython_lapack, scipy.linalg._cythonized_array_utils, scipy.linalg._solve_toeplitz, scipy.linalg._flinalg, scipy.linalg._decomp_lu_cython, scipy.linalg._matfuncs_sqrtm_triu, scipy.linalg.cython_blas, scipy.linalg._matfuncs_expm, scipy.linalg._decomp_update, scipy.sparse.linalg._dsolve._superlu, scipy.sparse.linalg._eigen.arpack._arpack, scipy.sparse.csgraph._tools, scipy.sparse.csgraph._shortest_path, scipy.sparse.csgraph._traversal, scipy.sparse.csgraph._min_spanning_tree, scipy.sparse.csgraph._flow, scipy.sparse.csgraph._matching, scipy.sparse.csgraph._reordering, scipy.spatial._ckdtree, scipy._lib.messagestream, scipy.spatial._qhull, scipy.spatial._voronoi, scipy.spatial._distance_wrap, scipy.spatial._hausdorff, scipy.special._ufuncs_cxx, scipy.special._ufuncs, scipy.special._specfun, scipy.special._comb, scipy.special._ellip_harm_2, scipy.spatial.transform._rotation, scipy.cluster._vq, scipy.cluster._hierarchy, scipy.cluster._optimal_leaf_ordering, yaml._yaml, numba.core.typeconv._typeconv, numba._helperlib, numba._dynfunc, numba._dispatcher, numba.core.runtime._nrt_python, numba.np.ufunc._internal, numba.experimental.jitclass._box, psutil._psutil_osx, psutil._psutil_posix, markupsafe._speedups, pandas._libs.tslibs.ccalendar, pandas._libs.tslibs.np_datetime, pandas._libs.tslibs.dtypes, pandas._libs.tslibs.base, pandas._libs.tslibs.nattype, pandas._libs.tslibs.timezones, pandas._libs.tslibs.fields, pandas._libs.tslibs.timedeltas, pandas._libs.tslibs.tzconversion, pandas._libs.tslibs.timestamps, pandas._libs.properties, pandas._libs.tslibs.offsets, pandas._libs.tslibs.strptime, pandas._libs.tslibs.parsing, pandas._libs.tslibs.conversion, pandas._libs.tslibs.period, pandas._libs.tslibs.vectorized, pandas._libs.ops_dispatch, pandas._libs.missing, pandas._libs.hashtable, pandas._libs.algos, pandas._libs.interval, pandas._libs.lib, pandas._libs.ops, pandas._libs.hashing, pandas._libs.arrays, pandas._libs.tslib, pandas._libs.sparse, pandas._libs.internals, pandas._libs.indexing, pandas._libs.index, pandas._libs.writers, pandas._libs.join, pandas._libs.window.aggregations, pandas._libs.window.indexers, pandas._libs.reshape, pandas._libs.groupby, pandas._libs.json, pandas._libs.parsers, pandas._libs.testing, matplotlib._c_internal_utils, PIL._imaging, matplotlib._path, kiwisolver._cext, matplotlib._image, rdkit.rdBase (total: 116)

A week ago I noticed that the tests were failing, and then a few days ago they were passing. Today they are not passing, due to new tests. This is caused by commit updates which don’t test on Apple Silicon.

This inconsistency is a problem for contributing to the project. As mentioned, @rusty1s has committed changes to the tests which accommodate Apple Silicon using his own device - but this is not a sustainable approach.

There’s not a good solution to variable pytest runs, since feature updates (commits) to the main line branch – which don’t build for Apple Silicon – are likely to break some test at least some of the time.

In some issues, @rusty1s mentioned a community effort to help test Apple Silicon builds, specifically the M1 architecture. I’d like to contribute, it’s just figuring out how to do so for the long term.

Else, the path of least resistance is to move to the cloud - I’ve recently received cloud credit from major providers.

More On Testing

These test related issues are not unique to the default testing for the package. The package has several kinds of tests, including: full, gpu, previous version, and nightly.

I looked over the github actions workflows, I noticed that full testing and gpu testing are not set up for apple silicon and also have to install undeclared dependencies, such as graphviz and bugfix pinned versions of packages (e.g. mpmath==1.3.0). These minor issues could be useful to work on - even with current default testing this seems to be the case.

I’m thinking of a Makefile to compose the different installation and testing steps into higher level portable grammars. This would be useful for the community, and also for me to use in the future.

Update, @rusty1s suggested more specific testing for mps (and by implication AS), starting with test decorators. That and test documentation could be a good place to start.

References