Use FP16 operands for UMA radial FC2 - #2150
Draft
mlazos wants to merge 1 commit into
Draft
Conversation
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 ```
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
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
This adds a default-off
fp16_radial_fc2_blockspath for the second unifiedUMA radial projection. Forward and frozen first-order backward use FP16 GEMM
operands with FP32 outputs, while layer normalization, activation, and
surrounding values remain FP32.
The implementation keeps nonpersistent FP16 weight caches stable across CUDA
graph replay, refreshes them after state or device changes, preserves state
dictionaries, and rejects non-FP32 base precision and Hessian inference.
Performance
With
torch.compile(dynamic=True), the option improved the full endpoint from17.01415 to 16.82399 ms with fresh caches and from 17.01992 to 16.82917 ms with
warm caches. The paired mean improvement is 0.19046 ms (1.12%).
A separate static-shape comparison improved 17.71949 to 17.55717 ms, a
0.16232 ms (0.92%) reduction. Maximum observed energy, force, and stress
differences were
0.006781,0.000299, and0.0000167.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 21 tests, and all pre-commit hooks passed. The added
GPU coverage exercises eager and compiled VJPs plus CUDA graph replay.