Pretrained GRACE Foundation Models#

Several GRACE models are pre-trained on large datasets.

NOTE! You should use the "Full Name" column to refer to the model in LAMMPS and ASE.

The UQ and Kokkos columns indicate, per model:

UQ — the distributed saved model ships with an uncertainty-quantification head (and the checkpoint includes gmm_artifacts.npz).
Kokkos — a kokkos.npz LAMMPS-Kokkos export is available. It is now distributed as a separate file (no longer bundled inside the model archive); fetch it with grace_models download <model_name> --kokkos.

Precision (fp32 / fp64). The current release defaults to fp32: the bare model name (e.g. GRACE-2L-OMAT-large-ft-E) resolves to the fp32 saved model and checkpoint — roughly 2× faster and half the memory, with a negligible accuracy gap (UQ gamma max-rel ≤ 5×10⁻⁵). A matching full-precision build is published non-default under the -fp64 suffix (e.g. GRACE-2L-OMAT-large-ft-E-fp64) for tight finite-difference or bit-reproducible work. Exceptions: -mx models are natively mixed precision, and the 3L models are natively fp32 (no -fp64 variant). The Kokkos export is precision-agnostic — one file serves both.

SMAX models#

Reference: arXiv

The SMAX (Maximum Entropy) models are trained on a chemistry-agnostic dataset generated via a multicomponent maximum information entropy structure generation protocol.

Unlike traditional datasets that focus on low-energy equilibrium structures, SMAX is constructed to deliberately sample broad and diverse regions of configurational space. This provides a robust physical prior for atomic interactions across the entire periodic table, enabling accurate modeling of large-strain phase transformations, defects in complex alloys, and reaction barriers in catalytic systems.

Custom Cutoffs & Interaction Ranges: SMAX models utilize a custom element-dependent cutoff radius ranging from 5.0 Å to 7.5 Å.

Recommendation: For most general-purpose applications, we recommend using the SMAX-OMAT models. They offer the best balance of structural robustness (from SMAX) and high-precision energy/force accuracy (from OMat24).

Two-layer (2L), semilocal models#

Model Name	Full Name	Size	κSRME	UQ	Kokkos	Description
GRACE-2L-SMAX-M	GRACE-2L-SMAX-medium	medium	0.469	✓	✓	Two-layer semi-local (SMAX)
GRACE-2L-SMAX-L	GRACE-2L-SMAX-large	large	0.444	✓	✓	Two-layer semi-local (SMAX)
GRACE-2L-SMAX-OMAT-M	GRACE-2L-SMAX-OMAT-medium	medium	0.197	✓	✓	Two-layer semi-local (SMAX + OMat24)
GRACE-2L-SMAX-OMAT-L	GRACE-2L-SMAX-OMAT-large	large	0.191	✓	✓	Two-layer semi-local (SMAX + OMat24)

Single-layer (1L), local models#

Model Name	Full Name	Size	κSRME	UQ	Kokkos	Description
GRACE-1L-SMAX-L	GRACE-1L-SMAX-large	large	0.696	✓	✓	Single-layer local (SMAX)
GRACE-1L-SMAX-OMAT-L	GRACE-1L-SMAX-OMAT-large	large	0.338	✓	✓	Single-layer local (SMAX + OMat24)

OMAT models#

Reference: npj Comp. Mat. (open access)

The base models (-OMAT) are trained on the OMat24 dataset. The fine-tuned versions (-OMAT-ft-E) are derived from these base models by fine-tuning with more emphasis on energies. All models listed use a fixed 6 Å cutoff.

Three-layer (3L) models#

Reference: npj Comp. Mat. (open access)

The three-layer (3L) GRACE models add a third message-passing layer for a larger effective receptive field. The 3L models use a fixed 6 Å cutoff and ship with UQ. They are natively fp32 (there is no -fp64 variant). For LAMMPS-Kokkos use pair_style grace/3l/kk (mixed precision) or pair_style grace/3l/kk/fp32.

Model Name	Full Name	Size	$\kappa_\mathrm{SRME}$	UQ	Kokkos	Description
GRACE-3L-OMAT-large	GRACE-3L-OMAT-large	large	0.141	✓	✓	Three-layer, trained on OMat24

Two-layer (2L), semilocal models#

