From 68bd0088177f5e220f2e5677b3d27372da671894 Mon Sep 17 00:00:00 2001 From: Michael Lazos Date: Tue, 4 Aug 2026 07:49:49 -0700 Subject: [PATCH] Reduce UMA radial input-gradient work UMA radial first layers receive distance features followed by source and target element embeddings. During force inference only the distance prefix depends on positions, but the ordinary frozen linear backward materializes all 288 input-gradient columns before the embedding columns are discarded. Add an opt-in fast-GPU inference path that keeps the forward unchanged while contracting the backward with only the distance-feature weight prefix. Derive and validate the prefix from model dimensions, preserve existing parameters and state dictionaries, and reject Hessian inference because this is a frozen first-order inference boundary. On H100 with UMA-S-1p2 at 1,000 atoms, the isolated option reduced dynamic full-CUDA-graph latency from 17.179976 ms to 16.826578 ms, a 0.353399 ms (2.06%) improvement averaged across independent fresh- and warm-cache processes. The configuration used energy, forces, and stress; TF32; merge_mole=True; external_graph_gen=False; internal graph v3; skin 0; compile_dynamic_shapes=True; and full CUDA graph replay. The optimization was also positive at every measured size from 32 through 1,000 atoms. Test Plan: ``` PYTHONPATH=src /home/mlazos/.conda/envs/pytorch-3.12/bin/python -m pytest -q -c packages/fairchem-core/pyproject.toml tests/core/models/uma/uma_fast/test_radial_prefix_grad.py -m 'not gpu' tests/core/models/uma/nn/test_unified_radial.py tests/core/units/mlip_unit/test_inference_settings.py pre-commit run --files src/fairchem/core/models/uma/nn/execution_backends.py src/fairchem/core/models/uma/nn/radial.py src/fairchem/core/models/uma/nn/unified_radial.py src/fairchem/core/units/mlip_unit/api/inference.py tests/core/models/uma/uma_fast/test_radial_prefix_grad.py tests/core/units/mlip_unit/test_inference_settings.py ``` --- .../core/models/uma/nn/execution_backends.py | 15 ++ src/fairchem/core/models/uma/nn/radial.py | 46 ++++- .../core/models/uma/nn/unified_radial.py | 21 ++- .../core/units/mlip_unit/api/inference.py | 13 ++ .../uma/uma_fast/test_radial_prefix_grad.py | 158 ++++++++++++++++++ .../mlip_unit/test_inference_settings.py | 24 ++- 6 files changed, 274 insertions(+), 3 deletions(-) create mode 100644 tests/core/models/uma/uma_fast/test_radial_prefix_grad.py 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