PaperlayerView exact revision on arXiv

Source-linked AI summary

86 PFLOPS Deep Potential Molecular Dynamics simulation of 100 million atoms with ab initio accuracy

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

arXiv:2004.11658v3physics.comp-phcs.CE

TL;DR

Large-scale ab initio molecular dynamics is constrained by the cost of first-principles calculations and by DeePMD-kit’s sub-optimal inference implementation. The paper reimplements and optimizes DeePMD-kit for GPU-accelerated heterogeneous supercomputers, achieving Summit-scale simulations with reported speedups and high peak performance.

  • Problem

    Density-functional-theory complexity limits the spatial and temporal scales of ab initio molecular dynamics, while DeePMD-kit inference remains sub-optimally implemented for large-scale runs.

  • Method

    The authors develop GPU-adapted algorithms and reimplement DeePMD-kit on Summit, optimizing data distribution and customized and standard TensorFlow operators.

  • Results

    86.2 PFLOPS, or 43% of Summit’s peak performance, is achieved while scaling to 99% of the machine; the GPU code is 7 times faster than CPU at equal power consumption.

  • Takeaways & Limitations

    The implementation enables nanosecond-timescale molecular dynamics with ab initio accuracy for systems reaching 113 million atoms on Summit.

  • Takeaways & Limitations

    GPU memory limits system size, while GPU and network latency remain bottlenecks for extreme-scale execution and future improvement.

Abstract

from arXiv · show

We present the GPU version of DeePMD-kit, which, upon training a deep neural network model using ab initio data, can drive extremely large-scale molecular dynamics (MD) simulation with ab initio accuracy. Our tests show that the GPU version is 7 times faster than the CPU version with the same power consumption. The code can scale up to the entire Summit supercomputer. For a copper system of 113, 246, 208 atoms, the code can perform one nanosecond MD simulation per day, reaching a peak performance of 86 PFLOPS (43% of the peak). Such unprecedented ability to perform MD simulation with ab initio accuracy opens up the possibility of studying many important issues in materials and molecules, such as heterogeneous catalysis, electrochemical cells, irradiation damage, crack propagation, and biochemical reactions.

NEW VERSION PROGRAM SUMMARY

DeePMD-kit is updated as a parallelized and optimized C++/Python/CUDA package for modern high-performance computers.

  • The new DeePMD-kit version parallelizes and optimizes the package for modern high-performance computers.It is distributed under the LGPL and uses C++, Python, and CUDA.

1. Introduction

The paper addresses the limited spatial and temporal scales of ab initio molecular dynamics by adapting and optimizing DeePMD-kit for heterogeneous supercomputers. It evaluates Summit scaling, GPU efficiency, bottlenecks, and speedups relative to the CPU version.

  • Motivation: Ab initio molecular dynamics provides atomic-level accuracy, but density-functional-theory complexity limits accessible spatial and temporal scales.Large practical problems often require thousands to millions of atoms or more, where empirical force fields are commonly used.
  • Motivation: Deep Potential models can approach ab initio molecular-dynamics accuracy while reducing computational cost by several orders of magnitude.The approach combines a smooth symmetry-preserving embedding network with adaptive data generation through DP-GEN.
  • Problem: DeePMD-kit performance is limited by sub-optimal implementation, especially during model inference of energies and forces within MD runs.Extensive optimization is required to extend accurate simulations to larger systems and longer timescales.
  • Problem: Summit exposes implementation challenges including CPU-only environment-matrix construction, limited multi-GPU assignment, and low efficiency for relatively small DP neural networks.These issues prevent the original code from fully using Summit’s heterogeneous architecture.
  • Research questions: The paper asks how to parallelize DeePMD-kit on heterogeneous supercomputers, improve customized and standard TensorFlow operators, and identify scaling bottlenecks.It also examines system-size and time-to-solution limits and GPU-versus-CPU speedups under matched nodes or power consumption.
  • Contributions: The work contributes GPU-adapted algorithms, optimized TensorFlow operators, and a heterogeneous Summit implementation.Its reported contributions include inherited LAMMPS data distribution and parallelization, 43% peak performance at 86 PFLOPS, and latency analysis.
  • Results: Weak scaling reaches the entire Summit supercomputer for a 113-million-atom copper system, while strong scaling reaches 110 MD steps per second for a 4-million-molecule water system.The GPU implementation is reported as 39 times faster than CPU at equal node count and 7 times faster at equal power consumption.

2. The Deep Potential model

