[Paper Review] Enhancing GPU-acceleration in the Python-based Simulations of Chemistry Framework

Xiaojie Wu, Qiming Sun+12 more|arXiv (Cornell University)|Apr 15, 2024

Machine Learning in Materials Science9 citations

TL;DR

GPU4PySCF is a GPU-accelerated Python package built on PySCF that delivers large speedups for DFT, SCF, gradient, and Hessian calculations, with substantial cost savings and ecosystem integration. It supports implicit solvent models and various quantum-chemistry workflows, validated against 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.

Motivation & Objective

• Motivate the need for GPU acceleration in quantum chemistry calculations to meet data-driven research demand.
• Develop a Python-based, GPU-optimized framework built on PySCF to improve performance for DFT/SCF workflows.
• Demonstrate integration with open-source ecosystems and applicability to industrial tasks such as PES exploration, solvation energies, and neural-network workflows.

Proposed method

• Prioritize density fitting algorithms to leverage GPU efficiency and provide limited direct SCF support.
• Adapt the codebase to modern GPU architectures, including tensor cores, to accelerate tensor contractions.
• Integrate a broad set of open-source quantum chemistry packages within GPU4PySCF for DFT methods, basis sets, solvent models, and property calculations.
• Cross-validate results with Q-Chem 6.1 to ensure reliability and accuracy.

Experimental results

Research questions

• RQ1 What performance gains (speedups) are achievable by GPU4PySCF for SCF, gradient, and Hessian calculations on modern GPUs compared to 32-core CPUs?
• RQ2 How do implicit solvent models (PCM/SMD) impact performance and accuracy in GPU-accelerated DFT calculations?
• RQ3 Can GPU4PySCF enable large-scale tasks such as torsion scans, dimer interaction energies, and solvation-energy computations within practical timeframes?

Key findings

• GPU4PySCF achieves approximately 20x speedup for SCF, around 50x for Hessian calculations, and more modest gains for gradients versus Q-Chem on 32 CPU cores.
• With NVIDIA A100-80G and density fitting, SCF/gradient/Hessian speedups in solvent models reach 40-80x, 20-40x, and 100-170x respectively, enabling large-scale vibrational analyses in a day for mid-sized molecules.
• Cost savings can be up to ~90% for many tasks on A100 GPUs compared to CPU-based execution, with gradient tasks in pure DFT showing around 70% savings.
• CHELPG charges speed improvements are about 10x faster on GPU4PySCF than CPU-based software.
• NMR shielding constants and other properties are implemented with GPU acceleration (GIAO-based) to achieve efficient calculations.
• Solvation free energies with implicit solvent models (PCM/SMD) generally achieve errors below 1 kcal/mol for neutrals and around 4 kcal/mol for ions relative to reference methods, with higher accuracy using larger basis sets and higher-level functionals.