编者按:2022年11月5日,“倚天开启云原生算力新时代”专场在杭州·云栖大会D馆云栖科创SHOW场举行,北京科学智能研究院高性能计算团队负责人蔡淳发表了《倚天助力AI for Science:ABACUS新实践》的主题分享。本文根据该演讲整理而成,主要分为三个部分:
1. AI for Science探索材料研发新范式
2. 倚天710的适配与调优
3. 基于E-HPC的万核级别算例验证
图:北京科学智能研究院高性能计算团队负责人 蔡淳
01 AI for Science探索材料研发新范式
首先,讲一讲AI for Science如何定义材料研发流程中的新范式。材料创新是药物设计、新能源等领域发展的源动力。材料研发的范式由传统的大量重复实验试错,升级为先通过理论模拟筛选出可能的材料,再进行验证研发的计算驱动流程。
但是,在进行理论模拟的过程中,维数灾难阻碍了高精度计算的进行。借用保罗·狄拉克所说:“有了量子力学之后,对于所有的化学问题和大部分的物理问题,所依照的基本物理定律都已经明确,困难仅在于这些定律的精确应用会导致数学方程过于复杂而无法求解。”
为了解决维数灾难的问题,科学家把最基础、最精确的物理模型层层抽象,在不同体系中,选择相对应的物理模型,以保证能在合理的时间内给出计算结果。但不同尺度的物理模型,在时间和空间上有多个数量级的差异,结果的精度也有不小差距。
AI for Science方法旨在解决上述问题:它能够通过机器学习的方式,将高精度方法下学习的特征应用到更大体系下,同时拥有高精度的计算结果,以及高效率的求解时间。
在原子尺度分子模拟方法方面,还存在一些挑战。传统的分子动力学方法,需要科学家提供力场经验参数,势函数开发周期非常长。
对于密度泛函理论(DFT)来说,DFT软件的代码分支非常庞杂,研发周期长;DFT算法使用了交换关联泛函近似,但精度越高的近似方式计算量越大。
深度势能方法是基于机器学习的分子动力学方法。它很好地耦合了科学计算、机器学习与高性能计算。
左侧的图展示了深度势能方法的训练流程。它通过DFT计算原子势,用神经网络学习势函数,最后将它应用到分子动力学中,从而实现高效率、高精度的计算。
上图中间的结果,展示了深度势能方法和传统DFT方法计算得出的分子径向分布函数对比,其结果非常一致。
DFT是一种通过求解薛定谔波函数方程,直接求解物质性质的算法。
假设给定一个晶体的晶胞参数等信息,通过DFT方法计算可以得到它的电导率、体系密度等基本的物理性质。它是几乎不需要经验参数的第一性原理方法。研究者也凭借着这项工作获得1998年的诺贝尔化学奖。
如上图公式所示,密度泛函理论的核心是,将系统的总能量E表示为电子密度ρ的泛函。
左下角展示了DFT的雅各布天梯,从最简单的局域密度近似开始,以计算量为代价,逐渐接近通过精确的量子力学方法计算得到的结果。DeePKS方法采用神经网络模型表示高精度方法与低精度方法之差;研究人员可以使用DeePKS软件,先使用低精度的泛函,用比较高的效率给出结果,再将结果加上DeePKS方法给出的修正值,从而让计算结果逼近高精度的泛函。右图是水分子中氧元素之间距离的径向分布函数:DeePKS方法和高精度DFT方法计算的结果能够很好的吻合;相比之下,使用PBE泛函的计算结果会出现比较大的偏差。
AI+材料科学范式需要依赖DFT软件为AI模型训练产生数据;训练中的模型会再次影响DFT软件的计算结果,这一过程需要反复迭代直至收敛,其中需要进行大量的DFT计算。
ABACUS(中文名“原子算筹”)是一款国产开源的密度泛函理论软件。ABACUS最早是由中科大的何力新老师课题组发展起来的,并且在2021年3月加入Deep Modeling开源社区。
作为一款科学计算软件,ABACUS研发突破了传统课题组的代码开发方式,把代码托管在GitHub平台,欢迎开源贡献者一起进行新功能的开发以及错误的修复。
加入DeepModeling社区之后,ABACUS迎来了来自中科大、北京大学、物理所以及北京科学智能研究院(AISI)等不同单位的贡献者。
值得一提的是,北京科学智能研究院(AISI)是鄂维南院士在2021年建立的首个以AI for Science为使命的科学研究机构。
02 倚天710的适配与调优
接下来,为大家介绍ABACUS如何迁移到倚天710云平台。倚天710的操作系统是阿里云提供的Ali Linux3,它能够很好地支持现有Linux的软件生态。用户在安装软件时,可以直接从包管理器下载需要的依赖,不再需要重新手动编译。
倚天710芯片是基于ARM v9架构开发的,支持SVE、INT8mm、BF16mm等指令集加速。除此之外,ARM平台为大家提供了高性能的数学库,它可以让研究者注重算法开发,不需要操心矩阵计算的实现方式。
上图是ABACUS自带算例进行的性能对比图。最右边一列是倚天710的计算时间,它和七代x86架构CPU以及六代高频实例,在计算时间上是一致的。
需要注意的是,上面的测试中倚天710使用的实例是4xlarge规格,仅为x86实例的一半。这个结果得益于倚天710的独立物理核心、独立cache和ALU性能,且无超线程损耗。
03 基于E-HPC的万核级别算例验证
研究人员在阿里云E-HPC上,进行了贴近实际算力的验证。阿里云团队提供的弹性高性能计算E-HPC服务,从ECS的虚拟机镜像创建计算节点,实现计算集群资源的弹性伸缩,保障计算在云上与本地超算类似的操作体验下高效进行。
E-HPC在创建ECS实例时,可以选用竞价实例,让时间要求不敏感的科学计算任务,在云平台资源使用的低谷期,以非常低廉的价格进行科学计算。
ABACUS针对350 eV的极端高温下的32个硼原子体系,在由倚天710实例构成集群上进行了最新发展的Stochastic DFT方法的计算。
如上图所示,研究人员使用了11008个倚天710的CPU核心,即86个128核节点。这是一个弱扩展的任务,每个核心数的运算量是一定的,计算资源消耗随着核心数增加而线性增加。
在ABACUS实现的轨道并行和K点并行两种不同的任务划分模式下,计算时间都在有限范围内增长,并行性能非常好。在软件精度方面,不同的核心数计算得到的能量和压强一致,计算结果是正确的。
每个展示的数据点都由十个随机种子初始化进行计算,以避免系统性的随机误差。随着计算量的增加,压强和能量的标准差在不断地收敛,直至理论的最优值。在实际运算时,研究员可以针对下游任务需要的精度,选用相应的核心数进行计算。
本次实验验证了倚天ECS实例的下列优势:
首先,倚天710有着稳定的主频,它能够保证在高密计算的科学计算场景下不降频,保持性能输出的一致性。
其次,倚天710实例的扩展性优秀,它能够在1万级别的核心规模上,实现接近线性的加速。
最后,倚天实例的性价比相较于x86非常高,研究院可以节省将近70%的成本。
倚天帮助我们在传统的科学计算任务由本地超算部署到云平台的过程中,实现降本增效。