BioNeMo Inference Runtime Python API

This is the public Python API for structure prediction with BioNeMo Inference Runtime (BioIR). It covers the two supported ways to run a model:

1. build_processor — parse sequences and MSAs, featurize, run inference, and write PDB/CIF. This is the production entry point.
2. Model constructor + forward — construct an nn.Module, load weights, and call it on a feature dict you already have.

A runnable wrapper around (1) lives at examples/folding/run_demo.py. Supported models, GPUs, and fused kernels: Support Matrix.

import bionemo_ir registers every model factory. Any import that pulls in bionemo_ir.registry or bionemo_ir.models.* does this transitively.

When to Use Which API

Tokenizer and feature-factory objects from the registry are pipeline specs, not callables. They are wired by build_processor. There is no tokenizer(request) / features.generate_features(...) helper on the public surface; going from an InputRequest to a feature dict is what the processor is for.

Input Requests

The processor consumes a list of row dicts. Each row must include record, an InputRequest:

1
from bionemo_ir.data.schemas import InputRequest, MSARecord, Polymer
2
3
SEQUENCE = (
4
    "ACKIENIKYKGKEVESKLGSQLIDIFNDLDRAKEEYDKLSSPEFIAKFGDWINDEVERNVNEDGEPLLIQDVRQDSSKHYFFILKNGERFDLLTR"
5
)
6
7
request = InputRequest(
8
    input_id="T1031",
9
    polymers=[
10
        Polymer(
11
            polymer_type="protein",
12
            chain_id=["A1"],
13
            sequence=SEQUENCE,
14
            msas=[MSARecord(content=f">T1031\n{SEQUENCE}\n")],
15
            paired_msas=[],
16
            templates=None,
17
        ),
18
    ],
19
)

Polymer fields:

MSARecord / Template take either path or inline content, plus format ("a3m" for MSAs; "cif" or "pdb" for templates). Template chain_id selects which chain of a multi-chain CIF or PDB to use; None auto-selects.

Paired MSAs are ordinary A3M (format="a3m"), not CSV and not a concatenated multi-chain alignment. Each protein polymer gets its own file covering that chain only. Row 0 is the query; row k on every chain is one pairing group, so the files must have the same number of records (AF2 multimer enforces this). Example (chain A; chain B has the same headers and row count, sequences aligned to B):

>query
SNAELFNLESRVEIEKSLTQMEDVLKALQMKLWEAESKLSFATCKS
>tr1
-DKELFNLESRVEIEKSLKQMEDVLKALQTKLWEVESKLSFTSCKS

Lowercase letters are deletions (standard A3M). Bundled files look like 7sfy_0_paired.a3m and 7sfy_1_paired.a3m (one paired A3M per chain, same row count).

The declarative JSON used under examples/data/samples/ is the same shape. A string msas path is accepted by examples/folding/run_demo.py and resolved relative to the JSON file; the Python schema wants list[MSARecord].

1
[
2
  {
3
    "input_id": "T1031",
4
    "polymers": [
5
      {
6
        "polymer_type": "protein",
7
        "chain_id": ["A1"],
8
        "sequence": "ACKIENIKYKGKEVESKLGSQLIDIFNDLDRAKEEYDKLSSPEFIAKFGDWINDEVERNVNEDGEPLLIQDVRQDSSKHYFFILKNGERFDLLTR",
9
        "msas": "msas/T1031.a3m",
10
        "paired_msas": null,
11
        "templates": null
12
      }
13
    ]
14
  }
15
]

Templates are protein-only. format is "cif" or "pdb". Pass hits you already have — BioIR does not run HHsearch / HMMsearch. chain_id selects which chain of a multi-chain CIF or PDB to use; omit it (or null) to auto-select. Bundled sample: T1047s1_with_template.json with 8wle_A.cif.

