Skip to navigation

Model Weights

View as Markdown

Checkpoints are not committed here. A model resolves them in __init__ when include_load_weights=True (default for every family except Protenix), so a miss fails at engine startup, not mid-inference.

Weights keep their upstream licences, and each upstream repo page is authoritative. Boltz also needs chemical-component metadata — refer to Metadata assets.

Notice and Disclaimer

The following notice applies to separately downloaded model checkpoints and metadata assets.

NOTICE AND DISCLAIMER: This software automatically retrieves, accesses or interacts with external materials. Those retrieved materials are not distributed with this software and are governed solely by separate terms, conditions and licenses. You are solely responsible for finding, reviewing and complying with all applicable terms, conditions, and licenses, and for verifying the security, integrity and suitability of any retrieved materials for your specific use case. This software is provided “AS IS”, without warranty of any kind. The author makes no representations or warranties regarding any retrieved materials, and assumes no liability for any losses, damages, liabilities or legal consequences from your use or inability to use this software or any retrieved materials. Use this software and the retrieved materials at your own risk.

Sources

HF_CHECKPOINTS in hubs/hf.py is the source of truth:

FamilyModel keysUpstream repoFile
Boltz-1boltz-1boltz-community/boltz-1boltz1_conf.ckpt
Boltz-2boltz-2boltz-community/boltz-2boltz2_conf.ckpt
Boltz-2boltz-2-affinityboltz-community/boltz-2boltz2_aff.ckpt
OpenFold2every openfold2_* keynz/OpenFoldfinetuning_*.pt (incl. no_templ / ptm)
OpenFold3openfold3OpenFold/OpenFold3checkpoints/of3-p2-155k.pt
Protenixprotenix-v2nonelocal file only
AlphaFold2alphafold2_*, alphafold2_multimer_*nonelocal file only
  • OpenFold2 Hub downloads come from a third-party mirror. nz/OpenFold is not aqlaboratory/openfold. scripts/fetch_weights.sh pulls the same parameters from OpenFold’s own S3 bucket and checks them against the same digests, or stage the files yourself.
  • AlphaFold2 and Protenix have no HF_CHECKPOINTS entry, so those keys resolve locally only. With nothing staged, load_hf_weights raises ValueError naming the <MODEL>_CKPT variable to set. Refer to AlphaFold2 parameters.
  • Protenix-v2: to get the weights, check <https://github.com/bytedance/Protenix> or contact the Protenix authors.
  • protenix-v2 has weights but no pipeline factory. Construct models/protenix and call load_weights() yourself (include_load_weights defaults to False).

Every Hub download is pinned to a commit (HF_REVISIONS in hubs/integrity.py), so an upstream re-upload changes nothing until BioIR moves the pin. Each file is also checked against its SHA-256 (CHECKPOINT_SHA256) before it is loaded, and a file with no recorded digest is refused. A repo_id you pass yourself resolves that repo’s default branch, and the digest check still applies. Local files are not checked.

Staging With fetch_weights.sh

scripts/fetch_weights.sh downloads checkpoints from their upstream publishers and symlinks them into the probe directory, so pytest and the examples resolve them with no *_CKPT set. This is the normal way to get weights; the manual route is Staging and overrides.

scripts/fetch_weights.sh # everything the source covers
scripts/fetch_weights.sh --model boltz # one family
scripts/fetch_weights.sh --model boltz-2 # one checkpoint
scripts/fetch_weights.sh --help

--source selects where they come from. The default auto uses NVIDIA’s internal NGC mirror when BIOIR_NGC_ORG, BIOIR_NGC_TEAM, and NGC credentials are all present, and the public upstreams otherwise — so outside NVIDIA it is exactly --source public and never needs credentials. --source ngc forces the mirror and fails when it is unusable; --source none stages nothing and only reports what already resolves.

What the public source covers, per family — the per-file list is the PUBLIC_URLS table in the script, not restated here to avoid drift:

FamilySourceLicence
OpenFold2openfold.s3.amazonaws.comCC BY 4.0
BoltzHuggingFace boltz-community, plus metadataMIT
OpenFold3HuggingFace OpenFold/OpenFold3 (gated)Apache-2.0
Protenixnone; stage your own fileApache-2.0
AlphaFold2conversion required, refer to the followingCC BY 4.0

