[논문 리뷰] Accelerating Density Fitting with Adaptive-precision and 8-bit Integer on AI Accelerators

Hua Huang, Wekai Shao외 1명|arXiv (Cornell University)|2026. 01. 12.

Machine Learning in Materials Science인용 수 0

한 줄 요약

이 논문은 Tensor Cores의 INT8 GEMM을 이용한 적응 정밀도 밀도 적합 알고리즘을 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 기 BasIS에서 INT8 GEMMs를 사용한 DF를 위한 적응 정밀도 밀도 적합 방법 개발.
• 지원되는 정확도에서 FP64 기준과 비교하여 수렴 정확도가 다양한 분자와 기 basis에 대해 유지되도록 보장.
• PySCF 구현 및 여러 NVIDIA GPU에서의 성능 평가.

제안 방법

• Coulomb 행렬 J를 계산하고 적응 정밀도로 INT8으로 시뮬레이션된 FP64 GEMMs를 사용하여 교환 행렬 K를 계산한다.
• SCF 반복 동안 상대적 SCF 에너지 변화(∆Erel)에 기반하여 INT8 에뮬레이션 레벨을 선택하는 적응 정밀도 스킴을 채용한다.
• ∆Erel이 임계값 아래로 떨어지면 FP64로 전환하여 같은 또는 추가 두 번의 SCF 반복 이내에 수렴을 보장한다.
• CuPy를 통한 GPU 가속 및 대 시스템에 대한 메모리 임계값 조정과 함께 PySCF에 메서드를 구현한다.
• 여러 기저집합 및 분자 시스템에서 RTX 4090, RTX 6000 Ada, H100의 성능을 평가한다.

실험 결과

연구 질문

• RQ1 INT8 에뮬레이션 FP64 GEMMs가 밀도 적합에서 K-매트릭스 구성을 가속하고 정확성을 희생하지 않는가?
• RQ2 적응 정밀도 전략이 다양한 분자 시스템에서 수렴 효율과 최종 에너지 정확도를 보존하는가?
• RQ3 완전 FP64 정밀도로 언제 되돌릴지에 대한 실용적인 성능 한계와 결정은 무엇인가?

주요 결과

• 적응 정밀도 DF는 RTX 4090에서 최대 3.04×, RTX 6000 Ada에서 4.64×의 속도향상을 DFT 계산에서 달성합니다.
• 적응 정밀도로 얻은 수렴 에너지는 모든 테스트 분자 및 기저에서 참조 FP64 결과와 1e-7 Ha 이내로 근접합니다.
• 적응 스키마는 FP64에 비해 보통 두 번의 추가 SCF 반복을 넘지 않는 경우가 많습니다.
• K-매트릭스 계산은 FLOP 수가 더 많아 INT8 에뮬레이션으로부터 가장 큰 이점을 얻어 정밀도 가속의 타깃이 됩니다.
• H100에서 적응 정밀도는 여전히 유익하지만 에뮬레이션 레벨과 수렴 동작에 따라 때때로 FP64가 더 빠를 수 있습니다.