diff --git a/docs/installation.md b/docs/installation.md index 8377bb08..ef33e0cf 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -492,6 +492,13 @@ Since JAX doesn't support running natively on Mac GPU as of 2026, you have to resort to running AlphaFold 3 in the slow CPU-only mode even though it has a GPU (`jax-metal` is unfinished as of July 2026). +> **Experimental Apple Silicon GPU (Metal / MPS) inference.** As an alternative +> to CPU-only mode, AlphaFold 3 inference can run on the Apple Silicon GPU with +> `--jax_backend=mps` via the community `jax-mps` Metal plugin. This is an +> unofficial, experimental feasibility path (not numerically certified), but it +> is substantially faster than CPU for small and medium structures. See +> [installation_apple_silicon_gpu.md](installation_apple_silicon_gpu.md). + 1. Download all required databases and AlphaFold 3 weights (see above). 2. Install the [HMMER Suite](http://hmmer.org/). See http://hmmer.org/documentation.html for installation instructions. diff --git a/docs/installation_apple_silicon_gpu.md b/docs/installation_apple_silicon_gpu.md new file mode 100644 index 00000000..26bec8d1 --- /dev/null +++ b/docs/installation_apple_silicon_gpu.md @@ -0,0 +1,190 @@ +# Running AlphaFold 3 on Apple Silicon GPU (Metal / MPS) + +> **Status: experimental feasibility path.** This mode runs AlphaFold 3 +> inference natively on the Apple Silicon integrated GPU through the community +> [`jax-mps`](https://pypi.org/project/jax-mps/) Metal plugin for JAX. It is +> **not** an officially supported or numerically certified configuration. The +> NVIDIA CUDA path is unchanged and remains the reference. Use this for research, +> local development, and prototyping on Macs. + +This guide describes a **native build with no Docker**. Docker Desktop on macOS +runs a Linux VM that cannot see the Mac GPU, so the containerized recipe in +[`installation.md`](installation.md) can only ever reach the CPU. To use the +Apple GPU you must build AlphaFold 3 natively, as described below. + +If you only want CPU-only inference on a Mac (much slower), follow the "Running +AlphaFold 3 without a GPU" section of [`installation.md`](installation.md) +instead. + +## What works and what to expect + +Validated in July 2026 across a 52-target campaign (all 12 molecule-type classes +plus token-scaling targets up to ~2,200 tokens) on an **M2 Ultra Mac Studio +(76-GPU-core, 192 GB unified memory)**, macOS 15. + +- **Correctness:** the full data pipeline (HMMER MSA + templates) and inference + run end to end; predictions match the CUDA reference closely for the cases + tested, and inference is **bit-reproducible run-to-run** with fixed seeds. +- **Attention:** only the portable `xla` implementation works on Metal; + `triton` and `cudnn` require NVIDIA GPUs and are rejected early. +- **Memory is the main limit:** peak GPU (unified) memory reaches ~167 GB at + ~2,200 tokens on a 192 GB machine — the practical single-structure ceiling on + this hardware (see [Performance and memory](#performance-and-memory)). +- **Large targets are slow:** per-seed inference scales roughly cubically with + tokens, with a steep Metal-specific penalty above ~1,500 tokens. + +## Prerequisites + +- Apple Silicon Mac (M1/M2/M3/M4 family). More unified memory = larger tractable + structures; 64 GB is a reasonable floor, 128–192 GB recommended for large + complexes. +- macOS 14 or newer. +- Xcode Command Line Tools (provides `clang`, `make`, `patch`): + + ```sh + xcode-select --install + ``` + +- [`uv`](https://docs.astral.sh/uv/) and a Python 3.11–3.13 toolchain (a + Conda/Miniforge environment works too — see + [Alternative: Conda](#alternative-conda-environment)). + +## Step 1 — Install the HMMER suite (for the MSA data pipeline) + +HMMER is only needed if you run the data pipeline (i.e. not with +`--norun_data_pipeline`). Prefer a user-local source build of HMMER 3.4 with the +sequence-truncation patch that ships in this repository +(`docker/jackhmmer_seq_limit.patch`): + +```sh +export AF3_REPO="$PWD" # path to your alphafold3 clone +export HMMER_PREFIX="$HOME/af3-tools/hmmer-3.4" +mkdir -p "$HOME/af3-tools/hmmer-source" && cd "$HOME/af3-tools/hmmer-source" + +curl -LO http://eddylab.org/software/hmmer/hmmer-3.4.tar.gz +echo "ca70d94fd0cf271bd7063423aabb116d42de533117343a9b27a65c17ff06fbf3 hmmer-3.4.tar.gz" | shasum -a 256 -c +tar -xzf hmmer-3.4.tar.gz +cd hmmer-3.4 +patch -p0 < "$AF3_REPO/docker/jackhmmer_seq_limit.patch" +./configure --prefix="$HMMER_PREFIX" +make -j"$(sysctl -n hw.perflevel0.logicalcpu)" +make install +"$HMMER_PREFIX/bin/jackhmmer" -h | head -1 # sanity check +``` + +Pass the binaries to `run_alphafold.py` in Step 4 via the +`--jackhmmer_binary_path`, `--nhmmer_binary_path`, `--hmmalign_binary_path`, +`--hmmsearch_binary_path`, and `--hmmbuild_binary_path` flags (or put +`$HMMER_PREFIX/bin` first on your `PATH`). + +## Step 2 — Get the code, model weights, and databases + +1. Clone the repository (if you have not already): + + ```sh + git clone https://github.com/google-deepmind/alphafold3.git + cd alphafold3 + ``` + +2. Obtain the model weights by following the process in the + [main README](https://github.com/google-deepmind/alphafold3). Keep the + weights **outside** the git checkout. +3. Download the genetic databases (see [`installation.md`](installation.md) and + `fetch_databases.py`) onto fast local storage. + +## Step 3 — Build and install AlphaFold 3 (with the MPS backend) + +On `darwin arm64`, `pyproject.toml` adds the `jax-mps` Metal backend (the base +install already provides `jax`/`jaxlib` 0.10.2; the CUDA plugin is excluded on +macOS), so a normal install pulls in everything needed for GPU inference: + +```sh +uv venv --python 3.12 +source .venv/bin/activate +uv sync # builds/installs alphafold3 AND the jax-mps Metal backend +uv run build_data # builds the CCD / chemical-component data used at runtime +``` + +Verify that JAX sees the Metal device: + +```sh +python -c "import jax; print([d.platform for d in jax.local_devices()])" +# -> ['mps'] (a one-time 'Platform mps is experimental' warning is expected) +``` + +## Step 4 — Run inference on the Apple GPU + +Select the Metal backend with `--jax_backend mps` and the portable `xla` +attention implementation: + +```sh +uv run python run_alphafold.py \ + --json_path=/path/to/fold_input.json \ + --model_dir=/path/to/weights \ + --db_dir=/path/to/databases \ + --output_dir=/path/to/output \ + --jax_backend=mps \ + --gpu_device=0 \ + --flash_attention_implementation=xla \ + --jackhmmer_binary_path="$HMMER_PREFIX/bin/jackhmmer" \ + --nhmmer_binary_path="$HMMER_PREFIX/bin/nhmmer" \ + --hmmalign_binary_path="$HMMER_PREFIX/bin/hmmalign" \ + --hmmsearch_binary_path="$HMMER_PREFIX/bin/hmmsearch" \ + --hmmbuild_binary_path="$HMMER_PREFIX/bin/hmmbuild" +``` + +For inference on a pre-computed input (skipping the MSA pipeline and HMMER), add +`--norun_data_pipeline` and drop the HMMER flags. + +## Performance and memory + +Measured on the M2 Ultra (192 GB). Inference time and peak memory are set by the +**padded token bucket**; MSA (CPU) is driven by chain count / RNA, not token +count. Treat as order-of-magnitude guidance; other Apple Silicon chips differ. + +| Padded tokens | MSA time¹ | Inference / seed | Peak unified memory (GB) | +|--:|--:|--:|--:| +| 256 | ~7 min | ~29 s | 36 | +| 512 | ~13 min | ~90 s | 50 | +| 768 | ~13 min | ~3.9 min | 73 | +| 1,024 | ~7 min | ~8.6 min | 101 | +| 1,536 | ~19 min | ~26.6 min | 155 | +| ~2,180 | ~20 min | ~3.9 h | 167 | + +¹ MSA / data pipeline is CPU-bound and driven by the number of unique chains and +RNA presence, not token count (observed range 7–35 min). + +Key points: + +- **Unified memory is the ceiling.** GPU tensors live in unified memory (macOS + "Memory Used") and are **not** reflected in process RSS (~7 GB). On a 192 GB + machine, ~2,200 tokens is the practical single-structure limit; larger + structures will run out of memory. Scale expectations to your Mac's memory. +- **No persistent JAX compilation cache on Metal.** `--jax_compilation_cache_dir` + does not persist across processes on MPS, so fold many inputs in one process + (e.g. `--input_dir`) to amortize compile cost. +- Set HMMER thread counts so that (concurrent searches × threads-per-search) ≈ + the number of performance cores (`sysctl hw.perflevel0.logicalcpu`). + +## Alternative: Conda environment + +`pyproject.toml` still selects `jax-mps` on `darwin arm64`, so a plain install +pulls in the Metal backend: + +```sh +conda create -n af3-mps python=3.13 cmake -y +conda activate af3-mps +pip install --no-build-isolation -e . # installs alphafold3 + jax-mps +build_data +``` + +Then run `run_alphafold.py` with the same MPS flags as in Step 4. + +## Troubleshooting + +- **`No local MPS device was found`** — `jax-mps` is not installed. Re-check the + install (`python -c "import jax; print(jax.local_devices())"`). +- **`--flash_attention_implementation must be set to "xla"`** — `triton`/`cudnn` + are NVIDIA-only; use `xla` on Metal. +- **Out-of-memory / heavy swapping on large targets** — reduce the structure + size or run on a Mac with more unified memory (see the table above). diff --git a/pyproject.toml b/pyproject.toml index 9b9e02dc..0ff0b106 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,8 @@ dependencies = [ "etils[epath]", "jax==0.10.2", "jax[cuda12]==0.10.2; sys_platform != 'darwin'", + # Apple Silicon (Metal/MPS) inference backend. + "jax-mps==0.10.9; sys_platform == 'darwin' and platform_machine == 'arm64'", "numpy", "rdkit==2025.9.4", "tokamax==0.0.12", diff --git a/run_alphafold.py b/run_alphafold.py index 1c665ddf..f95c3709 100644 --- a/run_alphafold.py +++ b/run_alphafold.py @@ -73,6 +73,7 @@ class JaxBackend(enum.StrEnum): CPU = enum.auto() GPU = enum.auto() + MPS = enum.auto() # Input and output paths. @@ -341,6 +342,8 @@ def _num_cpus_for_msa_tools() -> int: ' only for inference. This is much slower than using a GPU, but can be' ' useful for testing or running on systems without a GPU supported by' ' JAX. If you set this flag to "cpu", you must also set' + ' --flash_attention_implementation=xla. "mps" selects the experimental' + ' Apple Silicon (Metal) backend and also requires' ' --flash_attention_implementation=xla.' ), ) @@ -940,6 +943,12 @@ def main(_): 'For CPU-only inference, the --flash_attention_implementation must' ' be set to "xla".' ) + elif _JAX_BACKEND.value == JaxBackend.MPS: + if _FLASH_ATTENTION_IMPLEMENTATION.value != 'xla': + raise ValueError( + 'For Apple Silicon (MPS) inference, the' + ' --flash_attention_implementation must be set to "xla".' + ) elif _JAX_BACKEND.value == JaxBackend.GPU: gpu_devices = jax.local_devices(backend='gpu') if gpu_devices: @@ -1031,6 +1040,12 @@ def main(_): f'{_GPU_DEVICE.value}: {devices[_GPU_DEVICE.value]}' ) device = devices[_GPU_DEVICE.value] + elif _JAX_BACKEND.value == JaxBackend.MPS: + print( + f'Found local MPS devices: {devices}, using device ' + f'{_GPU_DEVICE.value}: {devices[_GPU_DEVICE.value]}' + ) + device = devices[_GPU_DEVICE.value] else: raise ValueError(f'Unsupported JAX backend: {_JAX_BACKEND.value}') diff --git a/uv.lock b/uv.lock index 35a75fab..1f74c275 100644 --- a/uv.lock +++ b/uv.lock @@ -115,6 +115,7 @@ dependencies = [ { name = "etils", extra = ["epath"], marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "jax", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "jax", extra = ["cuda12"], marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, + { name = "jax-mps", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" }, { name = "numpy", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "rdkit", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "tokamax", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -745,6 +746,18 @@ with-cuda = [ { name = "nvidia-nvshmem-cu12", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, ] +[[package]] +name = "jax-mps" +version = "0.10.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jax" }, + { name = "jaxlib" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/70/34/7f01c633b3ba61d0beaacf3bce238b1c7110bcfcf8b2e924ac3aba0d9be9/jax_mps-0.10.9-py3-none-macosx_14_0_arm64.whl", hash = "sha256:6b289db0bea58f9c11b7b09b73626d9ee0cb4b8268e2f88dca01286159fec3eb", size = 44034532, upload-time = "2026-07-10T03:50:16.636Z" }, +] + [[package]] name = "jaxlib" version = "0.10.2"