[论文解读] 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 Physics被引用 4

一句话总结

本文在开源的 ipie 框架中实现了 GPU 加速的多斯莱特行列式试函数相位无偏辅助场量子蒙特卡罗（MSD-ph-AFQMC）方法，实现了对强关联体系的高效、高精度模拟。该方法在 A100 GPU 上相比 32 核 CPU 节点实现了六倍速提升，并在使用超过 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.

研究动机与目标

• 通过利用 GPU 加速，克服多斯莱特行列式试函数在相位无偏 AFQMC 中的计算瓶颈。
• 实现对多参考体系的大规模、高精度模拟，实现静态与动态电子相关性的平衡处理。
• 利用 CUDA 和 cuPy 扩展 ipie AFQMC 框架在多斯莱特行列式试函数上的性能与可扩展性。
• 在需要大量行列式以描述动态相关性的基准过渡金属簇上展示该方法的能力。
• 为强关联分子的电子结构计算提供高性能、开源的 GPU 实现。

提出的方法

• 该方法采用 GPU 加速的 CUDA 内核与 cuPy 接口，对多斯莱特试波函数中的激发与行列式计算进行并行化。
• 采用相位无偏辅助场量子蒙特卡罗（ph-AFQMC）框架，并通过 Hubbard–Stratonovich 变换将两体相互作用分解为单体项。
• 该实现利用电子排斥积分的 Cholesky 分解与一阶 Trotter 分解进行时间演化。
• 原子操作（atomicAdd）用于解决多个行列式同时贡献于 CI 矩阵中同一行走器系数时的线程冲突问题。
• 通过协作线程组对行走器、行列式、激发顺序和辅助场进行并行化，并优化内存访问。
• 关键操作如行列式计算预先计算并传入内核，以避免重复的 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 上，MSD-ph-AFQMC 的 GPU 加速实现对 [Cu₂O₂]²⁺ 和 [Fe₂S₂(SCH₃)]²⁻ 的计算速度相比 32 核 CPU 节点至少提升了六倍。
• 对于 [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