大型科学研究领域的高性能运算,因其极高的吞吐量需求,都离不开超级计算机,如何优化也是一门学问。比如分子量子行为模拟,相互作用状态数量会指数级增长,因此要么借助量子计算机(但不成熟),要么简化处理以适配现有超级计算机(但太复杂)。
据报道,来自“神威”团队的中国科学家,成功将现有“海光lite”(Oceanlite)超级计算机与AI结合,实现了真实分子尺度下的复杂量子化学的建模运算。
可以说,这既是一项科学突破,也是一项技术突破。
在量子力学中,用来描述量子态的是波函数(Ψ),它决定了量子系统的所有可能构型及其对应概率,例如分子中电子等粒子的位置、自旋或能级。
对量子态进行建模极具挑战性,因为其状态空间会随粒子数量呈指数级急剧扩张,传统超级计算机很难甚至不可能完成。
为此,科学家们不得不采用多种近似方法,简化量子方程,同时保留描述分子结构、反应、能量所需要的精度,但是现有基于波函数近似的方法在规模上存在局限,只能用于小分子体系。
为了研究具有强电子关联的多体量子系统,比如包含数十个电子、超过100个自旋轨道等,物理学家提出了采用现代机器学习替代神经网络量子态(NNQS)等模型,近似模拟分子内电子的所有可能构型与运动。
这种方法有望将AI扩展性与量子精度结合,突破传统方法的研究领域。
来自中国“神威”团队的研究人员,自主研发了NNQS框架,模拟过程通过训练神经网络来近似分子波函数,进而确定电子最可能出现的位置。
对于每一种采样得到的电子排布,系统会计算局部能量,并调整网络参数,直至其预测结果与分子真实量子能量模式相符。
这款专属NNQS框架是为中国“海光lite”超级计算机量身定制的,它搭载了384核心的申威SW26010-Pro处理器,支持FP16半精度浮点、FP32单精度浮点、FP64双精度浮点数据格式,具备百亿亿次级算力。
不过,它是专为高性能计算(HPC)设计的独特架构,缺乏足够的AI设计,因此研究人员必须考虑负载并行处理方式,以及数据处理机制。
因此,研究人员专门设计了一种分层通信模型,由管理核心协调处理器与节点之间的通信,同时由数百万个轻量级的双路路计算处理单元(CPE)执行局部量子计算,这些单元都配备了512位向量引擎。
此外,他们还开发了动态负载均衡算法,以避免因计算负载不均匀而导致部分核心闲浪费置。
最终,团队在3700万个CPE核心上运行代码,实现了92%的强扩展性、98%的弱扩展性,实现了近乎完美的软硬件协同,对中国超级计算而言是一项重大成就。
截至目前,这项包含120个自旋轨道的分子系统模拟,是全球范围内在传统超级计算机上完成的规模最大的AI驱动量子化学计算,标志着中国在AI与量子科技领域取得了突破性进展。