PaperlayerView exact revision on arXiv

Source-linked AI summary

Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learning

Weile Jia, Han Wang, Mohan Chen, Denghui Lu, Lin Lin, Roberto Car, Weinan E, Linfeng Zhang

arXiv:2005.00223v3physics.comp-ph

TL;DR

AIMD offers first-principles accuracy but is limited to much smaller systems and shorter times than many atomistic problems require. This paper develops an optimized GPU implementation of Deep Potential Molecular Dynamics and demonstrates more than 100 million atoms with ab initio accuracy at Summit-scale performance. The result establishes nanosecond-per-day simulation at unprecedented system sizes.

  • Problem

    AIMD and conventional DFT-based methods remain too costly for many problems requiring thousands to hundreds of millions of atoms and long time scales.

  • Method

    The paper optimizes DeePMD-kit for heterogeneous GPUs by combining DP's symmetry-preserving neural representation with redesigned neighbor-list data structures and mixed-precision computation.

  • Results

    More than 100 million atoms were simulated with ab initio accuracy, reaching 91/162/275 PFLOPS in double, mixed-single, and mixed-half precision.

  • Takeaways & Limitations

    The work opens the door to simulating unprecedented size and time scales with ab initio accuracy.

Abstract

from arXiv · show

For 35 years, {\it ab initio} molecular dynamics (AIMD) has been the method of choice for modeling complex atomistic phenomena from first principles. However, most AIMD applications are limited by computational cost to systems with thousands of atoms at most. We report that a machine learning-based simulation protocol (Deep Potential Molecular Dynamics), while retaining {\it ab initio} accuracy, can simulate more than 1 nanosecond-long trajectory of over 100 million atoms per day, using a highly optimized code (GPU DeePMD-kit) on the Summit supercomputer. Our code can efficiently scale up to the entire Summit supercomputer, attaining $91$ PFLOPS in double precision ($45.5\%$ of the peak) and {$162$/$275$ PFLOPS in mixed-single/half precision}. The great accomplishment of this work is that it opens the door to simulating unprecedented size and time scales with {\it ab initio} accuracy. It also poses new challenges to the next-generation supercomputer for a better integration of machine learning and physical modeling.

I. JUSTIFICATION FOR PRIZE

AIMD provides first-principles atomic interactions but remains constrained by steep computational scaling, while many realistic problems require thousands to hundreds of millions of atoms and long trajectories. The work reports a >100-million-atom, ab initio-accuracy simulation at substantial Summit performance.

  • The reported record is a molecular dynamics simulation of more than 100 million atoms with ab initio accuracy.
  • AIMD computational cost generally scales cubically with the number of electronic degrees of freedom.Typical desktop-scale AIMD reaches approximately 100 atoms and 10 picoseconds.
  • Realistic problems including chemical reactions, electrochemical cells, nanocrystalline materials, radiation damage, and fracture can require thousands to hundreds of millions of atoms.Some also demand trajectories extending to microseconds or beyond, outside AIMD's usual scope.
  • Empirical force fields scale to very large systems, but their accuracy is often uncertain and suitable models remain difficult for multielement systems and bond formation or cleavage.The passage contrasts their scalability with the accuracy and coverage challenges motivating machine-learning-based molecular dynamics.

B. Deep Potential Molecular Dynamics

Deep Potential Molecular Dynamics uses symmetry-preserving local-environment descriptors and neural networks to predict atomic energies, offering AIMD-comparable accuracy with much greater efficiency. This paper focuses on overcoming the implementation and scaling limits that had prevented very large simulations.

  • The DP model combines deep neural networks, physical symmetry constraints, and training with ab initio data to provide AIMD-comparable accuracy and efficiency close to empirical-force-field molecular dynamics.
  • Deep Potential converts each atom's local environment into symmetry-preserving descriptors and passes them through a fitting network to obtain its atomic energy contribution.The total potential energy is constructed by summing atomic energy contributions.
  • Conventional DeePMD-kit could not fully exploit heterogeneous supercomputers because it was single-node oriented, retained CPU-only customized operators, and used relatively small computational kernels.The implementation was built on LAMMPS and TensorFlow, with TensorFlow providing GPU support for standard operators.
  • Efficient simulation of 100 million atoms with ab initio accuracy had not previously been demonstrated with AIMD or MLMD schemes.The authors identify integration of physics-based modeling, machine learning, and efficient implementation as necessary for routine use.
  • The paper adopts serial DeePMD-kit as a baseline and demonstrates how its performance can be substantially improved on Summit.

