Compile UMA energy output operations by mlazos · Pull Request #2125 · facebookresearch/fairchem

mlazos · 2026-07-31T10:18:09Z

Stack from ghstack (oldest at bottom):

The current end-to-end validation uses UMA-S-1p2 with energy/forces/stress,
external_graph_gen=False, internal_graph_gen_version=3, merge_mole=True,
execution_mode="umas_fast_gpu", and compile_dynamic_shapes=False. Energy
output processing is downstream of graph construction and is executed for this
checkpoint, so removing its explicit compile disable applies unchanged and removes
one graph break from this internal-graph configuration.

compute_energy carried an unconditional compile disable around the float64
per-system reduction. The disable was added for an older float64 index_add
accuracy issue, but current Inductor matches eager forward and backward results
byte-for-byte. Keeping it now guarantees a boundary on every inference:

@torch.compiler.disable
def compute_energy(...):
    ...

Remove the decorator after adding a regression test for the production-like
float32 node-energy to float64 system-reduction path. The test checks compiled
forward outputs and input gradients byte-for-byte for static and dynamic compilation.

Numerical validation

The following AI-assisted numerical analysis was reviewed for inclusion because it
bounds the nondeterminism relevant to this change:

With identical per-node inputs and deterministic algorithms disabled, two
arbitrary FP64 atomic or two-worker NCCL reduction orders satisfy
|E_a - E_b| <= 2 * gamma_(n-1) * sum_i |e_i|, where
gamma_k = k * 2^-53 / (1 - k * 2^-53). This assumes IEEE FP64 rounding and no
overflow or underflow.

A seeded recreation of the regression test's 257-node, four-system input has a
maximum bound of 5.57e-13; the compiled and eager outputs and gradients matched
exactly. In the 1000-atom two-worker endpoint experiment, the maximum observed
energy difference was 1.22e-5 eV. Explaining that difference through FP64
reduction order alone would require sum(abs(node_energy)) >= 5.5e7 eV, or about
55,000 eV/atom. The endpoint difference therefore reflects upstream
model-parallel FP32 computation rather than this final FP64 system reduction.

If upstream per-node values differ, their contribution is bounded separately by
sum_i |e_i^(1) - e_i^(2)|, in addition to the FP64 rounding terms.

Activation

No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled.

settings = InferenceSettings(compile=True)

The current benchmark additionally uses merge_mole=True, external_graph_gen=False, and execution_mode="umas_fast_gpu", but those settings are not required for this output-processing fix.

Test Plan:

PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile
ruff check src/fairchem/core/models/uma/outputs.py

Authored with assistance from Codex.

UMA disabled compilation around float64 energy reduction and used a determinant operation that split stress postprocessing. Keep the reduction in the compiled graph after validating strict float64 forward and backward accuracy, and express the fixed 3x3 determinant directly. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. [ghstack-poisoned]

Target configuration: the current uma-s-1p2 benchmark with `torch.compile(dynamic=True)`, energy/forces/stress enabled, `external_graph_gen=True`, `merge_mole=True`, and `umas_fast_gpu`. The `torch.compiler.disable` on energy is active in this exact run and accounts for one explicit graph break. Stress determinant handling is also active, but was not counted as a separate Dynamo graph-break occurrence. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a strict float64 compiled regression test. Stress postprocessing also used the general `torch.det` linalg operator for fixed 3x3 cells. Expressing that determinant with the ordinary 3x3 cofactor formula leaves only indexing, multiply, and add operations that can remain with the surrounding virial calculation. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. [ghstack-poisoned]

