UMA
Model Introduction
UMA (Universal Materials Interaction Model) is a general-purpose machine-learning interatomic potential for materials and catalytic systems. Built on an equivariant graph neural network, it predicts the energies and forces of atomic structures.
Model Description
UMA uses an equivariant graph neural network architecture trained on multiple materials and catalysis datasets, including OC20, OC22, OC25, OMat, OMOL, ODAC, and OMC. It supports energy and force prediction and structure optimization for catalytic adsorption, inorganic materials, molecular systems, and MOFs.
Use Cases
| Use case | Description |
|---|---|
| OC20 energy/force fine-tuning | Read OC20 fine-tuning data with the standard configuration and train a UMA model |
| OC22 oxide-catalysis fine-tuning | Fine-tune on OC22 with the standard configuration (1P2 only) |
| OC25 (electro)catalysis fine-tuning | Fine-tune on OC25 with the standard configuration (1P2 only) |
| OMat inorganic-material fine-tuning | Fine-tune on OMat with the standard configuration |
| OMOL molecule/polymer fine-tuning | Fine-tune on OMOL with the standard configuration |
| ODAC MOF fine-tuning | Fine-tune on ODAC with the standard configuration |
| OMC molecular-crystal fine-tuning | Fine-tune on OMC with the standard configuration |
| Training preflight | Check configurations, data paths, run scripts, and checkpoint locations |
| Catalytic adsorption modeling | Use OC20/OC22/OC25 tasks as references for training and inference on adsorption and catalytic-surface systems |
| Inference-script reference | Adapt upstream examples for crystal relaxation, adsorption-system relaxation, or molecular MD |
| Custom data migration | Convert ASE-readable structures into UMA fine-tuning data, then replace the training and validation paths |
Usage
1. Using OneCode
Try intelligent, one-click AI4S programming in the OneCode online environment:
Try intelligent, one-click AI4S programming
2. Manual Installation and Usage
Hardware requirements
- A GPU or DCU is recommended for fine-tuning.
- A CPU can be used to check configurations and data paths, but is not recommended for production UMA training.
- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
Download the Model Package
hf download --model OneScience-Sugon/UMA --local-dir ./UMA
cd UMA
Install the Runtime Environment
DCU environment
# Activate DTK and conda first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# uv installation is also supported
pip install onescience[matchem-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
GPU environment
# Activate conda first
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
# uv installation is also supported
pip install onescience[matchem-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
Training Datasets
Training data is not bundled with this repository. The following example uses OC20 fine-tuning. OC22, OC25, OMat, OMOL, ODAC, and OMC follow the same workflow; only --uma-task and the data paths need to be changed.
hf download --dataset OneScience-Sugon/oc20 --local-dir ./data
Data conversion
The downloaded raw data usually consists of .extxyz files. Convert them to ASE-LMDB format with scripts/create_uma_finetune_dataset.py, which also computes elem_refs and normalizer_rmsd. The script supports these tasks:
| Task | Description |
|---|---|
oc20 |
Catalysis (example) |
oc22 |
Oxide catalysis (1P2 only) |
oc25 |
(Electro)catalysis (1P2 only) |
omat |
Inorganic materials |
omol |
Molecules and polymers |
odac |
MOFs |
omc |
Molecular crystals |
python scripts/create_uma_finetune_dataset.py \
--train-dir data/oc20/s2ef_200k_uncompressed \
--val-dir data/oc20/s2ef_val_id_uncompressed \
--uma-task oc20 \
--regression-tasks ef \
--output-dir data/oc20_finetune \
--num-workers 8
The conversion produces:
data/oc20_finetune/
├── train/ # ASE-LMDB training data
├── val/ # ASE-LMDB validation data
└── data/ # Generated data configuration YAML
└── uma_conserving_data_task_energy_force.yaml
Next, use scripts/update_demo_config.py to copy the generated elem_refs, normalizer_rmsd, and data paths into the demo configuration:
python scripts/update_demo_config.py --demo-config demo/configs/oc20_ef_4dcu.yaml
demo/run.sh automatically sets the repository root as ONESCIENCE_DATASETS_DIR, so relative paths in the configuration resolve automatically.
Trained Weights
This repository includes the rotation-basis file weight/Jd.pt. Download a pretrained UMA checkpoint, such as uma-s-1p1_converted.pt, from the official fairchem repository, convert it to UMA format, and place it at the following location (upload pending):
weight/uma-s-1p1_converted.pt
- Official fairchem repository: https://github.com/facebookresearch/fairchem
demo/run.sh and the example scripts under inference/ automatically detect weight/Jd.pt and set ONESCIENCE_UMA_JD_PATH.
Fine-Tuning
bash demo/run.sh --config demo/configs/oc20_ef_4dcu.yaml
Inference
python inference/run_molecular_md.py
Specify the path to the pretrained weights before running inference.
Official OneScience Resources
| Platform | OneScience Main Repository | Skills Repository |
|---|---|---|
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
Citation and License
- The UMA example code comes from the MatChem examples in the OneScience project and refers to the upstream fairchem project (https://github.com/facebookresearch/fairchem). The upstream fairchem software is released under the MIT License. Individual fairchem model checkpoints and datasets may have separate licenses; follow the applicable terms.
- If you use UMA fine-tuning or inference results in research, please cite UMA or the relevant universal materials interaction model method, the fairchem and OneScience projects, and the datasets used.