Using the Universal Model for Atoms for single point calculations

David_R

Hello all,

This is post is intended to be an introductory guide for getting started with using the Universal Model for Atoms (UMA). I am not entirely sure this is in the right section of the forum, so I am happy to be corrected accordingly. Since the model is trained on DFT data, and the workflows I present here follow those conventially used with DFT methodogy, I figured it would reach the right audience here. I also hope that this post being in English does not present a barrier to its pedagogical value.

1. Introduction

UMA is a machine learning interatomic potential (MLIP) from Facebook’s AI research team. It is trained on more than 6 billion CPU core-hours of DFT data (at the ωB97M-V/def2-TZVPD level of theory) across a broad set of research domains. You can find more details in the corresponding paper here.

Having played around with it for a few weeks now, I believe this is a major innovation for computational chemistry. From benchmarks on my datasets, the accuracy is on par with high-level DFT calculations, at a fraction of the computational cost. Compared to DFT, the time complexity of running calculations with UMA is also very favourable, with very large systems (thousands of atoms) being computationally accessible. I have thoroughly tested and incorporated it into my workflows, which has made previously inaccessible research ideas now feasible on my hardware.

2. Hardware and software considerations

The UMA model is implemented in PyTorch, and runs on a GPU. The VRAM requirements are relatively modest, however, and one does not need exceptionally powerful (and expensive!) hardware to run the model on most systems. In fact, it runs just fine on my laptop with an RTX 4050 GPU. However, my main system specifications are as follows:

CPU: AMD Ryzen Threadripper 7980X: 64 cores/128 threads (running overclocked @ 4.9 GHz all-core)

RAM: 256 GB (4 × 64 GB) V-color DDR5 ECC RDIMM (running overclocked at 6000 MT/s)

GPU: NVIDIA RTX 4090

Motherboard: Asus Pro WS TRX50-SAGE

Storage: 3 × 2 TB Samsung 990 Pro

PSU: Seasonic Prime TX-1600

CPU cooler: Silverstone XE360-TR5

Case: Silverstone RM52

Ubuntu Desktop 24.04.2 LTS

Python version 3.12.3

Nvidia drivers version 570.133.07

CUDA version 12.8 (PyTorch using 12.4)

PyTorch version 2.6.0+cu124

3. Getting started with UMA and ASE

I shall assume some basic knowledge of installing and setting up Python, but I hope that this guide should point complete beginners in the right direction of getting their own calculations running. However, since I spent a rather long time troubleshooting problems when setting up this workflow, many of the more trivial steps may help reproducibility for more adept users. I am far from a master, so I shall welcome suggestions or improvements from others.

The easiest way to use the UMA model is with the Atomic Simulation Environment (ASE) library implemented in Python. For my workflows, I have found this library quite capable, with most functionalities well-implemented. The documentation for ASE can be found here.

3.1. Setting up a virtual environment

It is recommended to start a new virtual environment forusing UMA, particularly since specific versions of various Python libraries are required for intercompatibility. I find the Python venv module quite convenient to use.

• Create a new directory of your choosing: e.g. /home/david/PyEnv/UMA (note: it’s important to make sure the path does not contain spaces).
  

• Run the following command in a terminal window (replacing my example path with yours):
  

1. python3 -m venv “/home/david/PyEnv/UMA”

复制代码

• Whenever you want to use it, open a terminal window and navigate to the following path:
  

1. cd “/home/david/PyEnv/UMA/bin”

复制代码

then enter the command:

1. source activate

复制代码

...now you’re good to go and execute Python commands as required.

3.2. Obtaining the UMA model checkpoint

This is a difficult part, and, disappointingly, this appears to be due to political reasons. The uma_sm.pt (Edit: the latest version of the model weights now has the filename 'uma-s-1.pt') checkpoint is a 1.2 GB file that can be downloaded from here, but it requires approval from the FAIRChem team to be able to download it. I understand that it is difficult to get approval from China, but I have thoroughly read the permissive license agreement that came withmy copy of the model, and I didn’t find anything about geographical restrictions in there. I don’t particularly want to get myself in trouble, but I don’t see why it couldn’t be made more widely available, should there be sufficient will to do so.

3.3. Installing required packages
I found this process to be a real nuisance, and spent a considerable amount of time troubleshooting issues. The FAIRChem package underwent changes each time I went about installing it, with earlier versions having a lot of trouble with dependencies. The currently available version seems to be a lot more straightforward to get working. However, this also means I cannot guarantee that this guide will work with newer versions as they become available. Nonetheless, the following method is working for me as of this post date:

• I recommend using Jupyter notebook. If you wish to do so too, start with this step:
  

1. pip install jupyterlab

复制代码

Once you have this done, you can launch it by activating your virtual environment (see above) and entering the command:

