标题： 用于Fugaku上从头算量子化学模拟的大规模神经网络量子态 显示英文标题

标题： Large-scale Neural Network Quantum States for ab initio Quantum Chemistry Simulations on Fugaku

摘要： 解决量子多体问题是量子化学中的一个基本挑战。 虽然神经网络量子态（NQS）已成为一种有前途的计算工具，但其训练过程会带来指数级增长的计算需求，对于大规模分子系统来说变得极其昂贵，并为实际应用创造了根本性的可扩展性障碍。 为了解决上述挑战，我们提出了\ours ，一个用于\textit{从头计算}电子结构计算的高性能 NQS 训练框架。 首先，我们提出了一种可扩展的采样并行策略，具有多层工作负载划分和混合采样方案，打破了大规模 NQS 训练的可扩展性障碍。 然后，我们引入了多级并行局部能量并行，实现了更高效的局部能量计算。 最后，我们对基于变压器的\textit{假设}进行了以缓存为中心的优化，并将其与采样并行策略结合，进一步加速了 NQS 训练，并在大规模下实现了稳定的内存占用。 实验表明，\ours 在扩展到 1,536 个节点时，将 NQS 训练加速了高达 8.41 倍，并达到了高达 95.8% 的并行效率。 显示英文摘要

摘要： 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 above challenges, we present \ours, a high-performance NQS training framework for \textit{ab initio} electronic structure calculations. First, we propose a scalable sampling parallelism strategy with multi-layers workload division and hybrid sampling scheme, which break the scalability barriers for large-scale NQS training. Then, we introduce multi-level parallelism local energy parallelism, enabling more efficient local energy computation. Last, we employ cache-centric optimization for transformer-based \textit{ansatz} and incorporate it with sampling parallelism strategy, which further speedup up the NQS training and achieve stable memory footprint at scale. Experiments demonstrate that \ours accelerate NQS training with up to 8.41x speedup and attains a parallel efficiency up to 95.8\% when scaling to 1,536 nodes.