Target configuration: the current uma-s-1p2 benchmark with `torch.compile(dynamic=True)`, energy/forces/stress enabled, `external_graph_gen=True`, `merge_mole=True`, and `umas_fast_gpu`. The `torch.compiler.disable` on energy is active in this exact run and accounts for one explicit graph break. Stress determinant handling is also active, but was not counted as a separate Dynamo graph-break occurrence. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a strict float64 compiled regression test. Stress postprocessing also used the general `torch.det` linalg operator for fixed 3x3 cells. Expressing that determinant with the ordinary 3x3 cofactor formula leaves only indexing, multiply, and add operations that can remain with the surrounding virial calculation. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. The explicit 3x3 stress determinant is exercised when the requested task computes stress. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=True`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. [ghstack-poisoned]

Target configuration: the current uma-s-1p2 benchmark with `torch.compile(dynamic=True)`, energy/forces/stress enabled, `external_graph_gen=True`, `merge_mole=True`, and `umas_fast_gpu`. The `torch.compiler.disable` on energy is active in this exact run and accounts for one explicit graph break. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a strict float64 compiled regression test. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=True`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. [ghstack-poisoned]

Target configuration: the current uma-s-1p2 benchmark with `torch.compile(dynamic=True)`, energy/forces/stress enabled, `external_graph_gen=True`, `merge_mole=True`, and `umas_fast_gpu`. The `torch.compiler.disable` on energy is active in this exact run and accounts for one explicit graph break. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward and backward results at strict tolerance. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=True`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. [ghstack-poisoned]

> The current end-to-end validation uses UMA-S-1p2 with energy/forces/stress, > `external_graph_gen=False`, `internal_graph_gen_version=3`, `merge_mole=True`, > `execution_mode="umas_fast_gpu"`, and `compile_dynamic_shapes=False`. Energy > output processing is downstream of graph construction and is executed for this > checkpoint, so removing its explicit compile disable applies unchanged and removes > one graph break from this internal-graph configuration. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward and backward results at strict tolerance. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=True`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. [ghstack-poisoned]

Target configuration: the current uma-s-1p2 benchmark with `torch.compile(dynamic=True)`, energy/forces/stress enabled, `external_graph_gen=True`, `merge_mole=True`, and `umas_fast_gpu`. The `torch.compiler.disable` on energy is active in this exact run and accounts for one explicit graph break. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` @torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward and backward results at strict tolerance. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=True`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. ghstack-source-id: 1d6cb5c Pull Request resolved: #2125

…cebookresearch#2123) Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom): * facebookresearch#2125 * facebookresearch#2124 * __->__ facebookresearch#2123 > The change in this PR is specific to the precomputed/external multi-system graph > branch. Its isolation configuration uses `external_graph_gen=True`, but the branch > is not inherently restricted to that setting: a multi-system internal-v3 fallback > can also produce a precomputed graph. The current final-stack benchmark uses > `external_graph_gen=False`, `internal_graph_gen_version=3`, and one system, so it > takes the single-system batched-matmul branch and does not execute the changed > `repeat_interleave`. This remains a synchronization optimization, not a graph-break > fix and not a contributor to the reported one-system internal-graph result. For multiple systems, `repeat_interleave` expands each cell according to the number of edges in that system. Without `output_size`, PyTorch must obtain the sum of `nedges` to allocate the result. Supplying the already-known edge count from `cell_offsets.shape[0]` avoids discovering that output size from tensor values: ```python cell.repeat_interleave( nedges, dim=0, output_size=cell_offsets.shape[0] ) ``` On the current PyTorch 2.14 nightly, the original operation already compiles with `fullgraph=True` and accepts changing edge counts without recompilation. This is therefore a synchronization optimization, not a graph-break fix. The test runs two different two-system batches with `dynamic=True`, rejects recompilation on the second shape, and compares edge vectors and distances with eager execution. **Activation** No independent flag enables this optimization. It is used automatically when external/precomputed graph inference receives multiple systems in one batch. ```python settings = InferenceSettings( compile=True, external_graph_gen=True, ) ``` `merge_mole=True` and `execution_mode="umas_fast_gpu"` are part of the performance configuration studied here, but are not required for this change. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_compile.py -k test_compile_batched_external_graph ``` Authored with assistance from Codex.

