[论文解读] Accelerating Density Fitting with Adaptive-precision and 8-bit Integer on AI Accelerators
本文提出一种自适应精度密度拟合算法，在 Tensor Cores 上使用 INT8 GEMM 进行实现，基于 PySCF，相比 FP64 最高可实现 3.04×–4.64× 的加速，并保持收敛的 DFT 能量。
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.
研究动机与目标
- 在严格精度约束下，推动在量子化学中利用 AI 加速器张量核心。
- 开发基于 INT8 GEMMs 的自适应精度密度拟合方法用于高斯基组的 DF。
- 确保跨多种分子和基组的收敛精度与 FP64 基线相当。
- 在 PySCF 中实现并在多款 NVIDIA GPU 上评估性能。
提出的方法
- 使用 FP64 计算 Coulomb 矩阵 J，及 INT8 伪 FP64 GEMMs 来计算带自适应精度的交换矩阵 K。
- 采用自适应精度方案，在 SCF 迭代过程中根据相对 SCF 能量变化量 (∆Erel) 选择 INT8 伪装等级。
- 当 ∆Erel 降至阈值以下时切换回 FP64，以确保在相同或额外两次 SCF 迭代内收敛。
- 在 PySCF 中实现该方法，并使用 CuPy 进行 GPU 加速，调整大系统的内存阈值。
- 在 RTX 4090、RTX 6000 Ada 和 H100 上，覆盖多种基组和分子体系评估性能。
实验结果
研究问题
- RQ1INT8 伪 FP64 GEMMs 是否能在不牺牲精度的前提下加速密度拟合中的 K 矩阵构造？
- RQ2自适应精度策略是否能在多种分子体系中保持收敛效率和最终能量精度？
- RQ3何时回退到完全 FP64 精度的实际性能极限与决策？
主要发现
- 自适应精度密度拟合在 RTX 4090 上为 DFT 计算提供最高 3.04× 的加速，在 RTX 6000 Ada 上达到 4.64×。
- 在测试的分子和基组中，使用自适应精度得到的收敛能量与参考 FP64 结果相差不超过 1e-7 Ha。
- 与 FP64 相比，自适应方案通常不需要超过额外的两次 SCF 迭代。
- K 矩阵的计算最受益于 INT8 伪装，因为其 FLOP 数较高，因而成为目标精度加速的主要对象。
- 在 H100 上，自适应精度仍然有利，尽管有时取决于伪装等级和收敛行为 FP64 的速度更快。
更好的研究，从现在开始
从论文设计到论文写作，大幅缩短您的研究时间。
无需绑定信用卡
本解读由 AI 生成，并经人工编辑审核。