[论文解读] Scaling the leading accuracy of deep equivariant models to biomolecular simulations of realistic size

Albert Musaelian, Anders Johansson+2 more|arXiv (Cornell University)|Apr 20, 2023

Machine Learning in Materials Science被引用 14

一句话总结

Allegro 通过将可扩展的等变神经势能与大规模 GPU 并行性相结合，提供可扩展、量子级精确的生物分子模拟，在 Perlmutter 上实现纳秒级蛋白质动力学和 4400万原子 HIV 膜衣壳仿真，并具有强/弱规模化。

ABSTRACT

This work brings the leading accuracy, sample efficiency, and robustness of deep equivariant neural networks to the extreme computational scale. This is achieved through a combination of innovative model architecture, massive parallelization, and models and implementations optimized for efficient GPU utilization. The resulting Allegro architecture bridges the accuracy-speed tradeoff of atomistic simulations and enables description of dynamics in structures of unprecedented complexity at quantum fidelity. To illustrate the scalability of Allegro, we perform nanoseconds-long stable simulations of protein dynamics and scale up to a 44-million atom structure of a complete, all-atom, explicitly solvated HIV capsid on the Perlmutter supercomputer. We demonstrate excellent strong scaling up to 100 million atoms and 70% weak scaling to 5120 A100 GPUs.

研究动机与目标

• 展示一种可扩展、准确的等变深度学习势能，用于大规模生物分子系统。
• 将最先进的等变模型的精度与Exascale 计算能力结合起来。
• 展示前所未有系统规模下的稳定、长时间尺度生物分子模拟（蛋白质动力学）。

提出的方法

• 引入 Allegro，一种可扩展的等变 MLIP，将系统能量分解为原子级和逐序对分量。
• 采用带标量轨道和等变张量轨道的双轨架构，在提升容量的同时控制计算成本。
• 采用步幅内存布局和优化的张量乘积，实现高效的GPU执行和单次收缩运算。
• 应用混合精度计算（TF32 与 FP32 及偶发的 FP64），在速度与数值稳定性之间取得平衡。
• 引入每种物种对的截断半径，以在不牺牲精度的前提下减少邻域计算。
• 通过一个GPU驻留管线将 Allegro 与 LAMMPS 集成，以利用已建立的分子动力学空间分解。

实验结果

研究问题

• RQ1 Allegro 是否能够在扩大到数亿原子规模的同时达到最先进的等变精度？
• RQ2 在现代 GPU 上，大规模生物分子系统下 Allegro 的性能（steps/s，ns/day）是多少？
• RQ3 在真实生物分子结构上，Allegro 如何在长时间模拟中处理稳定性与精度？
• RQ4 哪些体系结构创新使得在分子动力学中的可扩展等变学习在不牺牲精度的前提下成为可能？
• RQ5 Allegro 在精度与可扩展性方面，与现有 MLIP 和经典力场相比，在生物分子系统上有何比较？

主要发现

• Allegro 在基准测试上达到高精度，并在 perlmutter 规模的硬件上扩展到 126.4 百万原子，性能为 100 timesteps/s。
• 展示可到达 100 百万原子强标度，以及在 5120 A100 GPU（1280 节点）上达到 70% 的弱尺度。
• 首次将最先进的机器学习原子间势应用于大规模生物分子仿真。
• 产生了蛋白质动力学的纳秒级稳定模拟（如 DHFR、Factor IX）以及一个 44M-atom 的 HIV capsid 溶剂化系统。
• 将标量轨道和张量轨道结合，采用步幅内存布局和优化的张量乘积，在保持精度的同时降低计算成本。
• 结合每种物种对的截断和混合精度技术，最大化 GPU 吞吐率。