Table of Contents

ESMFold2

ESMFold2 is installed in a conda environment (see usage Biohub/esm / model card).

conda activate esm

This points python at /opt/miniconda/envs/esm and sets HF_HOME to the shared offline cache (/mnt/scratch/esm/hf). Do not activate another conda env afterwards.

HF_HUB_OFFLINE=1 means from_pretrained will not download anything. Cached checkpoints: biohub/ESMFold2, biohub/ESMFold2-Fast, ESMC backbones (300M / 600M / 6B), and biohub/ESMC-6B-sae-k64-codebook16384. Deactivate the env if you need to fetch other Hugging Face models.

Do not copy Hugging Face model-card snippets that use transformers.AutoModel / AutoTokenizer. Darwin has transformers 4.57.6 (pinned by esm 3.4.1); that stack has no esmc architecture, so those snippets raise KeyError: esmc. Native ESMC support in Transformers landed only in 5.16. Use esm.models.* as below.

ESMC and SAE

The ESMC backbones can be used independently of ESMFold2. The cached ESMC-6B SAE is optional: it converts ESMC-6B hidden states into sparse, interpretable features. It is not needed for ordinary ESMC inference or ESMFold2.

Example using the SAE for layer 60 (the Atlas / paper layer for ESMC-6B):

import torch
from esm.models.esmc import EsmcForMaskedLM, EsmcSaeModel, EsmcTokenizer
 
sequence = 'MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG'
model = EsmcForMaskedLM.from_pretrained(
    'biohub/ESMC-6B',
    device='cuda',
    dtype=torch.bfloat16,
).eval()
tokenizer = EsmcTokenizer()
sae = EsmcSaeModel.from_pretrained(
    'biohub/ESMC-6B-sae-k64-codebook16384',
    allow_patterns=['config.json', 'layer_60.safetensors'],
    device=model.device,
)
sae.initialize_layers([60])
model.add_sae_models([sae.layers['60']])
 
inputs = tokenizer(sequence, return_tensors='pt')
inputs = {key: value.to(model.device) for key, value in inputs.items()}
with torch.inference_mode():
    output = model(**inputs)
 
features = output.sae_outputs['layer60']

Use an A6000 (48 GB) for ESMC-6B. The SAE repository contains all 81 backbone layers; initialize only the layers needed by the job. SAE inputs must not contain <mask> tokens.

Starting a Job

1. Create a working directory and a Python script, e.g. my_esmfold2.py. Example (ubiquitin):

from esm.models.esmfold2 import EsmFold2Model
 
sequence = 'MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG'
model = EsmFold2Model.from_pretrained('biohub/ESMFold2-Fast').cuda().eval()
open('ubq.pdb', 'w').write(model.infer_protein_as_pdb(sequence))

For complexes (protein + DNA/RNA + ligands) use ESMFold2InputBuilder and write mmCIF with result.complex.to_mmcif() — see the model card. ESMFold2-Fast is single-sequence only; full biohub/ESMFold2 can take an MSA.

Fused kernels are already in the env (Transformer Engine 2.19 + flash-attn 2.8.3, Ampere sm_86: A5000 / A6000 / 3080 / A4500). Do not module load cuda in the job; PyTorch already ships the CUDA 13 runtime. Prefer those GPUs — flash-attn does not support Turing (2080 Ti).

2. Create a Slurm job from the ESMFold2 template.

3. Adjust the job script (for more details see Slurm Guide). GPU: A5000 (24 GB) is enough for ESMFold2-Fast on short-to-medium proteins. Full ESMFold2 uses the ESMC-6B backbone — prefer A6000 (48 GB).

4. Start the job. Do not call the Biohub Platform API from Darwin; use the local weights.