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.

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:

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.