Enable Untrained Property Predictions Including Hessians (#1811)

## Summary

This PR extends FAIRChem's inference capabilities to compute derivative
properties (forces, stress, hessians) for
models trained only on energy, eliminating the need to retrain models
when derivative properties are needed.

Adds infrastructure for computing "untrained" properties—derivative
quantities that can be calculated via automatic
differentiation from trained energy predictions. Users can now request
forces, stress, and hessians at inference
time even if these properties were not included in the original training
task.

  ## Key Changes

  1. New Output Computation Modules (+463 lines)

src/fairchem/core/models/uma/outputs.py - UMA-specific gradient
computations:
- get_displacement_and_cell() - Prepares strain tensors for stress
computation
  - compute_energy() - Reduces node energies to system level
  - compute_forces() - Computes forces as -∇E/∇pos
  - compute_forces_and_stress() - Computes forces and stress from virial
- compute_hessian() - Second-order derivatives with vmap and loop
implementations
  - Helper functions: get_l_component(), reduce_node_to_system()

  src/fairchem/core/models/utils/outputs.py - General utilities:
- get_numerical_hessian() - Finite difference hessian computation for
validation

  2. Enhanced Inference Settings & Predict Unit (+240 lines)

  src/fairchem/core/units/mlip_unit/predict.py:
- _create_untrained_tasks() - Dynamically generates Task objects for
untrained properties
- _configure_head_gradients() - Enables autograd in model heads for
requested derivatives
- _validate_untrained_property_requests() - Validates inference settings
  - Supports per-dataset property requests via InferenceSettings:
- compute_untrained_forces: set[str] - Dataset names requiring forces
- compute_untrained_stress: set[str] - Dataset names requiring stress
- compute_untrained_hessian: set[str] - Dataset names requiring hessian

  3. Model Task Management (+72 lines)

  src/fairchem/core/models/base.py:
  - Added add_tasks() method to HydraModel and HydraModelV2
  - Allows runtime addition of inference-only tasks
  - Converted tasks to property with backing _tasks attribute
  - Rebuilds dataset_to_tasks map when tasks are added

  4. ESCN-MD Refactoring (+381/-209 lines)

  src/fairchem/core/models/uma/escn_md.py:
  - Refactored to use centralized gradient computation from outputs.py
- Cleaner separation between energy prediction and derivative
computation
  - Maintains backward compatibility
  
  ## Usage Example

```python
from fairchem.core.units.mlip_unit import InferenceSettings, MLIPPredictUnit

# Load energy-only checkpoint and enable untrained derivatives
settings = InferenceSettings(
    predict_untrained_forces={"omol"},
    predict_untrained_stress={"omol"},
    predict_untrained_hessian={"omol"}
)

predictor = MLIPPredictUnit(
    "uma-s.pt",
    device="cuda",
    inference_settings=settings
)

# Predictions now include all requested properties
preds = predictor.predict(batch)
# preds = {"energy": ..., "forces": ..., "stress": ..., "hessian": ...}
```

---------

Co-authored-by: Eric Yuan <87563575+ericyuan00000@users.noreply.github.com>
Co-authored-by: zulissimeta <122578103+zulissimeta@users.noreply.github.com>
Co-authored-by: ericyuan00000 <ericyuan@berkeley.edu>
Co-authored-by: misko <misko@meta.com>

Author: 4 people

docs/core/common_tasks/ase_calculator.md

@@ -108,6 +108,29 @@ predictor = pretrained_mlip.get_predict_unit(
 )
 ```
 
+## Enabling gradient stress or Hessian prediction
+
+Some tasks, for example omol, odac, or oc20/25, were not trained using stress labels. Similarly, no tasks were supervised to predict Hessians. However, predictions of untrained derivatives of energy, such as stress and Hessians, can be enabled by using the following inference settings flags,
+
+| Setting Flag  | Description |
+| ----- | ----- |
+| predict_untrained_forces | A set of task/dataset names (e.g., `{"omol", "oc20"}`) for which forces will be computed via autograd even though the checkpoint was not trained with a forces head for those tasks. |
+| predict_untrained_stress | A set of task/dataset names for which stress tensors will be computed via autograd even though the checkpoint was not trained with a stress head for those tasks. The default empty set disables this. |
+| predict_untrained_hessian | A set of task/dataset names for which the Hessian matrix will be computed via autograd. |
+
+For example, to enable stress and Hessian predictions with `omol` level of theory, the following settings can be used,
+
+```{code-cell} python3
+settings = InferenceSettings(
+    predict_untrained_stress={'omol'},
+    predict_untrained_hessian={'omol'}
+)
+
+predictor = pretrained_mlip.get_predict_unit(
+    "uma-s-1p1", device="cuda", inference_settings=settings
+)
+```
+
 ## Multi-GPU Inference
 
 UMA supports Graph Parallel inference natively. The graph is chunked into each rank and both the forward and backwards communication is handled by the built-in graph parallel algorithm with torch distributed. Because Multi-GPU inference requires special setup of communication protocols within a node and across nodes, we leverage [ray](https://www.ray.io/) to launch Ray Actors for each GPU-rank under the hood. This allows us to seamlessly scale to any infrastructure that can run Ray.

src/fairchem/core/calculate/__init__.py

@@ -12,5 +12,11 @@
     FAIRChemCalculator,
     FormationEnergyCalculator,
 )
+from fairchem.core.units.mlip_unit.api.inference import InferenceSettings
 
-__all__ = ["FAIRChemCalculator", "FormationEnergyCalculator", "InferenceBatcher"]
+__all__ = [
+    "FAIRChemCalculator",
+    "FormationEnergyCalculator",
+    "InferenceBatcher",
+    "InferenceSettings",
+]

src/fairchem/core/calculate/ase_calculator.py

@@ -14,7 +14,7 @@
 from typing import TYPE_CHECKING, Literal
 
 import numpy as np
-from ase.calculators.calculator import Calculator
+from ase.calculators.calculator import Calculator, PropertyNotImplementedError
 from ase.stress import full_3x3_to_voigt_6_stress
 
 from fairchem.core.calculate import pretrained_mlip
@@ -181,6 +181,22 @@ def check_state(self, atoms: Atoms, tol: float = 1e-15) -> list:
             state.append("info")
         return state
 
+    def get_property(self, name, atoms=None, allow_calculation=True):
+        try:
+            result = super().get_property(
+                name, atoms=atoms, allow_calculation=allow_calculation
+            )
+        except PropertyNotImplementedError as exc:
+            msg = str(exc)
+            if name in ("forces", "stress", "hessian"):
+                msg += (
+                    f"\n {name} prediction can be enabled by setting `predict_untrained_{name}=set('{self.task_name}')` "
+                    f"in the InferenceSettings."
+                )
+            raise PropertyNotImplementedError(msg) from exc
+
+        return result
+
     def calculate(
         self, atoms: Atoms, properties: list[str], system_changes: list[str]
     ) -> None:
@@ -236,6 +252,9 @@ def calculate(
                 stress = pred[calc_key].detach().cpu().numpy().reshape(3, 3)
                 stress_voigt = full_3x3_to_voigt_6_stress(stress)
                 self.results["stress"] = stress_voigt
+            if calc_key == "hessian":
+                hessian = pred[calc_key].detach().cpu().numpy().squeeze()
+                self.results["hessian"] = hessian
 
     def _check_atoms_pbc(self, atoms) -> None:
         """

src/fairchem/core/models/base.py

@@ -147,6 +147,8 @@ def __init__(
         # the old config system at some point, this will prevent the need to make major modifications to the trainer
         # because they all expect the name of the outputs directly instead of the head_name.property_name
         self.pass_through_head_outputs = pass_through_head_outputs
+        self._tasks = None
+        self._dataset_to_tasks = None
 
         # Does this model support inference on single atom systems
         self.supports_single_atoms = supports_single_atoms
@@ -255,6 +257,13 @@ def forward(self, data: AtomicData):
 
         return out
 
+    @property
+    def tasks(self) -> dict[str, Task]:
+        """
+        Mapping from task names to their associated Task objects.
+        """
+        return self._tasks
+
     @property
     def direct_forces(self) -> bool:
         """
@@ -317,18 +326,46 @@ def setup_tasks(self, tasks_config: list) -> None:
             tasks_config: List of task configurations from checkpoint
         """
         tasks = [hydra.utils.instantiate(task_config) for task_config in tasks_config]
-        self.tasks = {t.name: t for t in tasks}
+        self._tasks = {t.name: t for t in tasks}
         self._dataset_to_tasks = _get_dataset_to_tasks_map(tasks)
 
         # Let backbone validate tasks
         self.backbone.validate_tasks(self._dataset_to_tasks)
 
+    def add_tasks(self, tasks: Sequence[Task]) -> None:
+        """
+        Add additional tasks to the model.
+
+        This is useful for adding inference-only tasks that weren't in the
+        original checkpoint, such as untrained derivative properties.
+
+        Args:
+            tasks: List of Task objects to add
+        """
+        if not hasattr(self, "tasks"):
+            raise RuntimeError("setup_tasks() must be called before add_tasks()")
+
+        # Add new tasks to the tasks dict
+        for task in tasks:
+            if task.name in self.tasks:
+                logging.warning(
+                    f"Task '{task.name}' already exists, skipping adding as a new task."
+                )
+                continue
+            self.tasks[task.name] = task
+
+        # Rebuild dataset_to_tasks map
+        self._dataset_to_tasks = _get_dataset_to_tasks_map(self.tasks.values())
+
+        # Let backbone validate the updated task set
+        self.backbone.validate_tasks(self._dataset_to_tasks)
+
     @property
     def dataset_to_tasks(self) -> dict[str, list]:
         """
         Mapping from dataset names to their associated tasks.
         """
-        if not hasattr(self, "_dataset_to_tasks"):
+        if self._dataset_to_tasks is None:
             raise RuntimeError(
                 "setup_tasks() must be called before accessing dataset_to_tasks"
             )
@@ -346,6 +383,9 @@ def __init__(
         self.backbone = backbone
         self.output_heads = torch.nn.ModuleDict(heads)
         self.device = None
+        self._tasks = None
+        self._dataset_to_tasks = None
+
         if freeze_backbone:
             for param in self.backbone.parameters():
                 param.requires_grad = False
@@ -371,6 +411,13 @@ def forward(self, data):
                 out[k] = self.output_heads[k](data, emb)
         return out
 
+    @property
+    def tasks(self) -> dict[str, Task]:
+        """
+        Mapping from task names to their associated Task objects.
+        """
+        return self._tasks
+
     @property
     def direct_forces(self) -> bool:
         """
@@ -430,18 +477,41 @@ def setup_tasks(self, tasks_config: list) -> None:
             tasks_config: List of task configurations from checkpoint
         """
         tasks = [hydra.utils.instantiate(task_config) for task_config in tasks_config]
-        self.tasks = {t.name: t for t in tasks}
+        self._tasks = {t.name: t for t in tasks}
         self._dataset_to_tasks = _get_dataset_to_tasks_map(tasks)
 
         # Let backbone validate tasks
         self.backbone.validate_tasks(self._dataset_to_tasks)
 
+    def add_tasks(self, tasks: Sequence[Task]) -> None:
+        """
+        Add additional tasks to the model.
+
+        This is useful for adding inference-only tasks that weren't in the
+        original checkpoint, such as untrained derivative properties.
+
+        Args:
+            tasks: List of Task objects to add
+        """
+        # Add new tasks to the tasks dict
+        for task in tasks:
+            if task.name in self.tasks:
+                logging.warning(f"Task '{task.name}' already exists, skipping addition")
+                continue
+            self._tasks[task.name] = task
+
+        # Rebuild dataset_to_tasks map
+        self._dataset_to_tasks = _get_dataset_to_tasks_map(self.tasks.values())
+
+        # Let backbone validate the updated task set
+        self.backbone.validate_tasks(self._dataset_to_tasks)
+
     @property
     def dataset_to_tasks(self) -> dict[str, list]:
         """
         Mapping from dataset names to their associated tasks.
         """
-        if not hasattr(self, "_dataset_to_tasks"):
+        if self._dataset_to_tasks is None:
             raise RuntimeError(
                 "setup_tasks() must be called before accessing dataset_to_tasks"
             )

src/fairchem/core/models/uma/escn_md.py

@@ -1100,14 +1100,6 @@ def regress_forces(self) -> bool:
     def regress_stress(self) -> bool:
         return self.regress_config.stress
 
-    @property
-    def regress_hessian(self) -> bool:
-        return self.regress_config.hessian
-
-    @property
-    def hessian_vmap(self) -> bool:
-        return self.regress_config.hessian_vmap
-
     @conditional_grad(torch.enable_grad())
     def forward(
         self, data: AtomicData, emb: dict[str, torch.Tensor]
@@ -1139,7 +1131,7 @@ def forward(
 
         # Determine if we need create_graph for higher-order derivatives
         # Hessian computation requires second derivatives, so we need create_graph=True
-        create_graph = self.training or self.regress_hessian
+        create_graph = self.training or self.regress_config.hessian
 
         if self.regress_config.stress and not self.regress_config.direct_stress:
             forces, stress = compute_forces_and_stress(
@@ -1158,7 +1150,7 @@ def forward(
         else:
             forces = None
 
-        if self.regress_hessian:
+        if self.regress_config.hessian:
             if forces is None:
                 raise ValueError(
                     "Hessian computation requires forces. "
@@ -1171,7 +1163,10 @@ def forward(
                 )
 
             hessian = compute_hessian(
-                forces, data["pos"], vmap=self.hessian_vmap, training=self.training
+                forces,
+                data["pos"],
+                vmap=self.regress_config.hessian_vmap,
+                training=create_graph,
             )
             outputs[hessian_key] = (
                 {"hessian": hessian} if self.wrap_property else hessian

src/fairchem/core/models/uma/outputs.py

@@ -244,11 +244,13 @@ def compute_grad_component(vec):
         )[0]
 
     # Use vmap to compute all components in parallel
+    # autograd.grad returns shape [N, 3] (same as pos), vmap gives [N*3, N, 3]
     hessian = torch.vmap(compute_grad_component)(
-        torch.eye(forces_flat.numel(), device=forces_flat.device)
+        torch.eye(forces_flat.shape[0], device=forces_flat.device)
     )
 
-    return hessian
+    n_dof = forces_flat.numel()
+    return hessian.reshape(n_dof, n_dof)
 
 
 def compute_hessian_loop(
@@ -306,7 +308,8 @@ def compute_hessian(
         training: Whether to create graph for third-order derivatives.
 
     Returns:
-        Hessian matrix of shape [N*3, N*3].
+        Hessian matrix of shape [1, N*3, N*3] (batch dim always 1 since
+        hessian requires single-system batches).
 
     Note:
         Graph parallel (GP) mode is not fully supported. The Hessian should
@@ -325,4 +328,4 @@ def compute_hessian(
     else:
         hessian = compute_hessian_loop(forces_flat, pos, create_graph=training)
 
-    return hessian
+    return hessian.unsqueeze(0)

src/fairchem/core/units/mlip_unit/api/inference.py

@@ -7,7 +7,7 @@
 
 from __future__ import annotations
 
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, field
 
 import torch  # - needed at runtime for dataclass field type resolution
 
@@ -106,6 +106,17 @@ class InferenceSettings:
     # Set to "umas_fast_gpu" to enable highly optimized backend with triton kernels for maximum speed.
     execution_mode: str = "general"
 
+    # New fields for untrained derivative properties
+    # These flags request computation of properties NOT in the checkpoint's task list.
+    # If a property is already in the checkpoint (e.g., omol_forces task exists),
+    # it will be computed regardless of these flags.
+    # Specify datasets as a set of strings (e.g., {"omol", "oc20"}).
+    # Empty set means no untrained properties will be computed (default).
+    predict_untrained_forces: set[str] = field(default_factory=set)
+    predict_untrained_stress: set[str] = field(default_factory=set)
+    predict_untrained_hessian: set[str] = field(default_factory=set)
+    hessian_vmap: bool = True  # Use fast vmap vs memory-efficient loop
+
     def __post_init__(self):
         if isinstance(self.base_precision_dtype, str):
             self.base_precision_dtype = getattr(torch, self.base_precision_dtype)