Reduce UMA radial input-gradient work

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

@@ -263,6 +263,18 @@ def edge_degree_scatter(
         )
 
 
+def _configure_radial_first_linear_prefix_grad(model, unified_radial_mlp) -> None:
+    if model.regress_config.hessian:
+        raise ValueError("radial_first_linear_prefix_grad does not support Hessians")
+    expected_input_features = model.num_distance_basis + 2 * model.edge_channels
+    model.edge_degree_embedding.rad_func.configure_first_linear_grad_prefix(
+        model.num_distance_basis, expected_input_features
+    )
+    unified_radial_mlp.configure_first_linear_grad_prefix(
+        model.num_distance_basis, expected_input_features
+    )
+
+
 class UMASFastPytorchBackend(ExecutionBackend):
     """
     Optimized PyTorch backend using block-diagonal SO2 convolutions.
@@ -308,6 +320,9 @@ def prepare_model_for_inference(model: torch.nn.Module) -> None:
         # 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)
+        settings = getattr(model, "_inference_settings", None)
+        if getattr(settings, "radial_first_linear_prefix_grad", False):
+            _configure_radial_first_linear_prefix_grad(model, model._unified_radial_mlp)
 
     @staticmethod
     def get_layer_radial_emb(

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

@@ -13,6 +13,26 @@
 import torch.nn as nn
 
 
+class _FrozenLinearInputPrefixFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, detached_inputs, input_prefix, weight, bias):
+        ctx.save_for_backward(weight)
+        ctx.prefix = input_prefix.shape[1]
+        return torch.nn.functional.linear(detached_inputs, weight, bias)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (weight,) = ctx.saved_tensors
+        grad_prefix = torch.mm(grad_output, weight[:, : ctx.prefix])
+        return None, grad_prefix, None, None
+
+
+def _frozen_linear_input_prefix(inputs, weight, bias, prefix):
+    return _FrozenLinearInputPrefixFunction.apply(
+        inputs.detach(), inputs[:, :prefix], weight, bias
+    )
+
+
 @torch.jit.script
 def gaussian(x: torch.Tensor, mean, std) -> torch.Tensor:
     a = (2 * math.pi) ** 0.5
@@ -81,6 +101,30 @@ def __init__(self, channels_list) -> None:
             modules.append(torch.nn.SiLU())
 
         self.net = nn.Sequential(*modules)
+        self.first_linear_grad_prefix: int | None = None
+
+    def configure_first_linear_grad_prefix(
+        self, prefix: int, expected_input_features: int
+    ) -> None:
+        first_linear = self.net[0]
+        if not isinstance(first_linear, nn.Linear):
+            raise TypeError("radial first layer must be Linear")
+        if first_linear.in_features != expected_input_features:
+            raise ValueError("radial first-linear input width does not match x_edge")
+        if not 0 < prefix < expected_input_features:
+            raise ValueError("prefix must be between zero and the input width")
+        self.first_linear_grad_prefix = prefix
 
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:
-        return self.net(inputs)
+        if self.first_linear_grad_prefix is None:
+            return self.net(inputs)
+        first_linear = self.net[0]
+        hidden = _frozen_linear_input_prefix(
+            inputs,
+            first_linear.weight,
+            first_linear.bias,
+            self.first_linear_grad_prefix,
+        )
+        for index in range(1, len(self.net)):
+            hidden = self.net[index](hidden)
+        return hidden

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

@@ -22,6 +22,8 @@
 import torch
 import torch.nn as nn
 
+from .radial import _frozen_linear_input_prefix
+
 if TYPE_CHECKING:
     from .radial import RadialMLP
 
@@ -146,6 +148,18 @@ 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.first_linear_grad_prefix: int | None = None
+
+    def configure_first_linear_grad_prefix(
+        self, prefix: int, expected_input_features: int
+    ) -> None:
+        if self.W1_cat.shape[1] != expected_input_features:
+            raise ValueError(
+                "UnifiedRadialMLP first-linear input width does not match x_edge"
+            )
+        if not 0 < prefix < expected_input_features:
+            raise ValueError("prefix must be between zero and the input width")
+        self.first_linear_grad_prefix = prefix
 
     def umas_radial_mlp(self, h: torch.Tensor, i: int) -> torch.Tensor:
         """Apply layers 2+ (LN -> SiLU -> Linear -> LN -> SiLU -> Linear)."""
@@ -172,7 +186,12 @@ def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
             List of N tensors, each of shape [E, out_features]
         """
         # Single batched GEMM for first layer, then split into per-layer chunks
-        h_all = torch.nn.functional.linear(x, self.W1_cat, self.b1_cat)
+        if self.first_linear_grad_prefix is None:
+            h_all = torch.nn.functional.linear(x, self.W1_cat, self.b1_cat)
+        else:
+            h_all = _frozen_linear_input_prefix(
+                x, self.W1_cat, self.b1_cat, self.first_linear_grad_prefix
+            )
         h_per_layer = h_all.split(self.hidden_features, dim=1)
         return [self.umas_radial_mlp(h_per_layer[i], i) for i in range(self.num_layers)]
 

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
 
