Garden's Almanac of Matter Models | Checkpoint | Params | Polaris | Sophia | Perlmutter | Delta | Frontier | Della |
|---|---|---|---|---|---|---|---|
| uma-s-1p1 | 150M | ● | ● | ● | ● | ● | ● |
| uma-s-1p2p1 | 290M | — | — | — | ○ | — | — |
| uma-m-1p1 | 1.4B | ● | ● | ● | ● | ● | ● |
| uma:custom | ● | ● | — | — | ● | — |
● verified, last 30 days ○ installed, not recently verified — not installed
No cluster has installed any checkpoint of this model yet.
Running this model
This model has task heads. Select which output layer is active at inference time with
setup_kwargs={"task": "omol"} —
there is no default, and the calculator raises an error if no head is given.
| Task head | Trained against |
|---|---|
| omat | OMat24 — inorganic crystals |
| omol | OMol25 — molecules |
| oc20 | OC20 — surfaces/interfaces |
| odac | ODAC23 — surfaces/interfaces (MOF frameworks + adsorbates) |
| omc | OMC25 — molecular crystals |
1# from a job or interactive session on a supported cluster:
2from rootstock import RootstockCalculator
3
4# the omol head accepts charge and spin
5atoms.info["charge"] = -1
6atoms.info["spin"] = 2
7
8with RootstockCalculator(
9 cluster=YOUR_CLUSTER_ID, # eg, "polaris", "sophia"
10 checkpoint="uma-s-1p1",
11 device="cuda",
12 setup_kwargs={"task": "omol"}, # required: select a task head
13) as calc:
14 # model is now running in subprocess on compute node
15 atoms.calc = calc
16 atoms.get_potential_energy()
Environments
Rootstock runs each model family inside an isolated Python environment defined by a single file. This environment file includes the specific dependencies needed, plus a setup() function that loads the model and returns an ASE calculator. These files are usually almost identical for a given model family, but because of cluster-specific quirks (eg, an old CUDA driver) the dependencies and setup code can vary a bit.
uma_env.py
1# /// script
2# requires-python = ">=3.11"
3# dependencies = [
4# "fairchem-core>=2.20",
5# "ase>=3.22",
6# "torch>=2.4.0",
7# ]
8# ///
9"""UMA env — hosts Meta's UMA foundation model via FAIRChem.
10
11fairchem-core v2 dropped the torch-geometric / pyg-find-links install dance, so
12this env is a plain PyPI install. The original uma-s-1 had an extensivity bug
13and was removed from the fairchem 2.20 registry — use uma-s-1p1 or uma-s-1p2p1.
14"""
15
16CHECKPOINTS = {
17 "uma-s-1p1": "uma-s-1p1",
18 # uma-s-1p2 has a known major bug; uma-s-1p2p1 fixes it and is the
19 # upstream-recommended small model. 1p2 stays listed for reproducibility
20 # of existing runs.
21 "uma-s-1p2": "uma-s-1p2",
22 # uma-s-1p2p1 is in fairchem's registry on git main but NOT in any
23 # release yet (latest fairchem-core 2.21.0, 2026-06-08, lacks it — the
24 # 2026-07-30 sync failed on exactly this). Re-add when the next
25 # fairchem-core ships, and bump the dependency floor to that version.
26 # "uma-s-1p2p1": "uma-s-1p2p1",
27 "uma-m-1p1": "uma-m-1p1",
28 # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).
29 "uma:custom": None,
30}
31
32
33def _fairchem_device(device: str) -> str:
34 """Translate an indexed device ("cuda:2") into what fairchem v2 accepts.
35
36 MLIPPredictUnit._setup_device asserts `device in ["cpu", "cuda"]` and then
37 resolves the real GPU itself via get_device_for_local_rank(), which returns
38 f"cuda:{torch.cuda.current_device()}". So an index has to travel through
39 torch's current-device state, not the argument. Verifying several
40 checkpoints at once on a multi-GPU node hands each worker "cuda:N" — that
41 killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync
42 (4x A100, VERIFY_JOBS=4), while single-GPU Sophia never hit it.
43 """
44 if device.startswith("cuda:"):
45 import torch
46
47 torch.cuda.set_device(int(device.split(":", 1)[1]))
48 return "cuda"
49 return device
50
51
52def setup(checkpoint: str, device: str = "cuda", task: str = "omat"):
53 from fairchem.core import FAIRChemCalculator, pretrained_mlip
54
55 predictor = pretrained_mlip.get_predict_unit(
56 CHECKPOINTS[checkpoint], device=_fairchem_device(device)
57 )
58 return FAIRChemCalculator(predictor, task_name=task)
59
60
61def setup_from_path(path: str, device: str = "cuda", task: str = "omat"):
62 # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through
63 # load_predict_unit, not the registry-name lookup setup() uses.
64 from fairchem.core import FAIRChemCalculator
65 from fairchem.core.units.mlip_unit import load_predict_unit
66
67 predictor = load_predict_unit(path, device=_fairchem_device(device))
68 return FAIRChemCalculator(predictor, task_name=task)
69
Built on Polaris: 2026-08-06
Couldn't load the current environments from Rootstock.