🤖 AI

cuNNQS-SCI: A Fully GPU-Accelerated Framework for High-Performance Configuration Interaction Selection withNeural Network QQantum States

本文提出了 cuNNQS-SCI，一种完全基于 GPU 加速的框架，通过引入分布式全局去重、细粒度 CUDA 内核及 GPU 内存中心运行时，克服了传统 NNQS-SCI 方法中的 CPU-GPU 混合架构瓶颈，在保持化学精度的同时显著提升了大规模配置相互作用选择问题的求解效率与可扩展性。

原作者： Daran Sun, Bowen Kan, Haoquan Long, Hairui Zhao, Haoxu Li, Yicheng Liu, Pengyu Zhou, Ankang Feng, Wenjing Huang, Yida Gu, Zhenyu Li, Honghui Shang, Yunquan Zhang, Dingwen Tao, Ninghui Sun, Guangming T

发布于 2026-04-20

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Daran Sun, Bowen Kan, Haoquan Long, Hairui Zhao, Haoxu Li, Yicheng Liu, Pengyu Zhou, Ankang Feng, Wenjing Huang, Yida Gu, Zhenyu Li, Honghui Shang, Yunquan Zhang, Dingwen Tao, Ninghui Sun, Guangming Tan

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇论文介绍了一个名为 cuNNQS-SCI 的超级计算机程序，它的目标是解决一个困扰科学家多年的难题：如何更快速、更准确地模拟复杂分子和材料的微观世界。

为了让你轻松理解，我们可以把这项研究想象成**“在茫茫大海中寻找最珍贵的珍珠”**。

1. 背景：为什么要找珍珠？（量子化学的难题）

想象一下，科学家想要预测一种新材料（比如一种新药或一种电池）的性质。这需要计算原子和电子是如何相互作用的。

大海（配置空间）： 电子可能的排列组合就像大海里的每一滴水，数量多到无法想象（指数级增长）。
珍珠（重要状态）： 在这浩瀚的“大海”中，只有极少数特定的排列（配置）对最终结果最重要，就像大海里只有几颗真正的珍珠。
传统方法（全量搜索）： 以前的方法试图把大海里的每一滴水都检查一遍，这就像试图把整个太平洋的水都喝干，计算量太大，超级计算机也跑不动。
AI 辅助（智能筛选）： 后来，科学家引入了人工智能（神经网络），让它像“寻宝专家”一样，快速判断哪些地方可能有珍珠。这就是 NNQS-SCI 方法。

2. 旧方法的瓶颈：为什么之前的“寻宝队”跑不快？

虽然有了 AI 专家，但之前的系统（NNQS-SCI）在大规模运行时遇到了两个巨大的“拦路虎”：

拦路虎一：中央仓库的拥堵（CPU 去重瓶颈）
- 比喻： 想象有 64 个寻宝小队（GPU 显卡）在海里疯狂寻找珍珠。他们找到的珍珠都先扔给一个唯一的中央仓库管理员（CPU）。管理员要把所有珍珠倒出来，检查有没有重复的，然后再分回给小队。
- 问题： 当小队越来越多，珍珠越来越多时，管理员忙不过来了。所有的船都要排队等管理员，导致 64 个强大的寻宝小队有 30% 的时间都在干等着，非常浪费。
拦路虎二：搬运工太慢（CPU 生成瓶颈）
- 比喻： 在找珍珠之前，小队需要先制造一些“可能的珍珠位置”（耦合构型生成）。以前这个制造过程是由老式搬运工（CPU） 在岸上完成的，然后再把货搬给海里的寻宝小队（GPU）。
- 问题： 岸上的搬运工动作太慢，而海里的寻宝小队速度极快。结果就是，寻宝小队经常没货可找，只能停下来等搬运工，导致昂贵的显卡资源被闲置。

3. 新方案：cuNNQS-SCI 的“全速寻宝”策略

这篇论文提出的 cuNNQS-SCI，就是把整个寻宝过程完全搬到了海上（全 GPU 加速），并设计了全新的流程：

创新点一：分布式“自助去重”（解决拥堵）

做法： 不再把所有珍珠扔给一个管理员。每个寻宝小队（GPU）先在自己的小范围内把重复的珍珠扔掉。然后，大家通过一种**“智能排序交换”**的方式，像排队换座位一样，只把真正需要交换的珍珠传给邻居。
效果： 那个累死累活的“中央管理员”消失了。64 个小队各自忙碌，互不干扰，效率极高。

创新点二：海上“极速制造”（解决搬运慢）

做法： 把“制造可能位置”的工作也直接交给海里的寻宝小队（GPU）自己干。利用 GPU 强大的并行计算能力，把原本需要 CPU 做几小时的活，变成了几秒钟的“闪电战”。
效果： 寻宝小队不再需要等待岸上的搬运工，随时有货可找，显卡利用率拉满。

