[论文解读] Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learning

Weile Jia, Handong Wang+6 more|arXiv (Cornell University)|May 1, 2020

Machine Learning in Materials Science参考文献 72被引用 35

一句话总结

展示了一种基于机器学习的分子动力学协议（Deep Potential Molecular Dynamics），在系统超过1亿个原子时达到从头计算精度；在 Summit 上对 DeePMD-kit 进行了高度优化，达到高达 91 PFLOPS 双精度，并实现纳秒级每天的模拟。

ABSTRACT

For 35 years, {\it ab initio} molecular dynamics (AIMD) has been the method of choice for modeling complex atomistic phenomena from first principles. However, most AIMD applications are limited by computational cost to systems with thousands of atoms at most. We report that a machine learning-based simulation protocol (Deep Potential Molecular Dynamics), while retaining {\it ab initio} accuracy, can simulate more than 1 nanosecond-long trajectory of over 100 million atoms per day, using a highly optimized code (GPU DeePMD-kit) on the Summit supercomputer. Our code can efficiently scale up to the entire Summit supercomputer, attaining $91$ PFLOPS in double precision ($45.5\%$ of the peak) and {$162$/$275$ PFLOPS in mixed-single/half precision}. The great accomplishment of this work is that it opens the door to simulating unprecedented size and time scales with {\it ab initio} accuracy. It also poses new challenges to the next-generation supercomputer for a better integration of machine learning and physical modeling.

研究动机与目标

• 在不牺牲精度的前提下，将从头分子动力学（AIMD）提升到大尺度系统（数亿个原子级别）和较长时间尺度。
• 利用机器学习用一个代理模型替代昂贵的电子结构计算，同时保持从头保真度。
• 开发和优化面向高性能计算的实现，以充分利用现代超级计算机中异构 GPU 架构。

提出的方法

• 使用 Deep Potential (DP) 神经网络将原子间势能表示为原子能量之和。
• 在 from ab initio 数据训练 DP 模型，并在与 LAMMPS 集成的 DeePMD-kit 中实现以进行 MD 积分。
• 优化数据布局、邻居列表和自定义 TensorFlow 运算符，以在 Summit 上最大化 GPU 吞吐量。
• 采用混合精度方案（MIX-32 和 MIX-16）在加速计算的同时保持精度。
• 重构 TensorFlow 运算符（基于 GEMM 的替换、CUDA 内核）并减少 MPI 通信以提高可扩展性。

实验结果

研究问题

• RQ1 基于深度学习的原子间势能是否能够在非常大规模的系统（数千万到数亿原子）和较长的分子动力学轨迹中再现从头计算的精度？
• RQ2 在像 Summit 这样的异构超级计算机上扩展基于 DP 的 MD 需要哪些 HPC 优化和混合精度策略？
• RQ3 在将 DP-MD 推向极端系统规模时，实际的性能收益（FLOPS、求解时间）和精度权衡有哪些？
• RQ4 在材料与液体体系的可扩展性和保真度方面，DP-MD 与传统 AIMD 和经验力场相比如何？

主要发现

• DP-MD 在 Summit 上以 91 PFLOPS（双精度）实现每天超过 1 亿原子级 AIMD 风格轨迹，混合单/半精度在 4,560 个节点上达到 162 PFLOPS / 275 PFLOPS。
• 对于 127 百万原子铜系统，单步求解时间为 8.1e-10 s/step/atom（双精度），相当于约 0.8 ns/day，混合半精度为 2.5 ns/day。
• 优化邻居表数据布局并将邻居条目压缩为 64 位整数可消除分支并提高 GPU 效率，在单个 GPU 上，与基线 CPU 实现相比，定制 TensorFlow 运算符实现了 64.6x 的加速。
• 用 GEMM 替换 MATMUL 和 SUM，融合 TANH 与 TANHGrad，以及 CUDA 内核融合显著提升 TensorFlow 运算符效率，促成整体 MD 循环加速。
• 混合精度方案 MIX-32 和 MIX-16 在所测试的水配置中获得的能量/力精度可与双精度相当，其中 MIX-32 在能量和力误差方面与双精度相当。