Running ColabFold on Setonix

Overview

This guide explains how to run ColabFold batch jobs on Pawsey's Setonix supercomputer. ColabFold is a simplified and accelerated implementation of AlphaFold that can work with pre-computed Multiple Sequence Alignments (MSAs). Please note there is currently a memory limitation that restricts computations to proteins of approximately 3,500 amino acids or less. Attempting to process larger proteins will result in out-of-memory errors.

Prerequisites

A Pawsey account with GPU allocation
Pre-computed MSA file in A3M format unless you will get colabfold to generate the MSA for you
Basic familiarity with SLURM job submission

Job Script Template

Below is a template SLURM script for running ColabFold. Save this as run_colabfold.slurm:

#!/bin/bash -l
#SBATCH --job-name=colabbatch
#SBATCH --partition=gpu
#SBATCH --nodes=1
#SBATCH --gres=gpu:1
#SBATCH --time=01:00:00
#SBATCH --account=${PAWSEY_PROJECT}-gpu

# Load required module
module load singularity/3.11.4-nompi

# Set input and output paths
A3M=/path/to/your/msafile.a3m
OUT=$MYSCRATCH/colabfold/${SLURM_JOB_ID}
containerImage=docker://quay.io/pawsey/colabfold:1.6.1_rocm7.0.2

# Set JAX/XLA JIT compilation cache
export JAX_COMPILATION_CACHE_DIR=${MYSOFTWARE}/jax_cache

# Run ColabFold
srun -N 1 -n 1 -c 8 --gres=gpu:1 \
        singularity exec $containerImage \
        colabfold_batch \
        --data /scratch/references/colabfold_jun2024/database \
        --num-recycle 3 \
        --model-type alphafold2_multimer_v3 \
        --num-models 3 $A3M $OUT

Key Parameters and Settings

Resource Allocation:
- Uses the gpu partition
- Requests 1 GPU
- Default runtime is 1 hour
- Uses 8 CPU cores

Before Running

Modify the input path:
A3M=/path/to/your/msafile.a3m
Optional: Adjust the output directory:
OUT=$MYSCRATCH/colabfold/${SLURM_JOB_ID}
Replace ${PAWSEYPROJECT} with your project code.

Running Your Job

Submit your job:
sbatch run_colabfold.slurm
Monitor your job:
squeue -u $USER

Output Files

ColabFold will create a directory with the job ID containing:

Predicted structures in PDB format
Confidence scores
Ranking information
Log files

Common Issues and Solutions

Memory Issues: If you encounter memory errors:
- Check if your input sequence is too large i.e. >3500aa
JAX Compilation: JAX may take a long time to compile the first time it runs a new shape. Setting the cache dir means you will build up a library of compilation files over time that JAX can use instead of recompiling each time.

Using ColabFold on AMD GPUs