Garden's Almanac of Matter Models eSEN
January 2025
| Checkpoint | Params | Polaris | Sophia | Perlmutter | Delta | Frontier | Della |
|---|---|---|---|---|---|---|---|
| esen-30m-mptrj | 30M | — | — | — | — | — | — |
| esen-30m-oam | 30M | — | — | — | — | — | — |
| esen-30m-omat | 30M | — | — | — | — | — | — |
| esen-md-direct-all-omol | 50.7M | ● | ● | ● | ● | — | — |
| esen-sm-conserving-all-omol | 6.3M | ● | ● | ● | ● | — | — |
| esen-sm-direct-all-omol | 6.3M | ● | ● | ● | ● | — | — |
| esen:custom | ● | ● | — | — | — | — |
● verified, last 30 days ○ installed, not recently verified — not installed
No cluster has installed any checkpoint of this model yet.
Running this model
1# from a job or interactive session on a supported cluster:
2from rootstock import RootstockCalculator
3
4# this checkpoint 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="esen-md-direct-all-omol",
11 device="cuda",
12) as calc:
13 # model is now running in subprocess on compute node
14 atoms.calc = calc
15 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.
esen_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"""eSEN env — hosts FAIRChem eSEN single-task checkpoints.
10
11fairchem-core v2 dropped the torch-geometric / pyg-find-links install dance
12(simplified here 2026-07-30 to match uma.py/allscaip.py — the find-links were
13also blocking universal lockfile resolution).
14
15OMol checkpoints expect `charge` and `spin` in `atoms.info`.
16"""
17
18CHECKPOINTS = {
19 "esen-md-direct-all-omol": "esen-md-direct-all-omol",
20 "esen-sm-conserving-all-omol": "esen-sm-conserving-all-omol",
21 "esen-sm-direct-all-omol": "esen-sm-direct-all-omol",
22 # Your own fine-tuned weights: pair with weights= (loaded via setup_from_path).
23 "esen:custom": None,
24}
25
26
27def _fairchem_device(device: str) -> str:
28 """Translate an indexed device ("cuda:2") into what fairchem v2 accepts.
29
30 MLIPPredictUnit._setup_device asserts `device in ["cpu", "cuda"]` and then
31 resolves the real GPU itself via get_device_for_local_rank(), which returns
32 f"cuda:{torch.cuda.current_device()}". So an index has to travel through
33 torch's current-device state, not the argument. Verifying several
34 checkpoints at once on a multi-GPU node hands each worker "cuda:N" — that
35 killed all 8 fairchem-v2 checkpoints on the 2026-08-06 Polaris sync
36 (4x A100, VERIFY_JOBS=4), while single-GPU Sophia never hit it.
37 """
38 if device.startswith("cuda:"):
39 import torch
40
41 torch.cuda.set_device(int(device.split(":", 1)[1]))
42 return "cuda"
43 return device
44
45
46def setup(checkpoint: str, device: str = "cuda"):
47 from fairchem.core import FAIRChemCalculator, pretrained_mlip
48
49 predictor = pretrained_mlip.get_predict_unit(
50 CHECKPOINTS[checkpoint], device=_fairchem_device(device)
51 )
52 return FAIRChemCalculator(predictor)
53
54
55def setup_from_path(path: str, device: str = "cuda"):
56 # Custom checkpoints (`:custom` ids with user weights): a weights *file* loads through
57 # load_predict_unit, not the registry-name lookup setup() uses.
58 from fairchem.core import FAIRChemCalculator
59 from fairchem.core.units.mlip_unit import load_predict_unit
60
61 predictor = load_predict_unit(path, device=_fairchem_device(device))
62 return FAIRChemCalculator(predictor)
63
Built on Polaris: 2026-08-06
Couldn't load the current environments from Rootstock.