Reduce UMA radial input-gradient work - #2149
Draft
mlazos wants to merge 1 commit into
Draft
Conversation
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 ```
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_gradfast-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 from17.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 graphv3,
internal_graph_skin=0,compile=True,compile_dynamic_shapes=True, andfull CUDA graph replay. Steady-state execution had zero recompiles and zero
dynamic CUDA graph fallbacks.
Testing
The focused suite passed 25 tests, and all pre-commit hooks passed.