Model Name	Full Name	Size	$\kappa_\mathrm{SRME}$	UQ	Kokkos	Description
GRACE-2L-OMAT	GRACE-2L-OMAT	small	0.288	—	—	Two-layer semi-local
GRACE-2L-OMAT-M-base	GRACE-2L-OMAT-medium-base	medium	0.212	✓	✓	Two-layer semi-local (base)
GRACE-2L-OMAT-M	GRACE-2L-OMAT-medium-ft-E	medium	0.217	✓	✓	Two-layer semi-local (finetuned on energy)
GRACE-2L-OMAT-L-base	GRACE-2L-OMAT-large-base	large	0.165	✓	✓	Two-layer semi-local (base)
GRACE-2L-OMAT-L	GRACE-2L-OMAT-large-ft-E	large	0.186	✓	✓	Two-layer semi-local (finetuned on energy)
GRACE-2L-OMAT-L-mx	GRACE-2L-OMAT-large-mx	large	-	✓	✓	Two-layer semi-local (large, mixed precision)

Single-layer (1L), local models#

Model Name	Full Name	Size	$\kappa_\mathrm{SRME}$	UQ	Kokkos	Description
GRACE-1L-OMAT	GRACE-1L-OMAT	small	0.398	—	—	Single-layer local
GRACE-1L-OMAT-M-base	GRACE-1L-OMAT-medium-base	medium	0.380	✓	✓	Single-layer local (base)
GRACE-1L-OMAT-M	GRACE-1L-OMAT-medium-ft-E	medium	0.417	✓	✓	Single-layer local (finetuned on energy)
GRACE-1L-OMAT-L-base	GRACE-1L-OMAT-large-base	large	0.354	✓	✓	Single-layer local (base)
GRACE-1L-OMAT-L	GRACE-1L-OMAT-large-ft-E	large	0.383	✓	✓	Single-layer local (finetuned on energy)

OAM models#

Reference: npj Comp. Mat. (open access)

These models are first pre-trained on OMat24 and then fine-tuned on a combination of the sAlex dataset (10.4M structures) and the MPtraj dataset (1.58M structures).

Three-layer (3L) models#

The three-layer (3L) GRACE models add a third message-passing layer for a larger effective receptive field. The 3L models use a fixed 6 Å cutoff and ship with UQ. They are natively fp32 (there is no -fp64 variant). For LAMMPS-Kokkos use pair_style grace/3l/kk (mixed precision) or pair_style grace/3l/kk/fp32.

Model Name	Full Name	Size	F1	$\kappa_\mathrm{SRME}$	UQ	Kokkos	Description
GRACE-3L-OMAT-large-ft-AM	GRACE-3L-OMAT-large-ft-AM	large	0.925	0.121	✓	✓	Three-layer, fine-tuned on sAlex + MPtraj

Two-layer (2L), semilocal models#

Model Name	Full Name	Size	F1	$\kappa_\mathrm{SRME}$	UQ	Kokkos	Description
GRACE-2L-OAM	GRACE-2L-OAM	small	0.880	0.294	—	—	Two-layer semi-local
GRACE-2L-OAM-M	GRACE-2L-OMAT-medium-ft-AM	medium	0.881	0.200	✓	✓	Two-layer semi-local
GRACE-2L-OAM-L	GRACE-2L-OMAT-large-ft-AM	large	0.889	0.168	✓	✓	Two-layer semi-local
GRACE-2L-OAM-L-mx	GRACE-2L-OMAT-large-mx-ft-AM	large	-	-	✓	✓	Two-layer semi-local (large, mixed precision)

Single-layer (1L), local models#

Model Name	Full Name	Size	F1	$\kappa_\mathrm{SRME}$	UQ	Kokkos	Description
GRACE-1L-OAM	GRACE-1L-OAM	small	0.824	0.516	—	—	Single-layer local
GRACE-1L-OAM-M	GRACE-1L-OMAT-medium-ft-AM	medium	0.800	0.411	✓	✓	Single-layer local
GRACE-1L-OAM-L	GRACE-1L-OMAT-large-ft-AM	large	0.815	0.377	✓	✓	Single-layer local

Downloading Foundation Models#

You can use the grace_models utility to download and inspect pre-trained GRACE models:

Run grace_models list to view the list of available models (a compact capability table). Add -v/--verbose for a full per-model block (description, paths, license).
Run grace_models info <model_name> to show a model's declared capability flags (precision, UQ, parallel) and, if it is already downloaded, introspect the artifact.
Run grace_models download <model_name> to download a specific model. Add --kokkos to also fetch its LAMMPS-Kokkos export (kokkos.npz) into the model directory.
Run grace_models checkpoint <model_name> to download only the checkpoint (used for fine-tuning).

These models can be used for simulations within ASE and LAMMPS.

