Merge branch 'main' into dependabot/github_actions/actions/upload-artifact-7

Author: zulissimeta

README.md

@@ -42,6 +42,8 @@ demos, and application efforts for materials science and quantum chemistry.
 
 
 ## Latest news
+March 2026 - UMA-1.2 released! ~50% faster, ~40% more accurate on Open Molecules test set, and expanded data coverage for catalysts (oxides and interfaces), molecules, and polymers!
+
 Oct 2025 - [check out our seamless Multi-node, Multi-GPU and LAMMPs interfaces to run large scale dynamics!](#multi-gpu-inference-and-lammps)
 
 ## Read our latest release post!
@@ -87,14 +89,17 @@ Models are referenced by their name, below are the currently supported models:
 
 | Model Name | Description |
 |---|---|
-| uma-s-1p1 | Latest version of the UMA small model, fastest of the UMA models while still SOTA on most benchmarks (6.6M/150M active/total params) |
+| uma-s-1p2 | Latest version of the UMA small model, fastest of the UMA models while still SOTA on most benchmarks (6.6M/290M active/total params) |
+| uma-s-1p1 | Early version of the UMA small model while still SOTA on most benchmarks (6.6M/150M active/total params) |
 | uma-m-1p1 | Best in class UMA model across all metrics, but slower and more memory intensive than uma-s (50M/1.4B active/total params) |
 
 ### Set the task for your application and calculate
 
 - **oc20:** use this for catalysis
+- **oc22:** use this for oxide catalysis (1p2 only)
+- **oc25:** use this for (electro)catalysis (1p2 only)
 - **omat:** use this for inorganic materials
-- **omol:** use this for molecules
+- **omol:** use this for molecules+polymers
 - **odac:** use this for MOFs
 - **omc:** use this for molecular crystals
 
@@ -104,7 +109,7 @@ from ase.build import fcc100, add_adsorbate, molecule
 from ase.optimize import LBFGS
 from fairchem.core import pretrained_mlip, FAIRChemCalculator
 
-predictor = pretrained_mlip.get_predict_unit("uma-s-1p1", device="cuda")
+predictor = pretrained_mlip.get_predict_unit("uma-s-1p2", device="cuda")
 calc = FAIRChemCalculator(predictor, task_name="oc20")
 
 # Set up your system as an ASE atoms object
@@ -126,7 +131,7 @@ from ase.optimize import FIRE
 from ase.filters import FrechetCellFilter
 from fairchem.core import pretrained_mlip, FAIRChemCalculator
 
-predictor = pretrained_mlip.get_predict_unit("uma-s-1p1", device="cuda")
+predictor = pretrained_mlip.get_predict_unit("uma-s-1p2", device="cuda")
 calc = FAIRChemCalculator(predictor, task_name="omat")
 
 atoms = bulk("Fe")
@@ -144,7 +149,7 @@ from ase.md.langevin import Langevin
 from ase.build import molecule
 from fairchem.core import pretrained_mlip, FAIRChemCalculator
 
-predictor = pretrained_mlip.get_predict_unit("uma-s-1p1", device="cuda")
+predictor = pretrained_mlip.get_predict_unit("uma-s-1p2", device="cuda")
 calc = FAIRChemCalculator(predictor, task_name="omol")
 
 atoms = molecule("H2O")
@@ -166,7 +171,7 @@ dyn.run(steps=1000)
 from ase.build import molecule
 from fairchem.core import pretrained_mlip, FAIRChemCalculator
 
-predictor = pretrained_mlip.get_predict_unit("uma-s-1p1", device="cuda")
+predictor = pretrained_mlip.get_predict_unit("uma-s-1p2", device="cuda")
 
 #  singlet CH2
 singlet = molecule("CH2_s1A1d")
@@ -196,7 +201,7 @@ import time
 from fairchem.core.datasets.common_structures import get_fcc_crystal_by_num_atoms
 
 predictor = pretrained_mlip.get_predict_unit(
-    "uma-s-1p1", inference_settings="turbo", device="cuda", workers=8
+    "uma-s-1p2", inference_settings="turbo", device="cuda", workers=8
 )
 calc = FAIRChemCalculator(predictor, task_name="omat")
 

configs/uma/benchmark/uma-speed.yaml

@@ -0,0 +1,5 @@
+device_type: CUDA
+scheduler:
+  mode: LOCAL
+  ranks_per_node: 1
+run_dir: /checkpoint/ocp/rgao/speed

configs/uma/speed/job/local.yaml

@@ -0,0 +1,6 @@
+device_type: CUDA
+scheduler:
+  mode: LOCAL
+  ranks_per_node: 8
+graph_parallel_group_size: 8
+run_dir: /checkpoint/ocp/rgao/speed

configs/uma/speed/job/local_8gpu.yaml

@@ -0,0 +1,13 @@
+# SLURM 64 GPU job configuration
+run_dir: /checkpoint/ocp/rgao/speed
+run_name: test_speed
+scheduler:
+  mode: SLURM
+  num_nodes: 8
+  ranks_per_node: 8
+  slurm:
+    account: ocp
+    qos: h200_ocp_high
+    mem_gb: 0
+    cpus_per_task: 24
+graph_parallel_group_size: 64

configs/uma/speed/job/slurm.yaml

@@ -0,0 +1,35 @@
+# Usage:
+#   fairchem -c configs/uma/speed/uma-speed.yaml                    # default: local
+#   fairchem -c configs/uma/speed/uma-speed.yaml job=local_8gpu     # 8 GPU local
+#   fairchem -c configs/uma/speed/uma-speed.yaml job=slurm    # SLURM
+
+defaults:
+  - job: local
+  - _self_
+
+uma_s_1p2:
+  _target_: fairchem.core.calculate.pretrained_mlip.pretrained_checkpoint_path_from_name
+  model_name: "uma-s-1p2"
+
+runner:
+  _target_: fairchem.core.components.benchmark._single.uma_speed_benchmark.InferenceBenchRunner
+  device: "cuda"
+  timeiters: 100
+  repeats: 5
+  natoms_list: [1000]
+  model_checkpoints:
+    # uma_s_1p1: ${uma_s_1p1}
+    uma_s_1p2: ${uma_s_1p2}
+  # for large number atoms (ie: 1000+) turn off compile and turn on activation_checkpointing
+  # for profiling, we substract the graph generation time by using external_graph_gen=True
+  inference_settings:
+    _target_: fairchem.core.units.mlip_unit.api.inference.InferenceSettings
+    tf32: True
+    activation_checkpointing: False
+    merge_mole: True
+    compile: True
+    external_graph_gen: False
+    internal_graph_gen_version: 2
+    # edge_chunk_size: 32
+    execution_mode: "umas_fast_gpu"
+  # generate_traces: True

configs/uma/speed/uma-speed.yaml

@@ -83,6 +83,11 @@ The advanced user might quickly see that **default** mode and **turbo** mode are
 | merge_mole | This is useful in long rollout applications where the system composition stays constant. By pre-merge the MoLE weights, we can save both memory and compute. |
 | compile | This uses torch.compile to significantly speed up computation. Due to the way pytorch traces the internal graph, it requires a long compile time during the first iteration and can even recompile anytime it detected a significant change in input dimensions. It is not recommended if you are computing frequently on very different atomic systems. |
 | external_graph_gen | Only use this if you want to use an external graph generator. This should be rarely used except for development |
+| internal_graph_gen_version | currently we support v2[default], an internal implementation that is better suited for parallelism and v3 the neighborlist from Nvidia Alchemi library which is faster for single gpu operations. |
+| edge_chunk_size | Experimental. Used for padding edge sizes. This helps reduce re-compilations from torch compile, default to None |
+| use_quaternion_wigner | enable quaternion-based Wigner D matrix computation. If false we fall back to euler-angle based rotations. default True. |
+| base_precision_dtype | governs the main precision type of the computation, default to FP32, FP64 is also supported |
+| execution_mode | This allows manually toggling custom backends to maximize speed ups. default to "general". "umas-fast-gpu" will introduce 30-40% speedup for uma-s line of models. |
 
 For example, for an MD simulation use-case for a system of ~500 atoms, we can choose to use a custom mode like the following:
 

docs/core/common_tasks/ase_calculator.md

@@ -49,7 +49,7 @@ We separate the LAMMPS integration code into a standalone package (`fairchem-lam
 :::
 
 ```bash
-# first install conda and lammps following the instructions above
+# first install conda and lammps following the instructions above, ie: conda install lammps
 # then activate the environment and install fairchem
 conda activate lammps-env
 pip install fairchem-core[extras]
@@ -60,7 +60,7 @@ Assuming you have a classic LAMMPS .in script, make the following changes:
 
 1. Remove all other forces from your LAMMPS script (e.g., pair_style, etc.)
 2. Make sure the units are in "metal"
-3. Make sure there is only 1 run command at the bottom of the script
+3. Make sure there is only 1 run command at the bottom of the script, if you have multiple run segments, ie: NVT followed by NPT, you can separate them into separate scripts
 
 To run, use the Python entrypoint `lmp_fc` (shortcut name for the [python lammps_fc.py script](https://github.com/facebookresearch/fairchem/pull/1454)):
 

docs/core/common_tasks/lammps.md

@@ -17,7 +17,6 @@ kernelspec:
 
 ![UMA model architecture](uma.svg "UMA model architecture")
 
-
 ## The UMA Mixture-of-Linear-Experts routing function
 
 :::{note}
@@ -62,6 +61,10 @@ Choose your task based on the application domain. Each task corresponds to a spe
 | **omat** | [OMat24](https://arxiv.org/abs/2410.12771) | PBE/PBE+U as implemented in VASP using Materials Project suggested settings, except with VASP 54 pseudopotentials. No dispersion. | Inorganic materials discovery, solar photovoltaics, advanced alloys, superconductors, electronic materials, optical materials | UMA has not seen varying charge or spin multiplicity for the OMat task, and expects total_charge=0 and spin multiplicity=0 as model inputs. Spin polarization effects are included, but you can't select the magnetic state. Further, OMat24 did not fully sample possible spin states in the training data. |
 | **oc20** | [OC20*](https://arxiv.org/abs/2010.09990) | RPBE as implemented in VASP, with VASP5.4 pseudopotentials. No dispersion. | Renewable energy, catalysis, fuel cells, energy conversion, sustainable fertilizer production, chemical refining, plastics synthesis/upcycling | UMA has not seen varying charge or spin multiplicity for the OC20 task, and expects total_charge=0 and spin multiplicity=0 as model inputs. No oxides or explicit solvents are included in OC20. The model works surprisingly well for transition state searches given the nature of the training data, but you should be careful. RPBE works well for small molecules, but dispersion will be important for larger molecules on surfaces. |
 | **odac** | [ODAC23](https://arxiv.org/abs/2311.00341) | PBE+D3 as implemented in VASP, with VASP5.4 pseudopotentials. | Direct air capture, carbon capture and storage, CO2 conversion, catalysis | UMA has not seen varying charge or spin multiplicity for the ODAC task, and expects total_charge=0 and spin multiplicity=0 as model inputs. The ODAC23 dataset only contains CO2/H2O water absorption, so anything more than might be inaccurate (e.g. hydrocarbons in MOFs). Further, there is a limited number of bare-MOF structures in the training data, so you should be careful if you are using a new MOF structure. |
+| **oc25** | [OC25](https://arxiv.org/abs/2509.17862) | RPBE+D3 as implemented in VASP, with VASP6.4 pseudopotentials, and dipole corrections.
+ | Renewable energy, (electro)catalysis, fuel cells, energy conversion, sustainable fertilizer production, chemical refining, plastics synthesis/upcycling | UMA has not seen varying charge or spin multiplicity for the OC25 task, and expects total_charge=0 and spin multiplicity=0 as model inputs. The model works surprisingly well for charged systems despite not being explicitly provided that information, but one should be careful. Work functions are not provided by UMA, subsequent DFT calculations are required to extract such information, if desired. Only available in UMA-1.2 |
+ | **oc22** | [OC22](https://arxiv.org/abs/2206.08917) | PBE+U as implemented in VASP with VASP5.4 pseudopotentials and spin-polarization
+ | Direct air capture, carbon capture and storage, CO2 conversion, catalysis | UMA has not seen varying charge or spin multiplicity for the OC22 task, and expects total_charge=0 and spin multiplicity=0 as model inputs. No explicit solvents are included in OC22. Only available in UMA-1.2 |
 
 :::{note}
 *OC20 was updated from the original OC20 and recomputed to produce total energies instead of adsorption energies.

docs/core/uma.md

@@ -214,61 +214,28 @@ def calculate(
         # Standard call to check system_changes etc
         Calculator.calculate(self, atoms, properties, system_changes)
 
-        if len(atoms) == 1 and sum(atoms.pbc) == 0:
-            self.results = self._get_single_atom_energies(atoms)
-        else:
-            # Convert using the current a2g object
-            data = self.a2g(atoms)
+        # Convert using the current a2g object
+        data = self.a2g(atoms)
 
-            # Batch and predict
-            pred = self.predictor.predict(data)
+        # Batch and predict
+        pred = self.predictor.predict(data)
 
-            # Collect the results into self.results
-            self.results = {}
-            for calc_key in self.implemented_properties:
-                if calc_key == "energy":
-                    energy = float(pred[calc_key].detach().cpu().numpy()[0])
+        # Collect the results into self.results
+        self.results = {}
+        for calc_key in self.implemented_properties:
+            if calc_key == "energy":
+                energy = float(pred[calc_key].detach().cpu().numpy()[0])
 
-                    self.results["energy"] = self.results["free_energy"] = (
-                        energy  # Free energy is a copy of energy
-                    )
-                if calc_key == "forces":
-                    forces = pred[calc_key].detach().cpu().numpy()
-                    self.results["forces"] = forces
-                if calc_key == "stress":
-                    stress = pred[calc_key].detach().cpu().numpy().reshape(3, 3)
-                    stress_voigt = full_3x3_to_voigt_6_stress(stress)
-                    self.results["stress"] = stress_voigt
-
-    def _get_single_atom_energies(self, atoms) -> dict:
-        """
-        Populate output with single atom energies
-        """
-        if self.predictor.atom_refs is None:
-            raise ValueError(
-                "Single atom system but no atomic references present. "
-                "Please call fairchem.core.pretrained_mlip.get_predict_unit() "
-                "with an appropriate checkpoint name."
-            )
-        logging.warning(
-            "Single atom systems are not handled by the model; "
-            "the precomputed DFT result is returned. "
-            "Spin multiplicity is ignored for monoatomic systems."
-        )
-        elt = atoms.get_atomic_numbers()[0]
-        results = {}
-
-        atom_refs = self.predictor.atom_refs[self.task_name]
-        try:
-            energy = atom_refs.get(int(elt), {}).get(atoms.info["charge"])
-        except AttributeError:
-            energy = atom_refs[int(elt)]
-        if energy is None:
-            raise ValueError("This model has not stored this element with this charge.")
-        results["energy"] = energy
-        results["forces"] = np.array([[0.0] * 3])
-        results["stress"] = np.array([0.0] * 6)
-        return results
+                self.results["energy"] = self.results["free_energy"] = (
+                    energy  # Free energy is a copy of energy
+                )
+            if calc_key == "forces":
+                forces = pred[calc_key].detach().cpu().numpy()
+                self.results["forces"] = forces
+            if calc_key == "stress":
+                stress = pred[calc_key].detach().cpu().numpy().reshape(3, 3)
+                stress_voigt = full_3x3_to_voigt_6_stress(stress)
+                self.results["stress"] = stress_voigt
 
     def _check_atoms_pbc(self, atoms) -> None:
         """