Issue with pytorch 2.10.0

[stefdoerr]

Seems like a change in pytorch 2.10.0 is causing a crash in this library because: set_stride is not allowed on a Tensor created from .data or .detach().

self = <openmm.openmm.Context; proxy of <Swig Object of type 'OpenMM::Context *' at 0x72c58eec4830> >, positions = False, velocities = False, forces = False, energy = False
parameters = False, parameterDerivatives = False, integratorParameters = False, enforcePeriodicBox = False, groups = -1, getPositions = False, getVelocities = False, getForces = True
getEnergy = True, getParameters = False, getParameterDerivatives = False, getIntegratorParameters = False

    def getState(self, positions=False, velocities=False,
                 forces=False, energy=False, parameters=False,
                 parameterDerivatives=False, integratorParameters=False,
                 enforcePeriodicBox=False, groups=-1,
                 getPositions=False, getVelocities=False,
                 getForces=False, getEnergy=False, getParameters=False,
                 getParameterDerivatives=False, getIntegratorParameters=False):
        """Get a State object recording the current state information stored in this context.
    
        Parameters
        ----------
        positions : bool=False
            whether to store particle positions in the State
        velocities : bool=False
            whether to store particle velocities in the State
        forces : bool=False
            whether to store the forces acting on particles in the State
        energy : bool=False
            whether to store potential and kinetic energy in the State
        parameters : bool=False
            whether to store context parameters in the State
        parameterDerivatives : bool=False
            whether to store parameter derivatives in the State
        integratorParameters : bool=False
            whether to store integrator parameters in the State
        enforcePeriodicBox : bool=False
            if false, the position of each particle will be whatever position
            is stored in the Context, regardless of periodic boundary conditions.
            If true, particle positions will be translated so the center of
            every molecule lies in the same periodic box.
        groups : set={0,1,2,...,31}
            a set of indices for which force groups to include when computing
            forces and energies. The default value includes all groups. groups
            can also be passed as an unsigned integer interpreted as a bitmask,
            in which case group i will be included if (groups&(1<<i)) != 0.
        getPositions : bool=False
            Deprecated.  Use `positions` instead.
        getVelocities : bool=False
            Deprecated.  Use `velocities` instead.
        getForces : bool=False
            Deprecated.  Use `forces` instead.
        getEnergy : bool=False
            Deprecated.  Use `energy` instead.
        getParameters : bool=False
            Deprecated.  Use `parameters` instead.
        getParameterDerivatives : bool=False
            Deprecated.  Use `parameterDerivatives` instead.
        getIntegratorParameters : bool=False
            Deprecated.  Use `integratorParameters` instead.
        """
        try:
    # is the input integer-like?
            groups_mask = int(groups)
        except TypeError:
            if isinstance(groups, set):
    # nope, okay, then it should be an set
                groups_mask = functools.reduce(operator.or_,
                        ((1<<x) & 0xffffffff for x in groups))
            else:
                raise TypeError('%s is neither an int nor set' % groups)
        if groups_mask >= 0x80000000:
            groups_mask -= 0x100000000
        types = 0
        if positions or getPositions:
            types += State.Positions
        if velocities or getVelocities:
            types += State.Velocities
        if forces or getForces:
            types += State.Forces
        if energy or getEnergy:
            types += State.Energy
        if parameters or getParameters:
            types += State.Parameters
        if parameterDerivatives or getParameterDerivatives:
            types += State.ParameterDerivatives
        if integratorParameters or getIntegratorParameters:
            types += State.IntegratorParameters
