Load-Balanced Diffusion Monte Carlo Method with Lattice Regularization

Abstract

Ab initio quantum Monte Carlo (QMC) is a stochastic approach for solving the many-body Schr\"odinger equation without resorting to one-body approximations. QMC algorithms are readily parallelizable via ensembles of $N_w$ walkers, making them well suited to large-scale high-performance computing. Among the QMC techniques, Diffusion Monte Carlo (DMC) is widely regarded as the most reliable, since it provides the projection onto the ground state of a given Hamiltonian under the fixed-node approximation. One practical realization of DMC is the Lattice Regularized Diffusion Monte Carlo (LRDMC) method, which discretizes the Hamiltonian within the Green's Function Monte Carlo framework. DMC methods - including LRDMC - employ the so-called branching technique to stabilize walker weights and populations. At the branching step, walkers must be synchronized globally; any imbalance in per-walker workload can leave CPU or GPU cores idle, thereby degrading overall hardware utilization. The conventional LRDMC algorithm intrinsically suffers from such load imbalance, which grows as $\log(N_w)$, rendering it less efficient on modern parallel architectures. In this work, we present an LRDMC algorithm that inherently addresses the load imbalance issue and achieves significantly improved weak-scaling parallel efficiency. Using the binding energy calculation of a water-methane complex as a test case, we demonstrated that the conventional and load-balanced LRDMC algorithms yield consistent results. Furthermore, by utilizing the Leonardo supercomputer equipped with NVIDIA A100 GPUs, we demonstrated that the load-balanced LRDMC algorithm can maintain extremely high parallel efficiency ($\sim$98\%) up to 512 GPUs (corresponding to $N_{\rm w}= 51200$), together with a speedup of $\times~1.24$ if directly compared with the conventional LRDMC algorithm with the same number of walkers.