Does espaloma correctly work on CPUs without modification?

I am trying to craft a simple script to create and simulate a small molecule with espaloma via openmmforcefields.

There’s probably a much simpler way to do this with just espaloma, and I’d love to know what that is, but trying to run this on a CPU exposed a potential bug when running espaloma on a machine without a GPU:

$ python espaloma-aniline.py
Source espaloma-0.2.2.offxml could not be read. If this is a file, ensure that the path is correct.
If the file is present, ensure it is in a known SMIRNOFF encoding.
Valid formats are: ['XML']
Parsing failed with the following error:
syntax error: line 1, column 0

Warning: importing 'simtk.openmm' is deprecated.  Import 'openmm' instead.
/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/numpy/core/getlimits.py:499: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.
  setattr(self, word, getattr(machar, word).flat[0])
/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float32'> type is zero.
  return self._float_to_str(self.smallest_subnormal)
/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/numpy/core/getlimits.py:499: UserWarning: The value of the smallest subnormal for <class 'numpy.float64'> type is zero.
  setattr(self, word, getattr(machar, word).flat[0])
/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/numpy/core/getlimits.py:89: UserWarning: The value of the smallest subnormal for <class 'numpy.float64'> type is zero.
  return self._float_to_str(self.smallest_subnormal)
/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/heterograph.py:72: DGLWarning: Recommend creating graphs by `dgl.graph(data)` instead of `dgl.DGLGraph(data)`.
  dgl_warning('Recommend creating graphs by `dgl.graph(data)`'
/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/convert.py:997: DGLWarning: dgl.to_homo is deprecated. Please use dgl.to_homogeneous
  dgl_warning("dgl.to_homo is deprecated. Please use dgl.to_homogeneous")
Traceback (most recent call last):
  File "/lila/home/chodera/espaloma-aniline/espaloma-aniline.py", line 14, in <module>
    system = system_generator.create_system(openmm_topology, molecules=[molecule])
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/openmmforcefields/generators/system_generators.py", line 324, in create_system
    system = self.forcefield.createSystem(topology, **forcefield_kwargs)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/openmm/app/forcefield.py", line 1212, in createSystem
    templateForResidue = self._matchAllResiduesToTemplates(data, topology, residueTemplates, ignoreExternalBonds)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/openmm/app/forcefield.py", line 1417, in _matchAllResiduesToTemplates
    if generator(self, res):
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/openmmforcefields/generators/template_generators.py", line 304, in generator
    ffxml_contents = self.generate_residue_template(molecule)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/openmmforcefields/generators/template_generators.py", line 1630, in generate_residue_template
    self.espaloma_model(molecule_graph.heterograph)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/torch/nn/modules/container.py", line 141, in forward
    input = module(input)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/espaloma/nn/sequential.py", line 144, in forward
    x = self._sequential(g_, x)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/espaloma/nn/sequential.py", line 62, in forward
    x = getattr(self, exe)(g, x)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/espaloma/nn/layers/dgl_legacy.py", line 46, in forward
    return self.gn(g, x)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/nn/pytorch/conv/sageconv.py", line 235, in forward
    graph.update_all(msg_fn, fn.mean('m', 'neigh'))
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/heterograph.py", line 4895, in update_all
    ndata = core.message_passing(g, message_func, reduce_func, apply_node_func)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/core.py", line 357, in message_passing
    ndata = invoke_gspmm(g, mfunc, rfunc)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/core.py", line 332, in invoke_gspmm
    z = op(graph, x)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/ops/spmm.py", line 189, in func
    return gspmm(g, 'copy_lhs', reduce_op, x, None)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/ops/spmm.py", line 75, in gspmm
    ret = gspmm_internal(g._graph, op,
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/backend/pytorch/sparse.py", line 757, in gspmm
    return GSpMM.apply(gidx, op, reduce_op, lhs_data, rhs_data)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/torch/cuda/amp/autocast_mode.py", line 94, in decorate_fwd
    return fwd(*args, **kwargs)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/backend/pytorch/sparse.py", line 126, in forward
    out, (argX, argY) = _gspmm(gidx, op, reduce_op, X, Y)
  File "/lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/sparse.py", line 228, in _gspmm
    _CAPI_DGLKernelSpMM(gidx, op, reduce_op,
  File "dgl/_ffi/_cython/./function.pxi", line 287, in dgl._ffi._cy3.core.FunctionBase.__call__
  File "dgl/_ffi/_cython/./function.pxi", line 232, in dgl._ffi._cy3.core.FuncCall
  File "dgl/_ffi/_cython/./base.pxi", line 155, in dgl._ffi._cy3.core.CALL
dgl._ffi.base.DGLError: [11:21:00] /opt/dgl/src/array/cpu/./spmm_blocking_libxsmm.h:267: Failed to generate libxsmm kernel for the SpMM operation!
Stack trace:
  [bt] (0) /lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/libdgl.so(dmlc::LogMessageFatal::~LogMessageFatal()+0x4f) [0x2ad0bac7000f]
  [bt] (1) /lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/libdgl.so(void dgl::aten::cpu::SpMMRedopCsrOpt<long, float, dgl::aten::cpu::op::CopyLhs<float>, dgl::aten::cpu::op::Add<float> >(dgl::BcastOff const&, dgl::aten::CSRMatrix const&, dgl::runtime::NDArray, dgl::runtime::NDArray, dgl::runtime::NDArray, dgl::runtime::NDArray, dgl::runtime::NDArray)+0x3d4) [0x2ad0baeb07f4]
  [bt] (2) /lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/libdgl.so(void dgl::aten::cpu::SpMMSumCsrLibxsmm<long, float, dgl::aten::cpu::op::CopyLhs<float> >(dgl::BcastOff const&, dgl::aten::CSRMatrix const&, dgl::runtime::NDArray, dgl::runtime::NDArray, dgl::runtime::NDArray)+0x73) [0x2ad0baeb08a3]
  [bt] (3) /lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/libdgl.so(void dgl::aten::cpu::SpMMSumCsr<long, float, dgl::aten::cpu::op::CopyLhs<float> >(dgl::BcastOff const&, dgl::aten::CSRMatrix const&, dgl::runtime::NDArray, dgl::runtime::NDArray, dgl::runtime::NDArray)+0x12f) [0x2ad0baecbeaf]
  [bt] (4) /lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/libdgl.so(void dgl::aten::SpMMCsr<1, long, 32>(std::string const&, std::string const&, dgl::BcastOff const&, dgl::aten::CSRMatrix const&, dgl::runtime::NDArray, dgl::runtime::NDArray, dgl::runtime::NDArray, std::vector<dgl::runtime::NDArray, std::allocator<dgl::runtime::NDArray> >)+0xcd3) [0x2ad0baee2203]
  [bt] (5) /lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/libdgl.so(dgl::aten::SpMM(std::string const&, std::string const&, std::shared_ptr<dgl::BaseHeteroGraph>, dgl::runtime::NDArray, dgl::runtime::NDArray, dgl::runtime::NDArray, std::vector<dgl::runtime::NDArray, std::allocator<dgl::runtime::NDArray> >)+0x13d5) [0x2ad0baf14455]
  [bt] (6) /lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/libdgl.so(+0x46a8d8) [0x2ad0baf288d8]
  [bt] (7) /lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/libdgl.so(+0x46ae71) [0x2ad0baf28e71]
  [bt] (8) /lila/home/chodera/miniconda/envs/espaloma-perses/lib/python3.9/site-packages/dgl/libdgl.so(DGLFuncCall+0x48) [0x2ad0baf7c058]

Can we make sure we can handle CPU-only as well?

Here’s the script:

# Create OpenFF Molecule and conformers                                                                                                                                                                                            
from openff.toolkit.topology import Molecule
smiles = 'Nc1ccccc1' # aniline                                                                                                                                                                                                     
molecule = Molecule.from_smiles(smiles)
molecule.generate_conformers()

# Parameterize it                                                                                                                                                                                                                  
from openmmforcefields.generators import SystemGenerator
from openmm import app
from openmm import unit
openmm_topology = molecule.to_topology().to_openmm()
forcefield_kwargs = { 'constraints' : app.HBonds, 'rigidWater' : True, 'removeCMMotion' : False, 'hydrogenMass' : 4*unit.amu }
system_generator = SystemGenerator(small_molecule_forcefield='espaloma-0.2.2', forcefield_kwargs=forcefield_kwargs)
system = system_generator.create_system(openmm_topology, molecules=[molecule])

# Create OpenMM Context                                                                                                                                                                                                            
from simtk import unit
import openmm
temperature = 300 * unit.kelvin
collision_rate = 1.0 / unit.picoseconds
timestep = 4.0 * unit.femtoseconds
integrator = openmm.LangevinMiddleIntegrator(temperature, collision_rate, timestep)
context = openmm.Context(system, integrator)
context.setPositions(molecule.conformers[0])

# Minimize                                                                                                                                                                                                                         
openmm.LocalEnergyMinimizer.minimize(context)

# Write structure                                                                                                                                                                                                                  
from openmm.app import PDBFile
positions = context.getState(getPositions=True).getPositions(asNumpy=True)
with open('aniline-minimized.pdb', 'wt') as outfile:
    PDBFile.writeFile(openmm_topology, positions, outfile)

# Simulate                                                                                                                                                                                                                         
print('Simulating...')
integrator.step(25000)

# Write structure                                                                                                                                                                                                                  
from openmm.app import PDBFile
positions = context.getState(getPositions=True).getPositions(asNumpy=True)
with open('aniline-100ps.pdb', 'wt') as outfile:
    PDBFile.writeFile(openmm_topology, positions, outfile)

[jchodera]

This was on lilac, so it should be an x86 CPU.

[mikemhenry]

Okay cool, that might be the issue! Those virtual CPUs basically have the absolute minimum flags exposed so they can migrate a host from an intel server to an amd server without issues (could probably ask the admins to expose more since I doubt they use/need that functionality and users doing -march=native are going to end up with code as slow as molasses).