创新点三：超级“记忆背包”（突破内存限制）

做法： 以前，如果大海太大，单个寻宝小队的背包（显存）装不下所有数据，程序就崩溃了。现在，cuNNQS-SCI 设计了一个**“流动背包”**。
- 它把数据切成小块，像流水线一样：处理一块，扔掉一块，同时从岸上（CPU 内存）预取下一块。
- 它甚至让“搬运”和“计算”同时进行（就像一边做饭一边洗菜），完全隐藏了搬运的时间。
效果： 即使大海再大，单个小队的背包装不下，也能通过这种“流式”方法把整个大海装下，模拟以前根本算不动的超大分子（如铬二聚体 Cr2）。

4. 成果：快了多少？

在拥有 64 张顶级显卡（NVIDIA A100）的集群上测试：

速度提升： 整体速度比旧方法快了 2.32 倍。
效率： 在大规模并行计算时，效率保持在 90% 以上（意味着 64 个工人几乎都在干活，没人闲着）。
精度： 虽然速度飞快，但算出来的结果和“最精确的标准答案”几乎一模一样，没有因为快而牺牲准确性。

总结

cuNNQS-SCI 就像给量子化学计算装上了**“涡轮增压”**。它把原本依赖慢速 CPU 的“瓶颈”全部移除了，让昂贵的 GPU 显卡能 24 小时不间断地全速运转。

这使得科学家能够模拟以前完全无法想象的复杂分子系统，为设计新药物、新材料提供了强大的计算引擎。简单来说，它让“寻找微观世界的珍珠”这件事，从“大海捞针”变成了“高效捕捞”。

cuNNQS-SCI 技术总结

1. 研究背景与问题 (Problem)

背景：
基于神经网络的量子态（NNQS）方法，特别是结合了神经网络与选定的组态相互作用（SCI）的 NNQS-SCI 方法，已成为解决复杂多体系统薛定谔方程的先进手段。该方法在精度和可扩展性方面表现优异，能够处理传统方法难以企及的大规模量子系统。

核心问题：
尽管 NNQS-SCI 算法本身具有潜力，但其现有的实现严重受限于 混合 CPU-GPU 架构，导致在大规模系统上无法有效扩展。主要瓶颈包括：

集中式 CPU 去重（Global De-duplication）： 现有的去重过程依赖单个根节点（CPU）维护全局索引。随着 GPU 数量增加，通信开销呈 $O(N)$ 增长，且根节点的内存和计算能力成为严重瓶颈（在 64 GPU 配置下，去重可能占用 30% 的总运行时间）。
主机端耦合态生成（Host-resident Coupled-configuration Generation）： 耦合构型的生成主要在 CPU 上进行，导致频繁的主机 - 设备（Host-Device）数据传输，且无法充分利用 GPU 的并行计算能力。随着系统规模扩大，非 AI 组件（CPU 部分）的耗时占比从不到 10% 激增至 50% 以上。
显存墙（Memory Wall）： 完全 GPU 加速后，中间生成的构型数据量极易超过单张 GPU（如 A100 40GB）的显存容量，导致传统“全在显存内”的执行策略失效。

2. 方法论 (Methodology)

为了解决上述问题，论文提出了 cuNNQS-SCI，这是一个全 GPU 加速的 SCI 框架。其核心设计理念是将整个 SCI 工作流（包括构型生成、去重、推理和能量计算）完全迁移至 GPU 上，并引入以显存为中心的运行时机制。

2.1 分布式基于排序的全局去重算法 (Distributed Sort-based Global De-duplication)

策略： 摒弃传统的哈希去重和集中式 CPU 归约，采用 基于排序的定期采样（Sort-Based Regular Sampling） 策略。
流程：
1. 本地排序与采样： 每个 GPU 节点使用 GPU 优化的基数排序（Radix Sort）对本地构型进行排序，并定期采样（Pivots）以估计全局分布。
2. 全局划分： 根节点收集采样点，确定全局分割点（Splitters），广播给所有节点。
3. AlltoAll 交换： 各节点根据分割点将数据直接发送给目标节点（MPI_Alltoallv）。
4. 本地去重： 由于数据已全局有序，重复项仅出现在本地或相邻节点边界，通过本地合并和流压缩即可高效去重。
优势： 实现了确定性的负载均衡，消除了通信瓶颈，将去重任务从串行 CPU 转变为高吞吐的并行 GPU 操作。

2.2 细粒度 GPU 耦合构型生成内核 (Fine-grained GPU Kernel for Coupled Configurations)

