Using ESMFold on AMD GPUs

Overview

This guide explains how to run ESMFold batch jobs on Pawsey's Setonix supercomputer. ESMFold can only utilise the memory of a single GCD, which means there is an upper limit on amino acid size that it can process before running out of memory, around 2000aa. Changing the --chunk-size parameter may enable better memory performance.

Prerequisites

A Pawsey account with GPU allocation
Basic familiarity with SLURM job submission

Job Script Template

Below is a template SLURM script for running ESMFold. Save this as run_esm.slurm:

#!/bin/bash -l
#SBATCH --account=${PAWSEYPROJECT}-gpu
#SBATCH --partition=gpu
#SBATCH --nodes=1
#SBATCH --gres=gpu:1
#SBATCH --time=4:00:00
#SBATCH --job-name=esmfold

# Load required module
module load singularity/3.11.4-nompi 

# Set container
CONTAINER=docker://quay.io/pawsey/esmfold_openfold:rocm6.3.3

# Set reference directory
REF_DIR=/scratch/references/alphafold_feb2024/databases

# Set directory where models have been predownloaded for you
export TORCH_HOME=/scratch/references/esmfold/models

# Set input
INPUT=path/to/your/input

# Create results directory
RESULTS_DIR=$MYSCRATCH/esmfold/${SLURM_JOB_ID}/results/
mkdir -p $RESULTS_DIR

# Run the container
srun -N 1 -n 1 -c 8 --gres=gpu:1 \
  singularity exec \
  $CONTAINER \
  esm-fold -i $INPUT -o ${RESULTS_DIR} --chunk-size 16

Before Running

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

Running Your Job

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