Page 32 - 《软件学报》2025年第12期

沈莉 等: swJulia: 面向新一代神威超级计算机的 Julia 语言编译系统 5413

3 示范应用: 大规模并行量子化学模拟

在微观世界中, 描述粒子状态和运动的规律已不再遵循牛顿第二定律, 而是转向更为复杂的量子力学方程. 特
别是对于分子、原子等化学体系, 其状态和演化过程严格遵循薛定谔方程. 该方程可简化为:

ˆ H |Ψ⟩ = E |Ψ⟩ (1)
其中, ˆ H 的具体形式通常包含对坐标 x, y, z 的二阶导数算符等, 并可被解析地表达出来. 在此方程中, |Ψ⟩和 E 是待
求解的未知量. |Ψ⟩代表波函数, 是一个以所有电子和原子核的坐标为自变量的高维函数, 而 E 则代表该化学体系
的能量, 是一个具体的数值. 薛定谔方程提供了探索分子结构和光谱特性的理论框架, 在理解化学反应机制、开发
新型材料以及优化药物分子设计等领域具有不可替代的重要性. 总体来看, 这是一个二阶偏微分方程. 然而, 当化
学体系的规模超过一个电子和一个质子时, 我们就无法再获得解析解, 而只能通过数值方法进行求解.
求解薛定谔方程实质上等同于求解特征值方程 HC=CE, 其中 H 是一个维度为 2 ×2 N 的矩阵, 其特征值和特
N
征向量分别为 E 和 C. 依据量子力学的变分原理, 求解 E 和 C 的问题可以转化为优化问题, 即:

[ ]
T
E (θ) = min (C (θ)) HC (θ) (2)
θ
T
(C (θ)) HC (θ) 可以通过量子计算机或量子模拟器进行计算. 在此过程中, C (θ) 被编码到量
这个待优化的函数
θ 进行局部最小值的优化. 然而, 传统的
子线路中, 所得数值随后被提供给经典计算机, 进而使用经典算法对参数
量子化学方法, 如 Hartree-Fock 理论或密度泛函理论, 虽然能够提供精确的结果, 但在处理大分子系统时面临严重
的可扩展性和效率问题.
神经网络量子态 (neural network quantum state, NNQS) 方法是一种利用神经网络参数化方法描述量子系统的
新兴技术 [26,27] . 这种方法能够捕获指数级大规模编码希尔伯特空间内的非琐碎相关性, 为量子系统的模拟提供了
有效手段. NNQS 方法不仅具有强大的表示能力以捕获复杂的量子相关性, 而且可以利用神经网络的自学习特性,
通过训练持续优化波函数的表示, 从而提升模拟的准确性. 同时, 在评估 NNQS 定义的概率分布或生成样本时, 往
往需要进行大量采样. 这些采样任务是高度独立的, 非常适合在大规模的分布式环境中并行执行. 基于 Transformer
的 NNQS 方法是一种创新的量子计算方法. 该方法通过引入 Transformer 神经网络作为波函数拟设, 从而实现了
对量子多体波函数的高效表征. 其基本原理是利用 Transformer 网络的强大表达能力来模拟量子多体系统的基
态波函数, 并通过变分蒙特卡洛算法优化这些参数以逼近真实的量子态. 在该方法中, 波函数的振幅部分由
Transformer 网络处理, 而相位部分则由多层感知机来表征, 这种分离的处理方式使模型能更有效地捕捉波函数的
复杂结构.
作为量子计算领域的主流编程语言, Julia 不仅具备快速的原型设计能力, 还拥有丰富的科学计算库和基准测
试工具. 同时, Julia 提供的跨语言调用机制允许无缝集成第三方库, 能够显著增强应用的功能性和适应性. 基于
Julia 编程语言, 我们在新一代神威超级计算机上设计了可扩展的量子模拟器 NNQS-Transformer [28] , 以实现对大分
子化学体系能量的模拟求解. 其基本结构如后文图 12 所示. 除了第 2 节中介绍的 swJulia 编译系统外, NNQS-
Transformer 还包括深度学习框架 SWPyTorch、计算加速库、第三方 Julia 库以及底层文件系统等主要模块. 该模
拟器采用多种方式灵活使用众核计算资源, 并结合高效的策略以在高性能计算机上实现大规模分布式并行. 其不
仅能够克服传统量子化学方法的局限性, 还为大规模量子化学计算提供了一种新颖且强大的工具, 有望在量子化
学和其他相关领域发挥重要作用.
3.1 灵活的众核加速方式

基于 Julia 灵活高效的跨语言调用机制, NNQS-Transformer 得以支持多样化的众核加速方式. 首先, 通过
PyCall.jl 的 pyimport 方法, 我们可以实现对 SWPyTorch 框架的引用. SWPyTorch 是专为新一代神威超级计算机硬
件架构及其网络特性适配和优化的深度学习框架 [29] . 该框架在不改变上层 API 的前提下, 通过整合众核算子库
swDNN 及 swTENSOR, 实现了对 Transformer 等深度神经网络的并行加速计算, 其调用路径如图 12 中的①号箭
头所展示.