dvampire/autogenesis/esmfold2-skill

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:

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

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/

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:

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". Monkeypatch: redirect small batched SVDs to CPU:

_orig_svd = torch.linalg.svd
def _safe_svd(A, full_matrices=True, driver=None):
    if A.is_cuda and A.shape[-1] <= 4 and A.shape[-2] <= 4:
        Acpu = A.detach().float().cpu()
        if not torch.isfinite(Acpu).all():
            Acpu = torch.nan_to_num(Acpu, nan=0.0, posinf=1e6, neginf=-1e6)
        out = _orig_svd(Acpu, full_matrices=full_matrices)
        # torch.return_types.linalg_svd is a C structseq -> ctor takes ONE tuple.
        return type(out)(tuple(t.to(A.device, A.dtype) for t in out))
    return _orig_svd(A, full_matrices=full_matrices, driver=driver)
torch.linalg.svd = _safe_svd

Note type(out)(tuple(...)), not type(out)(*(...)) — torch.return_types.* are C structseqs whose constructor takes a single tuple argument.

With-MSA mode

ESMFold2 supports per-chain MSA input via ProteinInput(id, sequence, msa=MSA). The MSA object lives at esm.utils.msa.msa.MSA:

from esm.utils.msa.msa import MSA
# ProteinInput, StructurePredictionInput as imported above
msa_A = MSA.from_a3m("/path/chain_A.a3m", max_sequences=2048)
msa_B = MSA.from_a3m("/path/chain_B.a3m", max_sequences=2048)
inp = StructurePredictionInput(sequences=[
    ProteinInput(id="A", sequence=seq_A, msa=msa_A),
    ProteinInput(id="B", sequence=seq_B, msa=msa_B),
])

Gotchas:

• MSA.from_a3m(remove_insertions=True) asserts equal row lengths after

insertion removal. ColabFold a3m files often carry trailing null bytes and off-by-one rows vs the query — tr -d '\000' and force row 0 to the exact query sequence (or MSA.from_sequences on manually cleaned, query-length rows).

• ESMFold2-Fast does NOT support MSA (single-seq only).
• With-MSA mode improves AbAg interface pass-rate per the paper's evaluation.

Paper-matched inference configuration

The paper's FoldBench protocol (section A.2.11):

Training data cutoff

ESMFold2 and ESMFold2-Fast both use a Sept 2021 PDB training cutoff (HF biohub/ESMFold2 README).

ESMC language model

ESMC is the Biohub successor to ESM-2; three sizes: 300M (30L), 600M (36L), 6B (80L, d=2560). HF path: AutoModelForMaskedLM.from_pretrained("biohub/ESMC-6B").

Mask token is `<mask>` (id 32) — use tok.mask_token. The native-SDK _ convention does NOT apply to the HF tokenizer: _ is not in the vocab and encodes to <unk>, silently corrupting mutation scores.

Full API, mutation scoring, SAE features, contact prediction: see references/esmc.md.