fold-cp documentation and examples

cookbook/the_bionemo_contessa/README.md

@@ -0,0 +1,66 @@
+# The BioNeMo Contessa — ESMFold2 folding examples
+
+This cookbook shows how to run Biohub's ESMFold2 structure-prediction model with
+[NVIDIA BioNeMo libraries and methods](https://github.com/NVIDIA-BioNeMo). It
+covers single-GPU accelerated kernel backends for faster inference and context
+parallelism (CP) for inputs whose `L x L` pair representation does not fit on one
+GPU.
+
+Tested on 1x and 4x H100 80GB GPUs.
+
+| Script | Input | GPUs | Result |
+| --- | --- | --- | --- |
+| `esmfold2-none.py` | 7ysz (1 protein chain ×2) | 1 | Reference single-GPU path |
+| `esmfold2-cueq.py` | 7ysz (1 protein chain ×2) | 1 | cuEquivariance backend |
+| `esmfold2-fused.py` | 7ysz (1 protein chain ×2) | 1 | Triton fused backend |
+| `esmfold2-cp.py` | 7ysz (1 protein chain ×2) | 4 | Same fold API with CP setup |
+| `cuequivariance_cp.py` | 5xgo (1 protein chain ×12, CL ligand) | 4 | Larger input via CP |
+| `will_fail.py` | 5xgo (same as CP) | 1 | **OOMs** — too large for 1xH100 |
+
+## Environment setup
+
+Start from the NVIDIA PyTorch container, then install the dependencies:
+
+```bash
+docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:26.03-py3 bash -l
+
+# Main ESM package
+pip install "git+https://github.com/Biohub/esm.git@main"
+# fused dependency
+pip install "esm[fused] @ git+https://github.com/Biohub/esm.git@main"
+# cuEquivariance dependency (cueq12 also exists)
+pip install "esm[cueq13] @ git+https://github.com/Biohub/esm.git@main"
+# fold-cp dependency for larger inputs
+pip install "esm[fold-cp] @ git+https://github.com/Biohub/esm.git@main"
+```
+
+## Single-GPU accelerated backends
+
+[single-gpu.md](single-gpu.md) compares the three single-GPU backend choices:
+`None` for the pure PyTorch reference path, `"cuequivariance"` for NVIDIA
+cuEquivariance triangle-multiplication kernels, and `"fused"` for Triton kernels
+that fuse several ESMFold2 hot-path operations. These backends keep the same model
+weights and outputs while trading dependencies, warm-up behavior, and speed.
+
+```bash
+python esmfold2-none.py
+python esmfold2-cueq.py
+python esmfold2-fused.py
+```
+
+## Fold-CP for larger inputs
+
+[fold-cp.md](fold-cp.md) documents how `wrap_model_with_cp(model, dm, ...)`
+spreads one ESMFold2 fold across a square grid of GPUs using the [Fold-CP methodology](https://github.com/NVIDIA-BioNeMo/boltz-cp).
+CP shards the large `L x L` pair activations so longer proteins and complexes fit in memory; it is a
+capability feature rather than a general speedup.
+
+Launch CP examples with `torchrun` on a perfect-square number of GPUs.
+
+```bash
+torchrun --nproc-per-node=4 esmfold2-cp.py
+
+# Larger 5xgo example.
+PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
+  torchrun --nproc-per-node=4 cuequivariance_cp.py
+```

cookbook/the_bionemo_contessa/cuequivariance_cp.py

@@ -0,0 +1,85 @@
+import os, math, gc
+from collections import OrderedDict
+from time import time
+
+import torch
+from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model
+from transformers.models.esmfold2.distributed import (
+    DistributedManager,
+    wrap_model_with_cp,
+)
+from esm.models.esmfold2 import (
+    ESMFold2InputBuilder,
+    LigandInput,
+    ProteinInput,
+    StructurePredictionInput,
+)
+
+spi = StructurePredictionInput(  # 5xgo
+    sequences=[
+        ProteinInput(
+            id=["A1", "B1", "C1", "D1", "E1", "F1", "G1", "H1", "I1", "J1", "K1", "L1"],
+            sequence=(
+                "MAHHHHHHVDDDDKMSTAKLVKSKATNLLYTRNDVSDSEKKATVELLNRQVIQFIDLSLITKQAHWNMRG"
+                "ANFIAVHEMLDGFRTALIDHLDTMAERAVQLGGVALGTTQVINSKTPLKSYPLDIHNVQDHLKELADRYA"
+                "IVANDVRKAIGEAKDDDTADILTAASRDLDKFLWFIECNLDLIQKMGLQNYLQAQIREEG"
+            ),
+        ),
+        LigandInput(id=["M1", "N1"], ccd=["CL"]),
+    ]
+)
+
+# torchrun entrypoint:  torchrun --nproc_per_node=4 cuequivariance_cp.py
+# torchrun sets LOCAL_RANK / WORLD_SIZE / RANK / MASTER_ADDR / MASTER_PORT.
+# world_size must be a perfect square (the CP grid is n×n).
+local_rank = int(os.environ["LOCAL_RANK"])
+world_size = int(os.environ["WORLD_SIZE"])
+
+os.environ["RANK"] = str(local_rank)
+torch.cuda.set_device(local_rank)
+torch.cuda.reset_peak_memory_stats()
+
+n = math.isqrt(world_size)
+DistributedManager.initialize(
+    grid_group_sizes=OrderedDict([("dp", 1), ("cp", (n, n))]),
+    device_type="cuda",
+    backend="nccl",
+)
+dm = DistributedManager()
+
+model = ESMFold2Model.from_pretrained("biohub/ESMFold2", esmc_precision='bf16').cuda().eval()
+# Shard BOTH the folding trunk and the MSA encoder across the n×n CP grid.
+# (Trunk-only sharding leaves the MSA encoder running the full L×L pair on
+# every rank, which OOMs for large complexes.)
+#
+# Defaults applied here (override via args): bf16=True (bf16 distributed trunk)
+# and offload_esmc=True (move the ~12 GB ESM-C LM to CPU after its one-shot use
+# so the trunk reuses that memory). Together: ~2.14x faster, ~18.5 GB lower
+# peak vs fp32, quality-neutral. comm="gather" (bit-exact) | "ring" (n>=3).
+wrap_model_with_cp(model, dm, comm="ring")
+# Fused kernels accelerate the modules *outside* the CP region (notably the
+# diffusion sampler / structure head). The CP-wrapped trunk and MSA encoder
+# no-op this call, so it never touches the distributed ring math.
+model.set_kernel_backend("cuequivariance")
+
+start = time()
+with torch.inference_mode():
+    result = ESMFold2InputBuilder().fold(
+        model, spi, num_loops=10, num_sampling_steps=50,
+        num_diffusion_samples=1, seed=0,
+    )
+end = time()
+peak_mib = torch.cuda.max_memory_allocated() / (1024**2)
+
+if local_rank == 0:
+    print(f"pLDDT mean: {float(result.plddt.mean()):.3f}, "
+          f"pTM: {float(result.ptm):.3f}, ipTM: {float(result.iptm):.3f}")
+    print(f"Elapsed: {end - start} sec")
+    print(f"Max VRAM: {peak_mib} MB")
+    with open("esmfold2_cp_output.cif", "w") as f:
+        f.write(result.complex.to_mmcif())
+
+DistributedManager.cleanup()
+DistributedManager._state.clear()
+gc.collect()
+torch.cuda.empty_cache()

cookbook/the_bionemo_contessa/esmfold2-cp.py

@@ -0,0 +1,66 @@
+import math
+import os
+from collections import OrderedDict
+from time import time
+
+import torch
+from transformers.models.esmfold2.distributed import (
+    DistributedManager,
+    wrap_model_with_cp,
+)
+from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model
+from esm.models.esmfold2 import (
+    ESMFold2InputBuilder,
+    ProteinInput,
+    StructurePredictionInput,
+)
+
+spi = StructurePredictionInput(  # 7ysz
+    sequences=[
+        ProteinInput(
+            id=["A1", "B1"],
+            sequence=(
+                "MDNFDNYEQVASIKVIGIGGAGNNAVNRMIEAGVQGVEFIVANTDAQIISVSKSKNKIVLGKETSKGLGA"
+                "GANPDVGRQAAIESAEEIKDALKGADMVFVAAGMGGGTGTGAAPIIAKLAREQGALTVGIITTPFSFEGR"
+                "ARNSYAIQGTEELRKHVDSLIIISNDRLLEVIGGVPLKDSFKEADNILRQGVQTITDLIAVPSLINLDFA"
+                "DIKTVMKNKGNALFGIGIGSGKDKAIEAANKAIISPLLEASIRGARDAIINVTGGNTLTLNDANDAVDIV"
+                "KQAIGGEVNIIFGTAVNEHLDDEMIVTVIATGFDGSHHHHHH"
+            ),
+        ),
+    ]
+)
+
+local_rank = int(os.environ["LOCAL_RANK"])
+world_size = int(os.environ["WORLD_SIZE"])
+n = math.isqrt(world_size)
+assert n * n == world_size, "CP requires a square number of GPUs: 1, 4, 9, 16, ..."
+
+torch.cuda.set_device(local_rank)
+torch.cuda.reset_peak_memory_stats()
+
+DistributedManager.initialize(
+    grid_group_sizes=OrderedDict([("dp", 1), ("cp", (n, n))]),
+    device_type="cuda",
+    backend="nccl",
+)
+dm = DistributedManager()
+
+model = ESMFold2Model.from_pretrained("biohub/ESMFold2", esmc_precision="bf16").cuda().eval()
+wrap_model_with_cp(model, dm, comm="ring")
+
+start = time()
+with torch.inference_mode():
+    result = ESMFold2InputBuilder().fold(
+        model, spi, num_loops=10, num_sampling_steps=50,
+        num_diffusion_samples=1, seed=0,
+    )
+end = time()
+peak_mib = torch.cuda.max_memory_allocated() / (1024**2)
+
+if local_rank == 0:
+    print(f"pLDDT mean: {float(result.plddt.mean()):.3f}, "
+          f"pTM: {float(result.ptm):.3f}, ipTM: {float(result.iptm):.3f}")
+    print(f"Elapsed: {end - start:.2f} sec")
+    print(f"Max VRAM: {peak_mib:.1f} MB")
+
+DistributedManager.cleanup()

cookbook/the_bionemo_contessa/esmfold2-cueq.py

@@ -0,0 +1,44 @@
+import gc
+from time import time
+
+import torch
+from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model
+from esm.models.esmfold2 import (
+    ESMFold2InputBuilder,
+    ProteinInput,
+    StructurePredictionInput,
+)
+
+spi = StructurePredictionInput(  # 7ysz
+    sequences=[
+        ProteinInput(
+            id=["A1", "B1"],
+            sequence=(
+                "MDNFDNYEQVASIKVIGIGGAGNNAVNRMIEAGVQGVEFIVANTDAQIISVSKSKNKIVLGKETSKGLGA"
+                "GANPDVGRQAAIESAEEIKDALKGADMVFVAAGMGGGTGTGAAPIIAKLAREQGALTVGIITTPFSFEGR"
+                "ARNSYAIQGTEELRKHVDSLIIISNDRLLEVIGGVPLKDSFKEADNILRQGVQTITDLIAVPSLINLDFA"
+                "DIKTVMKNKGNALFGIGIGSGKDKAIEAANKAIISPLLEASIRGARDAIINVTGGNTLTLNDANDAVDIV"
+                "KQAIGGEVNIIFGTAVNEHLDDEMIVTVIATGFDGSHHHHHH"
+            ),
+        ),
+    ]
+)
+
+torch.cuda.reset_peak_memory_stats()
+
+model = ESMFold2Model.from_pretrained("biohub/ESMFold2", esmc_precision='bf16').cuda().eval()
+model.set_kernel_backend("cuequivariance")
+
+start = time()
+with torch.inference_mode():
+    result = ESMFold2InputBuilder().fold(
+        model, spi, num_loops=10, num_sampling_steps=50,
+        num_diffusion_samples=1, seed=0,
+    )
+end = time()
+peak_mib = torch.cuda.max_memory_allocated() / (1024**2)
+
+print(f"pLDDT mean: {float(result.plddt.mean()):.3f}, "
+      f"pTM: {float(result.ptm):.3f}, ipTM: {float(result.iptm):.3f}")
+print(f"Elapsed: {end - start:.2f} sec")
+print(f"Max VRAM: {peak_mib:.1f} MB")

cookbook/the_bionemo_contessa/esmfold2-fused.py

@@ -0,0 +1,44 @@
+import gc
+from time import time
+
+import torch
+from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model
+from esm.models.esmfold2 import (
+    ESMFold2InputBuilder,
+    ProteinInput,
+    StructurePredictionInput,
+)
+
+spi = StructurePredictionInput(  # 7ysz
+    sequences=[
+        ProteinInput(
+            id=["A1", "B1"],
+            sequence=(
+                "MDNFDNYEQVASIKVIGIGGAGNNAVNRMIEAGVQGVEFIVANTDAQIISVSKSKNKIVLGKETSKGLGA"
+                "GANPDVGRQAAIESAEEIKDALKGADMVFVAAGMGGGTGTGAAPIIAKLAREQGALTVGIITTPFSFEGR"
+                "ARNSYAIQGTEELRKHVDSLIIISNDRLLEVIGGVPLKDSFKEADNILRQGVQTITDLIAVPSLINLDFA"
+                "DIKTVMKNKGNALFGIGIGSGKDKAIEAANKAIISPLLEASIRGARDAIINVTGGNTLTLNDANDAVDIV"
+                "KQAIGGEVNIIFGTAVNEHLDDEMIVTVIATGFDGSHHHHHH"
+            ),
+        ),
+    ]
+)
+
+torch.cuda.reset_peak_memory_stats()
+
+model = ESMFold2Model.from_pretrained("biohub/ESMFold2", esmc_precision='bf16').cuda().eval()
+model.set_kernel_backend("fused")
+
+start = time()
+with torch.inference_mode():
+    result = ESMFold2InputBuilder().fold(
+        model, spi, num_loops=10, num_sampling_steps=50,
+        num_diffusion_samples=1, seed=0,
+    )
+end = time()
+peak_mib = torch.cuda.max_memory_allocated() / (1024**2)
+
+print(f"pLDDT mean: {float(result.plddt.mean()):.3f}, "
+      f"pTM: {float(result.ptm):.3f}, ipTM: {float(result.iptm):.3f}")
+print(f"Elapsed: {end - start:.2f} sec")
+print(f"Max VRAM: {peak_mib:.1f} MB")

cookbook/the_bionemo_contessa/esmfold2-none.py

@@ -0,0 +1,44 @@
+import gc
+from time import time
+
+import torch
+from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model
+from esm.models.esmfold2 import (
+    ESMFold2InputBuilder,
+    ProteinInput,
+    StructurePredictionInput,
+)
+
+spi = StructurePredictionInput(  # 7ysz
+    sequences=[
+        ProteinInput(
+            id=["A1", "B1"],
+            sequence=(
+                "MDNFDNYEQVASIKVIGIGGAGNNAVNRMIEAGVQGVEFIVANTDAQIISVSKSKNKIVLGKETSKGLGA"
+                "GANPDVGRQAAIESAEEIKDALKGADMVFVAAGMGGGTGTGAAPIIAKLAREQGALTVGIITTPFSFEGR"
+                "ARNSYAIQGTEELRKHVDSLIIISNDRLLEVIGGVPLKDSFKEADNILRQGVQTITDLIAVPSLINLDFA"
+                "DIKTVMKNKGNALFGIGIGSGKDKAIEAANKAIISPLLEASIRGARDAIINVTGGNTLTLNDANDAVDIV"
+                "KQAIGGEVNIIFGTAVNEHLDDEMIVTVIATGFDGSHHHHHH"
+            ),
+        ),
+    ]
+)
+
+torch.cuda.reset_peak_memory_stats()
+
+model = ESMFold2Model.from_pretrained("biohub/ESMFold2", esmc_precision='bf16').cuda().eval()
+model.set_kernel_backend(None)
+
+start = time()
+with torch.inference_mode():
+    result = ESMFold2InputBuilder().fold(
+        model, spi, num_loops=10, num_sampling_steps=50,
+        num_diffusion_samples=1, seed=0,
+    )
+end = time()
+peak_mib = torch.cuda.max_memory_allocated() / (1024**2)
+
+print(f"pLDDT mean: {float(result.plddt.mean()):.3f}, "
+      f"pTM: {float(result.ptm):.3f}, ipTM: {float(result.iptm):.3f}")
+print(f"Elapsed: {end - start:.2f} sec")
+print(f"Max VRAM: {peak_mib:.1f} MB")