GROMACS와 DeePMDkit 연동을 통한 AI 딥포텐셜 기반 고정밀 분자동역학

초록

본 연구는 최신 AI 딥포텐셜 모델(DPA2, DPA3)을 DeePMD‑kit을 통해 GROMACS에 통합하여, GPU 가속 환경에서 양자역학 수준의 정확도를 유지하면서도 전통적인 DFT 기반 AIMD보다 훨씬 높은 처리량을 달성한다. 네 가지 단백질‑수계 벤치마크(1YRF, 1UBQ, 3LZM, 2PTC)를 이용해 A100 및 GH200 GPU에서 성능을 평가했으며, DPA2가 DPA3보다 3‑4배 빠른 결과를 보였다. 또한 커널 런치 오버헤드와 도메인 분할 추론이 주요 병목임을 규명하고 향후 최적화 방향을 제시한다.

상세 분석

이 논문은 GROMACS라는 고성능 분자동역학 엔진에 DeePMD‑kit을 C++/CUDA 백엔드로 직접 연결함으로써, 기존에 Python‑중심이던 딥포텐셜 추론 파이프라인을 생산 환경에 맞게 재구성했다는 점에서 의미가 크다. 핵심 설계는 NNPot 모듈에 새로운 백엔드를 추가하고, MPI 기반 도메인 분할 시 발생하는 ‘ghost atom’ 문제를 해결하기 위해 서브도메인 간 원자 정보를 확대·재구성한 것이다. 특히 DPA2와 DPA3는 각각 어텐션 메커니즘과 그래프 신경망(GNN)을 기반으로 하는 대규모 원자 모델(LAM)이며, 이들 모델은 기존 딥포텐셜보다 더 복잡한 이웃 정보를 요구한다. 논문은 이를 위해 L × Rcut 영역까지 확장된 ghost 영역을 제공하고, 비선형 연산 특성상 양쪽 서브도메인 모두에서 완전한 원자 정보를 갖도록 통신을 추가함으로써 정확성을 유지한다.

성능 평가에서는 네 가지 단백질‑수계 시스템을 대상으로 A100과 차세대 GH200 GPU에서 실험했으며, DPA2가 DPA3보다 A100에서 최대 4.23배, GH200에서 3.18배 높은 스루풋을 기록했다. 메모리 사용량 측면에서도 DPA2가 상대적으로 효율적이었으며, 커널 레벨 프로파일링 결과 대부분의 실행 시간이 커널 런치 오버헤드와 작은 배치 크기로 인한 비효율적인 GPU 활용에 기인함을 확인했다. 또한 도메인‑분해 추론을 현재는 단일 MPI 랭크에서 수행하도록 설계했지만, 이는 확장성에 제한을 두는 요소로 지적한다.

이러한 분석을 토대로 저자들은 향후 최적화 과제로 (1) 커널 런치 수를 감소시키기 위한 배치 처리와 커스텀 CUDA 커널 설계, (2) 서브도메인 별 병렬 추론을 지원하는 분산 inference 프레임워크 구축, (3) 메모리 프리패치와 스트리밍을 통한 데이터 이동 최소화 등을 제시한다. 결과적으로 AI 딥포텐셜을 기존 MD 파이프라인에 원활히 통합함으로써, 양자역학 정확도를 유지하면서도 대규모 생물학적 시스템을 실시간에 가깝게 시뮬레이션할 수 있는 기반을 마련했다.