1
from bionemo_ir.data.schemas import Template
2
3
templated = InputRequest(
4
    input_id="T1047s1_with_template",
5
    polymers=[
6
        Polymer(
7
            polymer_type="protein",
8
            chain_id=["A1"],
9
            sequence="MQKNAAHTYAISSLLVLSLTGCAWIPSTPLVQGATSAQPVPGPTPVANGSIFQSAQPINYGYQPLFEDRRPRNIGDTLTIVLQENVSASKSSSANASRDGKTNFGFDTVPRYLQGLFGNARADVEASGGNTFNGKGGANASNTFSGTLTVTVDQVLVNGNLHVVGEKQIAINQGTEFIRFSGVVNPRTISGSNTVPSTQVADARIEYVGNGYINEAQNMGWLQRFFLNLSPM",
10
            msas=[MSARecord(path="msa.a3m", format="a3m")],
11
            templates=[
12
                Template(path="templates/8wle_A.cif", format="cif", chain_id="A"),
13
            ],
14
        ),
15
    ],
16
)

1
[
2
  {
3
    "input_id": "T1047s1_with_template",
4
    "polymers": [
5
      {
6
        "polymer_type": "protein",
7
        "chain_id": ["A1"],
8
        "sequence": "MQKNAAHTYAISSLLVLSLTGCAWIPSTPLVQGATSAQPVPGPTPVANGSIFQSAQPINYGYQPLFEDRRPRNIGDTLTIVLQENVSASKSSSANASRDGKTNFGFDTVPRYLQGLFGNARADVEASGGNTFNGKGGANASNTFSGTLTVTVDQVLVNGNLHVVGEKQIAINQGTEFIRFSGVVNPRTISGSNTVPSTQVADARIEYVGNGYINEAQNMGWLQRFFLNLSPM",
9
        "msas": "msas/T1047s1.a3m",
10
        "paired_msas": null,
11
        "templates": [
12
          {
13
            "path": "templates/8wle_A.cif",
14
            "format": "cif",
15
            "chain_id": "A"
16
          }
17
        ]
18
      }
19
    ]
20
  }
21
]

RNA, DNA, and ligands are Boltz-1/2 and OpenFold3 only (AF2 / OF2 are protein-only). Nucleic-acid and ligand chains carry no MSA. A CCD ligand uses polymer_type="ccd_ligand" and a CCD code in sequence ("ATP" or "ATP_FAD"). Bundled complexes: examples/data/samples/rna_dna_ligand/.

1
rna = InputRequest(
2
    input_id="rna_demo",
3
    polymers=[
4
        Polymer(
5
            polymer_type="rna",
6
            chain_id=["A"],
7
            sequence="UUGGGUUCCCUCACCCCAAUCAUAAAAA",
8
        ),
9
    ],
10
)
11
12
dna = InputRequest(
13
    input_id="dna_demo",
14
    polymers=[
15
        Polymer(
16
            polymer_type="dna",
17
            chain_id=["A"],
18
            sequence="CGTACGATCGTA",
19
        ),
20
    ],
21
)
22
23
# Protein + custom SMILES ligand. Protein still needs an unpaired MSA.
24
smiles = InputRequest(
25
    input_id="smiles_demo",
26
    polymers=[
27
        Polymer(
28
            polymer_type="protein",
29
            chain_id=["A"],
30
            sequence="MYTVKPGDTMWKIAVKYQIGISEIIAANPQIKNPNLIYPGQKINIPNILEHHHHHH",
31
            msas=[MSARecord(path="msa.a3m", format="a3m")],
32
        ),
33
        Polymer(
34
            polymer_type="smiles_ligand",
35
            chain_id=["B"],
36
            sequence="N[C@@H](Cc1ccc(O)cc1)C(=O)O",
37
        ),
38
    ],
39
)

Same shape in JSON (smiles_demo.json / R1117v2.json in that sample dir; there is no bundled DNA JSON — DNA is the RNA shape with ACGT):

