feat(models): add UMA (fairchem-core) interatomic-potential wrapper

[dallasfoster]

ALCHEMI Toolkit Pull Request

Description

Add UMAWrapper, a BaseModelMixin-compatible wrapper around fairchem-core's
UMA (Universal Models for Atoms) MLIPPredictUnit, so UMA foundation models can
drive nvalchemi dynamics and inference. UMA is multi-task: one checkpoint
(uma-s-1p1 / uma-s-1p2 / uma-m-1p1) ships heads for OMol, OMat, OC20, ODAC,
and OMC; the wrapper pins a single task at construction (one-wrapper-one-model,
drive nvalchemi dynamics and inference. UMA is multi-task: one checkpoint
(uma-s-1p1 / uma-s-1p2 / uma-m-1p1) ships heads for OMol, OMat, OC20, ODAC,
and OMC; the wrapper pins a single task at construction (one-wrapper-one-model,
matching MACEWrapper). The conversion is tensor-native (no ASE round trip),
energy is the differentiable primitive, and forces/stress come from autograd.

Type of Change

• New feature (non-breaking change that adds functionality)

Changes Made

• nvalchemi/models/uma.py: UMAWrapper (task-aware model_config, adapt_input/
  adapt_output, compute_embeddings, from_checkpoint). from_checkpoint exposes
  fairchem's native inference_settings (incl. "turbo"); forward routes the
  one-time lazy-init/MoLE-merge through CPU input to dodge a fairchem device-
  placement bug under turbo on GPU-resident first batches.
• nvalchemi/_optional.py + nvalchemi/models/init.py: register/export UMA.
• pyproject.toml: uma extra (fairchem-core>=2.0.0); declare uma conflicting
  with mace (e3nn pin) and cu12/cu13 (fairchem torch<2.9 vs toolkit-ops
  torch>=2.11); pin setuptools<81 for fairchem's torchtnt. uv.lock regenerated.
• examples/advanced/09_uma_nve.py: NVE/NVT/NPT MD example driven by built-in
  LoggingHook + EnergyDriftMonitorHook; NPT exercises the stress path; turbo
  selectable via inference_settings.
• test/models/test_uma.py: consolidated suite — structural (mock), forward
  equivalence vs FAIRChemCalculator, charged-input response, NVE drift (@slow),
  turbo/compile device path (@slow, CUDA).
• docs: userguide/models.md (supported-models table + UMA usage / HF-token /
  torch-environment notes), modules/models.rst, models/index.md, examples README.
• .github/workflows/ci.yml: install a dedicated .venv-uma and run the UMA tests
  from it (gated on UMA-file changes or full runs); optional HF_TOKEN secret.

Testing

• Unit tests pass locally (UMA suite: 36 passed with --slow against
  uma-s-1p1 on CUDA; 33 passed / 3 slow-skipped without --slow)
• Linting passes (make lint)
• New tests added for new functionality

Equivalence matches FAIRChemCalculator to 1e-4 (OMol energy/forces, OMat
energy/forces/stress); charged OMol runs match the calculator at charge -1;
small (uma-s-1p1/1p2) and medium (uma-m-1p1) checkpoints load and run.

Additional Notes

UMA's deps conflict with the mace/cuXX stack, so it must be installed in its
own environment (uv sync --extra uma); it brings its own CUDA-enabled torch
(2.8, cu12.8) and does not use the cuXX GPU stack. Checkpoints are gated on
HuggingFace (facebook/UMA) — CI structural tests run without a token; the
checkpoint-based tests skip unless an HF_TOKEN secret is provided.

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[dallasfoster]

/ok to test 8344178

[greptile-apps]

Greptile Summary

This PR introduces UMAWrapper, a BaseModelMixin-compatible wrapper around fairchem-core's MLIPPredictUnit, enabling UMA foundation model checkpoints to drive nvalchemi dynamics and inference. The integration is tensor-native (no ASE round-trip), task is pinned at construction, and forces/stress are derived via autograd.

• nvalchemi/models/uma.py: New UMAWrapper with adapt_input/adapt_output, task-aware ModelConfig, compute_embeddings, from_checkpoint, and a one-shot _cpu_route_first_forward flag to work around fairchem's turbo/MoLE device bug on the first call.
• Packaging & CI: Adds a uma extra with fairchem-core declared as conflicting with mace and cuXX; CI runs UMA tests from an isolated .venv-uma environment gated on UMA-file changes, with optional HF_TOKEN for checkpoint-based tests.
• Tests: Consolidated structural (mock) and checkpoint-based suite covering adapt_input/adapt_output, forward equivalence vs FAIRChemCalculator, charged input, NVE drift (@slow), and turbo/compile device path (@slow, CUDA).

Important Files Changed

Filename	Overview
nvalchemi/models/uma.py	New UMAWrapper for fairchem's MLIPPredictUnit; optional_inputs advertises "tags" but the adapter reads atom_categories — the two names don't match, so callers following the declared contract get silent zero-tags.
test/models/test_uma.py	Comprehensive structural and checkpoint test suite; tests atom_categories passthrough via test_passes_tags but doesn't test that optional_inputs accurately lists it.
.github/workflows/ci.yml	Adds isolated .venv-uma step gated on UMA file changes; coverage is correctly appended and HF_TOKEN secret is optional.
pyproject.toml	Adds uma extra with fairchem-core>=2.0.0, declares conflicts with mace and cuXX, and pins setuptools<81 for fairchem's torchtnt compatibility.
nvalchemi/models/init.py	Adds lazy-import and __all__ export for UMAWrapper following the existing pattern.
nvalchemi/_optional.py	Registers fairchem.core as an optional dependency under the UMA enum entry; straightforward addition.
examples/advanced/09_uma_nve.py	Well-documented NVE/NVT/NPT example; gracefully skips when checkpoint is unavailable and exercises the stress path via NPT.

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into dallasf/uma-wra..." | Re-trigger Greptile}

[dallasfoster]

/ok to test 00044fa

[laserkelvin]

Generally looks good to me; the spin default value needs to change though

[dallasfoster]

/ok to test 816b230