Using RFDiffusion on AMD GPUs

Overview

RFdiffusion is a powerful protein structure generation tool that uses diffusion models to create novel protein structures with desired properties. This guide details how to use RFdiffusion on the Pawsey Supercomputing Centre using our custom container for AMD GPUs.

Container Location

The RFdiffusion container is available for download from https://quay.io/pawseysc/rfdiffusion.

Batch Job Submission

This is a template SBATCH script to run RFDiffusion. It will also download the container for you if you haven’t already downloaded it.

#!/bin/bash -l
#SBATCH --job-name=rfdiffusion
#SBATCH -A ${PAWSEY_PROJECT}-gpu
#SBATCH --nodes=1
#SBATCH --partition=gpu
#SBATCH --time=10:00:00
#SBATCH --gres=gpu:1

# Load required modules
module load singularity/3.11.4-nompi

#Set up container variable
ContainerImage=docker://quay.io/pawsey/rfdiffusion:rocm6.4

# Set up working directory and output directory. This can be edited to suit you
WORKDIR=${MYSCRATCH}/rfdiffusion_runs/${SLURM_JOB_ID}
mkdir -p ${WORKDIR}
cd ${WORKDIR}
OUTDIR=${WORKDIR}/test_outputs/test

#Create schedules dir for RFDiffusion to use
mkdir ${WORKDIR}/schedules

# Run RFdiffusion
srun N 1 -n 1 -c 8 --gres=gpu:1 \
singularity exec \
-B schedules:-B schedules:/app/RFdiffusion/rfdiffusion/inference/../../schedules \
${ContainerImage} \
run_inference.py \
inference.model_directory_path=/app/RFdiffusion/models \
inference.output_prefix=${OUTDIR} \
<other parameters here as needed>

Save this script as run_rfdiffusion.sh and submit with:

sbatch run_rfdiffusion.sh