Difference between AceFF Energies Before and After Compilation with CUDA

[fjclark]

Summary

I recently got some very bad relative conformer energies using AceFF2 with CUDA and noticed that this was due to the first energies being systematically offset. This is due to the first energy coming from the uncompiled model and the later energies from the compiled model. This problem disappears if I load the model on CPU or disable compilation on GPU.

I've checked that this also happens with the torchnet-md ASE calculator, though here the first energy is discarded so you get the same energy for the same conformer evaluated twice.

To Reproduce

Here's an example of recalculating the energy of the same conformer (containing inputs, pixi toml, and lockfile): aceff_issue.tar.gz

Inside this directory, run pixi run python get_energies.py. I see:

...
Iteration 1
...
Energy: -369.6271265442011 kcal/mol
Iteration 2
Energy: -362.111691535089 kcal/mol

I'm running Ubuntu 22.04 and my driver version is 580.65.06 (please see lock file for other env details).

For convenience, the script is:

from openmmml import MLPotential
from openff.toolkit import Molecule
import openmm
import openmm.unit as unit

def main():
    conf1 = Molecule("mol.sdf")

    potential = MLPotential('aceff-2.0')
    topology = conf1.to_topology().to_openmm()
    system = potential.createSystem(topology, device="cuda")

    integrator = openmm.VerletIntegrator(0.001 * unit.picoseconds)
    platform   = openmm.Platform.getPlatformByName("CUDA")
    context = openmm.Context(system, integrator, platform)

    for i in range(2):
        print(f'Iteration {i+1}')

        positions = conf1.conformers[0].to_openmm()
        context.setPositions(positions)

        energy = context.getState(getEnergy=True).getPotentialEnergy()
        print(f'Energy: {energy.in_units_of(unit.kilocalories_per_mole)}')


if __name__ == '__main__':
    main()

and the pixi.toml is:

[workspace]
channels = ["conda-forge"]
name = "aceff_issue"
platforms = ["linux-64"]
version = "0.1.0"

[tasks]

[system-requirements]
cuda = "13"

[dependencies]
openmm-ml = ">=1.6,<2"
openff-toolkit-base = ">=0.18.0,<0.19"
torchmd-net = ">=3.0.3,<4"
huggingface_hub = ">=1.8.0,<2"
pytorch-gpu = ">=2.10"
rdkit = ">=2025.9.6,<2026"
cuda-cudart-dev = ">=13.2.51,<14"

Thanks.

[peastman]

@sef43 I think this is for you.

[sef43]

Thanks for reporting @fjclark !

This seems to be a torch.compile bug that occurs when the molecule has a number of atoms that is not a power of 2.
I have reported here: pytorch/pytorch#178871

It does some fusions that end up wrong. The backwards/forces are not affected. I also think the relative energies post compile end up correct (based on the fact I have not seen this in other benchmarks I have done).

[fjclark]

Thanks @sef43!

The relative energies post compile also looked reasonable to me (on a subset of https://onlinelibrary.wiley.com/doi/full/10.1002/qua.26381).

[sef43]

there is a workaround in torchmdnet here torchmd/torchmd-net#390