Configuration#

This page explains how to configure NIM for BGR at container startup, apply per-request overrides for geometry relaxation parameters, and manage model and GPU settings.

Startup Configuration#

Configure the model, periodic boundary conditions, batch sizing, optimization defaults, and DFT-D3(BJ) correction through environment variables when you launch the container. Refer to Environment Variables for the complete list, accepted values, and defaults.

For example, set the model architecture when launching the container:

docker run ... -e ALCHEMI_NIM_MODEL_TYPE="mace" ...

Periodic Boundary Conditions#

ALCHEMI_NIM_PBC selects the backend PBC mode when the container starts. In the default true mode, every structure in a request must include cell. In false mode, structures must omit cell. A container cannot accept a mixture of these two request forms; use separate containers when serving both periodic and non-periodic workloads.

The nested atoms[].pbc field has a different purpose. It selects the periodic axes for that structure’s supplied cell, in [x, y, z] order, and defaults to [true, true, true]. It does not change the backend PBC mode or relax the backend’s requirement that cell be present or absent.

DFT-D3 Parameters#

When ALCHEMI_NIM_DFTD3_ENABLED is true, you can provide custom damping parameters using the ALCHEMI_NIM_DFTD3_PARAM environment variable. The value must be a JSON string (for example, '{"s8": 0.3908, "a1": 0.5660, "a2": 3.1280}').

DFT-D3(BJ) can be combined with MACE, AIMNet2, or TensorNet. When ALCHEMI_NIM_MODEL_PATH is not set, ALCHEMI_NIM_MODEL_TYPE selects the following functional preset:

Default DFT-D3(BJ) Parameters by Model Type

Model type

Functional

a1

a2 (Bohr)

s6

s8

mace

PBE

0.4289

4.4407

1.0

0.7875

aimnet2

wB97M

0.5660

3.1280

1.0

0.3908

tensornet

PBE

0.4289

4.4407

1.0

0.7875

By default, DFT-D3(BJ) uses a 15 Å cutoff and smoothly switches the correction off from 12 Å to 15 Å.

When ALCHEMI_NIM_MODEL_PATH selects a user-provided checkpoint, the DFT-D3 configuration must specify a1, a2, and s8.

The JSON object supplied through ALCHEMI_NIM_DFTD3_PARAM accepts the following fields:

ALCHEMI_NIM_DFTD3_PARAM JSON Object Fields

Field

Type

Default

Description

a1

float

null

BJ damping parameter a1

a2

float

null

BJ damping parameter a2 (Bohr)

s6

float

null

Scaling factor for the C6 term

s8

float

null

Scaling factor for the C8 term

cutoff

float

15.0

Cutoff distance (Å)

smoothing_fraction

float

0.2

Fraction of the cutoff used for the smoothing region

For detailed information on DFT-D3(BJ) dispersion corrections, refer to the ALCHEMI Toolkit-Ops documentation.

Geometry Relaxation Configuration#

The startup configuration establishes the default force tolerance, cell-optimization setting, and pressure tolerance. An inference request can override these defaults with opttol, cellopt, and opttol_pressure. The maximum number of optimization steps and the optimizer preset apply to the backend as a whole.

Refer to API Reference for the request schema and to Tuning Parameters for an override example.

Model Configuration#

ALCHEMI_NIM_MODEL_TYPE selects the model architecture: MACE, AIMNet2, or TensorNet. ALCHEMI_NIM_MODEL_PATH points to the checkpoint that the selected architecture loads.

To load a checkpoint from the host, mount its model file or directory into the container and set ALCHEMI_NIM_MODEL_PATH to the mounted path. MACE and AIMNet2 use a single model file. TensorNet uses a directory containing model.json, model.pt, and state.pt.

For example, mount a TensorNet directory as follows:

docker run --rm -ti --name alchemi-bgr --gpus=all \
    -e NGC_API_KEY \
    -p 8000:8000 --shm-size=8g \
    -v /path/to/model-dir:/opt/nim/.cache/model-dir:ro \
    -e ALCHEMI_NIM_DISABLE_MODEL_DOWNLOAD=true \
    -e ALCHEMI_NIM_MODEL_TYPE="tensornet" \
    -e ALCHEMI_NIM_MODEL_PATH="/opt/nim/.cache/model-dir" \
    nvcr.io/nim/nvidia/alchemi-bgr:1.1.0

For model sources and checkpoint-specific requirements, refer to Custom Models.

GPU and Memory Configuration#

Configure memory allocation and GPU utilization when you launch the container:

  • Shared memory: Use --shm-size=8g to support large requests.

  • Batch size: The NIM estimates the optimal and maximum batch size for available GPU memory at startup. Set ALCHEMI_NIM_BATCH_SIZE to use a fixed batch size.

  • Multi-GPU: The NIM automatically uses all available GPUs. Each GPU runs an independent worker.