Breaking the Exascale Barrier for the Electronic Structure Problem in Ab-Initio Molecular Dynamics

arXiv:2205.12182 · physics.comp-ph, cond-mat.mtrl-sci, physics.chem-ph, q-bio.QM · Submitted 2022-05-24 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: Genomics Radio. Generated commentary on the latest computational biology and genomics papers.

Ines: I'm Ines, and with me are Marcus and Yuki, guest researcher.

Marcus: Today's paper: "Breaking the Exascale Barrier for the Electronic Structure Problem in Ab-Initio Molecular Dynamics".

Ines: The non-orthogonal local submatrix method applied to electronic-structure based molecular dynamics simulations demonstrates a sustained performance exceeding 1.1 EFLOP/s in mixed FP16/FP32 arithmetic,

Marcus: First, who's behind it and why it matters.

Paper summary: Ines: So, to get us started, we’re looking at this paper titled "Breaking the Exascale Barrier for the Electronic Structure Problem in Ab-Initio Molecular Dynamics" by Schade and his team. Basically, they are tackling a really tough problem in solid-state physics and chemistry: electronic structure based molecular dynamics simulations. The core idea of this paper is to see if they can solve these complex quantum mechanical problems on supercomputers at an exascale level, which is usually beyond what's currently feasible for this specific computational task.

Marcus: From a data science standpoint, I'm interested in what the authors claim regarding the performance metrics. They are proposing a technique called the non-orthogonal local submatrix method, or NOLSM, and they claim it manages to sustain more than one point one EFLOP/s when using mixed FP16/FP32 arithmetic on four thousand four hundred NVIDIA A100 GPUs on the Perlmutter system. That level of sustained performance is what makes this work significant because it pushes past the current limits for these kinds of calculations, which is a big deal for large-scale molecular dynamics.

Ines: That sounds intense; what exactly is the central thesis they are pushing here? They aren't just showing off a new trick; they are trying to solve the fundamental scaling issue where electronic structure calculations need to scale at most linearly with the number of atoms, which is crucial for large systems.

Yuki: From my perspective in population genetics, when you talk about scaling with the number of atoms, it really speaks to how we model complex biological systems. If a simulation can handle eighty-three million atoms, as they used for the SARS-CoV-two spike protein example, that implies we could potentially model much larger biological entities or more detailed molecular interactions without hitting insurmountable computational walls.

Marcus: Exactly, Yuki; the ability to run calculations on systems with up to eighty-three million atoms means we can test much more complex scenarios related to protein folding or drug binding effects in a way that was previously impossible because of those scaling constraints. The paper is focused on making this feasible using a method that avoids inter-node communication during the solution phase, which is another major hurdle for massive parallel systems.

Ines: And they’re claiming this modification to the original submatrix method pushes the sustained fraction of peak performance up to about eighty percent when combining submatrices. That detail about boosting the efficiency by that factor sounds like a very important piece of evidence supporting their claim regarding exascale potential for this problem.

Yuki: It's interesting how they are linking this computational technique directly to the complexity of the molecular structure being simulated; it shows that the mathematical approach is tailored to meet the physical requirements of electronic structure calculations, which ties into how we interpret genetic data at a structural level.

Conclusion: Ines: Considering the full scope of this work by Schade, Kenter, Elgabarty, and Lass, it seems the title "Breaking the Exascale Barrier for the Electronic Structure Problem in Ab-Initio Molecular Dynamics" accurately reflects their achievement in demonstrating sustained performance exceeding one point one EFLOP/s under mixed precision arithmetic on a large system. This isn't just an incremental improvement; it addresses a fundamental computational bottleneck that has kept these simulations from reaching the exascale level for this specific type of problem.

Marcus: I think what this work means, in simpler terms, is that we can now tackle much larger and more realistic electronic structure problems in molecular dynamics simulations than we could before because the computational engine can keep up with the physics required for those bigger systems. It validates a new way to solve the matrix function evaluations efficiently across thousands of GPUs without needing constant communication between them during that critical solving phase.

Yuki: From a wider view, this isn't just about protein simulations; it speaks to how computational power is unlocking our ability to understand the molecular basis of life at an even deeper level. If we can efficiently simulate these interactions on exascale systems, it opens doors for studying more intricate biological processes that might be driving evolution or disease mechanisms.

Ines: So, the authors are essentially showing that a novel submatrix combination heuristic, specifically one based on a cubic metric modified for GPU performance characteristics, allows them to sustain about eighty percent of peak performance when solving the required matrix functions for these AIMD simulations. This specific mechanism is what they believe unlocks that higher sustained fraction in their calculation involving up to eighty-three million atoms.

