用于Fugaku上ab initio量子化学模拟的大规模神经网络量子态
分布式、并行与集群计算
2025-07-01 v1
摘要
求解量子多体问题是量子化学中的基本挑战之一。虽然神经网络量子态(NQS)已成为有前景的计算工具，但其训练过程的计算需求呈指数增长，对于大规模分子系统和实际应用而言，已成为不可接受的瓶颈。为此，我们提出了\ours，一个用于\textit{ab initio}电子结构计算的高性能NQS训练框架。首先，我们提出了可扩展采样并行策略，采用多层工作负载划分和混合采样方案，突破了大规模NQS训练的可扩展性瓶颈。随后，我们引入多级并行本地能并行，实现更高效的本地能计算。最后，我们采用面向缓存的优化策略处理基于transformer的\textit{ansatz}，并将其与采样并行策略相结合，进一步加速NQS训练，在规模化时保持稳定的内存占用。实验表明，\ours在NQS训练方面可实现最高8.41倍加速，当扩展到1,536节点时，平行效率可达95.8%。
关键词
引用
@article{arxiv.2506.23809,
title = {Large-scale Neural Network Quantum States for ab initio Quantum Chemistry Simulations on Fugaku},
author = {Hongtao Xu and Zibo Wu and Mingzhen Li and Weile Jia},
journal= {arXiv preprint arXiv:2506.23809},
year = {2025}
}