Downloads are resumable and skipped when already present. Hugging Face URLs name the same commits as HF_REVISIONS, and every file is checked against its SHA-256 when it downloads and again each time it is staged. A download that does not match is deleted and fails the run; a cached file that no longer matches, or that has no pinned digest, fails staging until you delete it. Anything that cannot be fetched is reported and skipped, never fatal — so a run with no credentials completes, and the tests whose weights are missing skip themselves.

OpenFold3 is gated and needs HF_TOKEN; refer to Authentication. AlphaFold2 is not downloadable at all — refer to AlphaFold2 parameters, then point the script at your converted files:

ALPHAFOLD2_DIR=/path/to/converted scripts/fetch_weights.sh --model alphafold2

Resolution Order

load_weights (hubs/checkpoint.py) tries local, then Hugging Face. Local order (hubs/local.py):

  1. Explicit cache_path (Python-only; the pipeline never sets it).
  2. The model’s <MODEL>_CKPT env var, if it names an existing file. A bad path is warned and ignored — resolution continues (a typo can silently fall through to a download).
  3. A staged file under <checkpoint root>/<model key>/: *.pt first, then *.ckpt; sorted-first match wins (keep one file per directory).

A miss returns None and Hub is tried. hub="local" / hub="hf" restrict to one side; the pipeline uses the default (both). A key absent from LOCAL_CHECKPOINTS raises KeyError up front.

Cache Layout

Under BIOIR_CACHE (default ~/.cache/bionemo_ir):

PathPurpose
checkpoints/<model key>/staged checkpoints (BIOIR_CHECKPOINTS relocates the root)
metadata/<override env name>staged metadata (BIOIR_METADATA relocates)
model_cache/raw downloads before staging (MODEL_CACHE_DIR relocates)
<model key>/Hub-downloaded metadata for that model (and extracted archives)

Hub checkpoint downloads go to ~/.cache/hf (hardcoded cache_dir in hubs/hf.py) — HF_HOME / HF_HUB_CACHE do not move them.

Staging and Overrides

export BIOIR_CHECKPOINTS=/shared/bionemo-ir/checkpoints
mkdir -p "$BIOIR_CHECKPOINTS/boltz-2"
cp boltz2_conf.ckpt "$BIOIR_CHECKPOINTS/boltz-2/"

Model keys are the FoldingSupportMatrix strings — refer to support matrix. Extension matters; file name does not.

export BIOIR_METADATA=/shared/bionemo-ir/metadata
ln -s /shared/assets/ccd.pkl "$BIOIR_METADATA/BOLTZ_CCD_PATH"
ln -s /shared/assets/mols "$BIOIR_METADATA/BOLTZ_MOL_DIR"

Each key has one <MODEL>_CKPT in LOCAL_CHECKPOINTS (hubs/local.py) — do not invent names from the key. Uppercase + - → _ + _CKPT works for OpenFold2/3, AlphaFold2, and Protenix, but not Boltz: boltz-1 → BOLTZ1_CKPT, boltz-2 → BOLTZ2_CKPT, boltz-2-affinity → BOLTZ2_AFFINITY_CKPT.

BOLTZ2_CKPT=/data/boltz2_conf.ckpt python examples/folding/run_demo.py

Metadata overrides are BOLTZ_CCD_PATH and BOLTZ_MOL_DIR. Unlike checkpoint vars they are not validated — a bad path fails later inside the model.

What the Loader Does per Family

Local and Hub paths use the same family handlers:

FamilyHandling
Boltz-1/2torch.load under safe_globals (omegaconf), then ["state_dict"]
OpenFold3torch.load under safe_globals, then ["ema"]["params"] when present
Protenixtorch.load, then ["model"], strip a leading module. prefix
OpenFold2 / AlphaFold2plain torch.load(weights_only=True)

After load, each family’s convert.py remaps upstream names — refer to architecture.

Using a Checkpoint You Modified

Point <MODEL>_CKPT at your file (or pass a state dict to load_weights(weights) with include_load_weights=False). Layout must still match upstream key names and shapes. Local files are never digest-checked — verification is Hub-download only.

