Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Accelerating Density Fitting with Adaptive-precision and 8-bit Integer on AI Accelerators

Hua Huang, Wekai Shao, Jeff Hammond

TL;DR

The paper addresses the high computational cost of density fitting in Gaussian-basis DFT by introducing an adaptive-precision scheme that accelerates the K-matrix build using INT8 GEMM emulation on Tensor Cores, while keeping the J and XC-related steps in FP64. The method adaptively adjusts emulation precision based on SCF energy changes, achieving up to 364% speedup on high-end GPUs without sacrificing converged energies. Implemented in PySCF with CuPy, the approach demonstrates robustness across diverse molecules and basis sets and shows promising scalability with molecular size, particularly on GPUs with strong Tensor Core support. This work highlights a practical pathway to harness AI accelerator hardware for reliable quantum chemistry simulations, opening avenues for broader adoption and further optimization in scientific computing contexts.

Abstract

The emergence of artificial intelligence (AI) accelerators like NVIDIA Tensor Cores offers new opportunities to speed up tensor-heavy scientific computations. However, applying them to quantum chemistry is challenging due to strict accuracy demands and irregular data patterns. We propose an adaptive precision algorithm to accelerate the density fitting (DF) method with Gaussian basis sets on AI accelerators using 8-bit integer (INT8) arithmetics. Implemented in the GPU-accelerated PySCF package, the algorithm is tested on more than twenty molecular systems with different NVIDIA GPUs. Compared to the standard FP64 code, our algorithm is up to 204\% faster on a RTX 4090 gaming GPU and up to 364\% faster on a RTX 6000 Ada workstation GPU without compromising the converged energy. This work demonstrates a practical approach to use AI hardware for reliable quantum chemistry simulations.

Accelerating Density Fitting with Adaptive-precision and 8-bit Integer on AI Accelerators

TL;DR

The paper addresses the high computational cost of density fitting in Gaussian-basis DFT by introducing an adaptive-precision scheme that accelerates the K-matrix build using INT8 GEMM emulation on Tensor Cores, while keeping the J and XC-related steps in FP64. The method adaptively adjusts emulation precision based on SCF energy changes, achieving up to 364% speedup on high-end GPUs without sacrificing converged energies. Implemented in PySCF with CuPy, the approach demonstrates robustness across diverse molecules and basis sets and shows promising scalability with molecular size, particularly on GPUs with strong Tensor Core support. This work highlights a practical pathway to harness AI accelerator hardware for reliable quantum chemistry simulations, opening avenues for broader adoption and further optimization in scientific computing contexts.

Abstract

The emergence of artificial intelligence (AI) accelerators like NVIDIA Tensor Cores offers new opportunities to speed up tensor-heavy scientific computations. However, applying them to quantum chemistry is challenging due to strict accuracy demands and irregular data patterns. We propose an adaptive precision algorithm to accelerate the density fitting (DF) method with Gaussian basis sets on AI accelerators using 8-bit integer (INT8) arithmetics. Implemented in the GPU-accelerated PySCF package, the algorithm is tested on more than twenty molecular systems with different NVIDIA GPUs. Compared to the standard FP64 code, our algorithm is up to 204\% faster on a RTX 4090 gaming GPU and up to 364\% faster on a RTX 6000 Ada workstation GPU without compromising the converged energy. This work demonstrates a practical approach to use AI hardware for reliable quantum chemistry simulations.
Paper Structure (16 sections, 11 equations, 2 figures, 4 tables)

This paper contains 16 sections, 11 equations, 2 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Density Fitting
  4. Single- and Mixed-Precision Algorithms for Quantum Chemistry
  5. Floating-point Arithmetic and Low-precision Emulation
  6. Accelerate DF with Adaptive Precision and INT8 GEMMs on Tensor Cores
  7. Justification of Accelerating only DF $K$ Matrix Build
  8. Adaptive Precision DF $K$ Matrix Build with INT8 GEMM
  9. Implementation Details
  10. Numerical Experiments
  11. Experiment Setup
  12. Single SCF Iteration Performance Improvement
  13. DFT Calculation Convergence
  14. DFT Calculation Performance Improvement
  15. Implementing the Adaptive Precision Algorithm with Other Libraries
  16. ...and 1 more sections

Figures (2)

  • Figure 1: Running time breakdowns (in milliseconds) of the second SCF iteration on different GPUs and molecular systems.
  • Figure 2: Speedup of DFT calculations when using the adaptive precision method over using the standard FP64 algorithm. All computational steps, including the integrals $(ij|p)$ and the construction of the $B_{ij}^q$ tensor in Formula (\ref{['equ:cderi']}), are included. Missing data points indicate that the GPU did not have sufficient memory to complete the corresponding calculation.