>       state = _openmm.Context_getState(self, types, enforcePeriodicBox, groups_mask)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E       openmm.OpenMMException: set_stride is not allowed on a Tensor created from .data or .detach().
E       If your intent is to change the metadata of a Tensor (such as sizes / strides / storage / storage_offset)
E       without autograd tracking the change, remove the .data / .detach() call and wrap the change in a `with torch.no_grad():` block.
E       For example, change:
E           x.data.set_(y)
E       to:
E           with torch.no_grad():
E               x.set_(y)
E       Exception raised from set_stride at /pytorch/c10/core/TensorImpl.h:1819 (most recent call first):
E       frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >) + 0x9d (0x72c6c5b72fdd in /home/sdoerr/Work/acemd/.venv/lib/python3.11/site-packages/OpenMM.libs/lib/plugins/../../../torch/lib/libc10.so)
E       frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) + 0xc3 (0x72c6c5b08561 in /home/sdoerr/Work/acemd/.venv/lib/python3.11/site-packages/OpenMM.libs/lib/plugins/../../../torch/lib/libc10.so)
E       frame #2: <unknown function> + 0x803d0 (0x72c6c5b563d0 in /home/sdoerr/Work/acemd/.venv/lib/python3.11/site-packages/OpenMM.libs/lib/plugins/../../../torch/lib/libc10.so)
E       frame #3: torch::autograd::make_variable(at::Tensor, bool, bool) + 0x22e (0x72c6c10b60db in /home/sdoerr/Work/acemd/.venv/lib/python3.11/site-packages/OpenMM.libs/lib/plugins/libOpenMMTorchReference.so)
E       frame #4: torch::from_blob(void*, c10::ArrayRef<long>, c10::TensorOptions const&) + 0x81 (0x72c6c10b7643 in /home/sdoerr/Work/acemd/.venv/lib/python3.11/site-packages/OpenMM.libs/lib/plugins/libOpenMMTorchReference.so)
E       frame #5: TorchPlugin::ReferenceCalcTorchForceKernel::execute(OpenMM::ContextImpl&, bool, bool) + 0x16d (0x72c6c10b1f51 in /home/sdoerr/Work/acemd/.venv/lib/python3.11/site-packages/OpenMM.libs/lib/plugins/libOpenMMTorchReference.so)
E       frame #6: TorchPlugin::TorchForceImpl::calcForcesAndEnergy(OpenMM::ContextImpl&, bool, bool, int) + 0x6d (0x72c6fc7b3151 in /home/sdoerr/Work/acemd/.venv/lib/python3.11/site-packages/OpenMM.libs/lib/plugins/../libOpenMMTorch.so)
E       frame #7: OpenMM::ContextImpl::calcForcesAndEnergy(bool, bool, int) + 0xb9 (0x72c6c5d5aa59 in /home/sdoerr/Work/acemd/.venv/lib/python3.11/site-packages/openmm/../OpenMM.libs/lib/libOpenMM.so)
E       frame #8: OpenMM::Context::getState(int, bool, int) const + 0xbea (0x72c6c5d56b0a in /home/sdoerr/Work/acemd/.venv/lib/python3.11/site-packages/openmm/../OpenMM.libs/lib/libOpenMM.so)
E       frame #9: <unknown function> + 0x13c61d (0x72c6c633c61d in /home/sdoerr/Work/acemd/.venv/lib/python3.11/site-packages/openmm/_openmm.cpython-311-x86_64-linux-gnu.so)
E       frame #10: <unknown function> + 0x13cf9f (0x72c6c633cf9f in /home/sdoerr/Work/acemd/.venv/lib/python3.11/site-packages/openmm/_openmm.cpython-311-x86_64-linux-gnu.so)
E       frame #11: <unknown function> + 0x31e09f (0x6051bf7d809f in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #12: _PyEval_EvalFrameDefault + 0x34ebb (0x6051bf846bfb in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #13: _PyFunction_Vectorcall + 0x58b (0x6051bf7b365b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #14: _PyEval_EvalFrameDefault + 0x39429 (0x6051bf84b169 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #15: _PyFunction_Vectorcall + 0x58b (0x6051bf7b365b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #16: _PyEval_EvalFrameDefault + 0x3924b (0x6051bf84af8b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #17: _PyFunction_Vectorcall + 0x58b (0x6051bf7b365b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #18: _PyObject_Call_Prepend + 0x115 (0x6051bf8cfee5 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #19: <unknown function> + 0x447311 (0x6051bf901311 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #20: _PyEval_EvalFrameDefault + 0x34ebb (0x6051bf846bfb in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #21: _PyFunction_Vectorcall + 0x58b (0x6051bf7b365b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #22: _PyEval_EvalFrameDefault + 0x3924b (0x6051bf84af8b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #23: _PyFunction_Vectorcall + 0x58b (0x6051bf7b365b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #24: _PyObject_Call_Prepend + 0x115 (0x6051bf8cfee5 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #25: <unknown function> + 0x447311 (0x6051bf901311 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #26: _PyEval_EvalFrameDefault + 0x39327 (0x6051bf84b067 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #27: _PyFunction_Vectorcall + 0x58b (0x6051bf7b365b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #28: _PyEval_EvalFrameDefault + 0x3924b (0x6051bf84af8b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #29: _PyFunction_Vectorcall + 0x58b (0x6051bf7b365b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #30: _PyObject_Call_Prepend + 0x115 (0x6051bf8cfee5 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #31: <unknown function> + 0x447311 (0x6051bf901311 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #32: _PyEval_EvalFrameDefault + 0x34ebb (0x6051bf846bfb in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #33: _PyFunction_Vectorcall + 0x58b (0x6051bf7b365b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #34: _PyEval_EvalFrameDefault + 0x3924b (0x6051bf84af8b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #35: _PyFunction_Vectorcall + 0x58b (0x6051bf7b365b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #36: _PyObject_Call_Prepend + 0x115 (0x6051bf8cfee5 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #37: <unknown function> + 0x447311 (0x6051bf901311 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #38: _PyEval_EvalFrameDefault + 0x34ebb (0x6051bf846bfb in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #39: _PyFunction_Vectorcall + 0x58b (0x6051bf7b365b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #40: _PyEval_EvalFrameDefault + 0x3924b (0x6051bf84af8b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #41: _PyFunction_Vectorcall + 0x58b (0x6051bf7b365b in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #42: _PyObject_Call_Prepend + 0x115 (0x6051bf8cfee5 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #43: <unknown function> + 0x447311 (0x6051bf901311 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #44: _PyEval_EvalFrameDefault + 0x34ebb (0x6051bf846bfb in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #45: <unknown function> + 0x480052 (0x6051bf93a052 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #46: PyEval_EvalCode + 0xe8 (0x6051bf9395c8 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #47: <unknown function> + 0x4b41a1 (0x6051bf96e1a1 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #48: <unknown function> + 0x2332c2 (0x6051bf6ed2c2 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #49: _PyRun_SimpleFileObject + 0x109 (0x6051bf6ecdf9 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #50: _PyRun_AnyFileObject + 0xa5 (0x6051bf6ec815 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #51: <unknown function> + 0x23d469 (0x6051bf6f7469 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #52: <unknown function> + 0x23cfb0 (0x6051bf6f6fb0 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #53: Py_RunMain + 0x670 (0x6051bf97b1d0 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #54: <unknown function> + 0x1c3028 (0x6051bf67d028 in /home/sdoerr/Work/acemd/.venv/bin/python)
E       frame #55: <unknown function> + 0x2a1ca (0x72c6fec2a1ca in /lib/x86_64-linux-gnu/libc.so.6)
E       frame #56: __libc_start_main + 0x8b (0x72c6fec2a28b in /lib/x86_64-linux-gnu/libc.so.6)
E       frame #57: <unknown function> + 0x3e67bd (0x6051bf8a07bd in /home/sdoerr/Work/acemd/.venv/bin/python)