By default, all foundational models are stored into $HOME/.cache/grace, but you can overwrite it with GRACE_CACHE environment variable.

Fine-tuning foundation models#

Automatic input file generation#

You can generate an input.yaml file for fine-tuning a foundation model by running gracemaker -t.

Manual setup#

Fine-tuning foundation GRACE models can only be performed using checkpoints, not saved models. Run grace_models list -v to view the available models that include a CHECKPOINT: field (or grace_models info <MODEL-NAME> for a single model). You can either download this checkpoint manually using grace_models checkpoint <MODEL-NAME>, or it will be downloaded automatically when needed.

To fine-tune a foundation model, add the following to the potential section of your input.yaml:

potential:
  finetune_foundation_model: GRACE-1L-OAM
  shift: auto  # automatically align FM energies with your DFT reference data
# reduce_elements: True # if True - select from original models only elements present in the CURRENT dataset

Automatic energy shift correction (`shift: auto`)#

When shift: auto is set alongside finetune_foundation_model, GRACEmaker automatically computes optimal per-element energy shifts that minimize the difference between the foundation model's predictions and your reference DFT data. These shifts are injected directly into the model (via ConstantScaleShiftTarget) before training begins.

How it works:

Selects ~50 diverse structures from the training set covering all element combinations.
Predicts their energies with the foundation model (prior to training).
Solves a least-squares problem: A @ x = (E_DFT - E_FM), where A is the composition matrix and x represents per-element shifts.
Injects the computed shifts into the model's ConstantScaleShiftTarget instruction.

This is particularly useful when your DFT reference data uses different pseudopotentials or settings than the foundation model's training data, which usually results in a systematic energy offset. The correction minimizes the initial energy mismatch and leads to faster convergence.

Note: shift: auto is independent of data::reference_energy. Here is a typical usage example:

reference_energy: 0 (or dict) — set your DFT reference frame
shift: auto — align the FM output to your DFT reference frame

Learning rate reduction#

Additionally, it is recommended to set a small learning_rate (e.g., 1e-3 or 1e-4) and evaluate the initial metrics by setting eval_init_stats: True:

fit:

  # set a small learning rate
  opt_params: {learning_rate: 0.001,  ... }

  # evaluate initial metrics
  eval_init_stats: True

Frozen-weights fine-tuning#

By default, a naive fine-tuning approach is used, where all model parameters are updated. In some cases, this can lead to catastrophic forgetting of pre-trained knowledge. To mitigate this issue, we recommend using a frozen-weights fine-tuning approach.

For frozen-weights fine-tuning, you must specify the trainable_variable_names parameter in the fit section of your input.yaml.

Here are a few name pattern examples (for GRACE-2L models):

LinMLPOut2ScalarTarget_: Typical name for a linear+MLP energy readout function.
I2/reducing_: ACE expansion coefficients in the second layer.
rho/reducing_: ACE expansion coefficients in the first layer.
I1/reducing_: ACE expansion coefficients for the message being sent from the first layer to the second layer.

Here are the recommended trainable variables to preserve the pre-trained knowledge of unaffected chemical elements:

For GRACE-2L models:

fit:
  trainable_variable_names: ["I2/reducing_" ,"rho/reducing_","I1/reducing_"]

For GRACE-1L models:

fit:
  trainable_variable_names: ["rho/reducing_"]

You can find more variable names by running a model summary:

grace_utils -p ~/.cache/grace/checkpoints/SOME_FOUNDATIONAL_MODEL_NAME/model.yaml  summary -v 1

User-defined models#

You can add user-defined models to the local registry of foundation models. You need to create a $HOME/.cache/grace/models_registry.yaml file with a content:

FULL-MODEL-NAME:
  path: /path/to/saved_model
  checkpoint_path: /path/to/checkpoint/  # model.yaml and checkpoint.index should be in this folder
  description: some description
  license: some license

Pretrained GRACE Foundation Models#

SMAX models#

Two-layer (2L), semilocal models#

Single-layer (1L), local models#

OMAT models#

Three-layer (3L) models#

Two-layer (2L), semilocal models#

Single-layer (1L), local models#

OAM models#

Three-layer (3L) models#

Two-layer (2L), semilocal models#

Single-layer (1L), local models#

Downloading Foundation Models#

Fine-tuning foundation models#

Automatic input file generation#

Manual setup#

Automatic energy shift correction (shift: auto)#

Learning rate reduction#

Frozen-weights fine-tuning#

User-defined models#

Automatic energy shift correction (`shift: auto`)#