[Paper Review] Accelerating Density Fitting with Adaptive-precision and 8-bit Integer on AI Accelerators
The paper presents an adaptive-precision density fitting algorithm using INT8 GEMMs on Tensor Cores, implemented in PySCF, achieving up to 3.04×–4.64× speedups over FP64 and preserving converged DFT energies.
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.
Motivation & Objective
- Motivate leveraging AI accelerator tensor cores for quantum chemistry under strict accuracy constraints.
- Develop an adaptive-precision density fitting approach using INT8 GEMMs for DF in Gaussian basis sets.
- Ensure convergence accuracy remains on par with FP64 baseline across diverse molecules and bases.
- Demonstrate implementation in PySCF and assess performance across multiple NVIDIA GPUs.
Proposed method
- Use FP64 to compute the Coulomb matrix J and INT8-emulated FP64 GEMMs to compute the exchange matrix K with adaptive precision.
- Adopt an adaptive precision scheme that selects the INT8 emulation level based on the relative SCF energy change (∆Erel) during SCF iterations.
- Switch back to FP64 when ∆Erel falls below a threshold to guarantee convergence within the same or two extra SCF iterations.
- Implement the method in PySCF with CuPy for GPU acceleration and adjust memory thresholds for large systems.
- Evaluate performance on RTX 4090, RTX 6000 Ada, and H100 across multiple basis sets and molecular systems.
Experimental results
Research questions
- RQ1Can INT8-emulated FP64 GEMMs accelerate the K-matrix construction in density fitting without sacrificing accuracy?
- RQ2Does an adaptive precision strategy preserve convergence efficiency and final energy accuracy across diverse molecular systems?
- RQ3What are the practical performance limits and decisions for when to revert to full FP64 precision?
Key findings
- Adaptive precision DF achieves up to 3.04× speedup on RTX 4090 and 4.64× on RTX 6000 Ada for DFT calculations.
- Converged energies with adaptive precision are within 1e-7 Ha of reference FP64 results across tested molecules and bases.
- The adaptive scheme often requires no more than two extra SCF iterations compared to FP64.
- K-matrix computations benefit most from INT8 emulation due to their higher FLOP count, justifying targeted precision acceleration.
- On H100, adaptive precision remains beneficial, though sometimes FP64 is faster depending on emulation level and convergence behavior.
Better researchstarts right now
From reading papers to final review, dramatically reduce your research time.
No credit card · Free plan available
This review was created by AI and reviewed by human editors.