[论文解读] 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 Science被引用 3

一句话总结

本文介绍了 OpenMM-Python-Force 插件，该插件通过 PyTorch 的 torch.compile 和 CUDA Graph 技术，利用 CPython C-API 回调机制，实现了加速的基于 Python 的机器学习模型与分子动力学模拟的无缝集成。该方法在经典和从头算 MD 模拟中均保持了数值精度的同时，相较于基线实现最高实现了 8.2 倍的性能提升。

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.

研究动机与目标

• 解决 C 族 MD 引擎与基于 Python 的机器学习框架之间的技术差异，该差异限制了现代机器学习模型在分子动力学模拟中的集成。
• 克服现有工具（如 torch.jit.script）的局限性，后者因严格的语法限制导致约 50% 的真实世界模型无法使用。
• 在不牺牲可移植性或易用性的前提下，实现对优化机器学习模型（如通过 torch.compile 和 CUDA Graph）的高效、生产就绪部署，集成于 MD 模拟中。
• 提供一种通用且可扩展的解决方案，支持 PyTorch 和 NumPy 基础模型，并与 OpenMM 以外的其他 MD 引擎兼容。
• 在包括气相和从头算分子动力学（AIMD）模拟在内的多种 MD 工作负载中，端到端地展示性能与精度。

提出的方法

• 利用 CPython C-API 捕获可调用 Python 对象（如 PyTorch 模型）的 PyObject 指针，通过 Python 内置的 `id()` 函数实现，从而实现 C++ 层面对模型的低级访问。
• 在 C++ 中实现自定义的 `Callable` 类，用于存储模型的 PyObject 指针、输入/输出张量及执行参数，从而实现在 C++ 中直接调用 Python 函数。
• 使用 pybind11 将 C++ 回调机制绑定至 OpenMM 的力接口，使插件能够以最小代码修改的方式无缝集成到标准 OpenMM 模拟工作流中。
• 通过抽象化 `torch.Tensor` 或 `numpy.ndarray` 的数据交换，支持 PyTorch 和 NumPy 基础模型，实现与量子化学软件包（如 PySCF/GPU4PySCF）的灵活互操作。
• 通过 PyTorch 的 `torch.compile` 和 CUDA Graph 实现性能优化，减少内核启动开销并支持内核融合，显著提升推理速度。
• 设计插件时考虑可嵌入性，通过暴露最小化的 C-API 表面，支持在其他 MD 引擎（如 Tinker、LAMMPS）中使用，包括 Python 解释器初始化和函数调用。

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 ).

实验结果

研究问题

• RQ1 能否设计一种通用的回调机制，以实现高性能、生产就绪的基于 Python 的机器学习模型与基于 C 的分子动力学模拟的集成？
• RQ2 在不修改底层 MD 引擎的前提下，PyTorch 的 `torch.compile` 和 CUDA Graph 能在多大程度上提升 MD 模拟的性能？
• RQ3 与原生实现相比，基于回调的机制在使用高级编译和优化技术时，其数值精度如何？
• RQ4 该机制是否能够支持多种数据类型和后端（如 PyTorch 张量和 NumPy 数组），并适用于经典和从头算 MD 模拟？
• RQ5 所提出的架构是否具备可移植性和可扩展性，可适用于那些不原生支持 Python 嵌入的其他 MD 引擎？

主要发现

• 在乙醇模拟中，OpenMM-Python-Force 插件相较于基线实现（无优化）最高实现了 8.2 倍的性能提升，达到 0.486 ms/步，每日可计算 17.8 亿步。
• 使用 `torch.compile` 将推理时间从基线的 3.97 ms/步降低至 3.07 ms/步，实现 1.3 倍性能提升；进一步结合 CUDA Graph 和 AOT 编译后，总性能提升达 8.2 倍。
• 回调机制的数值精度经验证与原生实现相当，比较 `torch.compile` 与其他部署策略的结果时未发现显著差异。
• `torch.compile` 的性能优势在小系统（如单个乙醇分子）中最为显著，此时内核启动开销占主导地位，相较 `torch.jit.script` 实现了 3.5 倍的性能提升。
• 该机制成功通过 `NumPyForce` 插件实现了基于 PySCF/GPU4PySCF 的从头算分子动力学，证明了其与非可微分、基于量子化学的力计算的兼容性。
• 该架构具备可扩展性：通过 CPython C-API（如 `Py_Initialize`）嵌入 Python 解释器在其他 MD 引擎中是可行的，且仅需极少修改，类似于在 Tinker 中初始化 Fortran 运行时。