UMA Backbone
“UMA (Universal Models for Atoms) is a next-generation atomistic foundation model developed by Meta FAIRChem team. The UMA models are based on the eSEN architecture and trained on 5 billion unique 3D atomic structures — the largest training run to date — by compiling data from multiple chemical domains, including molecules, materials, catalysts, and metal–organic frameworks. In addition, UMA introduces a mixture-of-experts design, which scales the parameter size through the combination of multiple expert models. When performing structure optimization or molecular dynamics simulations, these expert models can be integrated to achieve the same speed as a single expert model.
Installation
To use the UMA series models, please first make sure you have obtained Hugging Face access at https://huggingface.co/facebook/UMA in order to access their checkpoints and created a Huggingface token at https://huggingface.co/settings/tokens/ with this permission:
Permissions: Read access to contents of all public gated repos you can access
Then, add the token as an environment variable
# pip install huggingface_hub
huggingface-cli login
It may be enough to use pip install fairchem-core to install UMA. This gets you the latest version on PyPi (https://pypi.org/project/fairchem-core/)
Or you can install from source code in edit mode:
git clone git@github.com:facebookresearch/fairchem.git
pip install -e fairchem/packages/fairchem-core
# or
# pip install -e fairchem/packages/fairchem-core[dev]
# for development
Key Features
We consider the key feature of the UMA model to be its unprecedented breadth of training domains among atomistic foundation models to date, encompassing molecules, materials, catalysts, and metal–organic frameworks. For detailed information about UMA, please see UMA: A Family of Universal Models for Atoms
License
The UMA backbone is available under MIT License
Usage Example
UMA Fine-tuning
from __future__ import annotations
import logging
from pathlib import Path
import rich
import os
from lightning.pytorch.strategies import DDPStrategy
import mattertune.configs as MC
from mattertune import MatterTuner
from mattertune.configs import WandbLoggerConfig
from mattertune.backbones.uma import UMABackboneModule
logging.basicConfig(level=logging.ERROR)
def main(args_dict: dict):
def hparams():
hparams = MC.MatterTunerConfig.draft()
hparams.model = MC.UMABackboneConfig.draft()
hparams.model.model_name = "uma-s-1.1"
hparams.model.task_name = "omat"
hparams.model.reset_output_heads = True
hparams.model.optimizer = MC.AdamWConfig(
lr=args_dict["lr"],
amsgrad=False,
betas=(0.9, 0.95),
eps=1.0e-8,
weight_decay=0.1,
)
hparams.model.lr_scheduler = MC.ReduceOnPlateauConfig(
mode="min",
monitor=f"val/forces_mae",
factor=0.8,
patience=5,
min_lr=1e-8,
)
# Add model properties
hparams.model.properties = []
energy = MC.EnergyPropertyConfig(
loss=MC.MSELossConfig(), loss_coefficient=1.0
)
hparams.model.properties.append(energy)
forces = MC.ForcesPropertyConfig(
loss=MC.MSELossConfig(), conservative=True, loss_coefficient=1.0
)
hparams.model.properties.append(forces)
stress = MC.StressesPropertyConfig(
loss=MC.MSELossConfig(), conservative=True, loss_coefficient=1.0
)
hparams.model.properties.append(stress)
## Data Hyperparameters
hparams.data = MC.ManualSplitDataModuleConfig.draft()
hparams.data.train = MC.XYZDatasetConfig.draft()
hparams.data.train.src = "./data/Li_electrode_finetune.xyz"
hparams.data.validation = MC.XYZDatasetConfig.draft()
hparams.data.validation.src = "./data/Li_electrode_val.xyz"
hparams.data.batch_size = args_dict["batch_size"]
## Add Normalization for Energy
hparams.model.normalizers = {
"energy": [
MC.PerAtomNormalizerConfig(),
]
}
## Configure EMA
hparams.trainer.ema = MC.EMAConfig(decay=0.99)
## Trainer Hyperparameters
hparams.trainer = MC.TrainerConfig.draft()
hparams.trainer.max_epochs = 20
hparams.trainer.accelerator = "gpu"
hparams.trainer.devices = args_dict["devices"]
hparams.trainer.strategy = DDPStrategy()
hparams.trainer.gradient_clip_algorithm = "norm"
hparams.trainer.gradient_clip_val = 1.0
hparams.trainer.precision = "32"
# Configure Early Stopping
hparams.trainer.early_stopping = MC.EarlyStoppingConfig(
monitor=f"val/forces_mae", patience=50, mode="min", min_delta=1e-4
)
# Configure Model Checkpoint
ckpt_name = "uma-s11-best"
if os.path.exists(f"./checkpoints/{ckpt_name}.ckpt"):
os.remove(f"./checkpoints/{ckpt_name}.ckpt")
hparams.trainer.checkpoint = MC.ModelCheckpointConfig(
monitor="val/forces_mae",
dirpath="./checkpoints",
filename=ckpt_name,
save_top_k=1,
mode="min",
)
# Configure Logger
hparams.trainer.loggers = [
WandbLoggerConfig(
project="MatterTune-UsageTest",
name="UMA",
)
]
# Additional trainer settings
hparams.trainer.additional_trainer_kwargs = {
"inference_mode": False,
}
hparams = hparams.finalize(strict=False)
return hparams
mt_config = hparams()
model, trainer = MatterTuner(mt_config).tune()
## Perform Evaluation
ckpt_path = "./checkpoints/uma-s11-best.ckpt"
model = UMABackboneModule.load_from_checkpoint(ckpt_path)
from ase.io import read
from ase import Atoms
import numpy as np
import torch
import wandb
from tqdm import tqdm
wandb.init(project="MatterTune-UsageTest", name="UMA", resume=True)
val_atoms_list:list[Atoms] = read("./data/Li_electrode_test.xyz", ":") # type: ignore
calc = model.ase_calculator(
device = f"cuda:{args_dict['devices'][0]}"
)
energies_per_atom = []
forces = []
stresses = []
pred_energies_per_atom = []
pred_forces = []
pred_stresses = []
for atoms in tqdm(val_atoms_list):
energies_per_atom.append(atoms.get_potential_energy() / len(atoms))
forces.extend(np.array(atoms.get_forces()).tolist())
stresses.extend(np.array(atoms.get_stress(voigt=False)).tolist())
atoms.set_calculator(calc)
pred_energies_per_atom.append(atoms.get_potential_energy() / len(atoms))
pred_forces.extend(np.array(atoms.get_forces()).tolist())
pred_stresses.extend(np.array(atoms.get_stress(voigt=False)).tolist())
e_mae = torch.nn.L1Loss()(torch.tensor(energies_per_atom), torch.tensor(pred_energies_per_atom))
f_mae = torch.nn.L1Loss()(torch.tensor(forces), torch.tensor(pred_forces))
s_mae = torch.nn.L1Loss()(torch.tensor(stresses), torch.tensor(pred_stresses))
rich.print(f"Energy MAE: {e_mae} eV/atom")
rich.print(f"Forces MAE: {f_mae} eV/Ang")
rich.print(f"Stresses MAE: {s_mae} eV/Ang^3")
wandb.finish()
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--batch_size", type=int, default=16)
parser.add_argument("--lr", type=float, default=1e-4)
parser.add_argument("--devices", type=int, nargs="+", default=[1, 2, 3])
args = parser.parse_args()
args_dict = vars(args)
main(args_dict)
Mixture of Experts Merging
If during your whole task, the overall charge and spin tag of your structure, as well as the element of each atom, are all fixed, you can accelerate UMA’s inferencing by model.merge_MOLE_model(atoms). An example is:
from ase.md.langevin import Langevin
from ase.io import read
from ase import Atoms
from mattertune.backbones.uma import UMABackboneModule
ckpt_path = "./checkpoints/uma-s11-best.ckpt"
model = UMABackboneModule.load_from_checkpoint(ckpt_path)
atoms: Atoms = read("./data/Li_electrode_test.xyz", index=0) # type: ignore
model.merge_MOLE_model(atoms)
calc = model.ase_calculator(
device = f"cuda:0"
)
atoms.set_calculator(calc)
dyn = Langevin(
atoms,
timestep=1.0,
temperature_K=300,
friction=0.02,
)
dyn.run(1000)