Search papers, labs, and topics across Lattice.
The paper introduces jQMC, a Python-based computational package for ab initio Quantum Monte Carlo simulations optimized for GPU-accelerated environments, implementing Variational Monte Carlo and lattice-regularized Diffusion Monte Carlo algorithms. By leveraging JAX for automatic differentiation and Just-In-Time compilation, jQMC achieves significant performance improvements, particularly in wall-time, outperforming the existing TurboRVB code on NVIDIA GPUs for system sizes up to 160 electrons. The results indicate that jQMC can provide speedups of up to an order of magnitude for VMC simulations, highlighting its potential for efficient quantum simulations in large-scale electronic systems.
jQMC achieves up to an order of magnitude speedup over TurboRVB for Variational Monte Carlo simulations on GPUs, revolutionizing the efficiency of quantum Monte Carlo methods.
We present jQMC, a Python-based computational package for {\it ab initio} Quantum Monte Carlo (QMC) simulations, designed for modern GPU-accelerated computing environments. jQMC implements two well-established QMC algorithms: Variational Monte Carlo (VMC) and the lattice-regularized variant of Diffusion Monte Carlo (LRDMC). The employed wave function is a Jastrow factor combined with the antisymmetrized geminal power with spin-singlet and spin-triplet pairings, which contains the single Slater determinant as its special lowest-rank case. The wave function can be initialized from external Hartree-Fock/Density Functional Theory calculations through the TREX-IO library (a common wave-function format across electronic-structure packages) and optimized by stochastic reconfiguration and linear-method energy minimization. One of the prominent features of jQMC is its use of JAX, which enables automatic differentiation for wave function optimization and atomic force calculations, and allows the main QMC algorithms to be Just-In-Time (JIT) compiled and portable across CPU and GPU. jQMC is vectorized over walkers at the top level of the QMC algorithms, providing efficient intra-GPU~(CPU) vectorization. The multi-GPU~(CPU) parallelization is also supported through MPI and JAX sharding. To assess the practical performance of this implementation, we benchmarked jQMC performance on NVIDIA GPUs (A100 and H100) and analyzed CUDA kernels. For the test cases analyzed here, with system sizes up to 160 electrons, the current version of jQMC is faster than TurboRVB, a Fortran90 code implementing the same algorithms and wave functions, once jQMC is run on GPUs. In terms of wall-time, the gain can reach an order of magnitude for VMC, while it is more moderate for LRDMC.