The Deep Potential model represents the potential energy surface as a sum of local atomic contributions, using symmetry-preserving descriptors built from cutoff-defined environments and neural networks. DeePMD-kit distributes these calculations spatially and evaluates energies and forces through a GPU-compatible workflow.

  • The Deep Potential model: The DP model expresses total energy as a sum of atomic contributions, with each contribution depending only on an atom’s local environment within cutoff radius rc.The environment comprises relative neighbor positions satisfying |r_ij| ≤ rc.
  • The Deep Potential model: Each local environment is mapped by an embedding net to a symmetry-preserving descriptor, then by a fitting network to the atomic energy E_i.The descriptor is required to preserve translational, rotational, and permutational symmetries.
  • The Deep Potential model: The environment matrix encodes each neighbor with a gated four-dimensional vector, while missing neighbor entries are padded with zeros.The gating function smoothly decays to zero at the cutoff radius, ensuring a smooth environment representation.
  • The Deep Potential model: The embedding net maps each neighbor scalar s(r_ij) to an M-dimensional vector, and the fitting network uses fully connected layers with skip connections and tanh activations.The embedding network’s hidden layers expand according to s_k = 2s_{k−1}, with final output size M.
  • The Deep Potential model: For multicomponent systems, separate fitting networks represent different central-atom species, while neighbor species are encoded through species-specific embedding nets.The species information is incorporated into the descriptor through embedding nets for possible neighboring-atom species combinations.
  • The Deep Potential model: Forces are obtained as negative gradients of total energy, while model evaluation costs O(N) for fitting and O(N × N_m) for embedding.Because N_m is typically 100–1000, embedding-net evaluation is roughly two to three orders of magnitude more expensive than fitting-net evaluation.

3. Implementation

DeePMD-kit integrates DP with LAMMPS and TensorFlow, then accelerates GPU execution through customized operators, optimized neighbor-list processing, and embedding-net transformations.

  • 3.1. Parallelization: DeePMD-kit replaces LAMMPS’s short-range empirical force field with energies and forces derived from the Deep Potential model.
  • 3.1. Parallelization: Each MD step updates neighbor lists, computes environment matrices, evaluates atomic energies and force increments, and updates atoms in each sub-region.
  • 3.2.2. Customized TensorFlow operators: The original Environment operator dominates post-TensorFlow GPU execution because it supports only CPUs and includes neighbor-list formatting and environment-matrix computation.
  • 3.2.2. Customized TensorFlow operators: 141 times faster than a single CPU core and 12.25 times faster than 14 CPU cores, neighbor-list formatting reaches these speedups after all GPU optimizations.
  • 3.2.3. Optimization of the embedding net: 55% reduction in CONCAT and SUM wall-clock time follows optimization that removes the separate SUM operation.
  • 3.2.3. Optimization of the embedding net: 1.18 overall speedup is achieved after the reported optimizations, while matrix-matrix multiplication rises from 30% to 61% of benchmark execution time.
  • 3.2.4. GPU memory accommodation: 4.5 copies of the embedding matrix are required, making GPU memory depend on atoms per GPU, maximum neighbors, and embedding-net width.

4. The physical system

The study benchmarks GPU DeePMD-kit on water and copper, selecting water for condensed-phase complexity and copper as a relatively simple metallic benchmark.

  • Water is tested despite its simple molecular structure because liquid-water behavior reflects hydrogen bonding, van der Waals dispersion, thermal effects, and nuclear quantum effects.
  • Copper is used as an important, relatively simple metallic benchmark system.
  • The water and copper models use maximal neighbor counts of 138 and 500, respectively, with accuracy treated as reasonably assured from prior benchmarks and theoretical studies.
  • The water strong-scaling system contains 12,582,912 atoms, while copper weak scaling uses 4,139 atoms per GPU card.

5. Machine configuration

The Summit tests use heterogeneous CPU–GPU nodes connected through NVLink, with MPI tasks bound to GPUs to exploit processor, memory, and network affinity.

  • Each Summit node contains two groups, each with one IBM POWER9 socket and three NVIDIA V100 GPUs connected by NVLink at 50 GB/s.
  • Each POWER9 socket has 22 physical CPU cores and 256 GB DDR4 memory, while each V100 GPU has 16 GB high-bandwidth memory.
  • GPU tests use the MPI+CUDA programming model with six MPI tasks per node and three tasks per socket.
  • Each MPI task is bound to an individual GPU to exploit CPU–GPU affinity and the network adapter.

6. Numerical results

