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:
If
MODEL_PATHis set, files are loaded directly from that directory.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 |
|
|
Affinity weights |
|
|
CCD components |
|
(not renamable) |
Molecule cache |
|
(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.pklandmols.tarassets 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.