[Paper Review] OpenMM-Python-Force: Deploying Accelerated Python Modules in Molecular Dynamics Simulation

Zhi Wang, Wen Yan|arXiv (Cornell University)|Dec 24, 2024

Machine Learning in Materials ScienceMaterials Science3 citations

TL;DR

This paper introduces OpenMM-Python-Force, a plugin that enables seamless integration of accelerated Python-based machine learning models—via PyTorch's torch.compile and CUDA Graph—into molecular dynamics simulations using a CPython C-API callback mechanism. The approach achieves up to 8.2× speedup over baseline implementations while maintaining numerical accuracy across classical and ab initio MD simulations.

ABSTRACT

We present OpenMM-Python-Force, a plugin designed to extend OpenMM's functionality by enabling integration of energy and force calculations from external Python programs via a callback mechanism. During molecular dynamics simulations, data exchange can be implemented through torch.Tensor or numpy.ndarray, depending on the specific use case. This enhancement significantly expands OpenMM's capabilities, facilitating seamless integration of accelerated Python modules within molecular dynamics simulations. This approach represents a general solution that can be adapted to other molecular dynamics engines beyond OpenMM. The source code is openly available at https://github.com/bytedance/OpenMM-Python-Force.

Motivation & Objective

• Address the technical divergence between C-family MD engines and Python-based ML frameworks, which limits integration of modern ML models into molecular dynamics simulations.
• Overcome the limitations of existing tools like torch.jit.script, which fail on ~50% of real-world models due to strict syntax restrictions.
• Enable efficient, production-ready deployment of optimized ML models (e.g., via torch.compile and CUDA Graph) within MD simulations without sacrificing portability or ease of use.
• Provide a general, extensible solution that supports both PyTorch and NumPy-based models, and is compatible with other MD engines beyond OpenMM.
• Demonstrate end-to-end performance and accuracy across diverse MD workloads, including gas-phase and ab initio molecular dynamics (AIMD) simulations.

Proposed method

• Leverage the CPython C-API to capture the PyObject pointer of a callable Python object (e.g., a PyTorch model) using Python’s built-in `id()` function, enabling low-level C++ access to the model.
• Implement a custom `Callable` class in C++ that stores the model’s PyObject pointer, input/output tensors, and execution parameters, enabling direct invocation of Python functions from C++.
• Use pybind11 to bind the C++ callback mechanism to OpenMM’s force interface, allowing the plugin to be integrated into standard OpenMM simulation workflows with minimal code changes.
• Support both PyTorch and NumPy-based models by abstracting data exchange through `torch.Tensor` or `numpy.ndarray`, enabling flexible interoperability with quantum chemistry packages like PySCF/GPU4PySCF.
• Enable performance optimization via PyTorch’s `torch.compile` and CUDA Graph, which reduce kernel launch overhead and enable kernel fusion, significantly improving inference speed.
• Design the plugin to be embeddable in other MD engines (e.g., Tinker, LAMMPS) by exposing a minimal C-API surface for Python interpreter initialization and function invocation.

Figure 1 : Illustration of the Python callback mechanism, demonstrating the translation between a Python function call and its corresponding pseudo C/C++ implementations using either the CPython API or pybind11 (with the C++ namespace pybind11 abbreviated as py ).

Experimental results

Research questions

• RQ1 Can a general-purpose callback mechanism be designed to enable high-performance, production-ready integration of Python-based ML models into C-based molecular dynamics simulations?
• RQ2 To what extent can PyTorch’s `torch.compile` and CUDA Graph improve performance in MD simulations without requiring changes to the underlying MD engine?
• RQ3 How does the numerical accuracy of the callback-based approach compare to native implementations, especially when using advanced compilation and optimization techniques?
• RQ4 Can the same mechanism support diverse data types and backends, such as PyTorch tensors and NumPy arrays, across classical and ab initio MD simulations?
• RQ5 Is the proposed architecture portable and extensible to other MD engines that do not natively support Python embedding?

Key findings

• The OpenMM-Python-Force plugin achieved up to 8.2× speedup over the baseline implementation (no optimization) in ethanol simulations, with 0.486 ms/step and 178 million steps per day.
• The use of `torch.compile` reduced inference time from 3.97 ms/step (baseline) to 3.07 ms/step, achieving a 1.3× speedup, and further optimization with CUDA Graph and AOT compilation led to 8.2× speedup in total.
• The numerical accuracy of the callback mechanism was validated to be on par with native implementations, with no significant discrepancies observed when comparing results from `torch.compile` and other deployment strategies.
• The performance advantage of `torch.compile` was most pronounced in small-scale systems (e.g., single ethanol molecule), where kernel launch overhead dominates, with a 3.5× speedup over `torch.jit.script`.
• The mechanism successfully enabled ab initio molecular dynamics using PySCF/GPU4PySCF via the `NumPyForce` plugin, demonstrating compatibility with non-differentiable, quantum chemistry-based force calculations.
• The architecture is extensible to other MD engines: embedding a Python interpreter via CPython C-API (e.g., `Py_Initialize`) is feasible and requires minimal modifications, similar to initializing Fortran runtimes in Tinker.