Breaking the Exascale Barrier for the Electronic Structure Problem in Ab-Initio Molecular Dynamics

为了理解物质在最基本层面的行为，科学家们经常转向一种被称为“从头算”分子动力学（ab-initio molecular dynamics）的技术。这种方法试图通过计算推动和拉动原子的力，来模拟分子、表面或固体中原子的运动。与依赖于简化经验规则的旧方法不同，该技术直接求解复杂的量子力学电子问题。它不将电子视为模糊的背景，而是将其视为决定原子如何相互作用的主要参与者。虽然这提供了高度精确的化学反应和材料属性图景，但也付出了沉重的代价：所需的计算量随着系统规模的增大而极速增长，以至于模拟大型、真实的结构长期以来被认为是不可能的。几十年来，研究人员被迫研究微小的物质碎片，无法看到由数十亿个原子共同运动构成的生命细胞或复杂材料的全貌。

来自德国帕德博恩大学的一个研究小组现在突破了这一障碍，展示了一种能够模拟包含高达 8300 万个原子的电子结构的方法。通过改进一种被称为非正交局部子矩阵（non-orthogonal local submatrix）的技术，并在配备了数千个专用图形处理器的超级计算机上运行，他们实现了超过 1.1 exaflops 的持续计算速度。这意味着该系统每秒执行了超过一百万亿亿次（one quintillion）浮点运算，这一里程碑式的成就使这项特定的科学应用成为首批突破“百亿亿级”（exascale）障碍的应用之一。研究人员不仅仅是运行了一次模拟；他们设计了一种新的数学工作组织方式，使得硬件能够以接近其最大理论容量的 80% 进行运行。他们的工作证明，通过适当的算法调整，计算整个蛋白质在溶液中的量子行为是可能的，这为以空前详尽的细节研究生物机器和复杂材料打开了大门。

这些模拟的核心挑战在于，每当一个原子移动哪怕极其微小的距离，都必须重新计算整个系统的电子结构，以确定作用在该原子上的新力。在传统方法中，随着原子数量的增加，这种重新计算会变得极其缓慢。研究人员使用了一种将庞大的数学问题分解为称为“子矩阵”的小型、可管理部分的方法。与其尝试一次性求解整个系统的方程，计算机会将数据隔离成较小的部分，独立求解这些部分，然后重新组合结果。这种“局部”方法避免了不同计算部分之间频繁的通信，而这通常是大型模拟中的瓶ль颈。团队将此方法应用于一个特定的生物靶点：SARS-CoV-2 病毒的刺突蛋白（spike protein），这是病毒附着在人类细胞上的结构。他们模拟了锚定在脂质层并被水包围的蛋白质，在初步测试中创建了一个包含约 170 万个原子的系统，随后通过扩展这些蛋白质的网格，达到了总计 8300 万个原子的规模。

为了达到这种性能水平，研究人员不仅需要使用更多的计算机，还必须从根本上重新思考软件如何与硬件进行交互。他们使用的超级计算机位于国家能源研究科学计算中心（National Energy Research Scientific Computing Center），配备了 4,400 个 NVIDIA A100 图形处理器。这些芯片旨在处理海量的并行计算，但它们在处理大型、稠密的数据块时效率最高。原始版本的算法创建的子矩阵往往太小，无法充分利用这些芯片的性能。团队引入了一种新的启发式规则（即决策规则），在处理之前将多个数据列合并为更大的子矩阵。这种调整并非微小的修补，而是基于图形处理器特定性能特征的战略性转变。通过测量芯片在处理不同大小矩阵时的速度，研究人员优化了数据的分组方式，以确保硬件在峰值效率下工作。这种修改使得系统能够维持一种此前被认为对于此类计算而言无法达到的性能水平。

这项工作的成果是通过在刺突蛋白网格上运行模拟来衡量的，这实际上创建了一个拥有 8300 万个原子的虚拟环境。团队追踪了完成每一步计算所需的时间以及执行的数学运算总数。他们发现，该系统持续提供 1.106 到 1.127 exaflops 的速度，保持了硬件理论峰值性能的约 80%。这是一个显著的成就，因为高性能计算系统在扩展到数千个处理器时，往往难以维持高效率；通常，随着通信开销的增加，效率会下降。在这种情况下，由于该方法的局部特性，处理器几乎将所有时间都花在了计算上，而不是等待数据。研究人员验证了模拟的准确性依然很高，确保了速度的提升并未以牺牲科学有效性为代价。

这一突破不仅仅是关于模拟一种病毒；它代表了计算科学的一种新能力。模拟拥有千万级原子电子结构的能力，意味着科学家现在可以研究以前无法触及的现象，例如大型生物分子在其天然水性环境中的行为，或者复杂材料在压力下的性质。该方法具有通用性，可以应用于任何需要对大型稀疏数据集进行数学函数求值的问题，而不局限于分子动力学。通过证明在现实世界的科学应用中实现百亿亿级性能是可行的，研究人员为未来的高性能计算提供了一个蓝图。他们证明了，通过将算法设计与硬件能力相结合，可以解决那些曾经被认为过于庞大而无法计算的问题，从而让原子和电子的量子力学世界在研究生命与物质时变得更加清晰。