[論文レビュー] 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

ひとこと要約

本論文では、PyTorchのtorch.compileおよびCUDA Graphを介して、加速されたPythonベースの機械学習モデルをCPython C-APIコールバック機構を用いて分子動力学シミュレーションにシームレスに統合できるプラグイン、OpenMM-Python-Forceを紹介する。この手法は、古典的およびアビ・イニオ・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ベースのMLフレームワークの間の技術的乖離を解消し、現代のMLモデルを分子動力学シミュレーションに統合するのを妨げる要因を除去すること。
• torch.jit.scriptのような既存のツールは、厳格な構文制限のため、実世界のモデルの約50％で失敗するという限界を克服すること。
• torch.compileおよびCUDA Graphによる最適化（例：AOTコンパイル）を活用した効率的で生産環境向けのMLモデルデプロイメントを可能にし、ポータビリティや使いやすさを損なわずにMDシミュレーション内に統合すること。
• PyTorchおよびNumPyベースのモデルを両方サポートする汎用的で拡張可能なソリューションを提供し、OpenMMに限らず他のMDエンジンとも互換性を持つこと。
• ガス相およびアビ・イニオ分子動力学（AIMD）シミュレーションを含む多様なMDワークロードにおいて、エンドツーエンドのパフォーマンスと正確性を実証すること。

提案手法

• CPython C-APIを活用し、Python組み込み関数`id()`を用いて、呼び出し可能なPythonオブジェクト（例：PyTorchモデル）のPyObjectポインタを取得することで、低レベルなC++からのモデルアクセスを可能にする。
• C++で独自の`Callable`クラスを実装し、モデルのPyObjectポインタ、入力/出力テンソル、実行パラメータを格納することで、C++から直接Python関数を呼び出せるようにする。
• pybind11を用いてC++コールバックメカニズムをOpenMMのフォースインターフェースにバインドし、標準的なOpenMMシミュレーションワークフローに最小限のコード変更で統合可能なプラグインを実現する。
• データ交換を`torch.Tensor`または`numpy.ndarray`を抽象化することで、PyTorchおよびNumPyベースのモデルを両方サポートし、量子化学パッケージ（例：PySCF/GPU4PySCF）との柔軟な相互運用性を実現する。
• PyTorchの`torch.compile`およびCUDA Graphを活用してパフォーマンス最適化を実現し、カーネル起動オーバーヘッドを低減し、カーネル統合を可能にすることで、推論速度を顕著に向上させる。
• Pythonインタプリタの初期化および関数呼び出しに最小限のC-API表面を公開することで、他のMDエンジン（例：Tinker、LAMMPS）への埋め込みが可能であり、同様にFortranランタイムの初期化と同様に最小限の変更で実現可能である。

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 Cベースの分子動力学シミュレーションに、PythonベースのMLモデルを高パフォーマンスかつ生産環境向けに統合できる汎用的なコールバックメカニズムを設計できるか？
• RQ2 PyTorchの`torch.compile`およびCUDA Graphは、下位のMDエンジンに変更を加えずに、MDシミュレーションにおけるパフォーマンスをどの程度向上できるか？
• RQ3 高度なコンパイルおよび最適化技術を用いた場合でも、コールバックベースのアプローチの数値的正確性は、ネイティブ実装と比べてどの程度の差があるか？
• RQ4 同じメカニズムが、古典的およびアビ・イニオMDシミュレーションにおいて、PyTorchテンソルやNumPy配列を含む多様なデータ型およびバックエンドをサポートできるか？
• RQ5 本提案アーキテクチャは、Python埋め込みをネイティブにサポートしない他のMDエンジンに対しても、ポータブルかつ拡張可能か？

主な発見

• OpenMM-Python-Forceプラグインは、エタノールシミュレーションにおいて、最適化なしのベースライン実装比で最大8.2倍の高速化を達成し、1ステップあたり0.486 ms、1日17800万ステップの性能を発揮した。
• `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を用いたアビ・イニオ分子動力学が成功裏に実行可能となり、非微分可能な量子化学ベースの力計算とも互換性があることが実証された。
• アーキテクチャは他のMDエンジンへも拡張可能である：CPython C-API（例：`Py_Initialize`）を用いたPythonインタプリタの埋め込みは実現可能で、TinkerにおけるFortranランタイム初期化と同様に、最小限の変更で実現可能である。