数据布局优化： 将庞大的哈密顿矩阵规则预压缩为两个高度压缩的“激发表”（Excitation Tables：单激发和双激发表），将内存占用减少了 15 个数量级（例如 N2 分子从 EB 级降至 KB 级）。
并行架构： 设计了两级并行内核：
- Block 级： 每个源构型对应一个 Block，利用共享内存（Shared Memory）协作。
- Thread 级： 每个线程处理虚拟的激发空间。
高效过滤： 采用并行流压缩（Parallel Stream Compaction）技术，在共享内存中对有效结果进行前缀和计算，避免原子操作带来的序列化开销，最后合并写入全局内存。

2.3 以显存为中心的执行模型 (GPU Memory-Centric Execution Model)

核心思想： 将 GPU 显存视为高速缓存（Cache），而非全量存储介质，通过分阶段流水线和异步数据流突破显存限制。
三阶段流水线：
1. 生成与去重阶段： 在 GPU 上分批生成构型，冷数据立即异步卸载（Offload）到主机内存，仅保留推理所需数据。
2. 推理与选择阶段： 将唯一构型分批流式加载回 GPU 进行神经网络推理。采用流式归约（Streaming Reduction）和堆结构实时维护 Top-K 候选，避免存储所有波函数振幅。
3. 能量计算与恢复阶段： 按需重建反向索引，计算局部能量。
延迟隐藏： 利用多 CUDA 流（Multi-stream）实现计算（Compute）、主机到设备传输（H2D）和设备到主机传输（D2H）的重叠执行，掩盖 PCIe 带宽瓶颈。

3. 主要贡献 (Key Contributions)

首个全 GPU 加速的 NNQS-SCI 框架： 成功将耦合构型生成、全局去重、神经网络推理和能量计算全部迁移至 GPU，消除了 CPU 侧的可扩展性瓶颈。
高性能内核与分布式去重算法：
- 设计了专用的 CUDA 内核进行精确的耦合计算，利用细粒度并行化。
- 提出了基于排序的分布式去重算法，实现了跨多 GPU 集群的确定性负载均衡。
以显存为中心的执行范式： 引入了包含 Mini-batch 流式处理和异步卸载的运行时模型，解决了大规模 SCI 模拟中的显存容量限制问题，使得在单卡显存不足的情况下也能运行。
广泛的实验评估： 在多种分子系统（从小分子到强关联的 Cr2 系统）上验证了框架的正确性、加速比和可扩展性。

4. 实验结果 (Results)

实验在 64 张 NVIDIA A100 GPU 的集群上进行：

端到端加速比：
- 在 Cr2 系统（强关联，84 量子位）上，相比高度优化的混合 CPU-GPU 基线，实现了 1.88 倍 的端到端加速。
- 在 N2 系统（大基组）上，实现了 2.32 倍 的端到端加速。
瓶颈转移： 优化后，系统瓶颈从 CPU 侧的通信和生成任务成功转移回 GPU 上的神经网络推理，使得 GPU 算力得到充分利用。
并行效率：
- 强扩展性（Strong Scaling）： 在 64 GPU 上，N2 系统的并行效率保持在 91% 以上。
- 弱扩展性（Weak Scaling）： 由于全局去重算法高效消除了冗余，随着 GPU 数量增加，有效计算负载反而减少，实现了 163.7% 的超线性加速（Parallel Efficiency > 100%）。
显存突破： 对于理论峰值显存需求超过 40GB（如 C2H4O 需 70GB）的系统，cuNNQS-SCI 通过分阶段流式处理，将实际峰值显存控制在 40GB 硬件限制以内，成功运行了原本无法在单卡上执行的模拟。
精度验证： 与精确的 NNQS-SCI 基线相比，cuNNQS-SCI 保持了相同的化学精度（Chemical Accuracy），误差在 $10^{-4}$ Hartree 级别，且无误差累积。

5. 意义与影响 (Significance)

突破规模限制： cuNNQS-SCI 从根本上扩展了可求解量子化学问题的规模，使得在现有 GPU 集群上模拟更大、更复杂的强关联电子系统成为可能。
系统级优化范式： 该工作展示了在 AI for Science 领域，仅仅优化 AI 模型是不够的，必须对非 AI 组件（如数据管理、去重、生成）进行系统级的重构，才能释放硬件的全部潜力。
通用性启示： 提出的“以显存为中心”的执行模型和分布式去重策略，对于其他需要处理大规模稀疏数据或状态空间的科学计算任务（如分子动力学、图神经网络等）具有重要的借鉴意义。
推动 AI 与 HPC 融合： 该框架为在超大规模集群上运行高精度量子化学模拟奠定了系统基础，推动了 AI 驱动的科学发现向更大规模迈进。