V. INNOVATIONS

The optimized DeePMD-kit combines scalable heterogeneous-GPU implementation with algorithmic changes that make DP computations more GPU-efficient. It reaches near-ideal scaling on a 127-million-atom copper system while retaining ab initio accuracy.

  • The optimized code scales almost perfectly to 4,560 Summit nodes for a 127,401,984-atom copper system.It reaches 91 PFLOPS in double precision and 162/275 PFLOPS in mixed-single/mixed-half precision.
  • The contribution combines DP's representation capability with a highly scalable, fine-tuned implementation for heterogeneous GPU architectures.
  • The implementation increases computational granularity by sorting and padding neighbor lists to avoid branching in embedding-matrix computation.
  • Neighbor-list elements are compressed into 64-bit integers to improve GPU optimization of customized TensorFlow operators.
  • Mixed-precision computation assigns single- or half-precision arithmetic to intensive tasks without reducing the accuracy of physical observables.

1) Increasing computational granularity:

The optimized DeePMD-kit increases GPU computational granularity by reorganizing and compressing neighbor-list data while migrating intensive calculations to GPU execution.

  • Sorting neighbors by atomic type and distance, then padding by type, avoids conditional branching in embedding-matrix computation.The layout increases computational granularity while preserving accuracy because descriptors are permutation-invariant.
  • 64-bit integer compression of each neighbor-list element improves GPU memory access compared with the original array-of-structures layout.The compressed representation encodes atomic type, distance, and index in one unsigned integer.
  • All customized TensorFlow operators, including environment, force, and virial calculations, are migrated and optimized on the GPU using fine-grained parallelism.
  • GPU memory is allocated once at initialization and reused throughout the simulation, while non-essential CPU-GPU transfers are eliminated.

3) Mixed-precision computation:

The optimized code uses mixed precision to reduce computational cost while retaining stability and accuracy, alongside operator-level GPU optimizations and kernel fusion.

  • Mixed-precision schemes retain double precision for atomic positions and environment-matrix construction while using lower precision in selected neural-network operations.Tensor Cores perform selected GEMM operations, with precision choices constrained by stability and accuracy considerations.
  • 5.2 × 10^-6 eV energy deviation and 2.5 × 10^-6 eV/Å force RMS deviation are observed for MIX-32 versus double precision.
  • Matrix-matrix multiplication dominates floating-point operations after customized TensorFlow operators are optimized.Bandwidth-intensive CONCAT and SUM operators perform few floating-point operations but remain optimization targets.
  • Merging multiplication with bias addition replaces separate CONCAT and SUM operations with a CUBLAS GEMM call.
  • Combining TANH and TANHGrad in one CUDA kernel saves computational time by reusing the forward-pass TANH result.The optimization trades GPU memory space for execution time.

D. Reducing MPI communication bottlenecks

DeePMD uses LAMMPS spatial decomposition and communication infrastructure, with GPU-aware MPI optimization for distributed simulations of water and copper.

  • DeePMD adopts LAMMPS’s parallelization scheme for EFFs while replacing EFF computation with DP and retaining spatial partitioning and communications.
  • CUDA-aware IBM Spectrum MPI optimizes ghost-region communication because the relevant data resides on the GPU.
  • Water and copper are selected as representative insulating and metallic benchmark systems for measuring optimized DeePMD-kit performance.Water is emphasized as difficult because of weak intermolecular interactions, thermal effects, and nuclear quantum effects.
  • The water benchmark uses a 6 Å cutoff and 144 neighbors, while copper uses an 8 Å cutoff and 512 neighbors.
  • 500-step velocity-Verlet tests use 0.5 fs time steps for water and 1.0 fs for copper, with velocities initialized at 330 K.
  • Weak scaling tests span 42,467,328 to 679,477,248 atoms on up to 4,560 Summit nodes, while copper requires 3.5 times more FLOPs than water at equal atom counts.