> The current end-to-end validation uses UMA-S-1p2 with energy/forces/stress, > `external_graph_gen=False`, `internal_graph_gen_version=3`, `merge_mole=True`, > `execution_mode="umas_fast_gpu"`, and `compile_dynamic_shapes=False`. Energy > output processing is downstream of graph construction and is executed for this > checkpoint, so removing its explicit compile disable applies unchanged and removes > one graph break from this internal-graph configuration. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward and backward results at strict tolerance. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=False`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. [ghstack-poisoned]

> The current end-to-end validation uses UMA-S-1p2 with energy/forces/stress, > `external_graph_gen=False`, `internal_graph_gen_version=3`, `merge_mole=True`, > `execution_mode="umas_fast_gpu"`, and `compile_dynamic_shapes=False`. Energy > output processing is downstream of graph construction and is executed for this > checkpoint, so removing its explicit compile disable applies unchanged and removes > one graph break from this internal-graph configuration. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` @torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward and backward results at strict tolerance. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=False`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. ghstack-source-id: 1d6cb5c Pull Request resolved: #2125

rayg1234 · 2026-08-06T18:26:12Z

    return reduced, system_values


-# Compile produces the wrong values using index_add with float64 precision :(


this needs to be tested VERY carefully on large number of atoms and multi-gpu reductions, the unit test is not enough, im not sure if anything has changed since this PR #1889

I actually think I may have fixed this in the compiler last half, I will run larger experiments from your PR.

that be would amazing!

so I validated the kernel/disabled region locally (compiled was within 1e-13 of eager) and when I enabled deterministic algorithms, we are exactly bitwise equivalent.

For distributed I ran with 2 workers (I actually only have a 2 H100 machine available) I can validate on more if you can share a bigger machine with me?

Let me know what you think!

rayg1234

See comment, this is a dangerous change

> The current end-to-end validation uses UMA-S-1p2 with energy/forces/stress, > `external_graph_gen=False`, `internal_graph_gen_version=3`, `merge_mole=True`, > `execution_mode="umas_fast_gpu"`, and `compile_dynamic_shapes=False`. Energy > output processing is downstream of graph construction and is executed for this > checkpoint, so removing its explicit compile disable applies unchanged and removes > one graph break from this internal-graph configuration. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward and backward results at strict tolerance. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=False`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. [ghstack-poisoned]

> The current end-to-end validation uses UMA-S-1p2 with energy/forces/stress, > `external_graph_gen=False`, `internal_graph_gen_version=3`, `merge_mole=True`, > `execution_mode="umas_fast_gpu"`, and `compile_dynamic_shapes=False`. Energy > output processing is downstream of graph construction and is executed for this > checkpoint, so removing its explicit compile disable applies unchanged and removes > one graph break from this internal-graph configuration. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` @torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward and backward results at strict tolerance. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=False`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. ghstack-source-id: 8af8878 Pull Request resolved: #2125

> The current end-to-end validation uses UMA-S-1p2 with energy/forces/stress, > `external_graph_gen=False`, `internal_graph_gen_version=3`, `merge_mole=True`, > `execution_mode="umas_fast_gpu"`, and `compile_dynamic_shapes=False`. Energy > output processing is downstream of graph construction and is executed for this > checkpoint, so removing its explicit compile disable applies unchanged and removes > one graph break from this internal-graph configuration. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward and backward results at strict tolerance. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=False`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. [ghstack-poisoned]

> The current end-to-end validation uses UMA-S-1p2 with energy/forces/stress, > `external_graph_gen=False`, `internal_graph_gen_version=3`, `merge_mole=True`, > `execution_mode="umas_fast_gpu"`, and `compile_dynamic_shapes=False`. Energy > output processing is downstream of graph construction and is executed for this > checkpoint, so removing its explicit compile disable applies unchanged and removes > one graph break from this internal-graph configuration. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` @torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward and backward results at strict tolerance. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=False`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. ghstack-source-id: 9006775 Pull Request resolved: #2125

> The current end-to-end validation uses UMA-S-1p2 with energy/forces/stress, > `external_graph_gen=False`, `internal_graph_gen_version=3`, `merge_mole=True`, > `execution_mode="umas_fast_gpu"`, and `compile_dynamic_shapes=False`. Energy > output processing is downstream of graph construction and is executed for this > checkpoint, so removing its explicit compile disable applies unchanged and removes > one graph break from this internal-graph configuration. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward and backward results at strict tolerance. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=False`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. [ghstack-poisoned]

