|
| 1 | +# Running AlphaFold 3 on Apple Silicon GPU (Metal / MPS) |
| 2 | + |
| 3 | +> **Status: experimental feasibility path.** This mode runs AlphaFold 3 |
| 4 | +> inference natively on the Apple Silicon integrated GPU through the community |
| 5 | +> [`jax-mps`](https://pypi.org/project/jax-mps/) Metal plugin for JAX. It is |
| 6 | +> **not** an officially supported or numerically certified configuration. The |
| 7 | +> NVIDIA CUDA path is unchanged and remains the reference. Use this for research, |
| 8 | +> local development, and prototyping on Macs. |
| 9 | +
|
| 10 | +This guide describes a **native build with no Docker**. Docker Desktop on macOS |
| 11 | +runs a Linux VM that cannot see the Mac GPU, so the containerized recipe in |
| 12 | +[`installation.md`](installation.md) can only ever reach the CPU. To use the |
| 13 | +Apple GPU you must build AlphaFold 3 natively, as described below. |
| 14 | + |
| 15 | +If you only want CPU-only inference on a Mac (much slower), follow the "Running |
| 16 | +AlphaFold 3 without a GPU" section of [`installation.md`](installation.md) |
| 17 | +instead. |
| 18 | + |
| 19 | +## What works and what to expect |
| 20 | + |
| 21 | +Validated in July 2026 across a 52-target campaign (all 12 molecule-type classes |
| 22 | +plus token-scaling targets up to ~2,200 tokens) on an **M2 Ultra Mac Studio |
| 23 | +(76-GPU-core, 192 GB unified memory)**, macOS 15. |
| 24 | + |
| 25 | +- **Correctness:** the full data pipeline (HMMER MSA + templates) and inference |
| 26 | + run end to end; predictions match the CUDA reference closely for the cases |
| 27 | + tested, and inference is **bit-reproducible run-to-run** with fixed seeds. |
| 28 | +- **Attention:** only the portable `xla` implementation works on Metal; |
| 29 | + `triton` and `cudnn` require NVIDIA GPUs and are rejected early. |
| 30 | +- **Memory is the main limit:** peak GPU (unified) memory reaches ~167 GB at |
| 31 | + ~2,200 tokens on a 192 GB machine — the practical single-structure ceiling on |
| 32 | + this hardware (see [Performance and memory](#performance-and-memory)). |
| 33 | +- **Large targets are slow:** per-seed inference scales roughly cubically with |
| 34 | + tokens, with a steep Metal-specific penalty above ~1,500 tokens. |
| 35 | + |
| 36 | +## Prerequisites |
| 37 | + |
| 38 | +- Apple Silicon Mac (M1/M2/M3/M4 family). More unified memory = larger tractable |
| 39 | + structures; 64 GB is a reasonable floor, 128–192 GB recommended for large |
| 40 | + complexes. |
| 41 | +- macOS 14 or newer. |
| 42 | +- Xcode Command Line Tools (provides `clang`, `make`, `patch`): |
| 43 | + |
| 44 | + ```sh |
| 45 | + xcode-select --install |
| 46 | + ``` |
| 47 | + |
| 48 | +- [`uv`](https://docs.astral.sh/uv/) and a Python 3.11–3.13 toolchain (a |
| 49 | + Conda/Miniforge environment works too — see |
| 50 | + [Alternative: Conda](#alternative-conda-environment)). |
| 51 | + |
| 52 | +## Step 1 — Install the HMMER suite (for the MSA data pipeline) |
| 53 | + |
| 54 | +HMMER is only needed if you run the data pipeline (i.e. not with |
| 55 | +`--norun_data_pipeline`). Prefer a user-local source build of HMMER 3.4 with the |
| 56 | +sequence-truncation patch that ships in this repository |
| 57 | +(`docker/jackhmmer_seq_limit.patch`): |
| 58 | + |
| 59 | +```sh |
| 60 | +export AF3_REPO="$PWD" # path to your alphafold3 clone |
| 61 | +export HMMER_PREFIX="$HOME/af3-tools/hmmer-3.4" |
| 62 | +mkdir -p "$HOME/af3-tools/hmmer-source" && cd "$HOME/af3-tools/hmmer-source" |
| 63 | + |
| 64 | +curl -LO http://eddylab.org/software/hmmer/hmmer-3.4.tar.gz |
| 65 | +echo "ca70d94fd0cf271bd7063423aabb116d42de533117343a9b27a65c17ff06fbf3 hmmer-3.4.tar.gz" | shasum -a 256 -c |
| 66 | +tar -xzf hmmer-3.4.tar.gz |
| 67 | +cd hmmer-3.4 |
| 68 | +patch -p0 < "$AF3_REPO/docker/jackhmmer_seq_limit.patch" |
| 69 | +./configure --prefix="$HMMER_PREFIX" |
| 70 | +make -j"$(sysctl -n hw.perflevel0.logicalcpu)" |
| 71 | +make install |
| 72 | +"$HMMER_PREFIX/bin/jackhmmer" -h | head -1 # sanity check |
| 73 | +``` |
| 74 | + |
| 75 | +Pass the binaries to `run_alphafold.py` in Step 4 via the |
| 76 | +`--jackhmmer_binary_path`, `--nhmmer_binary_path`, `--hmmalign_binary_path`, |
| 77 | +`--hmmsearch_binary_path`, and `--hmmbuild_binary_path` flags (or put |
| 78 | +`$HMMER_PREFIX/bin` first on your `PATH`). |
| 79 | + |
| 80 | +## Step 2 — Get the code, model weights, and databases |
| 81 | + |
| 82 | +1. Clone the repository (if you have not already): |
| 83 | + |
| 84 | + ```sh |
| 85 | + git clone https://github.com/google-deepmind/alphafold3.git |
| 86 | + cd alphafold3 |
| 87 | + ``` |
| 88 | + |
| 89 | +2. Obtain the model weights by following the process in the |
| 90 | + [main README](https://github.com/google-deepmind/alphafold3). Keep the |
| 91 | + weights **outside** the git checkout. |
| 92 | +3. Download the genetic databases (see [`installation.md`](installation.md) and |
| 93 | + `fetch_databases.py`) onto fast local storage. |
| 94 | + |
| 95 | +## Step 3 — Build and install AlphaFold 3 (with the MPS backend) |
| 96 | + |
| 97 | +On `darwin arm64`, `pyproject.toml` already selects the Metal backend |
| 98 | +(`jax==0.10.2`, `jaxlib==0.10.2`, `jax-mps==0.10.9`) in place of the CUDA stack, |
| 99 | +so a normal install pulls in everything needed for GPU inference: |
| 100 | + |
| 101 | +```sh |
| 102 | +uv venv --python 3.12 |
| 103 | +source .venv/bin/activate |
| 104 | +uv sync # builds/installs alphafold3 AND the jax-mps Metal backend |
| 105 | +uv run build_data # builds the CCD / chemical-component data used at runtime |
| 106 | +``` |
| 107 | + |
| 108 | +Verify that JAX sees the Metal device: |
| 109 | + |
| 110 | +```sh |
| 111 | +python -c "import jax; print([d.platform for d in jax.local_devices()])" |
| 112 | +# -> ['mps'] (a one-time 'Platform mps is experimental' warning is expected) |
| 113 | +``` |
| 114 | + |
| 115 | +## Step 4 — Run inference on the Apple GPU |
| 116 | + |
| 117 | +Select the Metal backend with `--jax_backend mps` and the portable `xla` |
| 118 | +attention implementation: |
| 119 | + |
| 120 | +```sh |
| 121 | +uv run python run_alphafold.py \ |
| 122 | + --json_path=/path/to/fold_input.json \ |
| 123 | + --model_dir=/path/to/weights \ |
| 124 | + --db_dir=/path/to/databases \ |
| 125 | + --output_dir=/path/to/output \ |
| 126 | + --jax_backend=mps \ |
| 127 | + --gpu_device=0 \ |
| 128 | + --flash_attention_implementation=xla \ |
| 129 | + --jackhmmer_binary_path="$HMMER_PREFIX/bin/jackhmmer" \ |
| 130 | + --nhmmer_binary_path="$HMMER_PREFIX/bin/nhmmer" \ |
| 131 | + --hmmalign_binary_path="$HMMER_PREFIX/bin/hmmalign" \ |
| 132 | + --hmmsearch_binary_path="$HMMER_PREFIX/bin/hmmsearch" \ |
| 133 | + --hmmbuild_binary_path="$HMMER_PREFIX/bin/hmmbuild" |
| 134 | +``` |
| 135 | + |
| 136 | +For inference on a pre-computed input (skipping the MSA pipeline and HMMER), add |
| 137 | +`--norun_data_pipeline` and drop the HMMER flags. |
| 138 | + |
| 139 | +## Performance and memory |
| 140 | + |
| 141 | +Measured on the M2 Ultra (192 GB). Inference time and peak memory are set by the |
| 142 | +**padded token bucket**; MSA (CPU) is driven by chain count / RNA, not token |
| 143 | +count. Treat as order-of-magnitude guidance; other Apple Silicon chips differ. |
| 144 | + |
| 145 | +| Padded tokens | MSA time¹ | Inference / seed | Peak unified memory (GB) | |
| 146 | +|--:|--:|--:|--:| |
| 147 | +| 256 | ~7 min | ~29 s | 36 | |
| 148 | +| 512 | ~13 min | ~90 s | 50 | |
| 149 | +| 768 | ~13 min | ~3.9 min | 73 | |
| 150 | +| 1,024 | ~7 min | ~8.6 min | 101 | |
| 151 | +| 1,536 | ~19 min | ~26.6 min | 155 | |
| 152 | +| ~2,180 | ~20 min | ~3.9 h | 167 | |
| 153 | + |
| 154 | +¹ MSA / data pipeline is CPU-bound and driven by the number of unique chains and |
| 155 | +RNA presence, not token count (observed range 7–35 min). |
| 156 | + |
| 157 | +Key points: |
| 158 | + |
| 159 | +- **Unified memory is the ceiling.** GPU tensors live in unified memory (macOS |
| 160 | + "Memory Used") and are **not** reflected in process RSS (~7 GB). On a 192 GB |
| 161 | + machine, ~2,200 tokens is the practical single-structure limit; larger |
| 162 | + structures will run out of memory. Scale expectations to your Mac's memory. |
| 163 | +- **No persistent JAX compilation cache on Metal.** `--jax_compilation_cache_dir` |
| 164 | + does not persist across processes on MPS, so fold many inputs in one process |
| 165 | + (e.g. `--input_dir`) to amortize compile cost. |
| 166 | +- Set HMMER thread counts so that (concurrent searches × threads-per-search) ≈ |
| 167 | + the number of performance cores (`sysctl hw.perflevel0.logicalcpu`). |
| 168 | + |
| 169 | +## Alternative: Conda environment |
| 170 | + |
| 171 | +`pyproject.toml` still selects `jax-mps` on `darwin arm64`, so a plain install |
| 172 | +pulls in the Metal backend: |
| 173 | + |
| 174 | +```sh |
| 175 | +conda create -n af3-mps python=3.13 cmake -y |
| 176 | +conda activate af3-mps |
| 177 | +pip install --no-build-isolation -e . # installs alphafold3 + jax-mps |
| 178 | +build_data |
| 179 | +``` |
| 180 | + |
| 181 | +Then run `run_alphafold.py` with the same MPS flags as in Step 4. |
| 182 | + |
| 183 | +## Troubleshooting |
| 184 | + |
| 185 | +- **`No local MPS device was found`** — `jax-mps` is not installed. Re-check the |
| 186 | + install (`python -c "import jax; print(jax.local_devices())"`). |
| 187 | +- **`--flash_attention_implementation must be set to "xla"`** — `triton`/`cudnn` |
| 188 | + are NVIDIA-only; use `xla` on Metal. |
| 189 | +- **Out-of-memory / heavy swapping on large targets** — reduce the structure |
| 190 | + size or run on a Mac with more unified memory (see the table above). |
0 commit comments