Skip to navigation

Custom Models

View as Markdown

The default quickstart uses the bundled MACE-MPA-0 model. For the default setup, follow the Quickstart Guide.

Use Custom or Downloaded Models

The bundled MACE model (auto-downloaded with an NGC API key) can be used. You can also provide MACE, AIMNet2, or TensorNet models by downloading them and mounting the model file or directory into the container. Choose the model type from the following tabs for download hints and the exact arguments to pass when launching the container.

MACE models are for periodic systems only (pbc: true). The default MACE-MPA-0 model is bundled and auto-downloaded when NGC_API_KEY is set. To use a different MACE model (for example, a custom-trained model), do the following:

Downloading a MACE model

Refer to the MACE foundation model GitHub releases for download links. For example:

curl -LO https://github.com/ACEsuit/mace-foundations/releases/download/mace_mpa_0/mace-mpa-0-medium.model
export ALCHEMI_MODEL_PATH=mace-mpa-0-medium.model

Launching the container with MACE

Run the following commands to launch the container:

export ALCHEMI_MODEL_PATH=/path/to/mace.model
export NGC_API_KEY=<Your NGC API Key>
docker run --rm -ti --name alchemi-bmd --gpus=all \
    -e NGC_API_KEY \
    -p 8000:8000 --shm-size=8g \
    -v "${ALCHEMI_MODEL_PATH}:/opt/nim/.cache/mace.model:ro" \
    -e NIM_DISABLE_MODEL_DOWNLOAD=true \
    -e ALCHEMI_NIM_MODEL_TYPE="mace" \
    -e ALCHEMI_NIM_MODEL_PATH=/opt/nim/.cache/mace.model \
    nvcr.io/nim/nvidia/alchemi-bmd:1.0.0

Example Configurations

This section is a quick reference for common configurations. For download steps and model-type-specific launch details, refer to Use Custom or Downloaded Models.

MACE for Periodic Systems (Default)

Run the following commands:

export NGC_API_KEY=<Your NGC API Key>
docker run --rm -ti --name alchemi-bmd --gpus=all \
    -e NGC_API_KEY \
    -p 8000:8000 \
    --shm-size=8g \
    nvcr.io/nim/nvidia/alchemi-bmd:1.0.0

AIMNet2 for Organic Molecular Systems

Run the following commands:

export ALCHEMI_MODEL_PATH=/path/to/aimnet-model.pt
docker run --rm -ti --name alchemi-bmd --gpus=all \
    -e NGC_API_KEY \
    -p 8000:8000 --shm-size=8g \
    -v "${ALCHEMI_MODEL_PATH}:/opt/nim/.cache/aimnet2.pt:ro" \
    -e NIM_DISABLE_MODEL_DOWNLOAD=true \
    -e ALCHEMI_NIM_MODEL_TYPE="aimnet2" \
    -e ALCHEMI_NIM_MODEL_PATH="/opt/nim/.cache/aimnet2.pt" \
    -e ALCHEMI_NIM_PBC=false \
    nvcr.io/nim/nvidia/alchemi-bmd:1.0.0

TensorNet for Periodic Materials

Run the following commands:

export ALCHEMI_MODEL_PATH=/path/to/model-dir
docker run --rm -ti --name alchemi-bmd --gpus=all \
    -p 8000:8000 --shm-size=8g \
    -v "${ALCHEMI_MODEL_PATH}:/opt/nim/.cache/TensorNet:ro" \
    -e NIM_DISABLE_MODEL_DOWNLOAD=true \
    -e ALCHEMI_NIM_MODEL_TYPE="tensornet" \
    -e ALCHEMI_NIM_MODEL_PATH="/opt/nim/.cache/TensorNet" \
    nvcr.io/nim/nvidia/alchemi-bmd:1.0.0