1
[
2
  {
3
    "input_id": "rna_demo",
4
    "polymers": [
5
      {
6
        "polymer_type": "rna",
7
        "chain_id": ["A"],
8
        "sequence": "UUGGGUUCCCUCACCCCAAUCAUAAAAA",
9
        "msas": null,
10
        "paired_msas": null,
11
        "templates": null
12
      }
13
    ]
14
  },
15
  {
16
    "input_id": "dna_demo",
17
    "polymers": [
18
      {
19
        "polymer_type": "dna",
20
        "chain_id": ["A"],
21
        "sequence": "CGTACGATCGTA",
22
        "msas": null,
23
        "paired_msas": null,
24
        "templates": null
25
      }
26
    ]
27
  },
28
  {
29
    "input_id": "smiles_demo",
30
    "polymers": [
31
      {
32
        "polymer_type": "protein",
33
        "chain_id": ["A"],
34
        "sequence": "MYTVKPGDTMWKIAVKYQIGISEIIAANPQIKNPNLIYPGQKINIPNILEHHHHHH",
35
        "msas": [{"path": "msas/T1152_0.a3m", "format": "a3m"}],
36
        "paired_msas": null,
37
        "templates": null
38
      },
39
      {
40
        "polymer_type": "smiles_ligand",
41
        "chain_id": ["B"],
42
        "sequence": "N[C@@H](Cc1ccc(O)cc1)C(=O)O",
43
        "msas": null,
44
        "paired_msas": null,
45
        "templates": null
46
      }
47
    ]
48
  }
49
]

Per-model coverage (monomer / MSA / templates / nucleic acids / ligands): support matrix — models and data pipeline.

build_processor

build_processor(config) in bionemo_ir.pipeline.processor.engine_proc builds a five-stage pipeline:

Parser → Tokenizer → Feature generator → Folding engine → Writer

It returns a SerialProcessor when config.executor_backend is None, or a Ray Processor when config.executor_backend == "ray".

Metadata (Boltz CCD + mols) and per-model runtime_args are filled in automatically if you omit them. User-supplied keys win over registry defaults.

Hello World (Serial)

The sequence is the bundled T1031 monomer (same as examples/folding/run_demo.py). The unpaired MSA is inlined as the query so the example runs without extra files; pass MSARecord(path=...) for a real a3m. executor_backend defaults to serial (None). Other Boltz-2 runtime_args come from the registry; only the sampling-step override is shown.

1
import json
2
3
from bionemo_ir.data.schemas import InputRequest, MSARecord, Polymer
4
from bionemo_ir.pipeline.processor.engine_proc import (
5
    EngineProcessorConfig,
6
    build_processor,
7
)
8
from bionemo_ir.pipeline.stages.configs import (
9
    FeatureGeneratorStageConfig,
10
    WriterStageConfig,
11
)
12
13
SEQUENCE = (
14
    "ACKIENIKYKGKEVESKLGSQLIDIFNDLDRAKEEYDKLSSPEFIAKFGDWINDEVERNVNEDGEPLLIQDVRQDSSKHYFFILKNGERFDLLTR"
15
)
16
17
request = InputRequest(
18
    input_id="T1031",
19
    polymers=[
20
        Polymer(
21
            chain_id=["A1"],
22
            sequence=SEQUENCE,
23
            msas=[MSARecord(content=f">T1031\n{SEQUENCE}\n")],
24
        )
25
    ],
26
)
27
rows = [{"record": request, "__record_id": request["input_id"]}]
28
29
config = EngineProcessorConfig(
30
    model_source="boltz-2",
31
    runtime_args={"num_sampling_steps": 50},
32
    feature_generator_stage=FeatureGeneratorStageConfig(
33
        init_context={"random_seed": 42},
34
    ),
35
    writer_stage=WriterStageConfig(output_path="output", format="cif"),
36
)
37
row = build_processor(config)(rows)[0]
38
scores = json.loads(row["scores"])
39
40
print(row["output_path"])  # output/T1031.cif
41
print(scores["ptm"], round(sum(scores["plddt"]) / len(scores["plddt"]), 2))

Each input row:

SerialProcessor.__call__ takes list[dict] and returns list[dict].

Ray (Multi-GPU Replicas)

Ray is the recommended executor for large inference on a GPU cluster. Staged map_batches overlaps parser / tokenizer / featurizer / writer with GPU forwards, so pre- and post-processing latency is hidden behind the engine. Serial (executor_backend=None) is for debugging or single-process measurements; it does not overlap those stages.

