esmfold2

# ESMFold2 (Biohub)

All-atom diffusion co-folding from the Biohub ESM release (2026). ESMFold2 =
48 pair layers with MSA support; ESMFold2-Fast = 24 layers, single-sequence
only, ~1.7x faster.

**License:** MIT (code github.com/Biohub/esm + weights HF `biohub/*`).
**Paper:** "Language Modeling Materializes a World Model of Protein Biology" (2026).

## Install

CUDA 12.x GPU (H100/A100-class); Python **3.12 only**. Fresh venv; needs
egress to HF Hub, GitHub, PyPI:

```bash
pip install --no-cache-dir uv
uv venv --python 3.12 /work/venv && source /work/venv/bin/activate
uv pip install \
  "torch>=2.5,<2.8" einops "biotite>=1.0" rdkit msgpack-numpy biopython \
  scikit-learn brotli attrs pandas cloudpathlib httpx tenacity zstd pydssp \
  pygtrie accelerate huggingface_hub safetensors "numpy<3" networkx \
  sentencepiece tokenizers regex packaging filelock pyyaml typing_extensions \
  "transformers @ git+https://github.com/Biohub/transformers.git@3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf"
uv pip install --no-deps "esm @ git+https://github.com/Biohub/esm.git@f652b471"
# OPTIONAL — only affects ESMC attention; trunk speedup comes from set_kernel_backend("fused")
uv pip install ninja packaging wheel setuptools
MAX_JOBS=8 uv pip install --no-deps --no-build-isolation "flash-attn<3"
# Do NOT install transformer-engine — RuntimeError (not ImportError) on import
# slips ESMC's guard and kills ESMFold2Model import.
```

The bundled `esmfold2_gpu` Modal env (remote-compute-modal skill) is the
canonical, version-pinned recipe.

**Gotchas:**
- **Default kernel backend is `None`** (reference PyTorch, ~12x slower than paper). Call `model.set_kernel_backend('fused')` after `from_pretrained()`. See section below.
- Match torch CUDA build to your driver; the pin `<2.8` targets CUDA 12.2.
- Weights via Xet bridge ~300 MB/s: ESMFold2 1.36 GB, ESMFold2-Fast 0.76 GB. Set `HF_HOME=/work/hf_cache`.

## Usage — local model

```python
from esm.models.esmfold2 import (
    ESMFold2InputBuilder, StructurePredictionInput,
    ProteinInput, DNAInput, RNAInput, LigandInput, Modification,
)
from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model

model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval()
# or "biohub/ESMFold2-Fast" (24 layers, no MSA, ~1.7x faster)
# or "biohub/ESMFold2-Experimental{,-Fast}{,-Cutoff2025}" (4 design-critic models)

spi = StructurePredictionInput(sequences=[
    ProteinInput(id="A", sequence=target_seq),
    ProteinInput(id="B", sequence=binder_seq),
    # DNAInput(id="C", sequence="ACGT", modifications=[Modification(position=5, ccd="C36")]),
    # RNAInput(id="D", sequence="ACGU"),
    # LigandInput(id="L", ccd=["SAH"]),  # or smiles="..."
])
# Homodimer: ProteinInput(id=["A","B"], sequence=seq)

results = ESMFold2InputBuilder().fold(
    model, spi,
    num_loops=10,             # paper FoldBench eval: 10; 20-loop variant: 20
    num_sampling_steps=68,    # paper eval: 68 (truncated EDM)
    num_diffusion_samples=5,  # paper eval: 5/seed
    seed=0,
)
# fold() returns list[Prediction], one per diffusion sample. Each carries
# .plddt [L], .ptm, .iptm, .pae [L,L], .pair_chains_iptm, .complex.to_mmcif().
# Rank by ipTM for complexes / mean pLDDT for monomers:
best = max(results, key=lambda r: float(r.iptm if r.iptm is not None
                                        else r.plddt.mean()))
open("pred.cif", "w").write(best.complex.to_mmcif())
```

**Paper-faithful FoldBench settings:** 10 loops, 68 sampling steps, 25 seeds
x 5 diffusion samples; rank by ipTM (complexes) or pLDDT (monomers); MSA mode
adds `msa_depth=1024` with 10% column masking and ESMC dropout 0.3.

## Model variants on HF `biohub/`

| repo | size | pair layers | MSA | use |
|---|---|---|---|---|
| `ESMFold2` | 0.94 GB + ccd.pkl 0.42 GB | 48 | yes | full eval |
| `ESMFold2-Fast` | 0.76 GB | 24 | no | fast single-seq |
| `ESMFold2-Experimental{,-Fast}` | 0.90 / 0.72 GB | 48 / 24 | — | design search (Alg 11) |
| `ESMFold2-Experimental{,-Fast}-Cutoff2025` | 0.90 / 0.72 GB | — | — | design search + critic |
| `ESMFold2-Experimental-Fast-base{300M,600M,6B}-step{250k..1500k}` | — | — | — | 15 critic ensemble |

## Throughput: `set_kernel_backend("fused")` is REQUIRED

**Default is the slow path.** `ESMFold2Model.from_pretrained(...)` loads with
`_kernel_backend=None` (reference PyTorch) and `chunk_size=64`. You MUST call:

```python
model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval()
model.set_kernel_backend("fused")   # vendored Triton TriMul/LN+SwiGLU/pair-bias kernels
model.set_chunk_size(None)           # optimal & OOM-safe L<=1024; use 256 above
```

`"fused"` gives ~1.5–6× trunk speedup over the reference backend, growing with
L; end-to-end `fold()` is diffusion-bound at short L so fused breaks even
around L≈300–400. Fused vs reference outputs are numerically consistent (pLDDT
within noise). `"fused"` (Triton, bundled with the GPU torch wheel) is
**inference-only** — auto-disables under backprop. Above ~L=1400
(`chunk_size=128`) it hits illegal memory access — fall back to
`set_kernel_backend(None)` + `set_chunk_size(64)`; validated through L=1024.

**Do NOT use** `set_kernel_backend("cuequivariance")`: the
`cuequivariance-torch==0.10.0` wheel lacks the compiled ops and **silently
falls back** to the reference path. **`apply_torch_compile()`** is an
alternative (NOT additive — call `set_kernel_backend(None)` first).

## ESMFold2-Experimental* — design hook

Experimental variants expose `res_type_soft` for gradient-guided design — see
`references/design-hook.md`. Do NOT use the fused backend with them (fp32/bf16
dtype crash; the reference path is correct).

## Gotcha: cusolver SVD poison + structseq constructor

The Kabsch alignment in `modeling_esmfold2_common.py` calls
`torch.linalg.svd(H32, driver="gesvd")` on batched 3x3 matrices. NaN/Inf inputs
(degenerate diffusion samples) corrupt the cusolver workspace — **all subsequent
CUDA calls fail with "illegal memory access"**. Mon