===== ESMFold2 =====

ESMFold2 is installed in a conda environment (see usage [[https://github.com/Biohub/esm|Biohub/esm]] / [[https://huggingface.co/biohub/ESMFold2|model card]]).

<code>conda activate esm</code>

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 [[https://huggingface.co/biohub/ESMC-6B-sae-k64-codebook16384|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):
<code python>
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']
</code>

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):
<code python>
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))
</code>

For complexes (protein + DNA/RNA + ligands) use ''ESMFold2InputBuilder'' and write mmCIF with ''result.complex.to_mmcif()'' — see the [[https://huggingface.co/biohub/ESMFold2|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** [[https://portal.darwin.bioeng.ru/pun/sys/myjobs/workflows/new|template]].

**3.** Adjust the job script (for more details see [[portal_guides:slurm|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.