.venv/lib/python3.11/site-packages/openmm/openmm.py:5991: OpenMMException

[peastman]

Do you know what tensors it's referring to? I don't think the plugin ever creates a tensor with .data() or .detach(), or ever calls set_stride() on anything.

[stefdoerr]

It's related to Context copying to a different device (here CUDA to CPU). I managed to create a minimal reproduction script.

import torch as pt
from openmmtorch import TorchForce
import openmm as mm


class ForceWithParameters(pt.nn.Module):

    def __init__(self):
        super(ForceWithParameters, self).__init__()

    def forward(self, positions: pt.Tensor, k: pt.Tensor) -> pt.Tensor:
        return k * pt.sum(positions**2)


numParticles = 10
system = mm.System()
for _ in range(numParticles):
    system.addParticle(1.0)

model = pt.jit.script(ForceWithParameters())
tforce = TorchForce(model)
tforce.setOutputsForces(False)
tforce.addGlobalParameter("k", 2.0)
tforce.addEnergyParameterDerivative("k")
system.addForce(tforce)
context = mm.Context(system, mm.VerletIntegrator(1.0))
context.setPositions(pt.rand(numParticles, 3).numpy())
state = context.getState(getParameterDerivatives=True)
dEdk = state.getEnergyParameterDerivatives()["k"]
print(dEdk)

