Large-Scale Neural Network Quantum States for ab initio Quantum Chemistry Simulations on Fugaku

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.