Using AlphaFold 3 on O2

Users must fill out an AF3 attestation form to load O2's AF3 module, which also provides access to the required weights and databases. After the form is submitted, HMS Research Computing will confirm when your attestation has been completed, after which you will have access to use AF3 on O2. Based on guidance from the Office of Research Administration, this attestation is required by the AlphaFold3 Terms of Use.

Do also note that in order to access the module, you need to have an active O2 account. Please do not submit the attestation form without first confirming you have an active O2 account - you will simply be instructed to submit the form to request an account anyway. To request an O2 account, please refer to How to request or retain an O2 account .

As the title states, this article is about using Alphafold 3 on O2. For information on accessing and using Alphafold 2 on O2, please refer to Using AlphaFold 2 on O2 .

1 About AlphaFold 3
2 How to load and Use the AlphaFold 3 Module
3 Example input and job submission
4 Other Important Details
5 General Troubleshooting

About AlphaFold 3

AlphaFold3 (GitHub - google-deepmind/alphafold3: AlphaFold 3 inference pipeline. ) (AF3) is the latest tool for predicting protein structures from DeepMind. This iteration can predict 3D protein structures from sequence along with other molecules, such as DNA, RNA, and chemical ligands. This advancement has significant implications for understanding biological processes and accelerating research in fields like drug discovery, disease treatment, and environmental science. This also makes significant changes to the input, runtime processes, and output over the previous version such that it is offered as a separate module rather than an update.

If you’re new to Slurm or O2, please see Using Slurm Basic for lots of information on submitting jobs.

If you have a relatively small number of structures to predict, consider using the AF3 server which is offered for non-commercial use but with a limit of 30 submissions per day.

Before starting: While AF3’s inference code is released under a Creative Commons license, the weights (files containing the model parameters) are under a strict non-commercial terms of use granted by sole discretion of Google. HMS RC has been granted permission to host these weights for use on O2 but can only provide access after users attest to Google’s terms of use. Please fill out the AF3 attestation form before running this module on O2. After the form is submitted, HMS Research Computing will confirm when your attestation has been completed, after which you will have access to use AF3 on O2.

How to load and Use the AlphaFold 3 Module

Alphafold 3 is installed on O2 as a module that loads a python environment. We have provided a python wrapper that includes the required databases and weights, but only users that have attested to terms-of-use and granted access by HMS RC will be able to run the script. It can be invoked as follows:

module load alphafold/3.0.1 
run_alphafold.py

The following flags are mandatory for invoking AlphaFold3:

--json_path or --input_dir specifies the location of the sequence inputs to predicted by alphafold3. Exactly one must be specified.
- if submitting a single JSON file, use --json_path, and provide the path to the JSON file itself (e.g., /path/to/input.json).
- if submitting multiple JSON files (within the same directory), use --input_dir, and point to the containing directory (e.g., /path/to/input).
--output_dir specifies the directory where your output will be written to.

The --db_dir flag is not mandatory, but it will point to /n/shared_db/alphafold-3.0 by default, where RC has centrally downloaded the databases. If you would rather use your own (not recommended due to requiring approximately 2T of free space), feel free to set this flag with the corresponding location.

You can invoke run_alphafold.py -h for more information about these, and other optional flags and their options.

Example input and job submission

The following example was graciously submitted by collaborators at the Core for Computational Biomedicine.

AlphaFold 3 uses a custom JSON input format. FASTA is not supported but we are working on implementing a tool that was developed at the NIH to convert fasta to the AF3 JSON format.

An example of this format for the PDB entry 7RCE:

{ 
  "name": "rcsb_pdb_7rce", 
  "modelSeeds": [ 
    1, 
    2, 
    3, 
    4, 
    5 
  ], 
  "dialect": "alphafold3", 
  "version": 1, 
  "sequences": [ 
    { 
      "protein": { 
        "id": "A",
        "sequence": "MASSRRESINPWILTGFADAEGSFGLSILNRNRGTARYHTRLSFTIMLHNKDKSILENIQSTWKVGSILNNGDHYVSLVVYRFEDLKVIIDHFEKYPLITQKLGDYKLFKQAFSVMENKEHLKENGIKELVRIKAKMNWGLNDELKKAFPENISKERPLINKNIPNFKWLAGFTSGDGSFFVRLRKSNVNARVRVQLVFEISQHIRDKNLMNSLITYLGCGHIYEGNKSERSWLQFRVEKFSDINDKIIPVFQENTLIGVKLEDFEDWCKVAKLIEEKKHLTESGLDEIKKIKLNMNKGR" 
      } 
    }, 
    { 
      "dna": { 
        "id": "B", 
        "sequence": "GGGGGCATGCAGATCCCACAGGCGCG" 
      } 
    }, 
    { 
      "dna": { 
        "id": "C", 
        "sequence": "CCGCGCCTGTGGGATCTGCATGCCCC" 
      } 
    } 
  ] 
}

It is generally recommended to use 5 seeds ("modelSeeds": [1,2,3,4,5]). Each seed can be any numeric value and will generate its own sampled structure in the output. Note that these values are the literal seed values, e.g, [83945, 348, 293, 238738] is also a list of 5 valid seeds.

Another example: the structure of a small homo-dimer (identical chains A and B) from Haemophilus influenzae.

