[論文レビュー] Accelerating Density Fitting with Adaptive-precision and 8-bit Integer on AI Accelerators

Hua Huang, Wekai Shao+1 more|arXiv (Cornell University)|Jan 12, 2026

Machine Learning in Materials Science被引用数 0

ひとこと要約

本論文は、Tensor Cores上のINT8 GEMMsを用いた適応精度密度適合アルゴリズムをPySCFで実装し、FP64に対して最大で3.04×–4.64×のスピードアップを実現し、収束したDFTエネルギーを保持する。

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.

研究の動機と目的

• AIアクセラレーターのテンソルコアを厳密な精度制約の下で量子化学に活用する動機付け。
• Gaussian基底集合に対するDFのためのINT8 GEMMsを用いた適応精度密度適合手法を開発。
• 様々な分子と基底集合に対してFP64ベースラインと同等の収束精度を維持。
• PySCFでの実装を実証し、複数のNVIDIA GPUで性能を評価。

提案手法

• Coulomb行列JをFP64で計算し、INT8エミュレート_FP64 GEMMsで交換行列Kを適応精度で計算。
• SCF反復中の相対SCFエネルギー変化量(∆Erel)に基づくINT8エミュレーションレベルを選択する適応精度スキームを採用。
• ∆Erelが閾値を下回るとFP64へ切り替え、同一回または追加の2回のSCF反復内で収束を保証。
• CuPyを用いたGPU加速と大規模系のメモリ閾値の調整を行い、PySCFに実装。
• RTX 4090、RTX 6000 Ada、H100で、複数の基底セットと分子系で性能を評価。

実験結果

リサーチクエスチョン

• RQ1 INT8エミュレート_FP64 GEMMsはDF密度適合のK行列構築を加速でき、精度を犠牲にしないか。
• RQ2 適応精度戦略は様々な分子系に対して収束効率と最終エネルギー精度を維持できるか。
• RQ3 エミュレーションレベルと収束挙動に応じて、実際にFP64精度へ戻すべき実用的な性能限界と判断基準は何か。

主な発見

• 適応精度DFは、DFT計算でRTX 4090で最大3.04×、RTX 6000 Adaで最大4.64×のスピードアップを達成。
• 適応精度による収束エネルギーは、検証対象の分子と基底集合全体で参照FP64結果と1e-7 Ha以内に収束。
• 適応スキームは、FP64と比べて通常は2回の追加SCF反復を超えない。
• K行列計算はFLOP数が多いためINT8エミュレーションの影響を最も受ける部分であり、対象精度加速の正当性を裏付け。
• H100では適応精度は引き続き有効だが、エミュレーションレベルと収束挙動次第でFP64の方が速い場合もある。