Protein Folding Inference Optimization (GPU MODE Hackathon)

Intro

Protein Folding models have been around for quite some time now. However, many of the models do not take advantage of new performance optimizations that have arrived in recent years. This repo speeds up inference for a BERT based protein folding model, that goes from 500 ms to 1 ms, a 500x speed up, using graph based optimization tools such as torch.compile

project structure

README.md profiler.py test.py bench.py model esm/ main.py ... esm_fast/ ...

steps to run

Download model weights

3B

curl https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t36_3B_UR50D-contact-regression.pt --output esm2_t36_3B_UR50D-contact-regression.pt
curl https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t36_3B_UR50D.pt --output esm2_t36_3B_UR50D.pt

15B

curl https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t48_15B_UR50D-contact-regression.pt --output esm2_t48_15B_UR50D-contact-regression.pt
curl https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t48_15B_UR50D.pt --output esm2_t48_15B_UR50D.pt

ESM FOLD

curl https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_3B_v1.pt --output esm_fold_v1.pt
curl https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_structure_module_only.pt --output esm_fold_structure.pt

uv venv --python 3.12
source ./.venv/bin/activate
uv pip install -r requirements.txt
cd esm
python inference.py

model todo list

ESM

• minimal example
• cool looking demo
• profile - runtime & memory per cuda operation
• op graph
• test.py
• bench.py
• display_autograd_graph.py