Skip to content

Reduce UMA radial input-gradient work - #2149

Draft
mlazos wants to merge 1 commit into
mainfrom
uma-radial-prefix-vjp
Draft

Reduce UMA radial input-gradient work#2149
mlazos wants to merge 1 commit into
mainfrom
uma-radial-prefix-vjp

Conversation

@mlazos

@mlazos mlazos commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Summary

UMA radial first layers consume distance features followed by source and target
element embeddings. Only the distance prefix depends on atomic positions, but
the ordinary frozen linear backward materializes all 288 input-gradient columns
before discarding the embedding columns.

This adds an opt-in radial_first_linear_prefix_grad fast-GPU inference path.
The forward remains unchanged, while the input VJP contracts only with the
distance-feature weight prefix. The implementation validates model dimensions,
preserves parameter and state-dict structure, and rejects Hessian inference.

Performance

With torch.compile(dynamic=True), the option improved the full endpoint from
17.17605 to 16.82399 ms with fresh caches and from 17.18390 to 16.82917 ms with
warm caches. The paired mean improvement is 0.35340 ms (2.06%).

The optimization was also positive at every measured size from 32 through
1,000 atoms.

Configuration: H100, TF32, UMA-S-1p2, 1,000 atoms, energy/forces/stress,
umas_fast_gpu, merge_mole=True, external_graph_gen=False, internal graph
v3, internal_graph_skin=0, compile=True, compile_dynamic_shapes=True, and
full CUDA graph replay. Steady-state execution had zero recompiles and zero
dynamic CUDA graph fallbacks.

Testing

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

The focused suite passed 25 tests, and all pre-commit hooks passed.

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
```
@meta-cla meta-cla Bot added the cla signed label Aug 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant