[論文レビュー] 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ノードと比較して6倍の高速化が達成され、[Fe₂S₂(SCH₃)]²⁻のような複雑な遷移金属クラスターに対しても10⁵個を超える行列式を用いて化学的精度を回復した。

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インタフェースを用いて、多スレイター試行波動関数における励起状態および行列式評価を並列化する。
• 2体相互作用を1体項に分解するためのハッバード＝ストラトニッチ変換を用いた位相なし補助場量子モンテカルロ（ph-AFQMC）フレームワークを採用する。
• 電子反発積分のチョルヘスキー分解と1次Trotter分解を用いて時間発展を実装する。
• 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