用全GPU加速突破量子化学计算规模瓶颈,速度提升2.32倍。
A Fully GPU-Accelerated Framework for High-Performance Configuration Interaction Selection with Neural Network Quantum States

- 全GPU架构重构去重与配置生成流程,消除CPU通信瓶颈。
- 在64张A100 GPU上实现2.32倍端到端提速,保持化学精度。
- 适合大规模量子多体系统模拟,尤其对高并发计算有需求者。
基于人工智能的方法在求解复杂多体系统的薛定谔方程方面展现出显著成效。在神经网络量子态(NNQS)方法中,NNQS-SCI(选配组态相互作用)因其高精度与可扩展性成为前沿技术。然而,其在更大系统上的应用受限于混合式CPU-GPU架构:集中式CPU全局去重导致严重通信瓶颈,主机端耦合组态生成带来巨大计算开销。本文提出QiankunNet-cuSCI,一种全GPU加速的SCI框架,通过分布式负载均衡的全局去重算法最小化冗余与通信开销;采用专用细粒度CUDA核实现精确耦合组态生成;并引入以GPU内存为中心的运行时机制,包括GPU侧池化、流式小批量处理和重叠卸载,突破单卡显存限制。该设计使可处理的组态空间大幅扩展,将性能瓶颈从主机侧重新转移至设备端推理。评估表明,该工作从根本上扩大了可求解问题的规模。在包含64张NVIDIA A100 GPU的集群上,相较高度优化的NNQS-SCI基线,实现了最高2.32倍的端到端加速,同时保持相同化学精度。此外,其分布式性能优异,在强缩放测试中维持超过90%的并行效率。
原文摘要 · Abstract (English)
AI-driven methods have demonstrated considerable success in tackling the central challenge of accurately solving the Schrödinger equation for complex many-body systems. Among neural network quantum state (NNQS) approaches, the NNQS-SCI (Selected Configuration Interaction) method stands out as a state-of-the-art technique, recognized for its high accuracy and scalability. However, its application to larger systems is severely constrained by a hybrid CPU-GPU architecture. Specifically, centralized CPU-based global de-duplication creates a severe scalability barrier due to communication bottlenecks, while host-resident coupled-configuration generation induces prohibitive computational overheads. We introduce QiankunNet-cuSCI, a fully GPU-accelerated SCI framework designed to overcome these bottlenecks. It first integrates a distributed, load-balanced global de-duplication algorithm to minimize redundancy and communication overhead at scale. To address compute limitations, it employs specialized, fine-grained CUDA kernels for exact coupled configuration generation. Finally, to break the single-GPU memory barrier exposed by this full acceleration, it incorporates a GPU memory-centric runtime featuring GPU-side pooling, streaming mini-batches, and overlapped offloading. This design enables much larger configuration spaces and shifts the bottleneck from host-side limitations back to on-device inference. Our evaluation demonstrates that our work fundamentally expands the scale of solvable problems. On an NVIDIA A100 cluster with 64 GPUs, our work achieves up to 2.32X end-to-end speedup over the highly-optimized NNQS-SCI baseline while preserving the same chemical accuracy. Furthermore, it demonstrates excellent distributed performance, maintaining over 90% parallel efficiency in strong scaling tests.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。