Marcus: And from a statistical viewpoint, the results showing node performance between one PFLOP/s and one point zero seven PFLOP/s, averaging about one point zero three PFLOP/s, represents roughly "eighty percent of the peak performance of one point two four eight PFLOP/s," which is a concrete measure of how well this method utilizes the available hardware resources on that Perlmutter system. That kind of quantification is what makes their conclusion solid for anyone analyzing genomic or structural data performance across different architectures.

Yuki: It really highlights how the mathematical innovations, like those submatrix combination heuristics, have a direct impact on the practical ability to model complex biological structures accurately, showing that computational methods can be tuned to match the physical scale of what we are studying.

Ines: So, for anyone listening who is interested in the computational biology side of this paper, it tells us that when you design algorithms for these quantum mechanics problems, focusing on how they handle matrix operations at scale and incorporating hardware specifics into your heuristics can lead to significant sustained gains in performance. This points toward a more robust framework for tackling the electronic structure challenge.

Marcus: It’s about showing that even with mixed precision FP16/FP32 arithmetic, this approach is powerful enough to achieve those high throughput figures, which is important because real-world simulations often have to balance accuracy and speed in different parts of the calculation. We need methods that handle both sides of that coin effectively for large datasets.

Yuki: And ultimately, the implication is that as we push these computational limits, we gain a more powerful tool for understanding structure from the ground up, allowing us to probe biological mechanisms with higher fidelity than previously possible.

Robert Schade, Tobias Kenter, Hossam Elgabarty, Michael Lass, Thomas D. Kuhne, Christian Plessl

Paderborn University

physics.comp-ph, cond-mat.mtrl-sci, physics.chem-ph, q-bio.QM

Submitted: 2022-05-24

Updated: 2022-06-07

Comments: 6 pages, 6 figures, 2 tables

Journal ref: Int. J. High Perform. Comput. Appl. 37, 530-538 (2023)

DOI: 10.1177/10943420231177631

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 86/100

The gist: The non-orthogonal local submatrix method applied to electronic-structure based molecular dynamics simulations demonstrates a sustained performance exceeding 1.1 EFLOP/s in mixed FP16/FP32

Key concepts

Electronic-structure based AIMD
These simulations use quantum mechanics to model atomic forces by solving the electronic structure problem at every time step. This is necessary because empirical models often fail to describe complex physical phenomena in solid-state chemistry and physics.
Non-Orthogonal Local Submatrix Method (NOLSM)
This is a massively parallel technique that approximates matrix functions needed for the density matrix calculation. It avoids inter-node communication during solving and scales well by intelligently combining submatrices derived from the input matrices.
Mixed Precision Arithmetic (FP16/FP32)
This refers to using different levels of floating-point precision (like 16-bit and 32-bit) simultaneously. The method is specifically designed to efficiently utilize these mixed precision tensor cores on GPUs to maximize computational throughput.
FLOPsNOLSM
This metric estimates the total floating-point operations performed by the NOLSM method, calculated as $2n^3$ for a gemm operation in FP16/FP32 mixed precision. It is used to quantify the computational effort required by this specific algorithm.

Terminology

Summary

The non-orthogonal local submatrix method applied to electronic-structure based molecular dynamics simulations demonstrates a sustained performance exceeding 1.1 EFLOP/s in mixed FP16/FP32 arithmetic, successfully breaking the exascale barrier for this computational problem.

Overview of the Problem

Electronic-structure based ab-initio molecular dynamics (AIMD) simulations are crucial tools in solid-state physics and chemistry because they explicitly treat quantum mechanical effects, which are necessary when empirical model potentials fail to describe relevant physical phenomena. To derive the forces acting on atoms in these simulations, the electronic structure problem must be solved at every time step by calculating forces as the negative gradient of the total energy: Fi = -∂E/∂Ri. The total energy is composed of electronic energy (Eelec), double counting terms (Edc), and nuclear Coulomb repulsion energy (Eion). The bulk of the computational effort lies in obtaining the electronic energy, which for very large systems requires methods that scale at most linearly with the number of atoms.

The Non-Orthogonal Local Submatrix Method (NOLSM)

The paper proposes the non-orthogonal local submatrix method (NOLSM) as a massively parallel technique to solve this electronic-structure problem by using an approximate solution of required matrix functions. This method is designed to avoid inter-node communication during the solution phase and scales extremely well, being shown to be capable of scaling to more than one thousand GPUs while efficiently utilizing mixed-precision tensor cores for linear algebra operations.

