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.

翻译：本研究将深度等变神经网络的领先精度、采样效率与鲁棒性推广至极端计算规模。通过创新模型架构、大规模并行化以及针对高效GPU利用率优化的模型与实现，我们构建了Allegro架构。该架构弥合了原子模拟中精度与速度的权衡，得以在量子保真度水平上描述空前复杂度结构的动力学过程。为验证Allegro的可扩展性，我们在Perlmutter超级计算机上实现了蛋白质动力学的纳秒级稳定模拟，并将模拟规模扩展至包含完整全原子显式溶剂化的4400万原子HIV衣壳结构。我们展示了在1亿原子规模下的优异强扩展性，以及在5120块A100 GPU上达到70%的弱扩展效率。