Quantum-Based Molecular Dynamics Simulations Using Tensor Cores

Abstract

Tensor cores, along with tensor processing units, represent a new form of hardware acceleration specifically designed for deep neural network calculations in artificial intelligence applications. Tensor cores provide extraordinary computational speed and energy efficiency but with the caveat that they were designed for tensor contractions (matrix–matrix multiplications) using only low-precision floating-point operations. Despite this perceived limitation, we demonstrate how tensor cores can be applied with high efficiency to the challenging and numerically sensitive problem of quantum-based Born–Oppenheimer molecular dynamics, which requires highly accurate electronic structure optimizations and conservative force evaluations. The interatomic forces are calculated on-the-fly from an electronic structure that is obtained from a generalized deep neural network, where the computational structure naturally takes advantage of the exceptional processing power of the tensor cores and allows for high performance in excess of 100 Tflops on a single Nvidia A100 GPU. Stable molecular dynamics trajectories are generated using the framework of extended Lagrangian Born–Oppenheimer molecular dynamics, which combines computational efficiency with long-term stability, even when using approximate charge relaxations and force evaluations that are limited in accuracy by the numerically noisy conditions caused by the low-precision tensor core floating-point operations. A canonical ensemble simulation scheme is also presented, where the additional numerical noise in the calculated forces is absorbed into a Langevin-like dynamics.

Quantum perturbation theory using Tensor cores and a deep neural network

“…It is therefore paramount that any numerical algorithms used with the Tensor core hardware are made robust to this introduction of lower precision arithmetic. We have previously shown how this is possible for quantum-based molecular dynamics simulations and density-matrix electronic structure calculations using the computational framework of a generalized convolutional deep neural network [24,27].…”

“…To further enhance accuracy, a refinement, or purification layer may be used as the final output layer of the network in Eq. (7), where here all matrix algebra is carried out in doubleprecision, FP64 [24,27]. This purification step is equivalent to a final double flip, e.g.…”

“…However, experience has shown this additional refinement step to be unnecessary in certain cases, e.g. for QMD simulations [27], and we will not consider any refinement layers in this article.…”

“…The recursive expansion of the density matrix response in Eqs. (25) to (27) forms the basis of DMPT in its DNN-SP2 form.…”

“…The development of accelerated quantum response calculations presented in this article continues a currently ongoing theme of using Tensor cores (or other similar machine learning inspired hardware) for more general scientific applications [24,[27][28][29][30][31][32][33][34]. This work is similar in spirit to the transition from CPUs to GPUs for scientific computations that started over a decade ago [35][36][37][38][39][40][41][42][43][44][45][46][47], and in some sense, represents the next phase of a Darwinianlike computational evolution driven by new hardware environments.…”