@@ -36,11 +36,115 @@ potentials:
3636| {py: class }` ~nvalchemi.models.demo.DemoModelWrapper ` | {py: class }` ~nvalchemi.models.demo.DemoModel ` | Non-invariant demo; useful for testing and tutorials |
3737| {py: class }` ~nvalchemi.models.aimnet2.AIMNet2Wrapper ` | {py: class }` ~aimnet.calculators.AIMNet2Calculator ` | Requires the ` aimnet2 ` optional dependency |
3838| {py: class }` ~nvalchemi.models.mace.MACEWrapper ` | Any MACE variant | Requires the ` mace ` optional dependency with a CUDA extra, such as ` cu13 ` or ` cu12 ` |
39+ | {py: class }` ~nvalchemi.models.uma.UMAWrapper ` | fairchem-core UMA (` MLIPPredictUnit ` ) | Requires the ` uma ` optional dependency; conflicts with ` mace ` (incompatible ` e3nn ` pins) |
3940
40- {py: class }` ~nvalchemi.models.aimnet2.AIMNet2Wrapper ` and {py: class }` ~nvalchemi.models.mace.MACEWrapper `
41+ {py: class }` ~nvalchemi.models.aimnet2.AIMNet2Wrapper ` , {py: class }` ~nvalchemi.models.mace.MACEWrapper ` ,
42+ and {py: class }` ~nvalchemi.models.uma.UMAWrapper `
4143are lazily imported --- they only load when accessed, so missing dependencies will not
4244break other imports.
4345
46+ ```` {note}
47+ **UMA resolves to a different torch than the `cu12` / `cu13` GPU stack.**
48+ `fairchem-core` caps torch below 2.9 (`fairchem-core>=2.8` requires
49+ `torch>=2.8,<2.9`), while the `cu12` / `cu13` extras pull
50+ `nvalchemi-toolkit-ops[torch-cuXX]`, which floors torch at `>=2.11`. The
51+ `uma` extra is therefore declared mutually exclusive with `cu12`, `cu13`,
52+ and `mace`, and `uv sync --extra uma` forks a standalone resolution that
53+ installs a PyPI CUDA torch wheel (~2.8) instead of the NVIDIA-indexed
54+ `cuXX` build. Keep UMA in its own environment, e.g.:
55+
56+ ```bash
57+ uv venv .venv-uma && uv sync --extra uma # UMA (fairchem's torch)
58+ uv venv .venv-mace && uv sync --extra cu13 --extra mace # MACE on the cu13 GPU stack
59+ ```
60+
61+ The core `nvalchemi-toolkit-ops` package (and its Warp kernels) is still
62+ installed in the UMA environment --- only the `cuXX` GPU-acceleration
63+ extras (cuEquivariance, cuML, the NVIDIA-indexed torch build) are dropped,
64+ none of which UMA uses, since it builds its neighbor graph inside
65+ fairchem. The toolkit-ops `torch-cuXX` extras pin `torch>=2.11`, but that
66+ tracks the `cuXX` wheel builds rather than a toolkit-ops API requirement
67+ --- the base package already declares `torch>=2.8`, so the Warp path is
68+ expected to work against the ~2.8 torch in the UMA environment.
69+ ````
70+
71+ ### Using UMA (fairchem-core)
72+
73+ UMA (Universal Models for Atoms) is a multi-task foundation model: one
74+ checkpoint ships task heads for molecules (` omol ` ), bulk crystals (` omat ` ),
75+ catalysis (` oc20 ` ), direct air capture (` odac ` ), and molecular crystals
76+ (` omc ` ). {py: class }` ~nvalchemi.models.uma.UMAWrapper ` pins a single task at
77+ construction; ` active_outputs ` is ` {energy, forces} ` for molecular tasks and
78+ ` {energy, forces, stress} ` for periodic ones.
79+
80+ ** 1. Install the optional dependency** (in its own environment, per the note
81+ above):
82+
83+ ``` bash
84+ uv venv .venv-uma && uv sync --extra uma
85+ # or, with pip: pip install 'nvalchemi-toolkit[uma]'
86+ ```
87+
88+ ** 2. Get HuggingFace access.** UMA checkpoints live in the ** gated**
89+ [ ` facebook/UMA ` ] ( https://huggingface.co/facebook/UMA ) repository, so a (free)
90+ HuggingFace account and a one-time access approval are required:
91+
92+ 1 . Sign in and click ** "Agree and access repository"** on the
93+ [ model page] ( https://huggingface.co/facebook/UMA ) .
94+ 2 . Create a ** read** token at
95+ < https://huggingface.co/settings/tokens > .
96+ 3 . Make the token available to your shell, either by logging in once (it is
97+ cached under ` ~/.cache/huggingface ` ):
98+
99+ ``` bash
100+ huggingface-cli login
101+ ```
102+
103+ or by exporting it for the session:
104+
105+ ``` bash
106+ export HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxxxxxx
107+ ```
108+
109+ ** 3. Load a checkpoint.** The first call downloads and caches the weights
110+ (under ` ~/.cache/fairchem ` ); later calls reuse the cache and need no network:
111+
112+ ``` python
113+ from nvalchemi.models.uma import UMAWrapper
114+
115+ # Molecular potential (OMol head)
116+ mol = UMAWrapper.from_checkpoint(" uma-s-1p1" task_name = " omol" device = " cuda"
117+
118+ # Bulk-crystal potential (OMat head, same checkpoint family)
119+ mat = UMAWrapper.from_checkpoint(" uma-m-1p1" task_name = " omat" device = " cuda"
120+ ```
121+
122+ Registered checkpoint names (see
123+ ` fairchem.core.calculate.pretrained_mlip.available_models ` for the full list):
124+
125+ | Checkpoint | Size | Notes |
126+ | ---| ---| ---|
127+ | ` uma-s-1p1 ` | small | Default in the examples / tests |
128+ | ` uma-s-1p2 ` | small | Updated small release |
129+ | ` uma-m-1p1 ` | medium | Higher accuracy, larger / slower |
130+
131+ ** ` torch.compile ` / turbo.** ` UMAWrapper ` does not add a ` compile_model `
132+ flag (unlike the MACE / AIMNet2 wrappers) because fairchem owns compilation
133+ internally as a field on its ` InferenceSettings ` . Reach it through
134+ ` from_checkpoint ` 's ` inference_settings ` argument --- pass ` "turbo" ` for
135+ fairchem's compiled preset (` torch.compile ` + TF32 + MoLE merge, for runs
136+ with fixed atomic composition), or an ` InferenceSettings ` instance for finer
137+ control:
138+
139+ ``` python
140+ fast = UMAWrapper.from_checkpoint(
141+ " uma-s-1p1" task_name = " omat" device = " cuda" inference_settings = " turbo"
142+ )
143+ ```
144+
145+ See {doc}` the UMA NVE/NVT example </auto_examples/advanced/09_uma_nve> ` for a
146+ runnable end-to-end molecular-dynamics walkthrough.
147+
44148## Architecture overview
45149
46150A wrapped model uses ** multiple inheritance** : your existing {py: class }` ~torch.nn.Module `
0 commit comments