Implementing a neural network potential for exascale molecular dynamics

1. Implementing a neural network potential
for exascale molecular dynamics
Saaketh Desai, Sam Reeve, Jim Belak
Lawrence Livermore National Laboratory
School of Materials Engineering, Purdue University
Computational Chemistry and Materials Science (CCMS) Summer Institute 2019

2. 2
Molecular dynamics and machine-learned potentials
Obtain updated
positions and velocities
Start with a set of
positions and velocities
Solve F = ma
Verlet
integration
Define atomic interactions
(interatomic potential)
Lennard-Jones
Ø DFT ~ O(n3
e)
Ø Interatomic potential ~ O(natoms)
Ø Can we get quantum mechanical
accuracy with linear scaling?
𝐸𝑖 = 𝜖 𝑐𝐹 𝛼 (
)*+
𝜌 𝛽 𝑟𝑖𝑗 +
1
2
(
)*+
𝜙 𝛼𝛽 𝑟𝑖𝑗
Embedded Atom Model
Spectral Neighbor Analysis Potential

3. 3
Caution: more than blind ab-initio sim (cf: yesterday)
Can ML be used to identify where new data is needed? How do
we decide if a funny feature is real or an artifact?

4. 4
Descriptor-based machine-learned potentials
Descriptor
Symmetry Functions
SOAP
Bispectrum coefficients
ML method
Neural network
Gaussian Process
Linear Regression
Atomic structure
GAP
SNAP
NNP
MTP
qSNAP
Figure from arXiv:1906.08888 [physics.comp-ph]
EAM
MEAM
Ø Can we improve implementation
to break the pareto front?

5. 5
A closer look at the kernel
traditional
𝑉 𝑟+) =
5
678
9:
−
5
678
<
𝐹𝑜𝑟 𝑎𝑡𝑜𝑚 𝑖 = 1: 𝑁
𝐹𝑜𝑟 𝑒𝑎𝑐ℎ 𝑛𝑒𝑖𝑔ℎ𝑏𝑜𝑟 𝑗 𝑜𝑓 𝑎𝑡𝑜𝑚 𝑖
unique
𝐸IJIKLM = (
+
𝐸+
𝑓+
N
= (
) * +
(
O
𝜕𝐸)
𝜕𝐺),O
∗
𝜕𝐺),O
𝜕𝑟+)
N
𝐺O
𝑟+) = 𝑒
TUV 678 T6WV
X
𝑓Y(𝑟+))
𝐹𝑜𝑟 𝑒𝑎𝑐ℎ 𝑑𝑒𝑠𝑐𝑟𝑖𝑝𝑡𝑜𝑟 𝑘
𝐸+ = 𝑁𝑁 (𝐺+)
𝐹𝑜𝑟 𝑎𝑡𝑜𝑚 𝑖 = 1: 𝑁
𝐹𝑜𝑟 𝑎𝑡𝑜𝑚 𝑖 = 1: 𝑁
1
G1k
G1j
G2j
G2k
atom i neighbor j neighbor k
G1(r)
G2(r)
G3(r)
G4(r)
G1(atom i) = G1j + G1k
G2(atom i) = G2j + G2k
2
3
~30
2

6. 6
Other neural network potentials
Weinan E. et. al., Computer Physics Communications 228 (2018) Schütt, Kristof, et al. Advances in Neural
Information Processing Systems. 2017
GAP SNAPHDNNP SchNet DeepMD
2007 2010 2015 2017 2018 2019
SimpleNN
wACSF
HIP-NN, AIMNet
ANI-1Ø Suggested improvements to HDNNPs
Ø Machine – learned descriptors vs
handcrafted descriptors?
SchNet DeepMD

7. 7
Implementation details
Singraber, Andreas, Jörg Behler, and Christoph Dellago, Journal
of chemical theory and computation 15.3 (2019): 1827-1840. https://github.com/CompPhysVienna/n2p2
Ø Evaluate scaling performance on water
Ø Compare with other potentials published in arXiv:1906.08888 [physics.comp-ph]
n2p2 – LAMMPS CabanaMD
Memory based model (large structs) Cabana AoSoAs/ Kokkos views
Minor explicit OpenMP parallelization Kokkos and Cabana parallelization
Ø Easy to read and modify C++ code
Ø Stand-alone package with no dependencies (Tensorflow etc)
Ø Neural network potential closest to traditional interatomic potentials
Why choose n2p2?
How to implement performance portable version of n2p2?

