[논문 리뷰] GPU-accelerated Auxiliary-field quantum Monte Carlo with multi-Slater determinant trial states

Yifei Huang, Zhen Guo외 2명|arXiv (Cornell University)|2024. 06. 12.

Theoretical and Computational Physics인용 수 4

한 줄 요약

이 논문은 오픈소스 ipie 프레임워크 내에서 GPU 가속을 통한 다중 슬래터 결정자 시도 상태 단계 없는 보조장 양자 몽테카를로(MSD-ph-AFQMC) 구현을 제시하며, 강한 상관관계를 가진 시스템의 효율적이고 고정밀 시뮬레이션을 가능하게 한다. 이 방법은 A100 GPU에서 32-CPU 노드 대비 6배 빠른 성능을 달성했고, 10⁵개 이상의 결정자를 사용하여 복잡한 전이 금속 클러스터인 [Fe₂S₂(SCH₃)]²⁻에서 화학적 정밀도를 회복했다.

ABSTRACT

The accuracy of phaseless auxiliary-field quantum Monte Carlo (ph-AFQMC) can be systematically improved with better trial states. Using multi-Slater determinant trial states, ph-AFQMC has the potential to faithfully treat strongly correlated systems, while balancing the static and dynamical correlations on an equal footing. This preprint presents an implementation and application of graphics processing unit-accelerated ph-AFQMC, for multi-Slater determinant trial wavefunctions (GPU-accelerated MSD-AFQMC), to enable efficient simulation of large-scale, strongly correlated systems. This approach allows for nearly-exact computation of ground state energies in multi-reference systems. Our GPU-accelerated MSD-AFQMC is implemented in the open-source code exttt{ipie}, a Python-based AFQMC package [ extit{J. Chem. Theory Comput.}, 2022, 19(1): 109-121]. We benchmark the performance of the GPU code on transition-metal clusters like [Cu$_2$O$_2$]$^{2+}$ and [Fe$_2$S$_2$(SCH$_3$)]$^{2-}$. The GPU code achieves at least sixfold speedup in both cases, comparing the timings of a single A100 GPU to that of a 32-CPU node. For [Fe$_2$S$_2$(SCH$_3$)]$^{2-}$, we demonstrate that our GPU MSD-AFQMC can recover the dynamical correlation necessary for chemical accuracy with an MSD trial, despite the large number of determinants required ($>10^5$). Our work significantly enhances the efficiency of MSD-AFQMC calculations for large, strongly correlated molecules by utilizing GPUs, offering a promising path for exploring the electronic structure of transition metal complexes.

연구 동기 및 목표

• 다중 슬래터 결정자 시도 상태의 단계 없는 AFQMC에서 발생하는 계산 병목 현상을 GPU 가속을 통해 해결하기 위해.
• 정적 및 동적 전자 상관관계를 균형 있게 다루는 다중 참고계 시스템의 대규모 고정밀 시뮬레이션을 가능하게 하기 위해.
• CUDA 및 cuPy를 활용하여 ipie AFQMC 프레임워크의 성능과 확장성을 다중 슬래터 결정자 시도 상태로 확장하기 위해.
• 동적 상관관계를 위해 큰 수의 결정게를 필요로 하는 기준 전이 금속 클러스터에서의 방법의 능력을 입증하기 위해.
• 강한 상관관계를 가진 분자의 전자 구조 계산을 위한 고성능이고 오픈소스의 GPU 구현을 제공하기 위해.

제안 방법

• 다중 슬래터 시도 파동함수에서의 도전 및 결정자 평가를 병렬화하기 위해 GPU 가속 CUDA 커널과 cuPy 인터페이스를 사용한다.
• 이중체 상호작용을 일차항으로 분리하기 위해 허버드-스트라토니시브 변환을 적용한 단계 없는 보조장 양자 몽테카를로(ph-AFQMC) 프레임워크를 사용한다.
• 전자 반발 적분의 콜레스키 분해와 시간 진동의 일阶 토르터 분해를 활용하여 구현한다.
• CI 행렬에서 동일한 워커 계수에 여러 결정자가 기여할 경우 발생하는 스레드 충돌 문제를 해결하기 위해 원자 연산(atomicAdd)을 사용한다.
• 협동 스레드 그룹과 최적화된 메모리 액세스를 통해 워커, 결정자, 도전 순서, 보조 장에 대해 병렬화한다.
• 결정자 계산과 같은 핵심 연산은 사전 계산하여 커널에 전달함으로써 반복적인 GPU 커널 실행을 방지한다.

Figure 1: (a) Comparison of time per MSD-AFQMC block with 10 walkers on a single CPU core and a single A100 GPU. The system we considered was [Cu 2 O 2 ] 2+ using the BS1 basis Mahajan, Lee, and Sharma ( 2022 ) , which was also a CPU time benchmark example in the original ipie release Malone et al.

실험 결과

연구 질문

• RQ1 GPU 가속이 큰 강한 상관관계를 가진 시스템에 대해 MSD-ph-AFQMC의 계산 비용을 크게 줄일 수 있는가?
• RQ2 ph-AFQMC에서 다중 슬래터 결정자 시도 상태가 복잡한 전자 상관관계를 가지는 전이 금속 클러스터에서 화학적 정밀도를 달성할 수 있는가?
• RQ3 대규모 시스템에서 GPU 가속 MSD-AFQMC의 성능은 CPU 기반 구현과 비교해 어떻게 되는가?
• RQ4 10⁵개 이상의 결정게를 필요로 하는 시스템을 유지하면서 정확성과 효율성을 유지할 수 있는가?
• RQ5 확장 가능한 고성능 MSD-AFQMC 시뮬레이션을 위해 필요한 주요 GPU 최적화 전략는 무엇인가?

주요 결과

• A100 GPU에서 [Cu₂O₂]²⁺ 및 [Fe₂S₂(SCH₃)]²⁻에 대해 32-CPU 노드 대비 최소 6배의 빠른 성능 향상을 달성하였다.
• [Fe₂S₂(SCH₃)]²⁻의 경우, 10⁵개 이상의 결정자를 사용하여 동적 상관관계를 성공적으로 복원하였고, 화학적 정밀도를 달성하였다.
• 동일한 워커 계수에 여러 결정자가 기여할 경우 발생하는 스레드 충돌 문제를 원자 연산(atomicAdd)을 통해 효과적으로 해결하였다.
• 결정자를 사전 계산하고 도전 및 행렬 연산을 위한 최적화된 요소별 CUDA 커널을 사용함으로써 효율적인 확장성이 확보되었다.
• 이 방법은 다중 참고계 시스템의 거의 정확한 기본 상태 에너지 계산을 가능하게 하였으며, 단일 결정자 AFQMC를 초월하고 MRCI+Q 정밀도에 도달하였다.
• 프레임워크는 다중 GPU 환경으로의 확장이 가능하며, 향후 성능 계산을 위한 자동 미분 통합도 지원한다.

Figure 2: Timing and absolute energy benchmarks on the [Fe 2 S 2 (SCH 3 )] 2- cluster. (a) Comparison of time per MSD-AFQMC block with 640 walkers on 32 CPU cores and a single A100 GPU. (b) Absolute energies derived using localized atomic orbitals and natural orbitals. The full configuration interac