{ 
  "name": "2PV7", 
  "sequences": [ 
    { 
      "protein": { 
        "id": ["A", "B"], 
        "sequence": "GMRESYANENQFGFKTINSDIHKIVIVGGYGKLGGLFARYLRASGYPISILDREDWAVAESILANADVVIVSVPINLTLETIERLKPYLTENMLLADLTSVKREPLAKMLEVHTGAVLGLHPMFGADIASMAKQVVVRCDGRFPERYEWLLEQIQIWGAKIYQTNATEHDHNMTYIQALRHFSTFANGLHLSKQPINLANLLALSSPIYRLELAMIGRLFAQDAELYADIIMDKSENLAVIETLKQTYDEALTFFENNDRQGFIDAFHKVRDWFGDYSEQFLKESRQLLQQANDLKQG" 
      } 
    } 
  ], 
  "modelSeeds": [1], 
  "dialect": "alphafold3", 
  "version": 1 
}

Once you have your sequences formatted, these can be put into `alphafold.py` in two ways:

Single input file: Use the --json_path flag followed by the path to a single JSON file.
Multiple input files: Use the --input_dir flag followed by the path to a directory of JSON files.

For efficiency, we recommend using separate jobs for the MSA search step on a cpu-only partition and the inference on a gpu partition. This necessitates 2 SBATCH scripts to submit to the cluster:

AF3_part1.sh

#!/bin/bash 
#SBATCH --job-name="af3_data_pipeline" 
#SBATCH --partition=short 
#SBATCH -c 10 
#SBATCH --time=4:00:00 
#SBATCH --mem=80G 

#load the module  
module load alphafold/3.0.1 

run_alphafold.py \ 
   --output_dir=path/to/output \ 
   --json_path=rcsb_pdb_7rce.json \ 
   --norun_inference

AF3_part1.sh runs the --norun_inference to only run the data processing pipeline in this first step. The script will run the sequences through jackhmmr to generate the MSAs needed for the inference. This requires no GPU resources but does require a lot of CPU and RAM.

Once the job has finished, the output directory should have a new json file (*_data.json) with the generated MSAs which can be used as an input. We can then proceed to the second SBATCH script.

AF3_part2.sh

#!/bin/bash 
#SBATCH --job-name="af3_inference_pipeline" 
#SBATCH --partition=gpu_quad 
#SBATCH -c 8 
#SBATCH --time=2:00:00 
#SBATCH --mem=20G 
#SBATCH --gres=gpu:l40s:1 

#load the module  
module load alphafold/3.0.1 

run_alphafold.py \ 
   --output_dir=path/to/output \ 
   --json_path=path/to/output/rcsb_pdb_7rce_data.json \ 
   --norun_data_pipeline

This script, AF3_part2.sh, has a request for an ‘l40s’ GPU on the ‘gpu_quad’ partition. A100 or newer GPUs are required. Though there are workarounds, there are known numerical issues with older CUDA Capability 7.x devices; this causes bad outputs thus we will not support them for AF3 on O2. In this step you must use the json(s) created in step1 as an input. In this example, we can submit two jobs with the second being the dependency of the first:

$ sbatch AF3_part1.sh 
Submitted batch job 150427 
$ sbatch --dependency=afterok:150427 AF3_part2.sh

After inference completion, your output directory should contain several .pdb files, of which the “best” one should be called *_model.cif.

Other Important Details

A standard AlphaFold 3 run with a single protein sequence launches 4 parallel jackhmmer processes requesting 8 worker threads for each plus a master thread for each, so a total of 36 threads. Each process searches one database with database sizes 17, 67, 102, and 120 GB (for bfd-first_nonconsensus_sequences.fasta, uniref90_2022_05.fa, uniprot_all_2021_04.fa, mgy_clusters_2022_05.fa). We thus recommend the max amount of CPU cores (20 per job on O2) and at least 80 GB of ram. Note that for sequences with deep MSAs, Jackhmmer may need a substantial amount of RAM beyond the recommended 80 GB of RAM.
- Depending on how many CPU cores you request in the sbatch job submission, the run_alphafold.py script will automatically enforce core/thread count matching your job allocation by setting the --nhmmer_n_cpu and --jackhmmer_n_cpu flags to equal the number of cores available to the SLURM scheduler (e.g., via the $SLURM_CPUS_ON_NODE environment variable). For example, if you request in your sbatch script -c 12, it automatically sets --nhmmer_n_cpu=12 and --jackhmmer_n_cpu=12, regardless of whether they were manually set to other values.
Predicted quality score measures (pTM and iPTM) are found in the output file *_summary_confidences.json. Predicted Aligned Error (PAE) for each residue is printed in the file *_confidences.json.

General Troubleshooting

ModuleNotFoundError: No module named 'absl' or similar error:
- This is likely due to the environment having a competing conda presence, as this module uses conda to manage its internal components. Check for whether you have a conda init block in your $HOME/.bashrc file, and comment it out (and re-login) and see if this error persists. If commenting the block out (and re-logging) does not fix it, check to see if you have any additional python packages, e.g., in your $HOME/.local directory that may conflict.
- The run_alphafold.py script will attempt to deactivate any existing conda or python environments prior to running, so this error should be mitigated in most cases. It also attempts to run in an isolated process (e.g., without being able to see places like .local), so it should hopefully mitigate the above situation in that case as well.
- If you continue to see this error even after manually removing personal conda presence from your environment, please contact HMS RC.