Inference using ASE and Predictor Interface

Inference is done using MLIPPredictUnit. The FairchemCalculator (an ASE calculator) is simply a convenience wrapper around the MLIPPredictUnit.

from __future__ import annotations

from fairchem.core import FAIRChemCalculator, pretrained_mlip

predictor = pretrained_mlip.get_predict_unit("uma-s-1p2p1", device="cuda")
calc = FAIRChemCalculator(predictor, task_name="oc20")

Warp 1.18.0 initialized:
   CUDA Toolkit 13.4, Driver 13.1
   Devices:
     "cpu"      : "x86_64"
     "cuda:0"   : "Tesla T4" (16 GiB, sm_75, mempool enabled)
   Kernel cache:
     /home/runner/.cache/warp/1.18.0

Adding a DFT-D3(BJ) dispersion correction¶

from fairchem.core import DFTD3Calculator, FAIRChemCalculator, pretrained_mlip

predictor = pretrained_mlip.get_predict_unit(
    "uma-s-1p2p1", device="cuda", inference_settings="turbo"
)
base_calc = FAIRChemCalculator(predictor, task_name="omat")

# Use "pbe" for PBE-trained models such as UMA's OMat head, or "r2scan"
# for models fine-tuned on r2SCAN data.
calc = DFTD3Calculator(base_calc, functional="pbe", device="cuda")
atoms.calc = calc

The wrapper adds the D3 energy, forces, and analytic stress to the base calculator. The pbe and r2scan presets both use Becke-Johnson damping, a 15 Å cutoff, and C5 smoothing over the outer 20% of the cutoff. Pass param_file and auto_download=False to use a local D3 parameter table without network access.

Both presets use k1=16.0 and k3=-4.0.

The PBE-D3(BJ) parameters are from:

S. Grimme, S. Ehrlich, and L. Goerigk, “Effect of the damping function in dispersion corrected density functional theory,” J. Comput. Chem. 32, 1456–1465 (2011). Grimme & others (2011)

The r2SCAN-D3(BJ) parameters were reported and benchmarked alongside D4 in:

S. Ehlert, U. Huniar, J. Ning, J. W. Furness, J. Sun, A. D. Kaplan, J. P. Perdew, and J. G. Brandenburg, “r²SCAN-D4: Dispersion corrected meta-generalized gradient approximation for general chemical applications,” J. Chem. Phys. 154, 061101 (2021). Extension and Evaluation of the D4 London-Dispersion Model for Periodic Systems (2021)

! pip install fairchem-core fairchem-data-oc fairchem-applications-cattsunami

# Login using the huggingface-cli utility
! huggingface-cli login

# alternatively,
import os
os.environ['HF_TOKEN'] = 'MY_TOKEN'

Default mode¶

UMA defaults to the merge_mole + compile fast mode with TF32 disabled. This fast path requires fixed composition, task, charge, and spin across repeated evaluations. If a later evaluation changes any of these, the calculator prints a warning and permanently falls back to the unmerged, uncompiled model. Batching is supported; a mixed batch across any of the same parameters triggers the same fallback.

Batch mode¶

Use batch mode for heterogeneous batches whose systems differ in composition, task, charge, or spin. It currently keeps MOLE unmerged and leaves compilation disabled. The named mode provides a stable entry point for future batch-specific optimizations, such as compilation without MOLE merging.

predictor = pretrained_mlip.get_predict_unit(
    "uma-s-1p2p1", device="cuda", inference_settings="batch"
)

Turbo mode¶

Turbo mode uses the same merge_mole + compile fast path as default mode and additionally enables TF32. TF32 can improve performance on compatible hardware at a small precision trade-off. Similar to default mode, any changes in composition, task, charge, and spin across different evaluations trigger a fallback to the unoptimized execution path.

predictor = pretrained_mlip.get_predict_unit(
    "uma-s-1p2p1", device="cuda", inference_settings="turbo"
)

Custom modes for advanced users¶

The advanced user might quickly see that default, batch, and turbo modes are special cases of our inference settings api. You can customize it for your application if you understand what you are doing. The following table provides more information.

For example, for an MD simulation use-case for a system of ~500 atoms, we can choose to use a custom mode like the following:

from fairchem.core.units.mlip_unit.api.inference import InferenceSettings

settings = InferenceSettings(
    tf32=True,
    activation_checkpointing=False,
    merge_mole=True,
    compile=True,
    external_graph_gen=False,
    internal_graph_gen_version=2,
)

predictor = pretrained_mlip.get_predict_unit(
    "uma-s-1p2p1", device="cuda", inference_settings=settings
)

Enabling gradient stress or Hessian prediction¶

Some tasks, for example omol, odac, or oc20/25, were not trained using stress labels. Similarly, no tasks were supervised to predict Hessians. However, predictions of untrained derivatives of energy, such as stress and Hessians, can be enabled by using the following inference settings flags,

For example, to enable stress and Hessian predictions with omol level of theory, the following settings can be used,

settings = InferenceSettings(
    predict_untrained_stress={'omol'},
    predict_untrained_hessian={'omol'}
)

predictor = pretrained_mlip.get_predict_unit(
    "uma-s-1p2p1", device="cuda", inference_settings=settings
)

Multi-GPU Inference¶

UMA supports Graph Parallel inference natively. The graph is chunked into each rank and both the forward and backwards communication is handled by the built-in graph parallel algorithm with torch distributed. Because Multi-GPU inference requires special setup of communication protocols within a node and across nodes, we leverage ray to launch Ray Actors for each GPU-rank under the hood. This allows us to seamlessly scale to any infrastructure that can run Ray.

To make things simple for the user that wants to run multi-gpu inference locally, we provide a drop-in replacement for MLIPPredictUnit, called ParallelMLIPPredictUnit

For example, we can create a predictor with 8 GPU workers in a very similar way to MLIPPredictUnit and perform an MD calculation with the ASE calculator. This mode of operation is also compatible with our LAMMPS integration.

from ase import units
from ase.md.langevin import Langevin
from fairchem.core import pretrained_mlip, FAIRChemCalculator
import time

from fairchem.core.datasets.common_structures import get_fcc_crystal_by_num_atoms

predictor = pretrained_mlip.get_predict_unit(
    "uma-s-1p2p1", inference_settings="turbo", device="cuda", workers=1
)
calc = FAIRChemCalculator(predictor, task_name="omat")

atoms = get_fcc_crystal_by_num_atoms(8000)
atoms.calc = calc

dyn = Langevin(
    atoms,
    timestep=0.1 * units.fs,
    temperature_K=400,
    friction=0.001 / units.fs,
)
# warmup 10 steps
dyn.run(steps=10)
start_time = time.time()
dyn.attach(
    lambda: print(
        f"Step: {dyn.get_number_of_steps()}, E: {atoms.get_potential_energy():.3f} eV, "
        f"QPS: {dyn.get_number_of_steps()/(time.time()-start_time):.2f}"
    ),
    interval=1,
)
dyn.run(steps=1000)