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-highmem
#SBATCH --nodes=1
#SBATCH --gres=gpu:2
#SBATCH --mem=0
#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}
CACHE_DIR=$MYSCRATCH/colabfold/cache
containerImage=docker://quay.io/pawsey/colabfold:rocm6.2.4

# Environment settings for optimal GPU performance
export SINGULARITYENV_XLA_PYTHON_CLIENT_PREALLOCATE="0"
export SINGULARITYENV_TF_FORCE_UNIFIED_MEMORY="0"
export SINGULARITYENV_XLA_PYTHON_CLIENT_MEM_FRACTION="4.0"
export SINGULARITYENV_TF_FORCE_GPU_ALLOW_GROWTH="true"
export SINGULARITYENV_JAX_PLATFORMS="rocm"
export XLA_FLAGS="--xla_gpu_autotune_level=0"

# Verify connection to ColabFold API (only needed if using colabfold server to generate MSAs)
# openssl s_client -connect api.colabfold.com:443 

# Create cachedir 
mkdir -p $CACHE_DIR

# Run ColabFold
srun -N 1 -n 1 -c 16 --gres=gpu:2 \
        singularity exec $containerImage \
        colabfold_batch \
        --data $CACHE_DIR \
        --num-recycle 5 \
        --pair-mode unpaired \
        --model-type alphafold2_multimer_v3 \
        --num-models 3 $A3M $OUT

Key Parameters and Settings

Resource Allocation:
- Uses the gpu partition
- Requests 1 GPU
- Sets unlimited memory (--mem=0)
- Default runtime is 1 hour
- Uses 8 CPU cores
Environment Variables:
export SINGULARITYENV_XLA_PYTHON_CLIENT_PREALLOCATE="0" export SINGULARITYENV_TF_FORCE_UNIFIED_MEMORY="0" export SINGULARITYENV_XLA_PYTHON_CLIENT_MEM_FRACTION="4.0" export SINGULARITYENV_TF_FORCE_GPU_ALLOW_GROWTH="true" export SINGULARITYENV_JAX_PLATFORMS="rocm" export XLA_FLAGS="--xla_gpu_autotune_level=0"
These settings optimize GPU memory usage and performance on Setonix's AMD GPUs.
ColabFold Settings:
- --num-recycle 5: Number of prediction refinement cycles
- --pair-mode unpaired: For single chain predictions
- --model-type alphafold2_multimer_v3: Uses the latest multimer model
- --num-models 3: Generates 3 prediction models

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

API Connection: The script checks the connection to the ColabFold API. If this fails, ensure you have internet connectivity from your compute node.
Memory Issues: If you encounter memory errors:
- Check if your input sequence is too large i.e. >3500aa

Using ColabFold on AMD GPUs