arrow
第十七卷, 第二十三期
机器学习模拟1亿原子:中美团队获2020「超算诺贝尔奖」戈登贝尔奖

机器之心公众号


有超算界诺贝尔奖美称的「戈登贝尔奖」今年颁给了一个中美合作研究团队,他们的研究被认为是当今计算科学中最令人兴奋领域的重大进展。
  
11月19日,在美国亚特兰大举行的国际超算大会SC2020上,美国计算机协会ACM公布了2020年戈登贝尔奖(Gordon Bell Prize)的颁奖结果。
  
来自中国和美国的一个研究小组获此奖项。他们通过机器学习将分子动力学极限从基线提升到了1亿原子的惊人数量,同时仍保证了「从头算(ab initio)」的高精度,效率是之前人类基线水平的1000倍。在基于物理模型的计算、机器学习和高性能GPU并行集群的共同助力下,研究人员已将超大系统的分子动力学模拟带进了一个全新时代。
  
研究团队在论文中写到:「这项工作的巨大成就在于,它在保持从头算精度的前提下,为分子模拟的空前数量和时间尺度打开了大门,对于更好地集成机器学习和物理建模的下一代超级计算机也提出了新的挑战。」
  
获奖团队成员包括来自 UC Berkeley 的博士后贾伟乐、副教授林霖,北京应用物理与计算数学研究所计算物理实验室的研究员王涵,北京大学助理教授陈默涵,北京大学研究生路登辉,普林斯顿大学化学系教授 Roberto Car、数学教授鄂维南以及普林斯顿大学数学系和化学系博士张林峰。
  
北京大学发来的祝贺中写到:在这项工作里,该研究团队在物理建模、机器学习与高性能计算的交叉学科领域取得突破,基于深度学习的分子动力学模拟通过机器学习和大规模并行的方法,将精确的物理建模带入了更大尺度的材料模拟中,有望在将来为力学、化学、材料、生物乃至工程领域解决实际问题发挥更大作用。
  
获奖团队做了什么?
  
著名物理学家理查德 $\cdot$ 费曼曾说过:「人类总是在不断地尝试理解生命,在这条道路上,如果非要我们说出一个最强大的假设,那就是:所有的东西都是由原子组成的,生物所做的一切都可以通过原子的振动和摆动来理解。」
  
分子动力学(MD)是一种计算机模拟方法,可以用来分析原子和分子在一段固定的时间内如何移动与交互。分子动力学模拟可以帮助科学家理解某个系统(从单个细胞到气体云)随时间流逝发生的演变。分子动力学的实际应用包括大分子的研究,如用于药物研发的蛋白质。
  
在模拟中,从头算分子动力学(Ab initio Molecular Dynamics,AIMD)用不同于标准分子动力学的方法计算原子间的力,所能达到的精度水平使其成为科学家35年来首选的模拟方法。但AIMD也有缺点,就是需要的计算量非常高,因此这方面的研究通常局限于小型系统(最多包含几千个原子的系统)。
  
在今年的戈登贝尔奖获奖论文中,研究团队引入了 Deep Potential 分子动力学(DPMD)。DPMD 是一种基于机器学习的新方法,每天能够模拟 1 亿个原子超过 1 纳秒的轨迹。虽然近年来分子动力学领域涌现了一些其他的基于机器学习的模拟方法,但获奖论文的作者表示,他们的研究首次在保持「从头算」精度的前提下实现了 1 亿个原子的高效分子动力学模拟。
  
具体来说,之前模拟的系统最大包含 100 万个硅原子(速度 = 4X 10^(-3) s/step/atom),模拟速度最快达到 1.3X 10^(-6) s/step/atom (9000 个水原子系统)。而这项研究模拟了6.79 亿个水分子,1.27 亿个铜原子。速度也比之前最快的模拟提高了几个数量级。
  
在算法的创新方面,研究团队写道:
  
为了有效利用 Summit 异构系统架构提供的算力,我们的目标是将几乎所有的计算任务和大量的通信任务迁移到 GPU。由于 DP 模型中计算粒度的大小相对有限,直接的 GPU 实现会遇到很多瓶颈,效率不高。因此,我们做出了以下算法方面的创新:
  
 (1)通过引入一种新的相邻列表数据分布,避免了在计算嵌入矩阵时出现分支,从而增加了 DeePMD的计算粒度;
  
 (2)相邻列表的新数据结构中的元素被压缩为 64 位整数,以便 GPU 对自定义 TensorFlow 算子进行更高效的优化;
  
 (3)为 DP 模型开发了混合精度计算,计算密集型任务以单精度或半精度执行,而不降低物理观测的精度。
  
由于戈登贝尔奖表彰在高性能计算方面的成就,决赛选手必须证明他们提出的算法可以在世界上最强大的超级计算机上高效运行。于是,获奖团队开发了一份高度优化的代码(GPU Deep MD-Kit)并在美国能源部橡树岭国家实验室的 Summit 超级计算机上顺利运行,在双精度下实现了 91 PFLOPS 的速度,在混合单 / 半精度下实现了 162/275 PFLOPS 的速度。Summit 超级计算机是 IBM 计划研发的一款超级计算机,目前在全球超级计算机 Top500 新榜单中排名第二。
  
戈登贝尔奖
  
戈登贝尔奖由美国计算机协会(ACM)设立于 1987 年,是国际高性能计算应用领域最高奖,每年颁发一次。该奖项主要颁发给高性能应用领域的最杰出成就。
  
戈登贝尔奖通常会由当年 TOP500 排行名列前茅的计算机系统应用获得,例如美国「泰坦」超级电脑、日本「京」超级计算机上的应用软件都曾获得此奖。奖项设立后近 30 年,都由美国和日本的软件获得此奖,直到 2016 年,中国研究团队打破了这一垄断。
  
2016 年,戈登贝尔奖授予了基于国家超级计算无锡中心的全球最快超算系统「神威$\cdot$太湖之光」运行的应用「千万核可扩展全球大气动力学全隐式模拟」,由中国科学院软件研究所、清华大学、北京师范大学、国家并行计算机工程技术研究中心,以及国家超级计算无锡中心的联合团队共同完成。