B. HPC Platforms and Software Environment

The experiments run on Summit’s heterogeneous CPU-GPU platform, using MPI and GPU binding, with performance assessed through time-to-solution, peak performance, and sustained performance.

  • B. HPC Platforms and Software Environment: Summit contains 4,608 nodes and provides a 200 PFLOPS peak performance platform with POWER9 CPUs and NVIDIA V100 GPUs linked by NVLink.
  • B. HPC Platforms and Software Environment: Six MPI tasks are used per computing node, with three tasks per socket and each task bound to an individual GPU.
  • B. HPC Platforms and Software Environment: FLOP counts are collected with NVIDIA CUDA NVPROF, which measures GPU operations; this is treated as reasonable because intensive calculations run on GPUs.
  • B. HPC Platforms and Software Environment: Time-to-solution measures average wall-clock time per MD step, including MD-loop I/O but excluding setup and MPI initialization or finalization.
  • B. HPC Platforms and Software Environment: Peak performance is total FLOPs divided by MD-loop time, whereas sustained performance divides total FLOPs by total wall-clock time including I/O.

VII. PERFORMANCE RESULTS

The optimized DeePMD-kit accelerates customized and standard TensorFlow operators while preserving the accuracy of mixed-precision models against AIMD.

  • 1) Customized TensorFlow operators: 64.6× overall speedup is achieved for customized TensorFlow operators, yielding a 6.2× reduction in MD loop time.For a 12,288-atom water system, customized operators account for about 85% of baseline MD loop time.
  • 2) Standard TensorFlow operators: 1.3×, 1.7×, and 1.6× acceleration is obtained for MATMUL+SUM, CONCAT+SUM, and TANH+TANHGrad, respectively.The optimizations use GEMM, GEMM, and merged TANH implementations.
  • 3) Mixed precision: MIX-32 is as accurate as double precision, while MIX-16 is slightly less accurate but usually sufficient for physical observables.Energy and force predictions are compared with AIMD for 100 water configurations containing 64 molecules.
  • 3) Mixed precision: MIX-32 and MIX-16 radial distribution functions agree with double precision and AIMD for gOO(r), gOH(r), and gHH(r).These RDFs characterize liquid-water structure at ambient conditions.
  • 3) Mixed precision: 1.7× and 2.6× speedups are obtained by MIX-32 and MIX-16 over double precision, while GPU memory costs fall by 50% and 75%.Combined optimizations make the total speedups over baseline 12.7× for MIX-32 and 19.5× for MIX-16.
  • Profiling: 67%, 62%, and 50% of water GPU time is spent in GEMM for double, MIX-32, and MIX-16, compared with 71%, 65%, and 54% for copper.Copper has larger neighbor counts and 3.5× more FLOPs than water.

B. Scaling