1
import ray
2
from bionemo_ir.pipeline.stages.configs import (
3
    EngineStageConfig,
4
    FeatureGeneratorStageConfig,
5
    ParallelismMode,
6
    ParserStageConfig,
7
    TokenizerStageConfig,
8
    WriterStageConfig,
9
)
10
11
config = EngineProcessorConfig(
12
    model_source="boltz-2",
13
    executor_backend="ray",
14
    parser_stage=ParserStageConfig(compute=4),
15
    tokenizer_stage=TokenizerStageConfig(compute=4, num_cpus=2),
16
    feature_generator_stage=FeatureGeneratorStageConfig(
17
        compute=8, num_cpus=4, init_context={"random_seed": 42}
18
    ),
19
    engine_stage=EngineStageConfig(
20
        parallelism_mode=ParallelismMode.REPLICA,
21
        compute=4,          # number of engine actors
22
        num_gpus=1.0,       # GPUs reserved per actor
23
        num_cpus=4,
24
    ),
25
    writer_stage=WriterStageConfig(
26
        compute=4, output_path="output", format="cif"
27
    ),
28
)
29
processor = build_processor(config)  # calls ray.init() if needed
30
31
ds = ray.data.from_items(rows)
32
out_rows = list(processor(ds).materialize().iter_rows())

EngineStageConfig.compute * num_gpus must not exceed visible GPUs, or build_processor raises ValueError.

One-replica-per-GPU helper:

1
config = EngineProcessorConfig.create_default_replica_mode_config(
2
    model_source="boltz-2",
3
    output_dir="output",
4
    output_format="cif",
5
)

That sets executor_backend="ray" and sizes CPU stages from torch.cuda.device_count().

EngineProcessorConfig

Inherits ProcessorConfig. Pass only documented fields.

engine_kwargs keys consumed by the folding engine:

Stage configs (ParserStageConfig, TokenizerStageConfig, FeatureGeneratorStageConfig, EngineStageConfig, WriterStageConfig) all share compute, num_cpus, memory, batch_size, drop_keys. Extra fields:

• Tokenizer / feature generator: init_context. Set init_context={"random_seed": N} on the feature-generator stage so the tokenizer can fall back to the same seed (RDKit ETKDG on OpenFold3 and MSA augmentation stay aligned). Setting it only on the tokenizer does not seed the feature stage.
• Writer: output_path, format ("pdb", "cif", or ["pdb", "cif"]).
• Engine: parallelism_mode=ParallelismMode.REPLICA, num_gpus (default 1.0).

All five stages always run. The enabled flag on a stage config is not a public way to skip a stage.

Runtime Args

build_processor starts from get_default_runtime_args(model_source) and overlays config.runtime_args. Only pass keys the model’s forward accepts.

Boltz-1 / Boltz-2

1
model(feed_dict, recycling_steps=3, num_sampling_steps=200,
2
      diffusion_samples=1, max_parallel_samples=None, steering_args=None)

OpenFold3

Same Boltz-style names, mapped inside forward:

You can also pin sample count at construction: OpenFold3(model_name="openfold3", diffusion_samples=N).

OpenFold2 / AlphaFold2

1
model(feed_dict, recycling_steps=None)

If recycling_steps is omitted, the recycle count is the last axis of aatype (sized max_recycling_iters + 1 by the feature factory). Pass runtime_args={"recycling_steps": N} to cap it. Do not pass Boltz sampling keys to OpenFold2.

CUDA Graphs (Boltz-1/2, OpenFold3, Protenix)

On Boltz-1/2, OpenFold3, and Protenix (protenix-v2) the diffusion module (including the token transformer) runs once per sampling step with a fixed shape. Capturing a CUDA graph of that module and replaying it removes per-kernel launch overhead — largest win on short sequences. OpenFold2 / AlphaFold2 have no CUDA-graph module; the same accelerated_configs entry is a no-op there. Protenix has no data pipeline; enable graphs with optimize() on the live module.

Wire it through engine_kwargs (this is what the engine’s optimize() call consumes):

1
from bionemo_ir.configs import AcceleratedConfig, BaseConfig
2
from bionemo_ir._torch.graph_optimization.config import (
3
    CUDAGraphOptimizationConfig,
4
    GraphOptimizationMode,
5
)
6
7
engine_kwargs = {
8
    "accelerated_configs": {
9
        "diffusion_module": AcceleratedConfig(
10
            backend="torch",
11
            default=BaseConfig(
12
                graph_optimization_config=CUDAGraphOptimizationConfig(
13
                    graph_optimization_mode=GraphOptimizationMode.CUDA_GRAPH_VIA_TORCH,
14
                )
15
            ),
16
        ),
17
    }
18
}