# Up to here it works fine.
# Copy the context

# Create a fake integrator for an auxilary context
integrator = mm.VerletIntegrator(0)
platform = mm.Platform.getPlatformByName("CPU")
properties = {"CpuThreads": "1"}

aux_context = mm.Context(context.getSystem(), integrator, platform, properties)
aux_context.setPositions(pt.rand(numParticles, 3).numpy())
state = aux_context.getState(getParameterDerivatives=True)  # Crashes here!
dEdk = state.getEnergyParameterDerivatives()["k"]
print(dEdk)

Error:

Traceback (most recent call last):
  File "/home/sdoerr/Work/acemd/test_pytorch_issue.py", line 40, in <module>
    state = aux_context.getState(getParameterDerivatives=True)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sdoerr/Work/acemd/.venv/lib/python3.11/site-packages/openmm/openmm.py", line 5991, in getState
    state = _openmm.Context_getState(self, types, enforcePeriodicBox, groups_mask)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
openmm.OpenMMException: set_stride is not allowed on a Tensor created from .data or .detach().
If your intent is to change the metadata of a Tensor (such as sizes / strides / storage / storage_offset)
without autograd tracking the change, remove the .data / .detach() call and wrap the change in a `with torch.no_grad():` block.
For example, change:
    x.data.set_(y)
to:
    with torch.no_grad():
        x.set_(y)
Exception raised from set_stride at /pytorch/c10/core/TensorImpl.h:1819 (most recent call first):

[peastman]

Is it related to having multiple contexts, or is it just that it happens on the CPU platform but not the CUDA platform? If you remove everything related to the first context and go straight to the CPU one, does the error still happen?

[stefdoerr]

Good point, that was a red herring. In reality it's an issue with using CPU platform.

import torch as pt
from openmmtorch import TorchForce
import openmm as mm


class ForceWithParameters(pt.nn.Module):
    def __init__(self):
        super(ForceWithParameters, self).__init__()

    def forward(self, positions: pt.Tensor, k: pt.Tensor) -> pt.Tensor:
        return k * pt.sum(positions**2)

numParticles = 10
system = mm.System()
for _ in range(numParticles):
    system.addParticle(1.0)

model = pt.jit.script(ForceWithParameters())
tforce = TorchForce(model)
tforce.setOutputsForces(False)
tforce.addGlobalParameter("k", 2.0)
tforce.addEnergyParameterDerivative("k")
system.addForce(tforce)
context = mm.Context(system, mm.VerletIntegrator(1.0), mm.Platform.getPlatformByName("CPU"))
context.setPositions(pt.rand(numParticles, 3).numpy())
state = context.getState(getEnergy=True)

This also crashes with the same error