Compile UMA energy output operations

> 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

Author: mlazos

src/fairchem/core/models/uma/outputs.py

@@ -81,8 +81,6 @@ def reduce_node_to_system(
     return reduced, system_values
 
 
-# Compile produces the wrong values using index_add with float64 precision :(
-@torch.compiler.disable
 def compute_energy(
     emb: dict[str, torch.Tensor],
     energy_block: torch.nn.Module,

tests/core/models/uma/test_outputs.py

@@ -256,6 +256,57 @@ def test_energy_part_for_gradients(self):
         assert node_embedding.grad is not None
         assert torch.allclose(node_embedding.grad, torch.ones_like(node_embedding))
 
+    @pytest.mark.gpu()
+    @pytest.mark.compile_gpu()
+    @pytest.mark.parametrize("dynamic", [False, True])
+    def test_float64_compile_index_add_regression(self, compile_reset_state, dynamic):
+        # Regression for pytorch/pytorch#108963.
+        def fn():
+            value = torch.zeros(1, dtype=torch.float64, device="cuda")
+            index = torch.tensor([0], dtype=torch.long, device="cuda")
+            source = torch.rand(1, dtype=torch.float64, device="cuda")
+            return source, value.index_add(0, index, source, alpha=2.0) / 2
+
+        torch.manual_seed(0)
+        source, output = torch.compile(fn, fullgraph=True, dynamic=dynamic)()
+
+        assert torch.equal(
+            output.contiguous().view(torch.uint8),
+            source.contiguous().view(torch.uint8),
+        )
+
+    @pytest.mark.gpu()
+    @pytest.mark.compile_gpu()
+    @pytest.mark.parametrize("dynamic", [False, True])
+    def test_float64_compile(self, compile_reset_state, dynamic):
+        energy_block = nn.Linear(8, 1).cuda()
+
+        def fn(node_embedding, batch):
+            return compute_energy(
+                {"node_embedding": node_embedding}, energy_block, batch, num_systems=4
+            )
+
+        node_embedding = torch.randn(
+            257, 9, 8, device="cuda", dtype=torch.float32, requires_grad=True
+        )
+        batch = torch.arange(257, device="cuda") % 4
+        expected = fn(node_embedding, batch)
+        actual = torch.compile(fn, fullgraph=True, dynamic=dynamic)(
+            node_embedding, batch
+        )
+
+        for actual_output, expected_output in zip(actual, expected):
+            assert torch.equal(
+                actual_output.contiguous().view(torch.uint8),
+                expected_output.contiguous().view(torch.uint8),
+            )
+        expected_grad = torch.autograd.grad(expected[1].sum(), node_embedding)[0]
+        actual_grad = torch.autograd.grad(actual[1].sum(), node_embedding)[0]
+        assert torch.equal(
+            actual_grad.contiguous().view(torch.uint8),
+            expected_grad.contiguous().view(torch.uint8),
+        )
+
     def test_reduce_mean(self):
         """Test that reduce='mean' divides energy by natoms per system."""
         emb, energy_block = _make_emb_and_block(