Use FP16 operands for UMA radial FC2

docs/core/common_tasks/ase_calculator.md

@@ -88,6 +88,7 @@ The advanced user might quickly see that **default**, **batch**, and **turbo** m
 
 | Setting Flag  | Description |
 | ----- | ----- |
+| `fp16_radial_fc2_blocks` | Optional UMA block indices whose second unified radial projection uses FP16 operands with an FP32 output. Its frozen backward similarly uses FP16 operands and returns an FP32 input gradient. This requires `execution_mode="umas_fast_gpu"` and FP32 base precision, and does not support Hessians. Validate energy, forces, and stress for the selected model and structures. The default empty selection preserves FP32 radial behavior. After manually changing `fc2_weight`, call `refresh_fp16_cache()` and recapture any CUDA graph if the refresh changes a cache pointer. |
 | tf32 | enables torch [tf32](https://docs.pytorch.org/docs/stable/notes/cuda.html) format for matrix multiplication. This will speed up inference at a slight trade-off for precision. In our tests, it makes minimal difference to most applications. It is able to preserve equivariance, energy conservation for long rollouts. However, if you are computing higher order derivatives such as Hessians or other calculations that requires strict numerical precision, we recommend turning this off |
 | activation_checkpointing | this uses a custom chunked activation checkpointing algorithm and allows significant savings in memory for a small inference speed penalty. If you are predicting on systems >1000 atoms, we recommend keeping this on. However, if you want the absolute fastest inference possible for small systems, you can turn this off |
 | 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. |

src/fairchem/core/models/uma/nn/execution_backends.py

@@ -284,6 +284,12 @@ def validate(
             raise ValueError(
                 "UMASFastPytorchBackend requires activation_checkpointing=False"
             )
+        if (
+            settings is not None
+            and settings.fp16_radial_fc2_blocks
+            and settings.predict_untrained_hessian
+        ):
+            raise ValueError("fp16_radial_fc2_blocks does not support Hessians")
 
     @staticmethod
     def prepare_model_for_inference(model: torch.nn.Module) -> None:
@@ -301,13 +307,23 @@ def prepare_model_for_inference(model: torch.nn.Module) -> None:
             convert_so2_conv2,
         )
 
+        settings = getattr(model, "_inference_settings", None)
+        fp16_radial_fc2_blocks = tuple(getattr(settings, "fp16_radial_fc2_blocks", ()))
+        if fp16_radial_fc2_blocks and model.regress_config.hessian:
+            raise ValueError("fp16_radial_fc2_blocks does not support Hessians")
+        if fp16_radial_fc2_blocks and max(fp16_radial_fc2_blocks) >= len(model.blocks):
+            raise ValueError(
+                "fp16_radial_fc2_blocks contains an index outside model.blocks"
+            )
+
         for block in model.blocks:
             block.edge_wise.so2_conv_1 = convert_so2_conv1(block.edge_wise.so2_conv_1)
             block.edge_wise.so2_conv_2 = convert_so2_conv2(block.edge_wise.so2_conv_2)
 
         # Create unified radial MLP for batched computation
         rad_funcs = [block.edge_wise.so2_conv_1.rad_func for block in model.blocks]
         model._unified_radial_mlp = UnifiedRadialMLP(rad_funcs)
+        model._unified_radial_mlp.configure_fp16_fc2(fp16_radial_fc2_blocks)
 
     @staticmethod
     def get_layer_radial_emb(

src/fairchem/core/models/uma/nn/unified_radial.py

@@ -21,6 +21,7 @@
 
 import torch
 import torch.nn as nn
+from torch.autograd.function import once_differentiable
 
 if TYPE_CHECKING:
     from .radial import RadialMLP
@@ -39,6 +40,20 @@
 )
 
 
+class _FrozenFP16RadialFC2Function(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, h, weight, bias):
+        ctx.save_for_backward(weight)
+        return torch.mm(h.half(), weight.T, out_dtype=torch.float32) + bias
+
+    @staticmethod
+    @once_differentiable
+    def backward(ctx, grad_output):
+        (weight,) = ctx.saved_tensors
+        grad_h = torch.mm(grad_output.half(), weight, out_dtype=torch.float32)
+        return grad_h, None, None
+
+
 def _validate_radial_mlp(mlp: RadialMLP, idx: int, reference: RadialMLP | None) -> None:
     """
     Validate a single RadialMLP has expected structure and matches reference.
@@ -146,6 +161,58 @@ def __init__(self, radial_mlps: list[RadialMLP]) -> None:
             "fc3_bias",
             torch.stack([mlp.net[6].bias.data for mlp in radial_mlps], dim=0),
         )
+        self.register_buffer("_fc2_weight_fp16", None, persistent=False)
+        self.fp16_fc2_blocks: tuple[int, ...] = ()
+
+    @staticmethod
+    def _update_buffer(current, value):
+        if (
+            current is not None
+            and current.shape == value.shape
+            and current.device == value.device
+            and current.dtype == value.dtype
+            and current.stride() == value.stride()
+        ):
+            current.copy_(value)
+            return current
+        return value
+
+    def configure_fp16_fc2(self, blocks: tuple[int, ...]) -> None:
+        blocks = tuple(blocks)
+        if len(set(blocks)) != len(blocks):
+            raise ValueError("fp16_radial_fc2_blocks must contain unique indices")
+        if any(type(index) is not int or index < 0 for index in blocks):
+            raise ValueError(
+                "fp16_radial_fc2_blocks must contain non-negative integers"
+            )
+        if blocks and max(blocks) >= self.num_layers:
+            raise ValueError(
+                "fp16_radial_fc2_blocks contains an index outside model.blocks"
+            )
+        self.fp16_fc2_blocks = blocks
+        self.refresh_fp16_cache()
+
+    def refresh_fp16_cache(self) -> None:
+        if not self.fp16_fc2_blocks:
+            self._fc2_weight_fp16 = None
+            return
+        self._fc2_weight_fp16 = self._update_buffer(
+            self._fc2_weight_fp16, self.fc2_weight.detach().half()
+        )
+
+    def _apply(self, fn, recurse=True):
+        result = super()._apply(fn, recurse)
+        self.refresh_fp16_cache()
+        return result
+
+    def _load_from_state_dict(self, *args, **kwargs):
+        super()._load_from_state_dict(*args, **kwargs)
+        self.refresh_fp16_cache()
+
+    def forward_fp16_fc2(self, h: torch.Tensor, i: int) -> torch.Tensor:
+        return _FrozenFP16RadialFC2Function.apply(
+            h, self._fc2_weight_fp16[i], self.fc2_bias[i]
+        )
 
     def umas_radial_mlp(self, h: torch.Tensor, i: int) -> torch.Tensor:
         """Apply layers 2+ (LN -> SiLU -> Linear -> LN -> SiLU -> Linear)."""
@@ -154,7 +221,10 @@ def umas_radial_mlp(self, h: torch.Tensor, i: int) -> torch.Tensor:
             h, (H,), self.ln1_weight[i], self.ln1_bias[i], self.ln_eps
         )
         h = torch.nn.functional.silu(h)
-        h = torch.nn.functional.linear(h, self.fc2_weight[i], self.fc2_bias[i])
+        if i in self.fp16_fc2_blocks:
+            h = self.forward_fp16_fc2(h, i)
+        else:
+            h = torch.nn.functional.linear(h, self.fc2_weight[i], self.fc2_bias[i])
         h = torch.nn.functional.layer_norm(
             h, (H,), self.ln2_weight[i], self.ln2_bias[i], self.ln_eps
         )

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

@@ -139,6 +139,9 @@ class InferenceSettings:
     # MLIPPredictUnit falls back to an unmerged model.
     merge_mole: bool = False
 
+    # SO2 block indices whose unified radial fc2 uses FP16 operands.
+    fp16_radial_fc2_blocks: tuple[int, ...] = ()
+
     # Flag to enable or disable the compilation of the inference model.
     compile: bool = False
 
@@ -207,6 +210,25 @@ def __post_init__(self):
             assert (
                 self.base_precision_dtype in ALLOWED_DTYPES
             ), f"base_precision_dtype must be one of {ALLOWED_DTYPES}, got {self.base_precision_dtype}"
+        self.fp16_radial_fc2_blocks = tuple(self.fp16_radial_fc2_blocks)
+        if len(set(self.fp16_radial_fc2_blocks)) != len(self.fp16_radial_fc2_blocks):
+            raise ValueError("fp16_radial_fc2_blocks must contain unique indices")
+        if any(
+            type(index) is not int or index < 0 for index in self.fp16_radial_fc2_blocks
+        ):
+            raise ValueError(
+                "fp16_radial_fc2_blocks must contain non-negative integers"
+            )
+        if self.fp16_radial_fc2_blocks and self.execution_mode != "umas_fast_gpu":
+            raise ValueError(
+                "fp16_radial_fc2_blocks requires execution_mode='umas_fast_gpu'"
+            )
+        if self.fp16_radial_fc2_blocks and self.predict_untrained_hessian:
+            raise ValueError("fp16_radial_fc2_blocks does not support Hessians")
+        if self.fp16_radial_fc2_blocks and self.base_precision_dtype != torch.float32:
+            raise ValueError(
+                "fp16_radial_fc2_blocks requires base_precision_dtype=torch.float32"
+            )
 
     def to_omegaconf(self) -> dict:
         """

tests/core/models/uma/nn/test_unified_radial.py

@@ -165,3 +165,79 @@ def test_unified_gradient_flow(self, radial_mlp_list):
         # Input should have gradients
         assert x_edge.grad is not None
         assert x_edge.grad.abs().sum() > 0
+
+    def test_fp16_fc2_cache_lifecycle(self, radial_mlp_list):
+        from fairchem.core.models.uma.nn.unified_radial import UnifiedRadialMLP
+
+        unified = UnifiedRadialMLP(radial_mlp_list)
+        original_keys = set(unified.state_dict())
+        unified.configure_fp16_fc2((0, 2))
+        assert set(unified.state_dict()) == original_keys
+        assert unified.fp16_fc2_blocks == (0, 2)
+        assert unified._fc2_weight_fp16.dtype is torch.float16
+        assert unified._fc2_weight_fp16.stride() == unified.fc2_weight.stride()
+        pointer = unified._fc2_weight_fp16.data_ptr()
+        state = {key: value.clone() for key, value in unified.state_dict().items()}
+        state["fc2_weight"].add_(1)
+        unified.load_state_dict(state)
+        assert unified._fc2_weight_fp16.data_ptr() == pointer
+        torch.testing.assert_close(unified._fc2_weight_fp16, state["fc2_weight"].half())
+
+        with pytest.raises(ValueError, match="outside model.blocks"):
+            unified.configure_fp16_fc2((len(radial_mlp_list),))
+        with pytest.raises(ValueError, match="unique indices"):
+            unified.configure_fp16_fc2((0, 0))
+        for blocks in ((-1,), (True,)):
+            with pytest.raises(ValueError, match="non-negative integers"):
+                unified.configure_fp16_fc2(blocks)
+
+    @pytest.mark.gpu()
+    @pytest.mark.compile_gpu()
+    def test_fp16_fc2_forward_and_vjp(self):
+        from fairchem.core.models.uma.nn.radial import RadialMLP
+        from fairchem.core.models.uma.nn.unified_radial import UnifiedRadialMLP
+
+        torch.manual_seed(42)
+        radial_mlps = [RadialMLP([64, 128, 128, 1536]) for _ in range(4)]
+        unified = UnifiedRadialMLP(radial_mlps).cuda()
+        unified.configure_fp16_fc2((0,))
+        h = torch.randn(17, 128, device="cuda", requires_grad=True)
+        grad_output = torch.randn_like(h)
+        expected = (
+            torch.mm(h.half(), unified.fc2_weight[0].half().T, out_dtype=torch.float32)
+            + unified.fc2_bias[0]
+        )
+        actual = unified.forward_fp16_fc2(h, 0)
+        expected_grad = torch.mm(
+            grad_output.half(),
+            unified.fc2_weight[0].half(),
+            out_dtype=torch.float32,
+        )
+        actual_grad = torch.autograd.grad(actual, h, grad_output)[0]
+        torch.testing.assert_close(actual, expected)
+        torch.testing.assert_close(actual_grad, expected_grad)
+
+        compiled = torch.compile(unified.forward_fp16_fc2, fullgraph=True)
+        warm_h = h.detach().clone().requires_grad_()
+        warm = compiled(warm_h, 0)
+        torch.autograd.grad(warm, warm_h, grad_output)
+
+        static_h = h.detach().clone().requires_grad_()
+        graph = torch.cuda.CUDAGraph()
+        torch.autograd.graph.set_override_stale_capture_stream(True)
+        try:
+            with torch.cuda.graph(graph):
+                captured = compiled(static_h, 0)
+                captured_grad = torch.autograd.grad(captured, static_h, grad_output)[0]
+        finally:
+            torch.autograd.graph.set_override_stale_capture_stream(False)
+        before = (captured.clone(), captured_grad.clone())
+        with torch.no_grad():
+            static_h.add_(0.125)
+        expected = unified.forward_fp16_fc2(static_h, 0)
+        expected_grad = torch.autograd.grad(expected, static_h, grad_output)[0]
+        graph.replay()
+        assert not torch.equal(captured, before[0])
+        assert not torch.equal(captured_grad, before[1])
+        torch.testing.assert_close(captured, expected)
+        torch.testing.assert_close(captured_grad, expected_grad)

tests/core/models/uma/uma_fast/test_execution_backends.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import os
+from types import SimpleNamespace
 
 import pytest
 import torch
@@ -62,6 +63,16 @@ def _mock_settings(
     )
 
 
+def test_fp16_radial_fc2_rejects_hessian_model():
+    model = torch.nn.Module()
+    model.regress_config = SimpleNamespace(hessian=True)
+    model._inference_settings = InferenceSettings(
+        execution_mode="umas_fast_gpu", fp16_radial_fc2_blocks=(0,)
+    )
+    with pytest.raises(ValueError, match="does not support Hessians"):
+        UMASFastGPUBackend.prepare_model_for_inference(model)
+
+
 @pytest.mark.gpu()
 def test_umas_fast_gpu_validation_requires_correct_lmax():
     """
@@ -679,35 +690,47 @@ def test_compiled_backends_match_baseline(pretrained_model_name, compile_reset_s
     )
     baseline_out = baseline_predictor.predict(batch.clone())
 
-    # Test configurations: (execution_mode, compile)
+    # Test configurations: (execution_mode, compile, FP16 radial FC2 blocks)
     test_configs = [
-        ("general", True),
-        ("umas_fast_gpu", True),
+        ("general", True, ()),
+        ("umas_fast_gpu", True, ()),
+        ("umas_fast_gpu", True, (0, 1, 2, 3)),
     ]
 
-    for test_mode, test_compile in test_configs:
+    for test_mode, test_compile, fp16_radial_fc2_blocks in test_configs:
         test_settings = InferenceSettings(
             activation_checkpointing=False,
             merge_mole=True,
             external_graph_gen=False,
             execution_mode=test_mode,
             compile=test_compile,
+            fp16_radial_fc2_blocks=fp16_radial_fc2_blocks,
         )
         test_predictor = MLIPPredictUnit(
             checkpoint_pt, "cuda", inference_settings=test_settings
         )
         test_out = test_predictor.predict(batch.clone())
+        if fp16_radial_fc2_blocks:
+            radial = test_predictor.model.module.backbone._unified_radial_mlp
+            assert radial.fp16_fc2_blocks == fp16_radial_fc2_blocks
+            assert radial._fc2_weight_fp16.dtype is torch.float16
 
         # Force comparison
         assert torch.allclose(
-            baseline_out["forces"], test_out["forces"], rtol=5e-4, atol=5e-5
+            baseline_out["forces"],
+            test_out["forces"],
+            rtol=(0 if fp16_radial_fc2_blocks else 5e-4),
+            atol=(1e-3 if fp16_radial_fc2_blocks else 5e-5),
         ), (
             f"{pretrained_model_name} {test_mode} compile={test_compile}: "
             f"force mismatch max diff = {(baseline_out['forces'] - test_out['forces']).abs().max()}"
         )
         # Energy comparison
         assert torch.allclose(
-            baseline_out["energy"], test_out["energy"], rtol=5e-4, atol=5e-5
+            baseline_out["energy"],
+            test_out["energy"],
+            rtol=(0 if fp16_radial_fc2_blocks else 5e-4),
+            atol=(3e-2 if fp16_radial_fc2_blocks else 5e-5),
         ), (
             f"{pretrained_model_name} {test_mode} compile={test_compile}: "
             f"energy mismatch {baseline_out['energy']} vs {test_out['energy']}"

tests/core/units/mlip_unit/test_inference_settings.py

@@ -71,6 +71,33 @@ def test_invalid_string_raises():
         InferenceSettings(base_precision_dtype="int8")
 
 
+def test_fp16_radial_fc2_blocks_validation():
+    with pytest.raises(ValueError, match="unique indices"):
+        InferenceSettings(fp16_radial_fc2_blocks=(0, 0))
+    with pytest.raises(ValueError, match="non-negative integers"):
+        InferenceSettings(fp16_radial_fc2_blocks=(-1,))
+    with pytest.raises(ValueError, match="execution_mode='umas_fast_gpu'"):
+        InferenceSettings(fp16_radial_fc2_blocks=(0,))
+    with pytest.raises(ValueError, match="base_precision_dtype=torch.float32"):
+        InferenceSettings(
+            base_precision_dtype=torch.float64,
+            execution_mode="umas_fast_gpu",
+            fp16_radial_fc2_blocks=(0,),
+        )
+    with pytest.raises(ValueError, match="does not support Hessians"):
+        InferenceSettings(
+            execution_mode="umas_fast_gpu",
+            fp16_radial_fc2_blocks=(0,),
+            predict_untrained_hessian={"omat"},
+        )
+
+    settings = InferenceSettings(
+        execution_mode="umas_fast_gpu",
+        fp16_radial_fc2_blocks=[0, 1, 2, 3],
+    )
+    assert settings.fp16_radial_fc2_blocks == (0, 1, 2, 3)
+
+
 # --- to_omegaconf ---
 
 
@@ -104,3 +131,14 @@ def test_to_omegaconf_roundtrip():
     assert isinstance(restored, InferenceSettings)
     assert restored.base_precision_dtype is torch.float64
     assert restored.tf32 is True
+
+
+def test_fp16_radial_fc2_blocks_omegaconf_roundtrip():
+    import hydra
+
+    original = InferenceSettings(
+        execution_mode="umas_fast_gpu",
+        fp16_radial_fc2_blocks=(0, 1, 2, 3),
+    )
+    restored = hydra.utils.instantiate(original.to_omegaconf())
+    assert restored.fp16_radial_fc2_blocks == (0, 1, 2, 3)