1. jupyter lab

复制代码

…and it’ll open in a new web browser window.

• Install ASE:
  

1. pip install ase

复制代码

• Install FAIRChem core:
  

1. pip install fairchem-core

复制代码

I found that the latest version of FAIRChem installs all of its dependencies without any trouble.

• Install Sella:
  

1. pip install sella

复制代码

This is a Python implementation of the very effective saddle point optimisation algorithm reported here (chemrxiv version here).

When running code, often lots of warnings pop up, though the code executes just fine. The following optional steps should fix the warnings:

• Revert to an earlier version of setuptools:
  

1. pip uninstall setuptools

复制代码

followed by:

1. pip install setuptools==80.7.1

复制代码

• Install the cuda version of jax:
  

1. pip install -U “jax[cuda12]”

复制代码

You’ll now want to put the UMA checkpoint file uma_sm.pt in your Python environment directory: e.g. /home/david/PyEnv/UMA/bin. In fact, you can put it wherever you like; you’ll just have to specify the path when setting up the calculator later.

I noticed that later versions of the FAIRChem package implemented the UMA model by downloading it with the huggingface_hub library using your account credentials, rather than letting you simply load it from your hard drive offline. I prefer the latter, so I had to go and modify some of the FAIRChem code as follows:

• Navigate to the following directory (replacing as necessary with your own Python virtual environment): /home/david/PyEnv/UMA/lib/python3.12/site-packages/fairchem/core/calculate
  

• Open pretrained_mlip.py (using a text editor will work just fine)
  

• Replace the code on lines 76-88:
  

1.     try:
   
2.         model_checkpoint = _MODEL_CKPTS.checkpoints[model_name]
   
3.     except:
   
4.         raise KeyError(
   
5.             f"Model '{model_name}' not found. Available models: {available_models}"
   
6.             ) from err
   
7.     checkpoint_path = hf_hub_download(
   
8.         filename=model_checkpoint.filename,
   
9.         repo_id=model_checkpoint.repo_id,
   
10.         subfolder=model_checkpoint.subfolder,
    
11.         revision=model_checkpoint.revision,
    
12.         cache_dir=CACHE_DIR,
    
13.     )

复制代码

With the following code:

1.     checkpoint_path=model_name

复制代码

Now we should be good to get started with some code. If one runs into problems with the above installation of the required Python packages, it is a good idea to first troubleshoot version mismatches between CUDA, PyTorch, and FAIRChem core. I am pretty sure this has been fixed now, but earlier versions of FAIRChem core installed a few packages that broke things—Jupyter in particular. If Jupyter Notebook fails to start properly, go into your python environment 'bin' directory and delete the ‘jupyter-config’ file.

4. Running a Geometry Optimisation with Example Code
OK, so now we should be good to get started with some code. I've provided the following code to give, for those that need it, an example of what can be done with ASE and UMA.

First, we'll need to create a new Notebook in Jupyter, by navigating to File > New > Notebook.

I found that ASE uses all available CPU cores by default, but this offers no increase in performance over using just one; a lot of heat, noise and electricity consumption for nothing. Accordingly, we want to stop ASE from using all CPU cores and limit to just one. I also found that it's better to ensure that jax runs on CPU and not the GPU, as it ends up allocating way more VRAM on the GPU than it needs, and actually performs worse.

1. #Set environment variables such that the number of CPU cores is just one.
   
2. import os
   
3. num_cpu_threads = '1'
   
4. os.environ['OMP_NUM_THREADS']=num_cpu_threads
   
5. os.environ['MKL_NUM_THREADS']=num_cpu_threads
   
6. os.environ['OPENBLAS_NUM_THREADS']=num_cpu_threads
   
7. os.environ['NUMEXPR_NUM_THREADS']=num_cpu_threads
   
8. 
   
9. #Using cuda with jax results in more VRAM usage and worse performance, so we'll make sure jax runs on CPU instead.
   
10. os.environ['JAX_PLATFORMS'] = 'cpu'

复制代码

I will note that it isn’t strictly necessary to specify all of these environment variables, but I use it as a convenient ‘catch all’ across different codebases and systems.

Now lets import the various packages we’ll need:

1. import time
   
2. from ase.io import Trajectory, read, write
   
3. from fairchem.core import pretrained_mlip, FAIRChemCalculator
   
4. from fairchem.core.units.mlip_unit.api.inference import InferenceSettings
   
5. from ase.optimize import BFGS
   
6. from sella import Sella

复制代码

Next, I spent some time optimising the model inference settings for UMA, in order to minimise VRAM usage and decrease inference time. If this causes problems, then it is recommended to revert to the default inference settings (see later). Note: you may get the error 'Cannot run on merged model on system. Embeddings seem different...' if you run consecutive calculations with the same calculator. In such a case, just restart the kernel between running calculations.