> The current end-to-end validation uses UMA-S-1p2 with energy/forces/stress, > `external_graph_gen=False`, `internal_graph_gen_version=3`, `merge_mole=True`, > `execution_mode="umas_fast_gpu"`, and `compile_dynamic_shapes=False`. Energy > output processing is downstream of graph construction and is executed for this > checkpoint, so removing its explicit compile disable applies unchanged and removes > one graph break from this internal-graph configuration. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results at strict tolerance. Keeping it now guarantees a boundary on every inference: ``` @torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward and backward results at strict tolerance. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=False`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. ghstack-source-id: 19b9ddf Pull Request resolved: #2125

> The current end-to-end validation uses UMA-S-1p2 with energy/forces/stress, > `external_graph_gen=False`, `internal_graph_gen_version=3`, `merge_mole=True`, > `execution_mode="umas_fast_gpu"`, and `compile_dynamic_shapes=False`. Energy > output processing is downstream of graph construction and is executed for this > checkpoint, so removing its explicit compile disable applies unchanged and removes > one graph break from this internal-graph configuration. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results byte-for-byte. Keeping it now guarantees a boundary on every inference: ``` torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward outputs and input gradients byte-for-byte for static and dynamic compilation. **Numerical validation** The following AI-assisted numerical analysis was reviewed for inclusion because it bounds the nondeterminism relevant to this change: > With identical per-node inputs and deterministic algorithms disabled, two > arbitrary FP64 atomic or two-worker NCCL reduction orders satisfy > `|E_a - E_b| <= 2 * gamma_(n-1) * sum_i |e_i|`, where > `gamma_k = k * 2^-53 / (1 - k * 2^-53)`. This assumes IEEE FP64 rounding and no > overflow or underflow. > > A seeded recreation of the regression test's 257-node, four-system input has a > maximum bound of `5.57e-13`; the compiled and eager outputs and gradients matched > exactly. In the 1000-atom two-worker endpoint experiment, the maximum observed > energy difference was `1.22e-5 eV`. Explaining that difference through FP64 > reduction order alone would require `sum(abs(node_energy)) >= 5.5e7 eV`, or about > `55,000 eV/atom`. The endpoint difference therefore reflects upstream > model-parallel FP32 computation rather than this final FP64 system reduction. > > If upstream per-node values differ, their contribution is bounded separately by > `sum_i |e_i^(1) - e_i^(2)|`, in addition to the FP64 rounding terms. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=False`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. [ghstack-poisoned]

