[Paper Review] GPU-accelerated Auxiliary-field quantum Monte Carlo with multi-Slater determinant trial states

Yifei Huang, Zhen Guo+2 more|arXiv (Cornell University)|Jun 12, 2024

Theoretical and Computational Physics4 citations

TL;DR

This paper presents a GPU-accelerated implementation of multi-Slater determinant trial state phaseless auxiliary-field quantum Monte Carlo (MSD-ph-AFQMC) in the open-source ipie framework, enabling efficient, high-accuracy simulations of strongly correlated systems. The method achieves a sixfold speedup on an A100 GPU versus a 32-CPU node and recovers chemical accuracy for complex transition-metal clusters like [Fe₂S₂(SCH₃)]²⁻ using over 10⁵ determinants.

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.

Motivation & Objective

• To overcome the computational bottleneck of multi-Slater determinant trial states in phaseless AFQMC by leveraging GPU acceleration.
• To enable large-scale, high-accuracy simulations of multi-reference systems with balanced treatment of static and dynamical electron correlation.
• To extend the performance and scalability of the ipie AFQMC framework to multi-Slater determinant trials using CUDA and cuPy.
• To demonstrate the method's capability on benchmark transition-metal clusters requiring large determinants for dynamical correlation.
• To provide a high-performance, open-source GPU implementation for electronic structure calculations of strongly correlated molecules.

Proposed method

• The method employs GPU-accelerated CUDA kernels with cuPy interface to parallelize excitation and determinant evaluation in multi-Slater trial wavefunctions.
• It uses the phaseless auxiliary-field quantum Monte Carlo (ph-AFQMC) framework with a Hubbard–Stratonovich transformation to decouple two-body interactions into one-body terms.
• The implementation leverages Cholesky decomposition of electron repulsion integrals and first-order Trotter decomposition for time evolution.
• Atomic operations (atomicAdd) resolve thread collisions when multiple determinants contribute to the same walker coefficient in the CI matrix.
• The code parallelizes over walkers, determinants, excitation orders, and auxiliary fields using cooperative thread groups and optimized memory access.
• Key operations like determinant calculations are precomputed and passed into kernels to avoid repeated GPU kernel launches.

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.

Experimental results

Research questions

• RQ1 Can GPU acceleration significantly reduce the computational cost of MSD-ph-AFQMC for large, strongly correlated systems?
• RQ2 Can a multi-Slater determinant trial state in ph-AFQMC achieve chemical accuracy for transition-metal clusters with complex electron correlation?
• RQ3 How does the performance of GPU-accelerated MSD-AFQMC compare to CPU-based implementations on large-scale systems?
• RQ4 Can the method handle systems requiring more than 10⁵ determinants while maintaining accuracy and efficiency?
• RQ5 What are the key GPU optimization strategies required for scalable, high-performance MSD-AFQMC simulations?

Key findings

• The GPU-accelerated MSD-ph-AFQMC implementation achieves at least a sixfold speedup on an A100 GPU compared to a 32-CPU node for both [Cu₂O₂]²⁺ and [Fe₂S₂(SCH₃)]²⁻.
• For [Fe₂S₂(SCH₃)]²⁻, the method successfully recovers dynamical correlation with over 10⁵ determinants, achieving chemical accuracy.
• The use of atomicAdd operations effectively resolves thread collision issues when multiple determinants contribute to the same walker coefficient.
• The implementation scales efficiently by precomputing determinants and using optimized, element-wise CUDA kernels for excitation and matrix operations.
• The method enables nearly exact ground state energy calculations for multi-reference systems, outperforming single-determinant AFQMC and matching MRCI+Q accuracy.
• The framework is extensible to multi-GPU setups and supports future integration of auto-differentiation for property calculations.

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