[논문 리뷰] Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learning

Weile Jia, Handong Wang외 6명|arXiv (Cornell University)|2020. 05. 01.

Machine Learning in Materials Science참고 문헌 72인용 수 35

한 줄 요약

제시된 기계학습 기반 MD 프로토콜 (Deep Potential Molecular Dynamics)이 1억 개가 넘는 원자 시스템에 대해 ab initio 정확도를 달성하고, Summit에서 highly optimized DeePMD-kit으로 최대 91 PFLOPS 이중 정밀도에 도달하며 나노초 규모의 시뮬레이션을 하루에 가능하게 한다는 것을 보여준다.

ABSTRACT

For 35 years, {\it ab initio} molecular dynamics (AIMD) has been the method of choice for modeling complex atomistic phenomena from first principles. However, most AIMD applications are limited by computational cost to systems with thousands of atoms at most. We report that a machine learning-based simulation protocol (Deep Potential Molecular Dynamics), while retaining {\it ab initio} accuracy, can simulate more than 1 nanosecond-long trajectory of over 100 million atoms per day, using a highly optimized code (GPU DeePMD-kit) on the Summit supercomputer. Our code can efficiently scale up to the entire Summit supercomputer, attaining $91$ PFLOPS in double precision ($45.5\%$ of the peak) and {$162$/$275$ PFLOPS in mixed-single/half precision}. The great accomplishment of this work is that it opens the door to simulating unprecedented size and time scales with {\it ab initio} accuracy. It also poses new challenges to the next-generation supercomputer for a better integration of machine learning and physical modeling.

연구 동기 및 목표

• ab initio 분자 동역학(AIMD)을 대규모 시스템 크기(수십~수백만 원자)로 확장하고 정확도를 희생하지 않으면서 긴 시간 척도를 달성한다.
• 비싼 전자구조 계산을 ab initio 충실성을 유지하는 대체 모델로 대체하기 위해 기계학습을 활용한다.
• 현대 슈퍼컴퓨터의 이종 GPU 아키텍처를 최대한 활용하도록 HPC 지향 구현을 개발·최적화한다.

제안 방법

• 원자 간 포텐셜을 원자 에너지의 합으로 표현하는 Deep Potential(DP) 신경망을 이용한다.
• ab initio 데이터로 DP 모델을 학습하고 LAMMPS와 통합된 DeePMD-kit에 구현하여 MD 적분을 수행한다.
• Summit에서 GPU 처리량을 극대화하기 위해 데이터 레이아웃, 이웃 리스트, 맞춤 TensorFlow 연산을 최적화한다.
• 정밀도 손실 없이 계산 속도를 높이기 위해 혼합 정밀도 스킴(MIX-32 및 MIX-16)을 사용한다.
• GEMM 기반 대체, TANH 및 TANHGrad의 융합, CUDA 커널 융합으로 TensorFlow 연산의 효율을 크게 향상시켜 MD 루프의 속도 향상을 달성한다.

실험 결과

연구 질문

• RQ1 깊은 학습 기반의 원자 간 포텐셜이 매우 큰 시스템(tens~hundreds of millions 원자)과 긴 MD 트레이젝토리에서 ab initio 정확성을 재현할 수 있는가?
• RQ2 Summit과 같은 이종 슈퍼컴퓨터에서 DP 기반 MD를 확장하기 위해 필요한 HPC 최적화 및 혼합 정밀도 전략은 무엇인가?
• RQ3 DP-MD를 극한의 시스템 크기로 확장할 때 실질적인 성능 이득(FLOPS, 해결 시간)과 정확도 간의 상충은 무엇인가?
• RQ4 재료 및 액체 시스템의 확장성 및 충실도 면에서 DP-MD가 전통적 AIMD 및 경험적 포스 필드와 비교하여 어떻게 되는가?

주요 결과

• DP-MD는 Summit에서 4,560 노드에서 이중 정밀도 91 PFLOPS 및 혼합 단일/반정밀도에서 162 PFLOPS / 275 PFLOPS로 초대형 AIMD 유사 궤적을 하루에 가능하게 한다.
• 127-million-atom 구리 시스템의 경우 한 스텝당 해결 시간은 이중 정밀도 기준 8.1e-10 s/스텝/원자에 해당하며, 하루에 약 0.8 ns에 해당하고 혼합-반정밀도에서는 하루에 2.5 ns이다.
• 이웃 리스트 데이터 레이아웃 최적화와 이웃 엔트리의 64비트 정수 압축은 분기 조건을 제거하고 GPU 효율을 향상시키며, 단일 GPU에서 커스텀 TensorFlow 연산에 대해 64.6배의 속도 향상을 제공한다(기본 CPU 구현 대비).
• MATMUL 및 SUM을 GEMM으로 대체하고 TANH 및 TANHGrad를 융합시키며 CUDA 커널 융합을 적용하면 TensorFlow 연산의 효율이 크게 향상되어 전체 MD 루프의 속도 증가에 기여한다.
• 혼합 정밀도 스킴 MIX-32 및 MIX-16은 테스트된 물 구성에서 이중 정밀도에 상응하는 에너지/힘 정확도를 달성하며, MIX-32는 에너지 및 힘 오차 면에서 이중 정밀도와 일치한다.