The GPU DeePMD-kit substantially outperforms the CPU version and scales from node-level comparisons to the full Summit supercomputer. Large water and copper simulations reach high throughput while preserving consistency with the baseline implementation.

  • CPU–GPU comparison: 39 times faster: the GPU version outperforms the CPU version on 80 Summit nodes, while the speedup decreases to 16 on 4,560 nodes.The comparison uses 480 GPUs versus 3,360 CPU cores at 80 nodes, and 27,360 GPUs versus 191,520 CPU cores at 4,560 nodes.
  • CPU–GPU comparison: GPU baseline performance remains 39 times faster than the CPU baseline with 3,360 CPU cores and exceeds CPU performance on 4,560 nodes.The reduced large-scale speedup reflects stronger CPU strong scaling than GPU strong scaling.
  • Strong scaling: 640 nodes: the GPU implementation scales perfectly for the 12,582,912-atom water system, then continues scaling to all 4,560 Summit nodes.The system reaches 3,840 GPUs at 640 nodes and 455 atoms per GPU at the full machine scale.
  • Strong scaling: 110 MD steps per second: the GPU code simulates the 12,582,912-atom water system for 4.8 ns per day at full Summit scale.The simulation uses 0.5 fs time steps and 4,194,304 water molecules.
  • Weak scaling: 86.2 PFLOPS: the GPU version achieves 43% of Summit’s peak performance for 113,246,208 copper atoms on 4,560 nodes.Each MD step takes 83 milliseconds, enabling one nanosecond of simulation per day.
  • Numerical consistency: Energy, force, and virial predictions remain consistent with the baseline implementation up to 15, 10, and 13 digits, respectively.The optimized GPU version uses the same floating-point precision as the baseline implementation.

7. Performance Analysis

Performance is dominated by atomic energy and force evaluation, especially the Pair component, while GPU and network latency limit efficiency at extreme scale. Workload size, neighbor counts, and communication patterns explain system-dependent performance.

  • Efficiency: 86 PFLOPS: the copper system reaches this double-precision peak on 4,560 Summit nodes, while the water system reaches 13% efficiency at 27,360 GPUs.The copper result corresponds to 43% of Summit’s peak performance for 113,246,208 atoms.
  • Operator analysis: 92% versus 64%: GEMM accounts for these shares of GPU time in copper and water, respectively.Copper has more average neighbors per atom, increasing GEMM’s proportion of computation; its mono-species structure also avoids extra embedding-matrix sorting and slicing.
  • Time breakdown: Less than 1%: the Others category contributes negligibly to total time, including I/O and computations invoked by fixes.The dominant costs therefore lie in force and energy evaluation and communication rather than miscellaneous operations.
  • Time breakdown: More than 93%: the Pair component dominates total time throughout strong scaling for the 12,582,912-atom water system.Pair includes CPU–GPU memory copies and atomic energy and force computations.
  • Communication: Communication time eventually stops scaling at 15,360 and 27,360 GPUs because ghost-region exchange becomes dominated by network latency.Ghost-region data decreases from 613 KB per GPU at 480 GPUs to 73 KB at 27,360 GPUs.
  • Communication: MPI Allreduce latency can become a bottleneck when global properties are collected every time step, but collecting them every 20 steps reduces its share below 1%.The authors identify asynchronous MPI Iallreduce as a possible further way to avoid implicit MPI barriers.

8. Conclusion

GPU DeePMD-kit combines GPU-adapted algorithms and implementation optimizations to enable large-scale molecular dynamics with ab initio accuracy. It scales across Summit and supports applications requiring systems from thousands to hundreds of millions of atoms.

  • 8. Conclusion: 86.2 PFLOPS (43% of the peak) demonstrates DeePMD-kit scaling to 99% of Summit.The implementation reaches this performance in weak-scaling tests.
  • 8. Conclusion: 16−39 times faster at the same number of nodes and 7 times faster at the same power consumption than the CPU version.These comparisons are reported for the GPU code relative to the CPU version.
  • 8. Conclusion: The GPU implementation addresses large molecular systems needed for heterogeneous catalysis, electrochemical cells, irradiation damage, crack propagation, and biochemical reactions.The motivating applications require systems ranging from thousands to hundreds of millions of atoms.
  • 8. Conclusion: Its performance comes from adapting data distribution and optimizing customized and standard TensorFlow operators on GPUs.The authors identify GPU and network latency as key targets for future exascale improvements.
  • 8. Conclusion: The optimization strategies can also be applied to other heterogeneous architectures, including conversion to HIP for AMD-based Frontier.The paper demonstrates these strategies on Summit but describes portability to other systems.
Loading 2004.11658v3…