jQMC: A JAX-based ab initio quantum Monte Carlo package designed for GPU-accelerated computing

Technical Summary of jQMC: A JAX-based ab initio quantum Monte Carlo package designed for GPU-accelerated computing

Problem Statement
Solving the many-body Schrödinger equation remains a central challenge in computational chemistry and physics. While Density Functional Theory (DFT) is widely used, its reliance on approximate exchange-correlation functionals limits its accuracy in strongly correlated systems. Ab initio Quantum Monte Carlo (QMC) methods, such as Variational Monte Carlo (VMC) and Diffusion Monte Carlo (DMC), offer a stochastic alternative with favorable scaling, but they are computationally demanding. The shift of High-Performance Computing (HPC) facilities toward GPU-dominated architectures presents a significant hurdle for legacy QMC codes, which are often written in Fortran or C and rely on CPU-centric designs. Refactoring these codes to exploit GPU parallelism is difficult, and existing GPU-accelerated QMC codes often lack the flexibility of modern frameworks or the specific algorithmic features required for high-precision ab initio simulations. Furthermore, there is a need for standardized wave function formats and efficient optimization techniques that can leverage automatic differentiation.

Methodology
The authors present jQMC, a Python-based computational package for ab initio QMC simulations built entirely on the JAX library. JAX enables Just-In-Time (JIT) compilation, automatic differentiation (AD), and seamless portability between CPU and GPU backends.

• Algorithms: jQMC implements two ground-state QMC algorithms:
  1. Variational Monte Carlo (VMC): Uses the Metropolis-Hastings algorithm to sample electron configurations. It employs the Attaccalite and Sorella (AS) reweighting scheme to handle singularities near nodal surfaces.
  2. Lattice-Regularized Diffusion Monte Carlo (LRDMC): A Green's Function Monte Carlo (GFMC) variant that discretizes the Hamiltonian on a lattice to improve numerical stability. jQMC implements both conventional and load-balanced branching schemes to mitigate load imbalance in parallel calculations.
• Wave Function Ansatz: The trial wave function is constructed as a product of a Jastrow factor and an antisymmetric part. The antisymmetric component supports:
  • Antisymmetrized Geminal Power (AGP): A resonating valence bond (RVB)-type wave function that includes spin-singlet and spin-triplet pairings.
  • AGP with constrained molecular orbitals (AGPn): A variant restricting the rank of the pairing matrix.
  • Single Slater Determinant (SD): As a special case of AGP.
    The wave function can be initialized from external Hartree-Fock (HF) or Density Functional Theory (DFT) calculations via the TREX-IO library.
• Optimization: The package implements Stochastic Reconfiguration (SR) and the Linear Method (LM) for optimizing variational parameters. These methods utilize the metric of the parameter space to update wave function amplitudes and nodal surfaces simultaneously.
• Differentiation and Forces: Leveraging JAX's automatic differentiation, jQMC efficiently computes energy derivatives. It implements atomic force calculations for both VMC (using Space-Warp Coordinate Transformation) and LRDMC (using the Reynolds approximation and Pathak-Wagner regularization).
• Architecture: The code separates data structures (implemented as JAX-compatible dataclasses) from algorithms. It utilizes vmap for vectorization over walkers (the top-level parallelism in QMC) and supports multi-GPU scaling via MPI (mpi4py) and JAX sharding. It also supports mixed-precision arithmetic (FP32 for atomic orbital evaluations, FP64 for determinants/geminials) to accelerate computation without sacrificing accuracy.

Key Contributions

1. JAX-Based Implementation: jQMC is the first QMC package to fully leverage JAX for both algorithmic differentiation and GPU acceleration, allowing for JIT-compiled, portable code across CPU and GPU architectures.
2. Advanced Wave Function Support: It provides a robust implementation of the Jastrow-Antisymmetrized Geminal Power (JAGP) ansatz, which generalizes the single Slater determinant and allows for the optimization of nodal surfaces.
3. Efficient Vectorization: By vectorizing over walkers at the top level, jQMC achieves high intra-GPU efficiency, saturating GPU resources more effectively than traditional CPU-parallelized codes.
4. Standardization: Integration with TREX-IO ensures compatibility with other electronic structure packages (e.g., PySCF, CP2K).
5. Force Evaluation: The implementation of atomic forces using automatic differentiation and specific regularization techniques (Pathak-Wagner) for LRDMC enables geometry optimization and phonon calculations within the QMC framework.

Results

• Performance Benchmarks:
  • GPU vs. CPU: On NVIDIA GPUs (A100 and H100), jQMC demonstrates speedups of over two orders of magnitude compared to its own CPU implementation.
  • Comparison with TurboRVB: When compared to TurboRVB (a Fortran90 code implementing the same algorithms and wave functions), jQMC on GPUs is significantly faster. For VMC, wall-time gains reach an order of magnitude. For LRDMC, the gain is more moderate but still substantial (approx. 3x faster for the tested system sizes).
  • Scalability: Weak-scaling benchmarks on the Miyabi (University of Tokyo) and Leonardo (CINECA) supercomputers show that jQMC maintains near-ideal efficiency (efficiency $\approx$ 1.0) up to 1024 GPUs for VMC and up to 512 GPUs for LRDMC (using the load-balanced algorithm).
• Accuracy Validation:
  • Reproducibility: jQMC results for the total energies and binding energies of methane, water, and the methane-water dimer agree with other major QMC codes (CASINO, QMCPACK, CMQMC, TurboRVB) within statistical error bars, validating the implementation of the DLA (Determinant Locality Approximation) and LRDMC algorithms.
  • Atomic Forces: The paper assesses the accuracy of atomic forces for ethanol and malonaldehyde against all-electron CCSD(T) calculations.
    • For ethanol, LRDMC forces achieve a Mean Absolute Error (MAE) of $\sim$1.5 kcal/mol/Å relative to CCSD(T), consistent with previous literature.
    • For malonaldehyde, a significant discrepancy is observed, with LRDMC and CCSD(T) forces differing by $\sim$5 kcal/mol/Å (MAE). The authors note that this discrepancy persists even with optimized nodal surfaces, suggesting the issue may not be solely due to the nodal surface approximation but requires further investigation.

Significance
The paper positions jQMC as a modern, versatile platform that bridges the gap between high-precision ab initio QMC methods and modern GPU-accelerated HPC architectures. By utilizing JAX, the package offers a flexible environment for rapid prototyping of new algorithms while delivering production-level performance. The authors claim that jQMC successfully addresses the challenges of porting legacy QMC codes to GPUs, providing a tool that is not only faster than existing Fortran-based implementations on GPU hardware but also easier to extend and maintain. The successful validation of forces and energies against other codes and high-level quantum chemistry methods demonstrates that jQMC is ready for production simulations of molecular and extended systems, particularly in regimes where strong correlation effects are critical.