Loading Custom / Fine-Tuned Boltz-2 Checkpoints#

The Boltz-2 NIM supports same-architecture checkpoint swapping: you can load your own Boltz-2 weights (for example, weights fine-tuned externally on synthetic data) instead of the default checkpoints downloaded from NGC. This enables iterative fine-tune → re-serve → screen workflows without changing the model architecture.

Note

Fine-tuning itself is performed outside the NIM. Supply .ckpt files that are weight-compatible with the stock Boltz-2 structure and affinity models.

How It Works#

At startup the NIM resolves each model file through a single model root:

  1. If MODEL_PATH is set, files are loaded directly from that directory.

  2. Otherwise files are loaded from the NGC cache under NIM_CACHE_PATH.

The following filenames are expected in the model root:

Purpose

Default filename

Override env var

Structure weights

boltz2_conf.ckpt

NIM_BOLTZ_CONF_CKPT_FILE

Affinity weights

boltz2_aff.ckpt

NIM_BOLTZ_AFFINITY_CKPT_FILE

CCD components

ccd.pkl

(not renamable)

Molecule cache

mols.tar

(not renamable)

If your checkpoint files use different names, set the corresponding override env var to the filename inside the mounted directory (not a full path).

Quickstart#

Mount your checkpoint directory, disable the default download, and point MODEL_PATH at the mount:

# Directory on the host holding your Boltz-2 assets
CUSTOM_MODEL_DIR=/data/boltz2-finetuned

docker run -it \
    --runtime=nvidia \
    --gpus '"device=0"' \
    -p 8000:8000 \
    --shm-size=16G \
    -v "$CUSTOM_MODEL_DIR":/opt/nim/custom-model:ro \
    -e MODEL_PATH=/opt/nim/custom-model \
    -e NIM_DISABLE_MODEL_DOWNLOAD=1 \
    nvcr.io/nim/mit/boltz2:1.9.0

Expected layout of $CUSTOM_MODEL_DIR (stock filenames):

/data/boltz2-finetuned/
├── boltz2_conf.ckpt
├── boltz2_aff.ckpt
├── ccd.pkl
└── mols.tar

Renaming Checkpoints#

If you keep multiple fine-tuned checkpoints in the same directory and want to select one by name, override the filename(s):

docker run -it \
    --runtime=nvidia \
    --gpus '"device=0"' \
    -p 8000:8000 \
    --shm-size=16G \
    -v "$CUSTOM_MODEL_DIR":/opt/nim/custom-model:ro \
    -e MODEL_PATH=/opt/nim/custom-model \
    -e NIM_DISABLE_MODEL_DOWNLOAD=1 \
    -e NIM_BOLTZ_CONF_CKPT_FILE=boltz2_conf_run42.ckpt \
    -e NIM_BOLTZ_AFFINITY_CKPT_FILE=boltz2_aff_run42.ckpt \
    nvcr.io/nim/mit/boltz2:1.9.0

With the above, the NIM loads boltz2_conf_run42.ckpt / boltz2_aff_run42.ckpt from MODEL_PATH; ccd.pkl and mols.tar are still read from the same directory with their default names.

Notes and limitations#

  • Same architecture only. Weights must match the stock Boltz-2 structure and affinity module shapes. Architecture changes are not supported.

  • The ccd.pkl and mols.tar assets must be present in the model root. If you are fine-tuning only the weights, copy these from the stock NGC download.

  • An empty override value is treated as unset and falls back to the default filename.