diff --git a/src/fairchem/core/models/uma/nn/execution_backends.py b/src/fairchem/core/models/uma/nn/execution_backends.py index a22188ee7f..827eefde2a 100644 --- a/src/fairchem/core/models/uma/nn/execution_backends.py +++ b/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( diff --git a/src/fairchem/core/models/uma/nn/radial.py b/src/fairchem/core/models/uma/nn/radial.py index 2e5dfcb3d4..3dc1c6090c 100644 --- a/src/fairchem/core/models/uma/nn/radial.py +++ b/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 diff --git a/src/fairchem/core/models/uma/nn/unified_radial.py b/src/fairchem/core/models/uma/nn/unified_radial.py index ea08e5e7e5..9d27e93adf 100644 --- a/src/fairchem/core/models/uma/nn/unified_radial.py +++ b/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)] diff --git a/src/fairchem/core/units/mlip_unit/api/inference.py b/src/fairchem/core/units/mlip_unit/api/inference.py index eadfe02f0c..72f2b16d7f 100644 --- a/src/fairchem/core/units/mlip_unit/api/inference.py +++ b/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 ( diff --git a/tests/core/models/uma/uma_fast/test_radial_prefix_grad.py b/tests/core/models/uma/uma_fast/test_radial_prefix_grad.py new file mode 100644 index 0000000000..13714a3b50 --- /dev/null +++ b/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:]) + ) diff --git a/tests/core/units/mlip_unit/test_inference_settings.py b/tests/core/units/mlip_unit/test_inference_settings.py index 19ebde2ed8..049c272337 100644 --- a/tests/core/units/mlip_unit/test_inference_settings.py +++ b/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