The evaluation of the density matrix, D, which is a matrix function of two matrices (H0 and S), is performed using the submatrix method. This involves three steps:

  1. In the first step, a submatrix Ti(A) is generated for every column i of the input matrix A by removing all rows and corresponding columns for which A has vanishing or negligible elements in column i, resulting in a smaller and much denser submatrix Ti(A).

  2. The matrix function is applied to this submatrix, i.e., f(Ti(A)) is evaluated.

  3. The matrix elements of f(Ti(A)) corresponding to column i are written to the result matrix B by applying the reverse submatrix construction in step 1).

For the non-orthogonal case, where D = D(H0,S), this combines sparsity patterns before submatrix construction, building two submatrices, Ti(H0) and Ti(S), and evaluating the function as: Ti(D) = 1/2 I - sign(Ti(S)−1Ti(H0) − µI).

Innovations in Submatrix Combination Heuristics

The paper introduces innovations beyond previous implementations like [9] by considering the matrix-size dependency of GPU performance when combining submatrices. The initial heuristic for combination, based on a cubic metric, is modified to include GPU performance characteristics. The new criterion compares predicted runtime:

p(ni+nj−ni∧j)×(ni + nj − ni∧j) < p(ni)×n3/i + p(nj)×n3/j. This criterion effectively increases the dimension of the submatrices and the achievable portion of peak performance.

Performance Measurement and Results

The benchmark system used is the full-length SARS-CoV-2 spike protein in aqueous solution, which includes approximately 1.7 mio. atoms (83 million atoms for one calculation). The simulations were performed on a 1,100-node Perlmutter system utilizing 4,400 NVIDIA A100 GPUs.

The main performance measurements include:

  1. Wall clock time of the NOLSM method (TNOLSM).

  2. Floating-point operations (FLOPsNOLSM), estimated as 2n3/ for a gemm-operation C = αA · B + βC in FP16/FP32 mixed precision.

  3. Node performance (PNOLSM,i), defined as FLOPsNOLSM,i/TNOLSM,i.

  4. Total performance (PNOLSM), the sum of node performances.

The results for the 7×7 grid (83 mio. atoms) show that nodes with 4 NVIDIA A100 GPUs mainly fall in the range between 1 PFLOP/s and 1.07 PFLOP/s, averaging about 1.03 PFLOP/s, which represents approximately 80% of the peak performance of 1.248 PFLOP/s. The floating-point throughput achieved with 4,400 NVIDIA A100 GPUs is between 1.106 to 1.127 EFLOP/s in FP16/FP32 mixed precision, achieving about "80% of the theoretical peak performance of the tensor cores.

Improvements for AI systems

Here are the specific improvements that can be made to AI systems, derived from the insights presented in this research paper:

  1. A significant increase in sustained peak performance for electronic-structure based molecular dynamics (AIMD) simulations, specifically for problems involving large atom counts (e.g., SARS-CoV-2 spike proteins with up to 83 million atoms).

  2. The development and application of a novel computational technique called the Non-Orthogonal Local Submatrix Method (NOLSM) combined with GPU acceleration that allows for sustained performance reaching approximately 80% of the theoretical peak performance on systems utilizing mixed-precision floating-point arithmetic (FP16/FP32).

  3. Implementation of advanced submatrix combination heuristics (Equation 10), which dynamically optimize the combination of smaller submatrices based on GPU matrix multiplication performance characteristics, leading to an estimated speedup improvement compared to traditional heuristics.

These improvements enable the following capabilities for AI systems:

  1. Ancillary or core simulations in materials science, chemistry, and drug discovery that require explicit treatment of quantum-mechanical effects (Ab-Initio Molecular Dynamics).

  2. The ability to accurately model complex biological structures, such as protein dynamics (e.g., SARS-CoV-2 spike proteins) in realistic aqueous environments with high atom counts, at a scale previously unattainable due to computational bottlenecks.

  3. More efficient and faster execution of linear-scaling electronic structure methods (like those based on reduced density matrices), allowing for the simulation of much larger systems while maintaining physical accuracy.

Abstract

The non-orthogonal local submatrix method applied to electronic-structure based molecular dynamics simulations is shown to exceed 1.1 EFLOP/s in FP16/FP32 mixed floating-point arithmetic when using 4,400 NVIDIA A100 GPUs of the Perlmutter system. This is enabled by a modification of the original method that pushes the sustained fraction of the peak performance to about 80%. Example calculations are performed for SARS-CoV-2 spike proteins with up to 83 million atoms.

Related papers