From 74e6650b95cf5a55e698b732c9decf010ac338e5 Mon Sep 17 00:00:00 2001 From: Michael Lazos Date: Wed, 5 Aug 2026 00:44:38 -0700 Subject: [PATCH] Use FP16 operands for UMA radial FC2 Add a default-off fast-inference path for the second unified radial projection. Forward and frozen first-order backward use FP16 GEMM operands with FP32 outputs while layer normalization, activation, and surrounding values remain FP32. Nonpersistent weight caches keep stable storage across CUDA graph replay and refresh after state or device changes. Reject non-FP32 base precision and Hessian inference. On H100 with UMA-S-1p2 at 1,000 atoms, the isolated option reduced dynamic full-CUDA-graph latency from 17.017035 ms to 16.826578 ms, a 0.190457 ms (1.12%) 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. A separate static measurement improved 17.719488 ms to 17.557168 ms. Maximum observed energy, force, and stress differences were 0.006781, 0.000299, and 0.0000167. 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/nn/test_unified_radial.py tests/core/models/uma/uma_fast/test_execution_backends.py tests/core/units/mlip_unit/test_inference_settings.py -m 'not gpu' pre-commit run --files docs/core/common_tasks/ase_calculator.md src/fairchem/core/models/uma/nn/execution_backends.py src/fairchem/core/models/uma/nn/unified_radial.py src/fairchem/core/units/mlip_unit/api/inference.py tests/core/models/uma/nn/test_unified_radial.py tests/core/models/uma/uma_fast/test_execution_backends.py tests/core/units/mlip_unit/test_inference_settings.py ``` --- docs/core/common_tasks/ase_calculator.md | 1 + .../core/models/uma/nn/execution_backends.py | 16 ++++ .../core/models/uma/nn/unified_radial.py | 72 +++++++++++++++++- .../core/units/mlip_unit/api/inference.py | 22 ++++++ .../core/models/uma/nn/test_unified_radial.py | 76 +++++++++++++++++++ .../uma/uma_fast/test_execution_backends.py | 35 +++++++-- .../mlip_unit/test_inference_settings.py | 38 ++++++++++ 7 files changed, 253 insertions(+), 7 deletions(-) diff --git a/docs/core/common_tasks/ase_calculator.md b/docs/core/common_tasks/ase_calculator.md index 4a9611e06e..6a14195e04 100644 --- a/docs/core/common_tasks/ase_calculator.md +++ b/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. | diff --git a/src/fairchem/core/models/uma/nn/execution_backends.py b/src/fairchem/core/models/uma/nn/execution_backends.py index a22188ee7f..4cb1fb6042 100644 --- a/src/fairchem/core/models/uma/nn/execution_backends.py +++ b/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,6 +307,15 @@ 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) @@ -308,6 +323,7 @@ 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) + model._unified_radial_mlp.configure_fp16_fc2(fp16_radial_fc2_blocks) @staticmethod def get_layer_radial_emb( diff --git a/src/fairchem/core/models/uma/nn/unified_radial.py b/src/fairchem/core/models/uma/nn/unified_radial.py index ea08e5e7e5..31bb06606b 100644 --- a/src/fairchem/core/models/uma/nn/unified_radial.py +++ b/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 ) diff --git a/src/fairchem/core/units/mlip_unit/api/inference.py b/src/fairchem/core/units/mlip_unit/api/inference.py index eadfe02f0c..659035f06d 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 + # 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: """ diff --git a/tests/core/models/uma/nn/test_unified_radial.py b/tests/core/models/uma/nn/test_unified_radial.py index 9b90e99e41..3708712a29 100644 --- a/tests/core/models/uma/nn/test_unified_radial.py +++ b/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) diff --git a/tests/core/models/uma/uma_fast/test_execution_backends.py b/tests/core/models/uma/uma_fast/test_execution_backends.py index d12e93611a..ffd824975a 100644 --- a/tests/core/models/uma/uma_fast/test_execution_backends.py +++ b/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']}" diff --git a/tests/core/units/mlip_unit/test_inference_settings.py b/tests/core/units/mlip_unit/test_inference_settings.py index 19ebde2ed8..958be62fbd 100644 --- a/tests/core/units/mlip_unit/test_inference_settings.py +++ b/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)