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

TL;DR

This paper introduces an adaptive-precision algorithm utilizing 8-bit integer arithmetic to accelerate density fitting in quantum chemistry on AI accelerators, achieving significant speedups without losing accuracy.

Contribution

It presents a novel adaptive-precision method implemented in GPU-accelerated PySCF, enabling faster density fitting on AI hardware while maintaining result accuracy.

Findings

Up to 204% speedup on RTX 4090 GPU.

Up to 364% speedup on RTX 6000 Ada GPU.

Maintains converged energy accuracy despite acceleration.

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.