+    # Restrict input gradients of the two radial first layers to distance features.
+    radial_first_linear_prefix_grad: bool = False
+
     # Flag to enable or disable the compilation of the inference model.
     compile: bool = False
 
@@ -202,6 +205,16 @@ class InferenceSettings:
     max_atoms: int | None = None
 
     def __post_init__(self):
+        if self.radial_first_linear_prefix_grad:
+            if self.execution_mode != "umas_fast_gpu":
+                raise ValueError(
+                    "radial_first_linear_prefix_grad requires "
+                    "execution_mode='umas_fast_gpu'"
+                )
+            if self.predict_untrained_hessian:
+                raise ValueError(
+                    "radial_first_linear_prefix_grad does not support Hessians"
+                )
         if isinstance(self.base_precision_dtype, str):
             self.base_precision_dtype = getattr(torch, self.base_precision_dtype)
             assert (

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

@@ -0,0 +1,158 @@
+"""Tests for restricting frozen radial input gradients to distance features."""
+
+from __future__ import annotations
+
+import copy
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from fairchem.core.models.uma.nn.execution_backends import (
+    _configure_radial_first_linear_prefix_grad,
+)
+from fairchem.core.models.uma.nn.radial import RadialMLP
+from fairchem.core.models.uma.nn.unified_radial import UnifiedRadialMLP
+
+
+def test_radial_first_linear_prefix_forward_and_vjp():
+    torch.manual_seed(0)
+    expected_model = RadialMLP([7, 5]).double()
+    actual_model = copy.deepcopy(expected_model)
+    actual_model.configure_first_linear_grad_prefix(3, 7)
+    reference_input = torch.randn(11, 7, dtype=torch.float64, requires_grad=True)
+    actual_input = reference_input.detach().clone().requires_grad_()
+    grad_output = torch.randn(11, 5, dtype=torch.float64)
+
+    expected = expected_model(reference_input)
+    actual = actual_model(actual_input)
+    expected_grad = torch.autograd.grad(expected, reference_input, grad_output)[0]
+    actual_grad, weight_grad, bias_grad = torch.autograd.grad(
+        actual,
+        (actual_input, actual_model.net[0].weight, actual_model.net[0].bias),
+        grad_output,
+        allow_unused=True,
+    )
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    torch.testing.assert_close(actual_grad[:, :3], expected_grad[:, :3])
+    torch.testing.assert_close(actual_grad[:, 3:], torch.zeros_like(actual_grad[:, 3:]))
+    assert weight_grad is None
+    assert bias_grad is None
+
+
+def test_radial_first_linear_prefix_double_backward():
+    torch.manual_seed(1)
+    expected_model = RadialMLP([7, 5]).double()
+    actual_model = copy.deepcopy(expected_model)
+    actual_model.configure_first_linear_grad_prefix(3, 7)
+    tail = torch.randn(4, dtype=torch.float64)
+
+    def expected(prefix):
+        inputs = torch.cat((prefix, tail)).unsqueeze(0)
+        return expected_model(inputs).sin().square().sum()
+
+    def actual(prefix):
+        inputs = torch.cat((prefix, tail)).unsqueeze(0)
+        return actual_model(inputs).sin().square().sum()
+
+    prefix = torch.randn(3, dtype=torch.float64)
+    expected_hessian = torch.autograd.functional.hessian(expected, prefix)
+    actual_hessian = torch.autograd.functional.hessian(actual, prefix)
+    torch.testing.assert_close(actual_hessian, expected_hessian, rtol=0, atol=0)
+
+
+def test_radial_first_linear_prefix_validation():
+    for prefix in (0, 7, 8):
+        with pytest.raises(ValueError, match="between zero and the input width"):
+            RadialMLP([7, 5]).configure_first_linear_grad_prefix(prefix, 7)
+
+
+def test_radial_first_linear_prefix_preserves_state_dict():
+    model = RadialMLP([7, 5, 5, 6])
+    state_dict_keys = tuple(model.state_dict())
+    weight = model.net[0].weight
+    model.configure_first_linear_grad_prefix(3, 7)
+    assert tuple(model.state_dict()) == state_dict_keys
+    assert model.net[0].weight is weight
+
+
+def test_unified_radial_prefix_grad_matches_full_input_grad():
+    torch.manual_seed(2)
+    radial_mlps = [RadialMLP([7, 5, 5, 6]), RadialMLP([7, 5, 5, 6])]
+    expected_model = UnifiedRadialMLP(radial_mlps).double()
+    actual_model = copy.deepcopy(expected_model)
+    actual_model.configure_first_linear_grad_prefix(3, 7)
+    expected_input = torch.randn(13, 7, dtype=torch.float64, requires_grad=True)
+    actual_input = expected_input.detach().clone().requires_grad_()
+
+    expected_outputs = expected_model(expected_input)
+    actual_outputs = actual_model(actual_input)
+    grads = [torch.randn_like(output) for output in expected_outputs]
+    expected_grad = torch.autograd.grad(expected_outputs, expected_input, grads)[0]
+    actual_grad = torch.autograd.grad(actual_outputs, actual_input, grads)[0]
+
+    for actual, expected in zip(actual_outputs, expected_outputs, strict=True):
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    torch.testing.assert_close(actual_grad[:, :3], expected_grad[:, :3])
+    torch.testing.assert_close(actual_grad[:, 3:], torch.zeros_like(actual_grad[:, 3:]))
+
+
+def _make_model(edge_input_features=7):
+    model = torch.nn.Module()
+    model.num_distance_basis = 3
+    model.edge_channels = 2
+    model.regress_config = SimpleNamespace(hessian=False)
+    model.edge_degree_embedding = torch.nn.Module()
+    model.edge_degree_embedding.rad_func = RadialMLP([edge_input_features, 5, 5, 6])
+    return model
+
+
+def test_configure_radial_first_linear_prefix_grad():
+    model = _make_model()
+    unified = UnifiedRadialMLP([RadialMLP([7, 5, 5, 6])])
+    _configure_radial_first_linear_prefix_grad(model, unified)
+    assert model.edge_degree_embedding.rad_func.first_linear_grad_prefix == 3
+    assert unified.first_linear_grad_prefix == 3
+
+
+def test_configure_radial_first_linear_prefix_grad_validates_widths():
+    with pytest.raises(ValueError, match="radial.*width"):
+        _configure_radial_first_linear_prefix_grad(
+            _make_model(edge_input_features=8),
+            UnifiedRadialMLP([RadialMLP([7, 5, 5, 6])]),
+        )
+
+    with pytest.raises(ValueError, match="UnifiedRadialMLP.*width"):
+        _configure_radial_first_linear_prefix_grad(
+            _make_model(), UnifiedRadialMLP([RadialMLP([8, 5, 5, 6])])
+        )
+
+    model = _make_model()
+    model.regress_config.hessian = True
+    with pytest.raises(ValueError, match="does not support Hessians"):
+        _configure_radial_first_linear_prefix_grad(
+            model, UnifiedRadialMLP([RadialMLP([7, 5, 5, 6])])
+        )
+
+
+@pytest.mark.gpu()
+@pytest.mark.compile_gpu()
+def test_radial_first_linear_prefix_compile_cuda(compile_reset_state):
+    torch.manual_seed(3)
+    model = RadialMLP([288, 128]).cuda()
+    model.configure_first_linear_grad_prefix(32, 288)
+    compiled = torch.compile(model, dynamic=True, fullgraph=True)
+    inputs = torch.randn(257, 288, device="cuda", requires_grad=True)
+    grad_output = torch.randn(257, 128, device="cuda")
+
+    actual = compiled(inputs)
+    actual_grad = torch.autograd.grad(actual, inputs, grad_output)[0]
+    expected = model.net[0](inputs)
+    expected_grad = torch.autograd.grad(expected, inputs, grad_output)[0]
+
+    torch.testing.assert_close(actual, expected)
+    torch.testing.assert_close(actual_grad[:, :32], expected_grad[:, :32])
+    torch.testing.assert_close(
+        actual_grad[:, 32:], torch.zeros_like(actual_grad[:, 32:])
+    )

tests/core/units/mlip_unit/test_inference_settings.py

@@ -71,6 +71,22 @@ def test_invalid_string_raises():
         InferenceSettings(base_precision_dtype="int8")
 
 
+def test_radial_first_linear_prefix_grad_validation():
+    with pytest.raises(ValueError, match="execution_mode='umas_fast_gpu'"):
+        InferenceSettings(radial_first_linear_prefix_grad=True)
+    with pytest.raises(ValueError, match="does not support Hessians"):
+        InferenceSettings(
+            execution_mode="umas_fast_gpu",
+            predict_untrained_hessian={"omat"},
+            radial_first_linear_prefix_grad=True,
+        )
+
+    settings = InferenceSettings(
+        execution_mode="umas_fast_gpu", radial_first_linear_prefix_grad=True
+    )
+    assert settings.radial_first_linear_prefix_grad
+
+
 # --- to_omegaconf ---
 
 
@@ -98,9 +114,15 @@ def test_to_omegaconf_roundtrip():
     """Hydra can reinstantiate InferenceSettings from to_omegaconf() output."""
     import hydra
 
-    original = InferenceSettings(base_precision_dtype=torch.float64, tf32=True)
+    original = InferenceSettings(
+        base_precision_dtype=torch.float64,
+        execution_mode="umas_fast_gpu",
+        radial_first_linear_prefix_grad=True,
+        tf32=True,
+    )
     config = original.to_omegaconf()
     restored = hydra.utils.instantiate(config)
     assert isinstance(restored, InferenceSettings)
     assert restored.base_precision_dtype is torch.float64
+    assert restored.radial_first_linear_prefix_grad is True
     assert restored.tf32 is True