The string form of the mode is "cuda_graph_via_torch". The first few calls for a given input shape run eager (kernel compile + allocator warmup); then the graph is captured. A shape mismatch or capture failure falls back to eager. token_transformer is nested inside diffusion_module; CUDA graphs cannot nest, so requesting both keeps the parent and drops the child. Refer to optimize().

Outputs

The writer is the terminal stage (update_row=False). Each output row:

scores always includes pLDDT / pTM / ipTM / PAE when the model produces them. Boltz-2 adds extras such as confidence_score, complex_plddt, ligand_iptm, protein_iptm, pde.

A sidecar {id}_scores.json is written next to the structure when output_path is set.

Errors

With the default should_continue_on_error=False, a failed forward raises FoldingPredictionError from bionemo_ir.pipeline.stages.engine_stage. The original exception is __cause__.

1
from bionemo_ir.pipeline.stages.engine_stage import FoldingPredictionError
2
3
try:
4
    outputs = processor(rows)
5
except FoldingPredictionError as exc:
6
    raise (exc.__cause__ or exc) from None

Model Constructor and forward

Use this at inference when you already have a feature dict (custom dataloader, composing models) and want a plain nn.Module. This is not a training API.

Registry

1
import bionemo_ir  # registers factories
2
from bionemo_ir.registry import (
3
    get_model_class,
4
    get_tokenizer,
5
    get_feature_factory,
6
    get_postprocessor,
7
    get_default_runtime_args,
8
    load_metadata,
9
)
10
11
ModelCls = get_model_class("boltz-2")

Unknown names raise ValueError listing registered keys.

Constructing a Model

Import the class (from bionemo_ir.models.boltz2 import Boltz2) or get it from get_model_class: get_model_class("boltz-2") is Boltz2. Then construct it.

All folding classes accept keyword arguments config, model_name, and include_load_weights (OpenFold3 also accepts diffusion_samples). Pass model_name= explicitly for AlphaFold2 / OpenFold2 variants: OpenFold2() defaults to openfold2_ptm_1, not to the key you looked up.

1
import os
2
from bionemo_ir.models.boltz2 import Boltz2
3
from bionemo_ir.models.openfold2 import OpenFold2
4
from bionemo_ir.models.openfold3 import OpenFold3
5
from bionemo_ir.models.protenix import Protenix
6
7
os.environ["ALPHAFOLD2_1_CKPT"] = "/checkpoints/alphafold2_1.pt"
8
af2 = OpenFold2(model_name="alphafold2_1").cuda().eval()
9
10
os.environ["ALPHAFOLD2_MULTIMER_1_CKPT"] = "/checkpoints/alphafold2_multimer_1.pt"
11
af2m = OpenFold2(model_name="alphafold2_multimer_1").cuda().eval()
12
13
b2 = Boltz2(model_name="boltz-2").cuda().eval()
14
15
of3 = OpenFold3(model_name="openfold3").cuda().eval()
16
17
# Protenix is not in the registry. include_load_weights defaults to False.
18
px = Protenix(model_name="protenix-v2", include_load_weights=True).cuda().eval()

from bionemo_ir.models.boltz1 import Boltz1 follows the same pattern as Boltz2.

include_load_weights=True (default on Boltz / OpenFold2 / OpenFold3) builds from ModelCls.get_pretrained_config(model_name) and loads weights through the hub resolver. Pass include_load_weights=False for an empty module you will load yourself (model.load_weights(state_dict)). On Protenix the default is False; pass True to load hub weights.

Pass config= to override dtypes, attention backends, recycle counts, and similar. Default triangle / pairwise backends: support matrix — fused kernels.

Calling forward

1
from bionemo_ir.registry import get_default_runtime_args, get_postprocessor
2
3
runtime_args = get_default_runtime_args("boltz-2")
4
# feats: dict[str, Tensor] already on CUDA, batch dim present
5
with torch.inference_mode():
6
    raw = model(feats, **runtime_args)
7
8
folding_output = get_postprocessor("boltz-2")()(feats, raw)

Post-processor signature is __call__(batch, raw_output) → FoldingOutput, not (raw, request, output_dir=...).

FoldingOutput

