skip to main content
US FlagAn official website of the United States government
dot gov icon
Official websites use .gov
A .gov website belongs to an official government organization in the United States.
https lock icon
Secure .gov websites use HTTPS
A lock ( lock ) or https:// means you've safely connected to the .gov website. Share sensitive information only on official, secure websites.

Attention:

The NSF Public Access Repository (PAR) system and access will be unavailable from 11:00 PM ET on Friday, November 14 until 2:00 AM ET on Saturday, November 15 due to maintenance. We apologize for the inconvenience.


Title: End-to-End Differentiable Reactive Molecular Dynamics Simulations Using JAX
The reactive force field (ReaxFF) interatomic potential is a powerful tool for simulating the behavior of molecules in a wide range of chemical and physical systems at the atomic level. Unlike traditional classical force fields, ReaxFF employs dynamic bonding and polarizability to enable the study of reactive systems. Over the past couple decades, highly optimized parallel implementations have been developed for ReaxFF to efficiently utilize modern hardware such as multi-core processors and graphics processing units (GPUs). However, the complexity of the ReaxFF potential poses challenges in terms of portability to new architectures (AMD and Intel GPUs, RISC-V processors, etc.), and limits the ability of computational scientists to tailor its functional form to their target systems. In this regard, the convergence of cyber-infrastructure for high performance computing (HPC) and machine learning (ML) presents new opportunities for customization, programmer productivity and performance portability. In this paper, we explore the benefits and limitations of JAX, a modern ML library in Python representing a prime example of the convergence of HPC and ML software, for implementing ReaxFF. We demonstrate that by leveraging auto-differentiation, just-in-time compilation, and vectorization capabilities of JAX, one can attain a portable, performant, and easy to maintain ReaxFF software. Beyond enabling MD simulations, end-to-end differentiability of trajectories produced by ReaxFF implemented with JAX makes it possible to perform related tasks such as force field parameter optimization and meta-analysis without requiring any significant software developments. We also discuss scalability limitations using the current version of JAX for ReaxFF simulations.  more » « less
Award ID(s):
2209717
PAR ID:
10432050
Author(s) / Creator(s):
; ; ; ; ;
Editor(s):
Bhatele, A.; Hammond, J.; Baboulin, M.; Kruse, C.
Date Published:
Journal Name:
Lecture notes in computer science
Volume:
13948
ISSN:
1611-3349
Page Range / eLocation ID:
202-219
Format(s):
Medium: X
Sponsoring Org:
National Science Foundation
More Like this
  1. null (Ed.)
    Reactive force elds provide an a ordable model for simulating chemical reactions at a fraction of the cost of quantum mechanical approaches. However classically accounting for chemical reactivity often comes at the expense of accuracy and transferability, while computational cost is still large relative to non-reactive force elds. In this Perspective we summarize recent e orts for improving the performance of reactive force elds in these three areas with a focus on the ReaxFF theoretical model. To improve accuracy we describe recent reformulations of charge equilibration schemes to overcome unphysical long-range charge transfer, new ReaxFF models that account for explicit electrons, and corrections for energy conservation issues of the ReaxFF model. To enhance transferability we also highlight new advances to include explicit treatment of electrons in the ReaxFF and hybrid non-reactive/reactive simulations that make it possible to model charge transfer, redox chemistry, and large systems such as reverse micelles within the framework of a reactive force eld. To address the computational cost we review recent work in extended Lagrangian schemes and matrix preconditioners for accelerating the charge equilibration method component of ReaxFF and improvements in its software performance in LAMMPs. 
    more » « less
  2. As FPGAs and GPUs continue to make inroads into high-performance computing (HPC), the need for languages and frameworks that offer performance, productivity, and portability across heterogeneous platforms, such as FPGAs and GPUs, continues to grow. OpenCL and SYCL have emerged as frameworks that offer cross-platform functional portability between FPGAs and GPUs. While functional portability across a diverse set of platforms is an important feature of portable frameworks, achieving performance portability often requires vendor and platform-specific optimizations. Achieving performance portability, therefore, comes at the expense of productivity. This paper presents a quantification of the tradeoffs between performance, portability, and productivity of OpenCL and SYCL. It extends and complements our prior work on quantifying performance-productivity tradeoffs between Verilog and OpenCL for the FPGA. In addition to evaluating the performance-productivity tradeoffs between OpenCL and SYCL, this work quantifies the performance portability (PP) of OpenCL and SYCL as well as their code convergence (CC), i.e., a measure of productivity across different platforms (e.g., FPGA and GPU). Using two applications as case studies (i.e., edge detection using the Sobel filter, and graph link prediction using the Jaccard similarity index), we characterize the tradeoffs between performance, portability, and productivity. Our results show that OpenCL and SYCL offer complementary tradeoffs. While OpenCL delivers better performance portability than SYCL, SYCL offers better code convergence and a 1.6× improvement in source lines of code over OpenCL. 
    more » « less
  3. Modern High Performance Computing (HPC) systems are built with innovative system architectures and novel programming models to further push the speed limit of computing. The increased complexity poses challenges for performance portability and performance evaluation. The Standard Performance Evaluation Corporation (SPEC) has a long history of producing industry-standard benchmarks for modern computer systems. SPEC’s newly released SPEChpc 2021 benchmark suites, developed by the High Performance Group, are a bold attempt to provide a fair and objective benchmarking tool designed for stateof-the-art HPC systems. With the support of multiple host and accelerator programming models, the suites are portable across both homogeneous and heterogeneous architectures. Different workloads are developed to fit system sizes ranging from a few compute nodes to a few hundred compute nodes. In this work we present our first experiences in performance benchmarking the new SPEChpc2021 suites and evaluate their portability and basic performance characteristics on various popular and emerging HPC architectures, including x86 CPU, NVIDIA GPU, and AMD GPU. This study provides a first-hand experience of executing the SPEChpc 2021 suites at scale on production HPC systems, discusses real-world use cases, and serves as an initial guideline for using the benchmark suites. 
    more » « less
  4. null (Ed.)
    Dense linear algebra (DLA) has historically been in the vanguard of software that must be adapted first to hardware changes. This is because DLA is both critical to the accuracy and performance of so many different types of applications, and because they have proved to be outstanding vehicles for finding and implementing solutions to the problems that novel architectures pose. Therefore, in this paper we investigate the portability of the MAGMA DLA library to the latest AMD GPUs.We use auto tools to convert the CUDA code in MAGMA to the Heterogeneous-Computing Interface for Portability (HIP) language. MAGMA provides LAPACK for GPUs and benchmarks for fundamental DLA routines ranging from BLAS to dense factorizations, linear systems and eigen-problem solvers. We port these routines to HIP and quantify currently achievable performance through the MAGMA benchmarks for the main workload algorithms on MI25 and MI50 AMD GPUs. Comparison with performance roofline models and theoretical expectations are used to identify current limitations and directions for future improvements. 
    more » « less
  5. null (Ed.)
    Abstract: Solid-state ion conduction (SSIC) is a mechanism of ionic current that has garnered increasing attention for applications in all-solid-state batteries and atomic switches. The Ag/S SSIC system in β-Ag S, possessing the highest ionic conductivity of any known material, provides a unique opportunity to better understand the fundamental nature of SSIC. β-Ag S is topographically similar to binary perovskites except that it is cubic, leading to isotropic SSIC exceeding 4 S/cm. The dynamic nature of SSIC makes it difficult to study by observational means, where inherent time-averaging obscures correlations among atomic transit routes.Molecular dynamics (MD) is a tool ideally suited for gaining insight into large atomic systems with subnanosecond time resolutions. However, traditional MD potentials lack a description of bond-breaking/forming reactions, which are an essential aspect of SSIC and related memristic properties. This limitation can be overcome by using a reactive force field (ReaxFF), which enables the simulation of bonding reactions with DFT-level accuracy. In this study, we present a ReaxFF force field for the Ag/S system, optimized for simulating SSIC in β-Ag S. Training data consisted of crystal structures, Bader partial charges, and energies of various Ag/S clusters calculated at the DFT-level. Energies were obtained with Gaussian 16, using the PBEh1PBE hybrid functional with a triple-zeta correlation-consistent basis set. Multiobjective parameter optimization was accomplished with an updated form of the Genetic Algorithm for Reactive Force Fields (GARFfield). The force field was validated with potential energy and ion conductivity calculations, along with relevant structural features. Results were compared with equivalent simulations from other established potentials. This new ReaxFF force field will enable modeling of realistic SSIC configurations for Ag/S-based materials and provides a viable approach for extending ReaxFF to other SSIC systems in the future. This work was supported by the National Science Foundation under grant #2025319. 
    more » « less