Accelerating Atomistic Simulations with GPUs and Machine Learning

Abstract

The field of atomistic simulations is rapidly evolving. The simulations are becoming more reliable and accurate through the invention of ever more advanced ML potentials such as uncertainty aware and equivariant methods. Simultaneously, the leading GPU-based supercomputers are finally reaching the exascale, which paves the way for larger and longer simulations. This thesis revolves around the intersection of these advances and the acceleration of state-of-the-art ML potentials with modern GPUs on the world's largest supercomputers.

More concretely, my focus has been the performance portability and scalability of FLARE and Allegro, two ML potentials near the opposite ends of the cost-accuracy trade-off. FLARE is a sparse Gaussian process potential that attempts to push the envelope for speed while maintaining reasonable accuracy. I have developed a Kokkos implementation of FLARE that outperformed previous state-of-the-art methods by 70\% on the second fastest supercomputer in the world, and recently reached one trillion atoms on Frontier, the fastest supercomputer. Allegro is an equivariant neural network implemented in PyTorch which sacrifices some speed to reach leading accuracy while maintaining scalability through its innovative architecture. Through the scalability and my LAMMPS interface, Allegro was able to efficiently utilize 5120 GPUs and reach relevant speeds for a wide range of biomoleular structures.