Verifying Resolution

python -c "
from bionemo_ir.hubs import load_metadata, load_weights
print(len(load_weights('boltz-2')), 'entries resolved')
print(load_metadata('boltz-2'))
"

Logs which branch ran at INFO (BIOIR_LOG_LEVEL). For an offline pre-flight, pass local_files_only=True (not exposed by the pipeline).

Downloads

Interrupted. Partial Hub transfers stay as .incomplete and are not linked into the snapshot; re-run to finish. Refer to the Hugging Face cache guide.

Authentication. BioNeMo Inference Runtime (BioIR) does no token handling — that is all huggingface_hub. OpenFold3 is gated and an anonymous request gets 401. Registering is free: create a HuggingFace account, accept the terms on OpenFold/OpenFold3, then export HF_TOKEN (or HUGGING_FACE_HUB_TOKEN / hf auth login).

export HF_TOKEN=hf_...

fetch_weights.sh passes it to HuggingFace downloads and the library’s own Hub fallback picks it up too. Without it the OpenFold3 tests skip with GatedRepoError in the skip reason; nothing else in the suite needs it. Alternatively stage the file and set OPENFOLD3_CKPT so resolution never hits the Hub. The grant is per account — CI needs its own.

Metadata Assets

Boltz needs a CCD pickle and a molecule directory, resolved separately from the checkpoint through HF_MODEL_METADATA in hubs/metadata.py (Boltz keys only; others return {}):

KeyAssetRepoFileOverride
ccd_pathCCD pickleboltz-community/boltz-1ccd.pklBOLTZ_CCD_PATH
mol_dirmoleculesboltz-community/boltz-2mols.tarBOLTZ_MOL_DIR

Identical for boltz-1, boltz-2, and boltz-2-affinity (CCD from Boltz-1, mols.tar from Boltz-2 — including for Boltz-1). Downloads use the pinned commits and are digest-checked; mols.tar is then extracted, which is skipped when the target already exists. Order: override → staged entry → download.

AlphaFold2 Parameters

BioIR distributes no AF2 weights. Converting DeepMind’s published parameters is manual:

  1. Obtain params_model_<N>[_ptm].npz from the DeepMind AlphaFold repository, under the terms stated there.
  2. Install upstream openfold — the converter imports openfold.config, openfold.model.model, and openfold.utils.import_weights. It is not a BioIR dependency, so use a throwaway environment.
  3. Run examples/folding/openfold2/jax_to_pt.py with --jax_path, --config_preset, --output_dir. The preset ↔ parameter mapping for the monomer presets is the table in models/openfold2/config.py — template and pTM variants are not interchangeable, so read it before choosing. Multimer presets are not in that table; they follow the model_<N>_multimer_v3 form.
  4. Set ALPHAFOLD2_<N>_CKPT or stage the result under the model key — or hand the output directory to fetch_weights.sh through ALPHAFOLD2_DIR, which recognises the converter’s params_model_*.pt names.

End to end, for one preset:

curl -O https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar
tar -xf alphafold_params_2022-12-06.tar -C params/
# In a throwaway env with upstream `openfold` installed. Repeat for
# params_model_{1..5}.npz and params_model_{1..5}_multimer_v3.npz, passing the
# matching --config_preset (model_N / model_N_multimer_v3):
python examples/folding/openfold2/jax_to_pt.py \
--jax_path params/params_model_1.npz --config_preset model_1 \
--output_dir converted/
ALPHAFOLD2_DIR=converted scripts/fetch_weights.sh --model alphafold2

That tarball is pinned to 2022-12-06 because that release is CC BY 4.0; the earlier ones were non-commercial. Review the licence before using these weights for anything.

Do not rename the .npz first. The converter derives the upstream weight version from the file name (params_model_1_ptm.npz → model_1_ptm) and passes it to import_jax_weights_; the output is named after the same stem (params_model_1_ptm.pt). A renamed input silently converts against the wrong version.

Multimer presets convert the same way: --config_preset model_<N>_multimer_v3 against params_model_<N>_multimer_v3.npz.