[논문 리뷰] OpenMM-Python-Force: Deploying Accelerated Python Modules in Molecular Dynamics Simulation

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

Machine Learning in Materials ScienceMaterials Science인용 수 3

한 줄 요약

이 논문은 CPython C-API 콜백 메커니즘을 통해 PyTorch의 torch.compile 및 CUDA Graph를 활용한 가속화된 파이썬 기반 기계학습 모델을 분자 동역학 시뮬레이션에 원활하게 통합할 수 있도록 해주는 플러그인인 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 엔진과 파이썬 기반 기계학습 프레임워크 간의 기술적 이질성 문제를 해결하여 현대적인 기계학습 모델을 분자 동역학 시뮬레이션에 통합하는 데 어려움을 해소한다.
• 엄격한 문법 제약 조건으로 인해 실제 모델의 약 50%에서 실패하는 기존 도구인 torch.jit.script의 한계를 극복한다.
• torch.compile 및 CUDA Graph를 통한 최적화된 기계학습 모델의 효율적이고 프로덕션 수준의 배포를 가능하게 하되, 이식성이나 사용 편의성을 희생하지 않는다.
• PyTorch 및 NumPy 기반 모델을 모두 지원하는 일반적이고 확장 가능한 솔루션을 제공하며, OpenMM 이외의 다른 MD 엔진과의 호환성도 보장한다.
• 가스상 및 아비니오 분자 동역학(AIMD) 시뮬레이션을 포함한 다양한 MD 워크로드에서 종단 간 성능과 정확도를 입증한다.

제안 방법

• 파이썬 내장 함수인 `id()`를 사용하여 가시성 있는 파이썬 객체(예: PyTorch 모델)의 PyObject 포인터를 CPython C-API를 통해 캡처함으로써, 저수준의 C++에서 모델에 접근할 수 있도록 한다.
• C++에서 모델의 PyObject 포인터, 입력/출력 텐서, 실행 매개변수를 저장하는 사용자 정의 `Callable` 클래스를 구현하여 C++에서 파이썬 함수를 직접 호출할 수 있도록 한다.
• pybind11를 활용해 C++ 콜백 메커니즘을 OpenMM의 힘 인터페이스에 바인딩함으로써, 표준 OpenMM 시뮬레이션 워크플로우에 최소한의 코드 변경으로 플러그인을 통합할 수 있도록 한다.
• 데이터 교환을 `torch.Tensor` 또는 `numpy.ndarray`를 통해 추상화함으로써 PyTorch 및 NumPy 기반 모델을 모두 지원하며, PySCF/GPU4PySCF와 같은 양자화학 패키지와의 탄력적 상호운용성을 확보한다.
• PyTorch의 `torch.compile` 및 CUDA Graph를 활용해 성능 최적화를 실현함으로써 커널 실행 오버헤드를 감소시키고 커널 융합을 가능하게 하여 추론 속도를 크게 향상시킨다.
• CPython C-API를 통한 파이썬 인터프리터 초기화 및 함수 호출 기능을 최소한의 인터페이스로 노출시켜 다른 MD 엔진(예: Tinker, LAMMPS)에 통합 가능한 임베딩 기능을 제공한다.

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기반 분자 동역학 시뮬레이션에 고성능, 프로덕션 수준의 파이썬 기반 기계학습 모델 통합을 가능하게 할 수 있는가?
• RQ2 기본 MD 엔진을 수정하지 않고도 PyTorch의 `torch.compile` 및 CUDA Graph가 MD 시뮬레이션 성능에 얼마나 기여할 수 있는가?
• RQ3 고급 컴파일 및 최적화 기법을 사용할 경우, 콜백 기반 접근 방식의 수치 정확도가 네이티브 구현과 비교해 어느 정도 유지되는가?
• RQ4 동일한 메커니즘이 고전적 및 아비니오 MD 시뮬레이션 전반에서 파이토치 텐서와 NumPy 배열을 포함한 다양한 데이터 유형 및 백엔드를 지원할 수 있는가?
• RQ5 제안된 아키텍처는 파이썬 임bedding을 네이티브로 지원하지 않는 다른 MD 엔진에 이식 가능하고 확장 가능한가?

주요 결과

• 에탄올 시뮬레이션에서 OpenMM-Python-Force 플러그인은 최대 8.2배의 성능 향상을 달성했으며, 스텝당 0.486ms, 일일 17800만 스텝의 성능을 기록했다.
• `torch.compile` 사용으로 추론 시간이 기준값인 3.97ms/step에서 3.07ms/step로 감소하여 1.3배의 성능 향상을 달성했으며, CUDA Graph 및 AOT 컴파일을 추가로 적용함으로써 총 8.2배의 성능 향상이 이루어졌다.
• 콜백 기반 접근 방식의 수치 정확도는 네이티브 구현과 유사한 수준임이 검증되었으며, `torch.compile` 및 기타 배포 전략 간 결과 비교 시 유의미한 차이가 관찰되지 않았다.
• kernel 실행 오버헤드가 지배적인 소규모 시스템(예: 단일 에탄올 분자)에서는 `torch.compile`의 성능 향상 효과가 두드러져 `torch.jit.script` 대비 3.5배의 성능 향상을 기록했다.
• PySCF/GPU4PySCF를 통한 `NumPyForce` 플러그인을 통해 비미분 가능이고 양자화학 기반의 힘 계산도 성공적으로 지원함으로써, 비차별적 힘 계산과의 호환성이 입증되었다.
• 아키텍처는 확장 가능하며, Tinker와 유사하게 Fortran 런타임 초기화와 유사한 방식으로 CPython C-API를 통해 파이썬 인터프리터를 임베딩하는 것이 가능하며, 최소한의 수정만으로도 다른 MD 엔진에 통합이 가능하다.