在新冠病毒的量子化学模拟中，超算首次实现百亿亿次双精度浮点计算
光子盒研究院出品
5月27日，德国帕德博恩大学（UPB）科学家Thomas D. Kühne教授、Christian Plessl教授领导团队宣布，他们打破了计算科学应用中混合精度计算的百亿亿次(Exascale)大关：成功在美国国家能源研究科学计算中心(NERSC) Perlmutter超级计算机上实现了量子化学应用程序CP2K，并达到了1.1 EFLOP/s(每秒百亿亿次)的应用级性能。
相关论文《打破从头算分子动力学中电子结构问题的百亿亿级障碍》[1]已发布在arXiv上。
目前，数据科学正处于“百亿亿次时代”（运算速度由1015提升到1018）的风口浪尖，人们普遍预计第一台突破百亿亿次双精度浮点计算门槛的超级计算机将在5月底汉堡的“ISC高性能会议（ISC）”上发布。
这一里程碑式事件将标志着通向百亿亿次级竞赛的结束，这场竞赛甚至可以被称为“21世纪的太空竞赛”。
HPL（High Performance Linpack）是并行计算集群的测试工具，通过调整测试文件的矩阵大小、CPU数量等测试参数，可以执行HPL测试程序来获取机器的最佳性能。衡量结果是每秒浮点运算(Flops)，可以并行对单个处理器或分布式处理器进行基准测试。用HPL基准对超级计算机进行排名的优点和缺点已经广为人知。提取接近HPL性能和效率的东西是雄心勃勃的，许多科学代码由于有限的并行性、不足的矢量化机会、通信开销、负载不平衡等原因，只利用了该理论性能的一小部分。因此，实际利用超大规模计算机在计算科学方面的能力将需要对已有的算法、数值库和应用代码进行调整。
如今，UPB团队已经克服了量子化学领域超大规模计算的挑战，并开发了一种新的方法。这种方法的最初变体以《用于矩阵函数近似计算的子矩阵方法》[2]为题，在SC'20会议上被提出；2021年，这一最初变体被改进为另一高度可扩展、高效的方法[3]，其核心是在一个非常大的稀疏矩阵上计算近似的矩阵函数，这也是量子力学线性扩展电子结构计算的关键操作。
子矩阵的步骤示意图。用于近似计算大稀疏矩阵A的矩阵函数f(A)，第一步是为矩阵A的每一列构建一个子矩阵Ti(A)。然后将矩阵函数应用于密集子矩阵，即f(Ti(A))，最后将相关结果列插入稀疏结果矩阵中。
该方法将最终产生的密度矩阵（一个巨大的稀疏矩阵）划分为许多更小但密集的子矩阵，在这些子矩阵上评估矩阵函数，并将这些中间解集合成一个全局解。由于子矩阵上的所有矩阵函数的评估都是独立的，该方法避免了产生交集并具有极大的可扩展性；由于子矩阵小而密集（几千行/列的数量级），在这些矩阵上执行线性代数达到了接近GPU的峰值性能。
因此，该方法引入了一个近似误差，其大小对该应用程序来说是可以接受的；此外，通过对原子运动的Langevin型方程补偿引入的误差，可以使该应用对低精度计算具有容错性。这为使用张量核心的混合精度计算开辟了道路：比双精度或单精度算术的性能高出一个数量级。
为了评估子矩阵方法，它被整合到流行的开源量子化学程序CP2K[4]中。在那里，它被用在解决电子结构问题的xTB方法中——迄今为止从头算分子动力学模拟中最大的计算时间部分。
2021年，帕德博恩的科学家们在当时欧洲最快的超级计算机（现在全球排名第8）——于利希超级计算中心的“JUWELS Booster”上对HIV病毒进行了多达1.02亿个原子的模拟，创造出了最大的基于电子结构的从头算分子动力学模拟记录[5]：在混合精度浮点运算中实现了324 petaflop/s(每秒千万亿次)的计算性能和67.7%的理论可用计算能力的效率，这在该应用领域是非常突出的。
自从在于利希创下模拟记录后，该方法得到了进一步优化：以提高GPU硬件加速器使用效率使GPU的运行更接近峰值性能。为了实际测试方法的超大规模能力，该团队运用位于美国国家能源研究科学计算中心（NERSC）、目前在Top500名单中排名第五的“Perlmutter”超级计算机，来突破混合精度运算的超大规模障碍。
水溶液中的新冠病毒刺突蛋白：完整细胞（左）和不含氢原子和氧原子（右）。
2022年4月，该团队报告具体完成情况：在对新冠病毒刺突蛋白的模拟中，使用4400个GPU加速器在实际科学计算应用中首次突破了百亿亿次障碍，在计算时间关键部分实现了1.1 exaflop/s的混合精度运算[6]。
为了对这一突破进行分类，我们可以考虑8300万个原子的单一模拟步骤需要42秒，执行大约42×1.127×1018=47×1018(47百亿亿)次浮点运算。如果不考虑内存要求，这样的计算步骤在2008年的第一个千万亿次级系统Roadrunner中需要花费约47000s或大约13个小时，在1997年的第一个百万亿次级系统ASCI Red中则需要大约1.5年。
尽管已经有了初步成功，对于相关团队来说这一课题还远远没有结束：该团队已经在研究下一步的工作。化学和固体物理学中的原子模拟的黄金标准是密度泛函理论方法，研究团队未来将试图运用子矩阵方法于密度泛函理论。
参考链接：
[1]https://arxiv.org/abs/2205.12182
[2]https://doi.org/10.5555/3433701.3433807
[3]https://arxiv.org/abs/2104.08245
[4]https://cp2k.org/
[5]https://doi.org/10.1016/j.parco.2022.102920
[6]https://doi.org/10.48550/arXiv.2205.12182