8. 8
Performance gains by re-implementation
8K atoms
36K atoms
121K atoms
saturating performance
memory-bound
(small systems)
Ø Our implementation scales better,
allowing large system sizes
Ø GPU implementation shows significant
improvement
Ø Intel Broadwell node / Nvidia V100 GPU
n2p2 (LAMMPS)
increasing
performance
n2p2 (LAMMPS)
CPU
GPU
CabanaMD
scales well
worker-bound
8K atoms
36K atoms
CabanaMD
121K atoms
562K atoms
1.5M atoms
4.5M atoms
8K atoms
36K atoms
121K atoms
saturating performance
memory-bound
(small systems)
n2p2 (LAMMPS)
increasing
performance

9. 9
Performance across various systems
GPU – H2O
GPU – Ni
Ni (~1.5M)
H2O (~36K)
strong scaling weak scaling
Ni (~40k per thread)
Ø Ni potential shows vastly superior
performance to H2O (CPU and GPU)
Ø Strong scaling shows parallel efficiency
of around 75%
Ø MPI scaling results coming soon…

10. 10
Data layout for better performance
Ø Neural network still a minor part of
computation
Ø Splitting 6 AoSoAs and changing
vector lengths show minor
improvements on the GPU
Ø Hierarchical parallelism shows
promise, subject to further testing
Ni (CPU)01
08
16
32
01
08
16
32
ForcesNeural
network
Symmetry
Functions
01
08
1632
01
08
16 32Ni (GPU)

11. Ø We implemented a neural network potential in CabanaMD
using Kokkos/Cabana constructs
Ø Kokkos/Cabana based implementation of neural network
potential offers significant on-node scalability (~4x speedup and
10x larger systems)
Ø First – ever GPU implementation of neural network potential
shows up to 10x improvement over CPU implementation (single
element performance ~ 1 million atomsteps per second)
Ø Data layout changes offer additional ~10% performance gain
Summary and future work

12. Seeds for Contemplation
12

13. • CALPHAD uses a power series expansion in temperature,
pressure and constitution
• By analogy to ML potentials, can ML be used to construct free
energies?
• Can ML augment CALPHAD and reveal features not possible
from an analytic fit?
• Are there correlations in multi-component systems revealed
using ML? (not possible using analytic models)
• Can ML be used to identify missing data and/or extend the
thermodynamic data used in CALPHAD assessments?
• How does one maintain thermodynamic consistency using ML?
• What about non-equilibrium free energies?
• What about transition / mixed regions (values of phi between 0
and 1)?
13
CALPHAD / PT is the low hanging fruit

14. • Liz Holm use machine vision in 2D, Peter’s students are
doing this in 3D.
• Can ML classify unique microstructures? (yes)
• Can ML be used to quantify the “distance” between
microstructure? How close are two microstructures?
Statistically?
• We want to quantitatively compare simulated
microstructure to experimental observation.
14
Classify Microstructures (CNN)
132 B.L. DeCost, E.A. Holm / Computational Materials Science 110 (2015) 126–133
DeCost and Holm, Comp. Mat. Sci. 110 (2015) 126-133

15. • Billion atom simulations of nucleation /
solidification are becoming routine
(Shibuta, Nature Comm 2017)
• Can ML by used to determine the best
model from simulation or experimental
data?
• Or create an ML model that uses a
know free energy? Generative Model.
15
Model determination (Supervised and GANs)
• Low fruit is determining optimal model parameters (e.g. mobilities)
• Can PFM datasets in a similar ML framework be used to determine
microstructure aware continuum models?
• Nucleation: Can ML be used to detect anomalies in PF simulations?
• Can ML be used to control / steer simulations? Drive design of
experiments (cf. Michael Tonks)?
• Is there an analogy to putting smart sensors into the simulation?