diff --git a/cookbook/the_bionemo_contessa/README.md b/cookbook/the_bionemo_contessa/README.md new file mode 100644 index 00000000..7f8f3d3e --- /dev/null +++ b/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 +``` diff --git a/cookbook/the_bionemo_contessa/cuequivariance_cp.py b/cookbook/the_bionemo_contessa/cuequivariance_cp.py new file mode 100644 index 00000000..f9970bec --- /dev/null +++ b/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() diff --git a/cookbook/the_bionemo_contessa/esmfold2-cp.py b/cookbook/the_bionemo_contessa/esmfold2-cp.py new file mode 100644 index 00000000..3d0ad30e --- /dev/null +++ b/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() \ No newline at end of file diff --git a/cookbook/the_bionemo_contessa/esmfold2-cueq.py b/cookbook/the_bionemo_contessa/esmfold2-cueq.py new file mode 100644 index 00000000..09ea35cd --- /dev/null +++ b/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") diff --git a/cookbook/the_bionemo_contessa/esmfold2-fused.py b/cookbook/the_bionemo_contessa/esmfold2-fused.py new file mode 100644 index 00000000..ef0643f3 --- /dev/null +++ b/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") diff --git a/cookbook/the_bionemo_contessa/esmfold2-none.py b/cookbook/the_bionemo_contessa/esmfold2-none.py new file mode 100644 index 00000000..aacf07aa --- /dev/null +++ b/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") diff --git a/cookbook/the_bionemo_contessa/fast_runtime_and_vram.png b/cookbook/the_bionemo_contessa/fast_runtime_and_vram.png new file mode 100644 index 00000000..e8eda16f Binary files /dev/null and b/cookbook/the_bionemo_contessa/fast_runtime_and_vram.png differ diff --git a/cookbook/the_bionemo_contessa/fold-cp.md b/cookbook/the_bionemo_contessa/fold-cp.md new file mode 100644 index 00000000..fb3bcf5a --- /dev/null +++ b/cookbook/the_bionemo_contessa/fold-cp.md @@ -0,0 +1,182 @@ +# Fold larger inputs by running ESMFold2 across several GPUs + +`wrap_model_with_cp(model, dm, …)` takes a normal `ESMFold2Model` and rewires it, to spread one fold across several GPUs using the [Fold-CP methodology](https://github.com/NVIDIA-BioNeMo/boltz-cp) so you can fold longer +proteins than fit on a single GPU. Wrapping results in the same object with a few internal pieces swapped for versions that split their work across the GPUs. + +## Table of contents + +- [Quickstart](#quickstart) +- [Requirements](#requirements) +- [API: wrap_model_with_cp](#api-wrap_model_with_cp) +- [Behavior and differences from single-GPU implementation](#behavior-and-differences-from-single-gpu-implementation) +- [Larger example](#larger-example) +- [How context parallelism works](#how-context-parallelism-works) + +## Quickstart + +Install the fold-cp dependency + +```bash +pip install "esm[fold-cp] @ git+https://github.com/Biohub/esm.git@main" +``` + +Use [esmfold2-cp.py](esmfold2-cp.py) as the CP version of the single-GPU +[esmfold2-none.py](esmfold2-none.py) example. Launch it with `torchrun`, using a +perfect-square number of GPUs: + +```bash +# Launch on 4 GPUs +torchrun --nproc-per-node=4 esmfold2-cp.py +``` + +Compared with `esmfold2-none.py`, the `-cp.py` script changes setup, not the input or the `fold()` call: + +- Assigns each process to its own CUDA device using `LOCAL_RANK` and `WORLD_SIZE` (set by `torchrun`) +- Initializes `DistributedManager` with `("cp", (n, n))`, where n is the `sqrt(WORLD_SIZE)` +- Wraps the normal `ESMFold2Model` with CP using `wrap_model_with_cp(model, dm, comm="ring")` +- It keeps `num_diffusion_samples=1`, which is required by the distributed diffusion path. + +The wrapped model is still an `ESMFold2Model`; `fold()` and the result object are +unchanged. + +## Requirements + +- Python environment: + - This ESM package + - transformer-engine 2 (installed by `esm[fold-cp]` dependency) +- Use a perfect-square number of GPUs: 1, 4, 9, 16, ... +- Keep `num_diffusion_samples=1` in `ESMFold2InputBuilder().fold(...)`. The CP + diffusion path currently expects one diffusion sample. +- If you enable `tp_esmc=True`, ESM-C's MLP hidden size must divide across the CP + ranks. The current ESM-C `ffn_hidden=6912` works for 4, 9, and 16 GPUs. + +> Use a fast GPU interconnect for good throughput. NVLink within a node and + InfiniBand between nodes are the intended setup; PCIe or Ethernet work and still save + memory, but communication may dominate runtime. + +## API: wrap_model_with_cp + +```python +replaced = wrap_model_with_cp( + model, + dm, + comm="gather", + bf16=True, + offload_esmc=True, + wrap_structure=True, + tp_esmc=False, +) +``` + +`wrap_model_with_cp` mutates the existing `ESMFold2Model` in place and returns a +list of module paths that were replaced or augmented. The model type, `fold()` API, +and output object stay the same. + +| Argument | Default | What it controls | +|---|---:|---| +| `model` | required | The `ESMFold2Model` to rewire. Load it normally with `from_pretrained(...)` first. | +| `dm` | required | The initialized `DistributedManager`; its CP mesh must be square. | +| `comm` | `"gather"` | MSA pair-weighted-averaging communication. Use `"gather"` for exact all-gather behavior, or `"ring"` for a cheaper online-softmax ring path on larger grids. | +| `bf16` | `True` | Runs the distributed trunk/MSA path in bf16. This is the practical default for lower memory and faster inference; use `False` only for tight fp32 parity checks. | +| `offload_esmc` | `True` | Moves the ESM-C language model back to CPU after its one-shot use, freeing GPU memory for the trunk and diffusion stages. | +| `wrap_structure` | `True` | Wraps the diffusion structure head so the conditioned pair representation can stay sharded through structure sampling. | +| `tp_esmc` | `False` | Tensor-parallelizes ESM-C's MLP across the CP ranks. Useful when ESM-C's replicated weights are the remaining memory floor. | + +If you enable `tp_esmc=True`, ESM-C's MLP hidden size must divide across the CP ranks. +The current ESM-C `ffn_hidden=6912` works for `4`, `9`, and `16` GPUs. + +Common choices: + +```python +# Recommended large-input path. +wrap_model_with_cp(model, dm, comm="ring") + +# Exact communication path; useful for parity checks. +wrap_model_with_cp(model, dm, comm="gather") + +# Also shard ESM-C's MLP when ESM-C memory is the bottleneck. +wrap_model_with_cp(model, dm, comm="ring", tp_esmc=True) +``` + +After wrapping, keep `num_diffusion_samples=1` in the `fold()` call. + +## Behavior and differences from single-GPU implementation + +The API and outputs are identical (`from_pretrained`, `fold()`, output keys, `type(model)`), and almost every step matches the single-GPU result. +Two differences are expected, both small: + +- **LM dropout** picks a different (still valid) random pattern than the single-GPU + run, because the sharded table can't cheaply reproduce the exact full-table draw. + This is the main reason the pTM/ipTM scores can wiggle slightly. +- **Rounding** differs at about 1e-3, because numbers are summed across GPUs in a + different order. + +Single-GPU kernel backends such as `"fused"` and `"cuequivariance"` are ignored by the CP-wrapped stages; those stages use the distributed CP implementations instead. + +## Larger example + +If you have a large 12-chain complex like [5xgo](https://www.rcsb.org/structure/5XGO), attempting to fold it on a single H100 80GB will fail with OOM + +```bash +python will_fail.py # expected to OOM +``` + +Context parallelism sharding will split the memory and allow this protein to be folded on 4x H100 SXM GPUs. +Launch with torchrun on a perfect-square number of GPUs: + +```bash +PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True torchrun --nproc-per-node=4 cuequivariance_cp.py +``` + +## How context parallelism works + +A protein is a chain of building blocks called residues (`L` = how many). +To predict its shape, ESMFold2 keeps a big matrix with one cell for every pair of residues. +This is the pair representation (or just `z`), and it's what +dominates memory, because `L` residues means `L x L` cells: + +- 500 residues → 250,000 cells +- 2,000 residues → 4,000,000 cells (16x bigger) + +This means that the memory requirement scales quadratically with the input size. + +Context parallelism (CP) shards tensors (vectors and matrices) across GPUs. +More GPUs means more memory for each tensor, allowing for larger inputs. + +CP arranges the GPUs into a square grid so the GPU count must be a perfect square (e.g. 1, 4, 9, 16, …). +Each GPU process is a **rank**, and the total count is the **world size**. +The `L x L` matrix is split into blocks owned by each rank with a [PyTorch DTensor](https://docs.pytorch.org/docs/2.12/distributed.tensor.html). + +The model weights are kept replicated, where each GPU has a copy. +Only the big activations get sharded in this method so per-GPU memory drops as GPUs are added to the inference pool. + +![memory usage](fast_runtime_and_vram.png) + +### Stage-by-stage: what runs where + +The base model runs every stage at full length `L` on every GPU, with the pair table +full-size. The table below walks that same pipeline; **a ✅ means +`wrap_model_with_cp` splits that stage across the grid** (everything else stays +replicated). + +| Stage (base ESMFold2Model.forward) | Split across GPUs? | What wrapping does | +|---|---|---| +| `inputs_embedder` (atom features) | — (replicated) | unchanged; its cost grows only **linearly** with size (it looks at a small window of atoms at a time), so it stays a small, fixed floor | +| ESM-C 6B language model → embeddings | ✅ MLP only, if `tp_esmc=True` | splits ESM-C's big MLP across GPUs; the rest stays replicated; ESM-C is moved to CPU right after it's used, to free room | +| `language_model` → `lm_z` (`LxL`) | ✅ | `_cp_language_model` builds this `LxL` table already sharded (the full thing is never assembled on any GPU) | +| pair init (`z_init`, rel_pos, token_bonds, initial `z`, pair mask) | ✅ | `_cp_pair_init` builds each `LxL` table one block per GPU — no full copy anywhere. | +| recycle loop (`_run_one_loop`) | ✅ | `_cp_recycle_engine` keeps the pair table sharded the whole time — it never gathers the full table between passes | +| `parcae_readout` + `parcae_coda` | ✅ | runs on each GPU's block plus a shared trunk | +| `distogram_head(z + zᵀ)` | ✅ | works on the sharded table and gathers only the small result | +| `structure_head.sample` (diffusion → 3D coords) | ✅ (if `wrap_structure`) | runs the diffusion with the pair table kept sharded | +| `confidence_head` (pLDDT/PAE/PDE/pTM/ipTM scores) | ✅ | every `LxL` input (pair, rel_pos, token_bonds, distance bins, mask) stays sharded — nothing full-size is rebuilt; only the small final scores are gathered | + +### Summary: what shrinks, what doesn't + +- **Grows with the `LxL` pair table → now sharded:** pair init, recycle loop, + parcae, distogram, structure, confidence. Each is built one block per GPU and never + assembled full, so **adding GPUs raises the longest protein you can fold**. +- **Grows only with `L` (atom-level work):** `inputs_embedder` — a slow-growing, + replicated floor. +- **Roughly fixed:** ESM-C's weights (split by `tp_esmc`); ESM-C's own activations are + the remaining floor. diff --git a/cookbook/the_bionemo_contessa/single-gpu.md b/cookbook/the_bionemo_contessa/single-gpu.md new file mode 100644 index 00000000..49d74105 --- /dev/null +++ b/cookbook/the_bionemo_contessa/single-gpu.md @@ -0,0 +1,180 @@ +# Single-GPU ESMFold2 — kernel backends + +The reference implementation in pure PyTorch is accurate but not the fastest way to run ESMFold2. +Most of the single-GPU runtime sits in a few repeated pair-representation operations, and there are alternate **kernel backends** that keep the same weights and outputs while swapping in faster implementations of those hot paths. +The result is the same ESMFold2 model, but with different speed, warm-up, and dependency trade-offs. + +ESMFold2 runs the expensive Pairformer/attention math through a selectable +**kernel backend**, chosen with: + +```python +model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval() +model.set_kernel_backend(None | "fused" | "cuequivariance") +``` + +This propagates out to every module that runs Pairformer-style blocks (`folding_trunk`, +`lm_encoder`, `parcae_coda`, `confidence_head`, `structure_head`). It only swaps the +*implementation* of a few hot ops. Weights are untouched and the three +backends are numerically equivalent to bf16 rounding (same pLDDT/pTM). + +> This page covers the single-GPU backends. For multi-GPU context parallelism see +> `fold-cp.md` (`wrap_model_with_cp`), which is an orthogonal choice. + +## Table of contents + +- [The three backends](#the-three-backends) +- [None (default experience)](#none-default-experience) +- [cuEquivariance](#cuequivariance) +- [Fused](#fused) +- [Performance summary](#performance-summary) +- [Decision guide](#decision-guide) + +## The three backends + +| | None | fused | cuequivariance | +|---|---|---|---| +| Implementation | PyTorch | Triton kernels | [cuEquivariance](https://github.com/nvidia/cuequivariance) | +| Accelerated Operations | none (reference) | tri-mul + LN+SwiGLU + dropout-residual + pair-bias | tri-mul | +| Extra dependency | none | `triton>=3` | `cuequivariance-torch` (CUDA-matched build) | +| Training / autograd | yes | inference-only | yes | +| First-call compilation | none | yes — Triton JIT + autotune | no — precompiled kernels | + +## None (default experience) + +### Dependencies + +None beyond this ESM package + +### Usage + +```python +model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval() +model.set_kernel_backend(None) # optional +``` + +### Performance + +Fold a 644-residue protein using the default backend with [esmfold2-none.py](esmfold2-none.py) on a single H100 SXM + +```shell +python esmfold2-none.py + +/usr/local/lib/python3.12/dist-packages/torch/jit/_script.py:1487: DeprecationWarning: `torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`. + warnings.warn( +🚨 No checkpoint found for ESMCForSequenceClassification.forward. Please add a `checkpoint` arg to `auto_docstring` or add one in ESMCConfig's docstring +🚨 No checkpoint found for ESMCForTokenClassification.forward. Please add a `checkpoint` arg to `auto_docstring` or add one in ESMCConfig's docstring +Loading checkpoint shards: 100%|█████████████████████████████████| 6/6 [00:00<00:00, 151.51it/s] +Loading CCD dictionary from /root/.cache/huggingface/hub/models--biohub--ESMFold2/snapshots/1ebf0e3481a5184eb6171d40615c79e384b48796/ccd.pkl +pLDDT mean: 0.751, pTM: 0.458, ipTM: 0.096 +Elapsed: 58.02 sec +Max VRAM: 19682.0 MB +``` + +## cuEquivariance + +[cuEquivariance](https://github.com/nvidia/cuequivariance) is NVIDIA’s precompiled CUDA kernel backend for ESMFold2’s triangle-multiplication hot path, giving a large single-GPU speedup without Triton JIT or accuracy changes. + +### Dependencies + +* `cuequivariance-torch` +* Depending on your CUDA version + * `cuequivariance-ops-torch-cu13` + * `cuequivariance-ops-torch-cu12` + +Install the CUDA-matched extra: + +```bash +pip install "esm[cueq12]" # CUDA 12 build +``` +or +```bash +pip install "esm[cueq13]" # CUDA 13 build +``` + +You can figure out which one you need with `nvidia-smi` + +```shell +$ nvidia-smi + +Thu Jul 16 13:03:23 2026 ++---------------------------------------------------------------------------------------+ +| NVIDIA-SMI 535.216.03 Driver Version: 535.216.03 CUDA Version: 13.2 | +|-----------------------------------------+----------------------+----------------------+ +``` + +In this case, the CUDA Version is 13.2, so you would need to install `esm[cueq13]`. + +### Usage + +```python +model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval() +model.set_kernel_backend("cuequivariance") +``` + +### Performance + +Fold a 644-residue protein using the cuEquivariance backend with [esmfold2-cueq.py](esmfold2-cueq.py) on a single H100 SXM + +```shell +python esmfold2-cueq.py + +pLDDT mean: 0.751, pTM: 0.459, ipTM: 0.096 +Elapsed: 19.90 sec +Max VRAM: 19680.1 MB +``` + +## Fused + +The "fused" backend uses Triton, a Python-based language for writing custom GPU kernels, to fuse several ESMFold2 hot-path operations into fewer CUDA launches for the fastest single-GPU inference while preserving the same model weights and outputs. + +### Dependencies + +The "fused" backend needs Triton 3. It's GPU-only (Triton JIT-compiles to PTX) and **inference-only** (falls back to reference under autograd). + +```bash +pip install "esm[fused]" # triton>=3,<4 +``` + +> If Triton is not importable, `set_kernel_backend("fused")` silently runs the reference path. Check that Triton can be imported if `"fused"` is unexpectedly slow. + +### Usage + +```python +model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval() +model.set_kernel_backend("fused") +``` + +### Performance + +Fold a 644-residue protein using the "fused" backend with [esmfold2-fused.py](esmfold2-fused.py) on a single H100 SXM + +```shell +python esmfold2-fused.py + +pLDDT mean: 0.750, pTM: 0.458, ipTM: 0.096 +Elapsed: 16.56 sec +Max VRAM: 19795.1 MB +``` + +## Performance summary + +These results fold the same 644-residue `7ysz` input used in the per-backend +examples above. + +| backend | elapsed | max VRAM | pLDDT | pTM | ipTM | speedup vs `None` | +|---|---:|---:|---:|---:|---:|---:| +| `None` | 58.02 s | 19682.0 MB | 0.751 | 0.458 | 0.096 | 1.0× | +| "cuequivariance" | 19.90 s | 19680.1 MB | 0.751 | 0.459 | 0.096 | 2.9× | +| "fused" | 16.56 s | 19795.1 MB | 0.750 | 0.458 | 0.096 | 3.5× | + +## Decision guide + +If you want: + +- **Fastest throughput →** `"fused"`. + - Fastest steady state, and the incremental per-new-length compile is only ~0.5 s +- **Huge length variety →** `"cuequivariance"` + — precompiled, steady state a bit slower than fused. +- **Bit-exact reference for debugging →** `None`. + +> Tip: For any benchmarking, include a warm-up fold at startup to pay the ~10s global cost diff --git a/cookbook/the_bionemo_contessa/will_fail.py b/cookbook/the_bionemo_contessa/will_fail.py new file mode 100644 index 00000000..cd986a9c --- /dev/null +++ b/cookbook/the_bionemo_contessa/will_fail.py @@ -0,0 +1,46 @@ +import gc +from time import time + +import torch +from transformers.models.esmfold2.modeling_esmfold2 import ESMFold2Model +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"]), + ] +) + +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} sec") +print(f"Max VRAM: {peak_mib} MB") +with open("esmfold2_output.cif", "w") as f: + f.write(result.complex.to_mmcif()) diff --git a/pyproject.toml b/pyproject.toml index 7ab6d609..a3bb3da0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,6 +45,17 @@ dependencies = [ "dna_features_viewer", "accelerate", ] + +[project.optional-dependencies] +# ESMFold2 "fused" kernel backend (set_kernel_backend("fused")) +# inference kernels for tri-mul / LN+SwiGLU / dropout-residual. +fused = ["triton>=3,<4"] +# ESMFold2 context-parallel ESM-C tensor parallelism (wrap_model_with_cp(tp_esmc=True)). +fold-cp = ["transformer-engine[pytorch]>=2,<3"] +# "cuequivariance" kernel backend — pick the build matching your CUDA toolkit. +cueq12 = ["cuequivariance-torch", "cuequivariance-ops-torch-cu12"] +cueq13 = ["cuequivariance-torch", "cuequivariance-ops-torch-cu13"] + # Pytest [tool.pytest.ini_options] addopts = """