1. #Custom inference settings designed to speed up inference and minimise VRAM usage.
   
2. def inference_settings_custom():
   
3.     return InferenceSettings(
   
4.         tf32=True,
   
5.         activation_checkpointing=True,
   
6.         merge_mole=True,
   
7.         compile=False,
   
8.         wigner_cuda=False,
   
9.         internal_graph_gen_version=2,
   
10.         external_graph_gen=False
    
11.     )

复制代码

Now, we can set up UMA as a calculator in ASE:

1. # Set up the UMA calculator.
   
2. predictor = pretrained_mlip.get_predict_unit(model_name="uma_sm.pt", device="cuda", inference_settings=inference_settings_custom())
   
3. calc = FAIRChemCalculator(predictor, task_name="omol")

复制代码

Here, we specify the model checkpoint (you can add full path to uma_sm.pt if you have it placed elsewhere on your system), ensure it runs with CUDA on the GPU, and set the inference settings to the custom ones we specified earlier. If you have trouble with these inference settings, you can change it to inference_settings=inference_settings_default(), or remove this argument altogether.

Additionally, in the FAIRChemCalculator, we specify the 'task_name', which depends on what type of system you want to model. Here, we'll be doing a geometry optimisation for a small molecule, for which "omol" is the corresponding task. Refer to the original paper for other available tasks.

Next, we can write a function that will perform a geometry optimisation in ASE:

1. def ground_state_minimisation(inputxyz, outputtrajectory, fmax=0.0025, maxsteps=1000, logfile='-'):
   
2.     GSatoms = read(inputxyz, format='xyz')
   
3. 
   
4.     GSatoms.calc = calc
   
5.     GSatoms.info = {'charge': 0, 'spin': 1}
   
6. 
   
7.     geo_opt = BFGS(GSatoms, alpha=50, logfile=logfile, trajectory=outputtrajectory)
   
8. 
   
9.     geo_opt.run(fmax=fmax, steps=maxsteps)
   
10. 
    
11.     return GSatoms

复制代码

To use this function, we specify a .xyz file correponding to our starting geometry, and a filename for the trajectory to be output to. fmax is the maximum per-atom force, provided as the convergence criteria for the geometry optimisation algorithm. maxsteps specifies the maximum number of iterations that the optimisation algorithm runs for. You can specify how the logs from the optimisation routine are output. 'None' suppresses this output, while '-' shows the logs in the cell output as the calculation runs. Alternatively, one can output the logs to a text file.

Let's give this a go on a test system now. To start, you'll need to select a suitable system, and move the corresponding .xyz file into the same directory as the Jupyter Notebook. I've called mine 'input_geometry.xyz', which corresponds to an organometallic copper complex with 58 atoms.

To run the geometry optimisation:

1. %%time
   
2. 
   
3. atoms = ground_state_minimisation('input_geometry.xyz', 'output_trajectory.xyz')

复制代码

I get the following output:

1.       Step     Time          Energy          fmax
   
2. BFGS:    0 13:12:55   -94732.728217        0.390381
   
3. BFGS:    1 13:12:55   -94732.800124        0.259838
   
4. BFGS:    2 13:12:55   -94732.857109        0.435664
   
5. BFGS:    3 13:12:55   -94732.867930        0.306413
   
6. BFGS:    4 13:12:55   -94732.876524        0.148563
   
7. BFGS:    5 13:12:55   -94732.888666        0.117449
   
8. .
   
9. .
   
10. .
    
11. BFGS:  386 13:13:09   -94733.020315        0.002951
    
12. BFGS:  387 13:13:09   -94733.020156        0.006825
    
13. BFGS:  388 13:13:09   -94733.020140        0.003175
    
14. BFGS:  389 13:13:09   -94733.019475        0.004464
    
15. BFGS:  390 13:13:09   -94733.019988        0.002441
    
16. CPU times: user 14.3 s, sys: 45.7 ms, total: 14.3 s
    
17. Wall time: 14.2 s

复制代码

It is pretty fast, less then 30 ms per iteration!—and the speed becomes even more impressive with larger systems.

From nvidia-smi, the VRAM usage was 742 MB, which is well within the capabilities of pretty much any GPU these days. I have seen larger systems of 500 atoms take up more than 3 GB.

The output_trajectory.traj file can be conveniently visualised with the ASE GUI, which can be started by activating your Python virtual environment, then entering the following command into the terminal:

1. ase gui

复制代码

The geometry optimisation function returns an ASE atoms object corresponding to the optimised geometry, from which we can get various properties we're interested in, such as the energy (in eV):

1. dE = atoms.get_potential_energy()

复制代码