The optimized DeePMD-kit scales to thousands of Summit nodes and systems from millions to hundreds of millions of atoms, with mixed precision improving throughput.

  • 1) Strong Scaling: 87.3%, 71.9%, and 61.9% parallel efficiency is reached for copper when scaling from 570 to 4,560 nodes in double, MIX-32, and MIX-16 precision.Peak performance reaches 78.3, 112.3, and 171.8 PFLOPS, respectively.
  • 1) Strong Scaling: 81.7%, 81%, and 77% efficiency is maintained for water to 640 nodes, decreasing to 38.3%, 24.9%, and 18.7% at 4,560 nodes.The decline is mainly attributed to decreasing data size per GPU; all versions still scale to 4,560 nodes.
  • 1) Strong Scaling: 679 million water atoms and 127 million copper atoms are simulated, more than three orders of magnitude beyond state-of-the-art ab initio MD system sizes.Both systems show perfect weak scaling with the number of nodes.
  • 1) Strong Scaling: 91 PFLOPS is achieved in double precision, with 162 and 275 PFLOPS in MIX-32 and MIX-16 precision for copper.The copper result uses 4,560 nodes and reaches 45.5% of peak in double precision.
  • 1) Strong Scaling: 8.1/4.6/2.7 × 10^-10 second/step/atom is attained for copper in double/MIX-32/MIX-16 precision.A 127-million-atom copper system can complete one nanosecond in 29/16/9.5 hours.
  • 1) Strong Scaling: 3.0/1.7/1.1 × 10^-10 second/step/atom is attained for water in double/MIX-32/MIX-16 precision.A 679-million-atom water system can complete one nanosecond in 112/64/42 hours.

C. Sustained performance

The sustained-performance measurements reduce one-time initialization and model-loading overheads while maintaining high double-precision throughput on Summit.

  • C. Sustained performance: Setup time falls from more than 263 seconds to less than 5 seconds for all tests.The optimization constructs atomic structures across MPI tasks and stages model data through one rank before broadcasting.
  • C. Sustained performance: 90.3 PFLOPS, or 45% of peak, is sustained for 127,401,984 copper atoms over 5,000 MD steps.The run corresponds to a 5-ps simulation on 4,560 Summit nodes.

D. Network size

Network size affects precision-dependent throughput: larger embedding matrices benefit MIX-16, although the 64×128 network already achieves ab initio accuracy.

  • D. Network size: 1.17 EFLOPS is reached by MIX-16 with a 1024 × 2048 embedding-net matrix.Double and MIX-32 FLOPS curves flatten after 256 × 512, while MIX-16 continues increasing.
  • D. Network size: 64×128 is sufficient to achieve ab initio accuracy despite higher FLOPS from larger networks.The tested embedding net is characterized by the largest hidden-layer parameter matrix.
  • Conclusion: The work combines physics-based modeling, machine learning, and efficient implementation to open possibilities in materials science, chemistry, and biology.It also poses challenges for integrating machine learning and physical modeling on next-generation supercomputers.

A. Applications of Optimized DeePMD-kit

Optimized DeePMD-kit extends ab initio-accuracy molecular dynamics to nanocrystalline metals and large-scale water and materials simulations. Its GPU-oriented redesign supports scaling toward systems with billions of atoms while retaining broad applicability across complex materials and molecules.

  • Nanocrystalline copper: Nanocrystalline metals typically contain tens to hundreds of millions of atoms, beyond ab initio methods, while DP addresses key EFF accuracy deficiencies for surface and stacking-fault energies.EFFs can provide strain-stress curves, but their accuracy is limited for certain properties.
  • Nanocrystalline copper: 10,401,218-atom nanocrystalline copper was tensile-deformed by 10% over 50,000 MD steps, revealing deformation-associated stacking faults.The simulation used a 50 × 50 × 50 nm3 cell, 300 K annealing, and a strain rate of 5×10^8 s−1.
  • Broader applications: DP applications can extend beyond copper and water to a wide spectrum of complex materials and molecules because the model combines machine learning with physics.The passage states that DP inherits accuracy from first-principles methods.
  • Scalable implementation: Redesigning DeePMD-kit data layout and MPI communication at the algorithmic level supports efficient execution on many-core architectures.The approach increases task granularity rather than simply offloading computationally intensive tasks.
  • Scalable implementation: MIX-16 performance is mainly bounded by GPU memory bandwidth, while expanded high-bandwidth memory could increase the computationally feasible system size.The passage also projects potential performance benefits on architectures with a lower FLOP/Byte ratio.
  • Scalable implementation: Scaling results show no intrinsic obstacle to running DeePMD-kit on exascale systems for simulations involving billions of atoms.Deep Potential's resilience to numerical noise may also reduce work needed for fault-tolerant treatments.
Loading 2005.00223v3…