> The current end-to-end validation uses UMA-S-1p2 with energy/forces/stress, > `external_graph_gen=False`, `internal_graph_gen_version=3`, `merge_mole=True`, > `execution_mode="umas_fast_gpu"`, and `compile_dynamic_shapes=False`. Energy > output processing is downstream of graph construction and is executed for this > checkpoint, so removing its explicit compile disable applies unchanged and removes > one graph break from this internal-graph configuration. `compute_energy` carried an unconditional compile disable around the float64 per-system reduction. The disable was added for an older float64 `index_add` accuracy issue, but current Inductor matches eager forward and backward results byte-for-byte. Keeping it now guarantees a boundary on every inference: ``` @torch.compiler.disable def compute_energy(...): ... ``` Remove the decorator after adding a regression test for the production-like float32 node-energy to float64 system-reduction path. The test checks compiled forward outputs and input gradients byte-for-byte for static and dynamic compilation. **Numerical validation** The following AI-assisted numerical analysis was reviewed for inclusion because it bounds the nondeterminism relevant to this change: > With identical per-node inputs and deterministic algorithms disabled, two > arbitrary FP64 atomic or two-worker NCCL reduction orders satisfy > `|E_a - E_b| <= 2 * gamma_(n-1) * sum_i |e_i|`, where > `gamma_k = k * 2^-53 / (1 - k * 2^-53)`. This assumes IEEE FP64 rounding and no > overflow or underflow. > > A seeded recreation of the regression test's 257-node, four-system input has a > maximum bound of `5.57e-13`; the compiled and eager outputs and gradients matched > exactly. In the 1000-atom two-worker endpoint experiment, the maximum observed > energy difference was `1.22e-5 eV`. Explaining that difference through FP64 > reduction order alone would require `sum(abs(node_energy)) >= 5.5e7 eV`, or about > `55,000 eV/atom`. The endpoint difference therefore reflects upstream > model-parallel FP32 computation rather than this final FP64 system reduction. > > If upstream per-node values differ, their contribution is bounded separately by > `sum_i |e_i^(1) - e_i^(2)|`, in addition to the FP64 rounding terms. **Activation** No independent flag enables this fix. Energy output processing remains in the captured graph whenever UMA inference is compiled. ```python settings = InferenceSettings(compile=True) ``` The current benchmark additionally uses `merge_mole=True`, `external_graph_gen=False`, and `execution_mode="umas_fast_gpu"`, but those settings are not required for this output-processing fix. Test Plan: ``` PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/test_outputs.py -k float64_compile ruff check src/fairchem/core/models/uma/outputs.py ``` Authored with assistance from Codex. ghstack-source-id: 19b9ddf Pull Request resolved: #2125

meta-cla Bot added the cla signed label Jul 31, 2026

mlazos added 4 commits July 31, 2026 12:01

mlazos changed the title ~~Compile UMA energy and stress output operations~~ Compile UMA energy output operations Jul 31, 2026

mlazos requested review from lbluque and rayg1234 and removed request for rayg1234 July 31, 2026 20:33

mlazos marked this pull request as ready for review July 31, 2026 20:33

mlazos requested a review from rayg1234 July 31, 2026 20:34

mlazos added the patch Patch version release label Jul 31, 2026

mlazos added 2 commits July 31, 2026 13:50

rayg1234 reviewed Aug 6, 2026

View reviewed changes

rayg1234 requested changes Aug 6, 2026

View reviewed changes

mlazos requested a review from rayg1234 August 12, 2026 21:07

Navigation Menu

Uh oh!

Compile UMA energy output operations - #2125

Compile UMA energy output operations#2125
mlazos wants to merge 12 commits into
gh/mlazos/3/basefrom
gh/mlazos/3/head

mlazos commented Jul 31, 2026 •

edited

Loading

Uh oh!

rayg1234 Aug 6, 2026 •

edited

Loading

Uh oh!

mlazos Aug 6, 2026 •

edited

Loading

Uh oh!

rayg1234 Aug 6, 2026

Uh oh!

mlazos Aug 12, 2026 •

edited

Loading

Uh oh!

rayg1234 left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		return reduced, system_values


		# Compile produces the wrong values using index_add with float64 precision :(

Uh oh!

Conversation

mlazos commented Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rayg1234 Aug 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mlazos Aug 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rayg1234 Aug 6, 2026

Choose a reason for hiding this comment

Uh oh!

mlazos Aug 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rayg1234 left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

mlazos commented Jul 31, 2026 •

edited

Loading

rayg1234 Aug 6, 2026 •

edited

Loading

mlazos Aug 6, 2026 •

edited

Loading

mlazos Aug 12, 2026 •

edited

Loading