Now, let's use Sella to optimise a geometry to a first order saddle point (transition state). I'll be using the starting geometry file 'input_geometry_TS.xyz'—a larger organometallic copper species with 147 atoms, placed in the same directory as my Jupyter notebook.

1. def TS_optimisation(inputxyz, outputtrajectory, fmax=0.0025, maxsteps=1000, logfile='-'):
   
2.     TSatoms = read(inputxyz, format='xyz')
   
3. 
   
4.     TSatoms.calc = calc
   
5.     TSatoms.info = {'charge': 0, 'spin': 1}
   
6. 
   
7.     TS_opt = Sella(TSatoms, trajectory = outputtrajectory, eta=2e-2, gamma=0.0001, delta0=0.02, logfile=logfile)
   
8. 
   
9.     TS_opt.run(fmax, maxsteps)
   
10. 
    
11.     return TSatoms

复制代码

The function is similar to the ground state optimisation one, but calls the Sella algorithm instead. The Sella algorithm is specified with a number of hyperparameters explained in the original publication, and here. I found that the default values for these hyperparameters did not work with UMA at all, so I recommend starting with the ones I have found. However, they may need adjusting to work with different systems.

Let's run the function on our test system.

1. %%time
   
2. 
   
3. TSatoms = TS_optimisation('input_geometry_TS.xyz', 'output_trajectory_TS.traj')

复制代码

I get the following output:

1.      Step     Time          Energy         fmax         cmax       rtrust          rho
   
2. Sella   0 15:58:47  -126516.018063       0.6620       0.0000       0.0200       1.0000
   
3. Sella   1 15:58:50  -126516.080038       0.4894       0.0000       0.0230       0.9984
   
4. Sella   2 15:58:51  -126516.141463       0.2988       0.0000       0.0264       1.0156
   
5. Sella   3 15:58:52  -126516.242900       0.1168       0.0000       0.0304       0.9913
   
6. Sella   4 15:58:53  -126516.252680       0.1975       0.0000       0.0304       0.7696
   
7. Sella   5 15:58:54  -126516.257909       0.1511       0.0000       0.0304       0.5737
   
8. .
   
9. .
   
10. .
    
11. Sella 758 16:08:14  -126516.310874       0.0030       0.0000       0.0200    -134.6560
    
12. Sella 759 16:08:15  -126516.311217       0.0065       0.0000       0.0200      77.1546
    
13. Sella 760 16:08:15  -126516.311980       0.0033       0.0000       0.0200     197.2380
    
14. Sella 761 16:08:16  -126516.311240       0.0034       0.0000       0.0200    -228.6339
    
15. Sella 762 16:08:17  -126516.311059       0.0024       0.0000       0.0200    -107.6856
    
16. CPU times: user 12min 8s, sys: 15.4 s, total: 12min 23s
    
17. Wall time: 9min 30s

复制代码

In this case, the calculation took substantially longer, but the bottleneck is mostly the approximate Hessian calculation that Sella carries out, and the large number of geometry iterations. However, for comparison, a single point DFT calculation (ωB97M-V/def2-TZVPD) on this system using ORCA takes over 2 hours on 16 CPU cores, compared to less than a second for the same calculation using UMA. The maximum VRAM usage for this system was 1570 MB.

I hope that others will have the confidence to try out this model for themselves, and share their experiences, findings, problems and suggestions in this thread. If the instructions in this guide do not work, I shall be happy to test things and correct them as appropriate.

Cheers!

Edit: Fixed various formatting and typographical errors.

ShangChien

https://huggingface.co/facebook/UMA
有没有好心人能够下载权重呀，个人申请access，然后被无情reject...
(小扎不厚道)

David_R

在中国，不少人提到获取模型权重有困难。如果你需要指点，可以私信我。

wxyhgk

这个模型实际上我五月的时候就关注到了，一直申请拒绝,我手里有大量的数据可以测试。

Illuminatia

用US或者其他地方的IP和填写其他地区申请会好很多, OMol25 就是这么申请的

senpoa

这个模型的效果还是很惊艳的，力的预测精度很好，能量稍微差一些

adong

senpoa 发表于 2025-7-15 19:35
这个模型的效果还是很惊艳的，力的预测精度很好，能量稍微差一些

您测试的是什么体系？最近UMA更新了middle的模型有近12G，不清楚最低得多大显存的GPU能顺利运行？精度相比small模型提升有多少

David_R

I've tested the UMA medium model and it is very impressive. For my datasets, the energy is very accurate when compared with ωB97M-V/def2-TZVPD calculations.

The VRAM usage and calculation time are approximately 2× that of the UMA small model, so still very usable on modest hardware.

小白风

UMA权重申请被拒绝了，还有什么办法可以获得权重吗

dishhhy

Hello, my application for UMA was rejected now. Do you have any suggestions? Thank you！