Scaling the leading accuracy of deep equivariant models to biomolecular simulations of realistic size
This enables quantum-fidelity simulations of large biomolecular systems, such as complete HIV capsids, for researchers in computational biology and drug discovery, representing a significant scaling advance rather than an incremental improvement.
The authors tackled the challenge of scaling deep equivariant neural networks to large biomolecular simulations, achieving stable simulations of a 44-million atom HIV capsid structure with excellent strong scaling up to 100 million atoms and 70% weak scaling to 5120 GPUs.
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.