FoldingOutput (bionemo_ir.data.schemas) is a dict the post-processor returns. Access fields as folding_output["atom_positions"]. Coordinates use the 37-atom protein layout the PDB/CIF writers expect. Confidence keys are None when the model does not produce them.

get_scores() returns JSON-able plddt / ptm / iptm / pae / max_pae (the writer’s scores payload). Boltz-2 also stores extras such as confidence_score and complex_plddt as additional dict keys; they are not constructor arguments.

To write a file from a FoldingOutput without the processor:

1
from bionemo_ir.data.utils import get_all_atom_types, get_all_residue_types
2
from bionemo_ir.data.writers import CIFWriter
3
4
res_types = get_all_residue_types("boltz-2")
5
atom_types = get_all_atom_types("boltz-2")
6
writer = CIFWriter(
7
    res_type_mapping=dict(enumerate(res_types)),
8
    atom_type_mapping=dict(enumerate(atom_types)),
9
    output_path="output/demo.cif",
10
)
11
writer.write(folding_output)

optimize() on a Live Module

Same CUDA-graph config as in the processor, applied yourself:

1
from bionemo_ir.configs import AcceleratedConfig, BaseConfig
2
from bionemo_ir.models.boltz2 import Boltz2
3
from bionemo_ir._torch.graph_optimization.config import (
4
    CUDAGraphOptimizationConfig,
5
    GraphOptimizationMode,
6
)
7
8
model = Boltz2(model_name="boltz-2").cuda().eval()
9
model.optimize({
10
    "diffusion_module": AcceleratedConfig(
11
        backend="torch",
12
        default=BaseConfig(
13
            graph_optimization_config=CUDAGraphOptimizationConfig(
14
                graph_optimization_mode=GraphOptimizationMode.CUDA_GRAPH_VIA_TORCH,
15
            )
16
        ),
17
    ),
18
})

optimize mutates the module in place and returns self. Unknown module names are warned and skipped. OpenFold2 has no graph-optimization modules, so this is a no-op.

token_transformer lives inside diffusion_module. CUDA graphs cannot be nested: if both are requested, optimize() keeps the parent and skips the child (Module 'token_transformer' is nested inside another requested module). Graph token_transformer alone if you only want that submodule captured. Unrelated modules (for example OpenFold3 structure_pairformer) are not nested and can be requested together.

Custom Architectures

If you already have a trained PyTorch model and want BioIR’s optimized Pairformer, diffusion transformer, or Evoformer in place of your module — not the full folding pipeline — construct the layer, remap weights, and swap it in. That path does not use build_processor.

The playbook is the module-onboard skill. Worked RF3 conversions (config, adapter, weight remap, swap) live under samples/.

The same custom-module path can take the pairwise memory optimizations already used in BioIR (Boltz, OpenFold, Protenix): bf16 pair tensors, shorter [N,N,*] lifetimes, never-materialize, and row-chunking. The playbook is the scan-mem-opt-patterns skill. Use it when the swapped layer still OOMs at large N or diffusion_samples > 1.

Layers (under bionemo_ir._torch.layers.transformers):

A typical conversion:

1. Map your hyperparameters onto the matching BioIR *Config (PairformerConfig, DiffusionTransformerConfig, EvoformerStackConfig from bionemo_ir.configs).
2. Remap state_dict keys into the BioIR layout (QKV / KV fusion, AdaLN gain+bias fusion, gate+input fusion, name renames such as tri_mul_outgoing → tri_mul_out).
3. Write a thin nn.Module adapter if signatures differ (mask polarity, extra sample/batch axes, bool vs float valid-masks).
4. Replace the original submodule on a live model.
5. Compare block-level then stack-level numerics against the original.
6. Optionally call model.optimize(...) for CUDA graphs on modules that declare graph optimization.

Fused kernels on supported SKUs: support matrix — fused kernels.

Related

• Config architecture (model tree vs pipeline stages): Config Architecture
• Support matrix (models, GPUs, fused kernels): Support Matrix
• Demo CLI: examples/folding/run_demo.py
• Sample JSON / MSA: examples/data/samples/
• Module onboarding (swap Pairformer / DiT / Evoformer into your model): module-onboard skill
• Pairwise memory optimizations (port BioIR patterns onto your module): scan-mem-opt-patterns skill