Large-Scale Neural Network Quantum States for ab initio Quantum Chemistry Simulations on Fugaku
Hongtao Xu, Zibo Wu, Mingzhen Li +1 more
Solving quantum many-body problems is one of the fundamental challenges in quantum chemistry. While neural network quantum states (NQS) have emerged as a promising computational tool, its training process incurs exponentially growing computational demands, becoming prohibitively expensive for large-scale molecular systems and creating fundamental scalability barriers for real-world applications. To address the above challenges, we present QChem-Trainer, a highperformance NQS training framework for ab initio electronic structure calculations. First, we propose a scalable sampling parallelism strategy with multi-layer workload division and hybrid sampling scheme, which breaks the scalability barriers in large-scale NQS training. Second, we introduce multi-level energy calculation parallelism, enabling more efficient local energy computation in NQS training. Last, we employ cachecentric optimization for transformer-based ansatz and incorporate it with sampling parallelism strategy, which further accelerates the NQS training and achieves stable memory footprints at scale. Experiments demonstrate that QChem-Trainer provides up to$8.41 \times$speedup in NQS training and attains a parallel efficiency of up to 95.8% when scaling to 1,536 nodes.
arXiv:2506.23809•IEEE International Symposium on Parallel and Distributed Processing with Applications (ISPA)•Oct 10, 2025•0 citations