[논문 리뷰] Enhancing GPU-acceleration in the Python-based Simulations of Chemistry Framework

Xiaojie Wu, Qiming Sun외 12명|arXiv (Cornell University)|2024. 04. 15.

Machine Learning in Materials Science인용 수 9

한 줄 요약

GPU4PySCF는 PySCF 기반의 GPU 가속 파이썬 패키지로, DFT, SCF, 그래디언트 및 해시안 계산에서 큰 속도향상과 상당한 비용 절감 및 생태계 통합을 제공합니다. 암시적 용매 모델과 다양한 양자화학 워크플로우를 지원하며 Q-Chem과의 검증을 거쳤습니다.

ABSTRACT

We describe our contribution as industrial stakeholders to the existing open-source GPU4PySCF project (https: //github.com/pyscf/gpu4pyscf), a GPU-accelerated Python quantum chemistry package. We have integrated GPU acceleration into other PySCF functionality including Density Functional Theory (DFT), geometry optimization, frequency analysis, solvent models, and density fitting technique. Through these contributions, GPU4PySCF v1.0 can now be regarded as a fully functional and industrially relevant platform which we demonstrate in this work through a range of tests. When performing DFT calculations on modern GPU platforms, GPU4PySCF delivers 30 times speedup over a 32-core CPU node, resulting in approximately 90% cost savings for most DFT tasks. The performance advantages and productivity improvements have been found in multiple industrial applications, such as generating potential energy surfaces, analyzing molecular properties, calculating solvation free energy, identifying chemical reactions in lithium-ion batteries, and accelerating neural-network methods. With the improved design that makes it easy to integrate with the Python and PySCF ecosystem, GPU4PySCF is natural choice that we can now recommend for many industrial quantum chemistry applications.

연구 동기 및 목표

• 데이터 주도 연구 수요를 충족하기 위해 양자 화학 계산에서 GPU 가속의 필요성을 제시합니다.
• DFT/SCF 워크플로의 성능 향상을 위해 PySCF를 기반으로 한 Python 기반의 GPU 최적화 프레임워크를 개발합니다.
• 개방형 소스 생태계와의 통합 및 PES 탐색, 용매화 에너지, 신경망 워크플로우와 같은 산업적 작업에의 적용 가능성을 입증합니다.

제안 방법

• GPU 효율을 활용하기 위해 밀도 적합(density fitting) 알고리즘을 우선적으로 사용하고, 직접 SCF 지원은 제한적으로 제공합니다.
• 텐서 코어를 포함한 현대 GPU 아키텍처에 코드베이스를 적응시키고 텐서 연산(텐서 수축)을 가속합니다.
• DFT 방법, 기저집합, 용매 모델 및 성질 계산을 위한 광범위한 오픈소스 양자화학 패키지를 GPU4PySCF에 통합합니다.
• 신뢰성과 정확성을 보장하기 위해 Q-Chem 6.1과의 교차 검증을 수행합니다.

실험 결과

연구 질문

• RQ1 GPU4PySCF가 32코어 CPU에 비해 현대 GPU에서 SCF, 그래디언트, 해시안 계산에서 어떤 속도향상을 달성할 수 있는가?
• RQ2 암시적 용매 모델(PCM/SMD)이 GPU 가속 DFT 계산의 성능 및 정확도에 어떤 영향을 미치는가?
• RQ3 GPU4PySCF가 큰 규모의 작업(터션 스캔, 이합체 상호작용 에너지, 용매화 에너지 계산)을 실용적인 시간 프레임 내에 가능하게 할 수 있는가?

주요 결과

분자	C-PCM_SCf	C-PCM_Gradient	C-PCM_Hessian	IEF-PCM_SCf	IEF-PCM_Gradient	IEF-PCM_Hessian
Vitamin C	41.1	28.4	106.4	41.6	27.3	104.9
Inosine	65.5	38.9	147.4	65.6	40.2	146.2
Bisphenol A	76.2	43.3	162.7	76.7	45.4	161.3
Mg Porphin	84.0	49.0	162.2	87.3	49.8	160.1
Penicillin V	83.2	44.1	170.5	84.0	43.4	169.4

• GPU4PySCF는 SCF에 대해 약 20배, 해시안 계산에는 약 50배의 속도향상을 달성하고, 그래디언트의 경우 32 CPU 코어의 Q-Chem 대비 더 보수적인 이득을 제공합니다.
• NVIDIA A100-80G 및 밀도 적합을 사용하면 용매 모델에서 SCF/그래디언트/해시안 속도향상은 각각 40-80x, 20-40x, 100-170x에 도달하여 중간 크기 분자에 대해 하루 내 대규모 진동 분석이 가능해집니다.
• CPU 기반 실행에 비해 A100 GPU에서 많은 작업에 대해 비용 절감이 최대 약 90%까지 가능하며, 순수 DFT의 그래디언트 작업에서 약 70%의 절감을 보입니다.
• CHELPG 전하 속도 향상은 GPU4PySCF가 CPU 기반 소프트웨어보다 약 10배 빠릅니다.
• NMR 차폐 상수 및 기타 특성은 GPU 가속(GIAO 기반)으로 구현되어 효율적인 계산을 달성합니다.
• 암시적 용매 모델(PCM/SMD)을 사용하는 용액화 자유 에너지는 중성에서 참조 방법 대비 일반적으로 오차가 1 kcal/mol 미만, 이온은 약 4 kcal/mol 정도로 더 높으며, 더 큰 기저집합과 더 높은 수준의 함수를 사용하면 정확도가 향상됩니다.