Hard scaling challenges for ab-initio molecular dynamics capabilities in NWChem: Using 100,000 CPUs per second

Size: px

Start display at page:

Download "Hard scaling challenges for ab-initio molecular dynamics capabilities in NWChem: Using 100,000 CPUs per second"

Justina Wilkerson
5 years ago
Views:

1 Hard scaling challenges for ab-initio molecular dynamics capabilities in NWChem: Using 100,000 CPUs per second Eric J. Bylaska, Kevin Glass, and Doug Baxter (PNNL) Scott B. Baden and John H. Weare (UCSD)

2 Outline Ab-Initio Molecular Dynamics Example simulations Plane-wave methods and their cost Performance of Parallel AIMD simulations Distribution of basis Extensions based on 2d processor geometry New algorithms for exact exchange in AIMD Data replication algorithm Incomplete butterfly algorithm 2

Ab-initio molecular dynamics For problems beyond

Environments Computational Biology Systems with

defects Metallic and semiconductor liquids

chemical transformations involving changes in

in liquid phase Solvation in polar liquids

biology, material science, hydrogen economy,

3 Ab-initio molecular dynamics For problems beyond classical MD simulations Chemistry in Extreme Environments Computational Biology Systems with unusual chemical bonding Clusters, surfaces and defects Metallic and semiconductor liquids Diffusion of impurities and defects Complex chemical transformations involving changes in electronic structure Band gap of semiconductors in liquid phase Solvation in polar liquids Chemical reactions Critical to computational biology, material science, hydrogen economy, environmental science,. Hydrogen Economy Polymers Environmental Science Material Science

0, the structures are less clear. Many species, such as AlOH (aq) 2+, exist and are traditionally assumed to be hexa-coordinate.

4 Computational geochemistry Trivalent aluminum ions 1 st coordination shell of Al 3+ is an unsolved problem. ph < 3.0, exists almost entirely as the octahedral Al(H 2 O) 6 3+ ion (ph > 7), has a tetrahedral Al(OH) 4 structure. In the biochemically and geochemically critical ph range of 4.3 to 7.0, the structures are less clear. Many species, such as AlOH (aq) 2+, exist and are traditionally assumed to be hexa-coordinate. The kinetics of proton- and water exchange on aqueous Al 3+, coupled with AIMD simulations, support a five-coordinate Al(H 2 O) 4 OH 2+ ion as the predominant form of AlOH (aq) 2+ under ambient conditions. This result contrasts Al 3+ with other trivalent metal aqua ions, for which there is no evidence for stable penta-coordinate hydrolysis products. Thomas W. Swaddle, Jörgen Rosenqvist, Ping Yu, Eric Bylaska, Brian L. Philiips, and William H. Casey (2005) Kinetic evidence for fivecoordination in AlOHaq2+ ion, Science, vol. 308, pages

5 Computational actinide chemistry Uranyl (UO 2 2+ ) ions Experimental and simulated extended x-ray absorption fine structure spectra of UO 2 2+ in water are in perfect agreement. Proper description of dynamics needed Static ab initio misses H 2 O in the apical regions For MD, UO water interaction difficult to parameterize Patrick Nichols, Eric J. Bylaska, G. Schenter, and Wibe A. de Jong (2008), Structured and Unstructured Solvent shells of the UO 2 2+ Ion, Journal of Chemical Physics, vol. 128(12):

6 Conventional versus ab-initio molecular dynamics Conventional molecular dynamics Empirical, usually two-body potentials Difficult to treat reactions Ab-initio molecular dynamics Potential obtained from Schrodinger equation, includes allbody and electronic behavior Empirical potentials parameterized for a small range of PT Equally applicable under all conditions 10 5 particles no problem 600 particles with significant dynamics 10 3 ps no problem 10 s of ps difficult 6

Basic features of ab-initio molecular dynamics H i r + ( 1 )V x DFT Equations H i = i i ()= 1 2 2 + V

7 Basic features of ab-initio molecular dynamics H i r + ( 1 )V x DFT Equations H i = i i ()= V l ()+ r V ˆ NL + V H []r () []r ()+ V c []r () i r () K ij () r j () r j CP dynamics: Ion and wavefunction motion coupled. Ground state energy μ=0 μ i = H i M I R I = F I N e i=1 ij j H F I = i i R I i=1 Want to do this in ~1second per step N e Plane-wave basis sets, pseudopotentials are used to solve PDE 7

8 Why do we need a second per step? Current ab-initio molecular dynamics simulations for 10 to 100 picoseconds can take several months to complete The step length in ab initio molecular dynamics simulation is on the order of fs/step 20 ps of simulation time 200,000 steps At 1 second per step 2-3 days At 10 seconds per step 23 days At 30 seconds per step 70 days 1 ns of simulation time 10,000,000 steps at 1 second per step 115 days of computing time At 10 seconds per step 3 years At 30 seconds per step 9 years At 0.1 seconds per step 11.5 days

2e12 Remember we want to do this 100,000+ times For hybrid-dft: A

9 Cost of AIMD step Na=500, Ne=500, Ng=256^3 Ne*Ng=8.4e9 Ne*Ng*Log(Ng)=2.0e11 Na*Ne*Ng=4.2e12, Ne*Ne*Ng=4.2e12 Remember we want to do this 100,000+ times For hybrid-dft: A day of computation on the PNNL Chinook system $16K/ 9 Hybrid-DFT: Ne*(Ne+1)*Ng*Log(Ng) = 1.0e14

10 Parallel strategies for plane-wave programs (w/o exact exchange) Ne N g basis N e molecular orbitals Three types of data distributions Distribute molecular orbitals (one eigenvector per cpu) < i j > requires communication Distribute basis(i.e. FFT grid) FFT requires communication proc=1 proc=2 proc=3 proc=... Minimal load balancing slab decomposition

Column decomposition Real Real and and k-space representation of of each each wavefunction is is distributed over over all all processors Requires parallel FFT FFT Column decomposition of of each

11 Column decomposition Real Real and and k-space representation of of each each wavefunction is is distributed over over all all processors Requires parallel FFT FFT Column decomposition of of each each wavefunction across across parallel processor 11 Processor 1 Processor 5 Processor 2 Processor Processor 6 Processor 3 Processor Processor 1 Processor 7 Processor 4 Processor Processor Processor 8 Processor Processor Processor 1 Processor 5Processor N Processor 2 Processor Processor 6 Processor 3 Processor Processor Processor 7 Processor 4 Processor Processor Processor 8 Processor Processor 1Processor L N e Processor 5Processor N Processor 2 Processor Processor 6 Processor 3 Processor Processor Processor 7 Processor 4 Processor Processor Processor 8 Processor Processor Processor N

Parallel timings for UO 2 2+ +122H 2 O - distributing the basis Bottlenecks Parallel 3d FFT stalls by 100

12 Parallel timings for UO H 2 O - distributing the basis Bottlenecks Parallel 3d FFT stalls by 100 cpus, (stalls by N g 1/3 ) Non-local pseudopotential stalls by 1000 cpus because of a broadcast across all cpus 12

13 2d processor geometry N orbitals N orbitals N grid =N 1 xn 2 xn 3 N grid =N 1 xn 2 xn 3 Old parallel distribution New parallel distribution (Gygi et al. Supercomputing 06) 13

14 2d processor geometry timings 14

Why do we need higher levels of solution to the electronic Schrödinger equation? Localized excess electron calculated at the hybrid DFT level with full surface hydration.

15 Why do we need higher levels of solution to the electronic Schrödinger equation? Localized excess electron calculated at the hybrid DFT level with full surface hydration. The charge is localized at an Fe 3+ Fe 2+ center as a result of exact exchange. DFT often fails for Elements with f and d valence electrons Reaction barriers Localized states Such failures often corrected by including exact exchange DFT & all independent particle methods do not describe long range forces e.g. interaction of water with (soft) surfaces and species Can result in too low of solvent coordination

16 Calculating exact exchange Algorithm 1: Serial algorithm for calculating exact exchange in a planewave basis Input: - N g x N e array Output: K - N g x N e array for m=1,n e for n=1,m (:) FFT_rc( (:,m)*. (:,n))) V(:) FFT_cr(f cutoff (:)*. (:)) K (:,m) -= V(:)*. (:,n); if m<>n K (:,n) -= V(:)*. (:,m) end for end for For Ne= (500+1) = 250,500 three-dimension FFTs per step For Ng=200x200x200 calculation can readily be run on leadership class machines (e.g cores, 4 to each of 9660 processing nodes) 16

Computation ~Npj^2, Message passing ~Npj Use replicated space to

17 Simple parallel algorithm for exact exchange Log Npj Multicast N e N e N e N e N e N e to big H big to The algorithm is scalable because: Computation ~Npj^2, Message passing ~Npj Use replicated space to compute exchange using 3d parallel FFTs along columns (load balanced) 17

Exact exchange timings 80 atom cell of hematite Exchange term is dominant Previous algorithm simple to implement Requires lots of workspace

18 Exact exchange timings 80 atom cell of hematite Exchange term is dominant Previous algorithm simple to implement Requires lots of workspace Sends approximately twice as much data as necessary Basic algorithm works fairly well - stalls at 7 seconds by 4096 cpus Can be improved 18

19 Combinations needed to compute exchange P0 P1 P2 P3 P4 P5 P6 P Data replicated unique pairs 19

Combinations needed to compute exchange P0 P1 P2 P3 P4 P5 P6 P7 01 02 03 04 05 06 07 10 12 13 14 15 16 17 20 21 23 24 25 26 27 30 31 32 34 35 36 37 40

20 Combinations needed to compute exchange P0 P1 P2 P3 P4 P5 P6 P N 2 Unique pairs need to be computed Pairs contained in two groups 20

Combinations needed to compute exchange P0 P1 P2 P3 P4 P5 P6 P7 01 02 03 04 05 06 07 12 13 14 15 16 17 23

21 Combinations needed to compute exchange P0 P1 P2 P3 P4 P5 P6 P Load balance by transposing the lower corner 21

22 Combinations needed to compute exchange P0 P1 P2 P3 P4 P5 P6 P Share the computations along the diagonal Pairs contained in two groups 22

23 Incomplete butterfly ( ) Only ( Floor (log ) 2 N ) N blocks are propagated in the last butterfly step 23

24 Incomplete butterfly timings 80 (160) atom cell of hematite 24

Incomplete butterfly algorithm timings Fitting performance model from 80 atom hematite PERFORMANCE MODEL t exchange = N e ( Ne + 1) t Npj + 4 (1/ 2) + 4 Log 2 N g 1 fft N Npi ( Npj) e Using this

25 Incomplete butterfly algorithm timings Fitting performance model from 80 atom hematite PERFORMANCE MODEL t exchange = N e ( Ne + 1) t Npj + 4 (1/ 2) + 4 Log 2 N g 1 fft N Npi ( Npj) e Using this formula we find that the maximum number of CPUs that can be gainfully used in the 80 and 160 atom Fe 2 O 3 Hybrid-DFT calculations to be 14,977 CPUs and 48,032 CPUs respectively and we plan to validate these results in the future. 25

26 Conclusions An incomplete butterfly implementation of exact exchange has been implemented that makes judicious use of data replication Overall performances of our AIMD are reasonable Serial and threaded optimization on-going Communications overheads are still and issue we are currently exploring latency hiding techniques via run-time substrates that implement a dataflow execution model (Taragon) Significant progress has been made in terms of accuracy, efficiency, and scalability of AIMD methods in recent years. However the algorithms and implementations of these methods need to be constantly upgraded to capture the performance of emerging computers Acknowledgements ASCR Multiscale Mathematics program, BES Geosciences program, BES Heavy element program, and EMSL of the U.S. Department of Energy, Office of Science ~ DE- AC06-76RLO EMSL operations are supported by the DOE's Office of Biological and Environmental Research. We wish to thank the Scientific Computing Staff, Office of Energy Research, and the U. S. Department of Energy for a grant of computer time at the National Energy Research Scientific Computing Center (Berkeley, CA). Some of the calculations were performed on the Chinook computing systems at the Molecular Science Computing Facility in the William R. Wiley Environmental Molecular Sciences Laboratory (EMSL) at PNNL. 26

27 Why Wannier? In general, a complete set of Bloch orbitals (complete sampling of the Brillouin zone) is needed to evaluate the complete set of Wannier orbitals. However, for non-metallic systems with sufficiently large unit cells, it turns out that one can obtain at least one of the complete sets of Wannier orbitals contained in the manifold of the sets of Wannier orbitals from having only the -point Bloch functions, of all the occupied bands. The strategy for doing this is quite simple. The trick is to apply a Marzari-Vanderbilt localization unitary transformation (which is the counterpart of the Foster-Boys transformation for finite systems) over n to produce a new set of -point Bloch functions to. These new orbitals, which are maximally localized, are extremely localized within each cell for non-metallic systems with a sufficiently large unit cells. If this is the case, then can be represented as a sum of piece-wise localized functions by Since this is just the defn of Wannier with k=0, complete set of Wannier orbitals contained in the manifold of the sets of Wannier functions is formed. With this new set of localized orbitals the exchange term per unit cell can be written as 27

Calculating E Exch with Wannier Functions Wannier functions constructed from single -point Kohn- Sham Bloch states, have infinitely many artificial periodic replicas.

28 Calculating E Exch with Wannier Functions Wannier functions constructed from single -point Kohn- Sham Bloch states, have infinitely many artificial periodic replicas. To use the plane-wave method to calculate E Exch, one has to limit integration to a single unit cell. N N * Exch nm() mn( ) n m ( ) E = drdr A r A r f r r To handle extended systems, we replace the Coulomb interaction by a screened Coulomb interaction: f () r N N ( 1 exp{ ( r / R) }) 1 = r This kernel is nearly equal to 1/r for r<r and rapidly decays to 0 for r>r. 28

29 Ab Initio Molecular Dynamics For Problems Beyond Classical MD Simulations Systems with unusual chemical bonding Clusters, surfaces and defects Metallic and semiconductor liquids Diffusion of impurities and defects Phenomena involving changes in electronic structure Band gap of semiconductors in liquid phase Solvation in polar liquids Chemical reactions 29

Incomplete Butterfly 0 1 2 3 4 5 6 7 07 10 21 32 43 54 65 76 0765 1076 2107 3210 4321 5432 6543 7654 07654 10765 21076

30 Incomplete Butterfly ( ) Only ( Floor (log ) 2 N ) N blocks are propagated in the last butterfly step 30

31 Parallel Strategy: One dimensional slab decomposition for N e orbitals Processor N Processor N Processor N Processor N Processor Processor Processor Processor Processor 2 Processor 2 Processor 2 Processor 2 Processor 1 Processor 1 Processor 1 Processor L Ne Real Real and and k-space representation of of each each wavefunction is is distributed over over all all processors Requires parallel FFT FFT Parallel scaling restricted to to number of of points along a single dimension 31

32 Conventional Versus ab-initio Molecular Dynamics Needs for geochemical/materials applications Accurately capture many-body interactions, reactions, polarization Extrapolates to difficult PT conditions Treat sufficient number of particles to capture chemistry. Thermodynamics? Conventional molecular dynamics Empirical usually twobody potentials, reactions difficult to treat Empirical potentials parameterized for a small range of PT 10 5 particles no problem 1 st principles molecular dynamics Potential from Schrödinger eq. All - body and electronic behavior included Equally applicable to all conditions 600 particles with significant dynamics Treat long time scales 10 3 ps no problem 10s of ps difficult

33 10 ps AIMD simulation of UO H 2 O (UO H 2 O 22ps) Average UO H 2 O (64 H 2 O) AIMD Exp. r U-O 1.77 (1.77) U-O 0.06 (0.06) 0.04 < OUO (174.1) < HOH 3.2 (3.2) r U-OI 2.46 (2.44) U-OI 0.11 (0.10) st shell tilt 32.5 (32.9) --- #H 2 O 1st 5 (5) r U-OII 4.59 (4.59) U-OII 0.11 (0.12) r OI-OII 2.74 (2.72) < OI-H-OII (163.5) # 2 2 nd 16.2 (14.8) 14 *Conventional molecular dynamics (2 and 33 3-body potential) lead to incorrect results.

New Capabilities Developed AIMD/MM internal MD potential MM potential is

but recent results shown it can model solvated cations very well.

per processor (where the total number of processors, P, can be written as

34 New Capabilities Developed AIMD/MM internal MD potential MM potential is fairly general also interfaced to NWMD (CHARM, ) Useful for equilibration, but recent results shown it can model solvated cations very well. Tera Petascale Parallelization parallel decomposition which the orbitals are distributed in a 2d processor grid. this decomposition only requires (O(log(p 1 )+O(log(p 2 )) communications per processor (where the total number of processors, P, can be written as P=p 1 p 2 ). uses of parallel matrix libraries, and MPI communicators. New algorithm developed for Exact-Exchange (MP2)

35 Molecular Dynamics Loop (1) Compute Forces on atoms, F I (t) for current atomic configuration, R I (t) F I (t) calculate using classical potentials (can do large systems and long simulation times) calculate directly from first principles by solving many-electron Schrödinger equations (can treat very complex chemistry, but simulations times are very long) (2) Update atom positions using Newtons laws R I (t+ t) 2*R I (t) R I (t- t) + t 2 /(M I )*F I (t) 35

Geochemistry 5 coordinated AlOH AIMD simulations of

cluster models with COSMO fail to predict 5-fold

Swaddle, Jörgen Rosenqvist, Ping Yu, Eric Bylaska,

36 Geochemistry 5 coordinated AlOH AIMD simulations of first hydrolysis species of Al 3+ (aq) performed for Casey et al. AlOH(H 2 O) 2+ 5 AlOH(H 2 O) 2+ 4 MD and ab initio cluster models with COSMO fail to predict 5-fold species! Thomas W. Swaddle, Jörgen Rosenqvist, Ping Yu, Eric Bylaska, Brian L. Philiips, and William H. Casey (2005) Kinetic evidence for five-coordination in AlOHaq2+ ion, Science, vol. 308, pages

37 10 ps AIMD simulation of UO H 2 O (UO H 2 O 22ps) Average UO H 2 O (64 H 2 O) AIMD Exp. r U-O 1.77 (1.77) U-O 0.06 (0.06) 0.04 < OUO (174.1) < HOH 3.2 (3.2) r U-OI 2.46 (2.44) U-OI 0.11 (0.10) st shell tilt 32.5 (32.9) --- #H 2 O 1st 5 (5) r U-OII 4.59 (4.59) U-OII 0.11 (0.12) r OI-OII 2.74 (2.72) < OI-H-OII (163.5) # 2 2 nd 16.2 (14.8) 14 37

38 Simple Parallel Algorithm for Exact Exchange Algorithm 2: Parallel algorithm for calculating exact exchange in a plane-wave basis using a two-dimensional processor grid Input: - (N g /Npi) x (N e

38 38 Simple Parallel Algorithm for Exact Exchange Algorithm 2: Parallel algorithm for calculating exact exchange in a plane-wave basis using a two-dimensional processor grid Input: - (N g /Npi) x (N e /Npj) array Output: K -(N g /Npi) x (N e /Npj) array Work Space: (N g /Npi) x N e array K (N g /Npi) x N e array N e = total number of orbitals Np = total number of processors, where Np=Npi*Npj Npi/ Npj = size of column (row) processor group taskid_i/j = rank along the column (row) processor group (:,:), (:,:) = 0 s=taskid_j*(n e /Npj); e = s+(n e /Npj); (:,s:e) = Row_Reduction( ) counter = 0 for m=1,n e for n=1,m if counter==taskid_j then (:) Column_FFT_rc( ( (:,m) *. (:,n)) ) V(:) Column_FFT_cr( (f screened (:) *. (:)) ) K (:,m ) -=V(:)*. (:,n) if m!= n then K (:,n) -=V(:)*. (:,m) end if end if counter = mod(counter + 1,Npj) end for end for Row_Reduce(K ) K =H (:,s:e)

39 Parallel FFT for Column Decomposed Data Need to reduce communication costs Parallel 2D data mapping dimension 2 Dimension 2 Dimension 3 dimension 3 i 39

40 Uranyl EXAFS The EXAFS spectra for the UO H 2 O and UO H 2 O simulations were simulated using MD-EXAFS procedure of Palmer et al.. As input to this procedure an ensemble of UO 2+ 2 (H 2 O) n clusters were generated from every 60 fs of the AIMD simulation in which the 100 nearest atoms to the uranium atom were extracted subject to periodic boundary conditions. Each cluster of the ensemble was then used as input into the FEFF electron multiple scattering code to generate the EXAFS fine structure factor and its windowed transform, i.e. 1 () R = k () k W() k e i kr dk Analysis performed by G. Schenter

Self-Consistent Implementation of Self-Interaction Corrected DFT and of the Exact Exchange Functionals in Plane-Wave DFT

Self-Consistent Implementation of Self-Interaction Corrected DFT and of the Exact Exchange Functionals in Plane-Wave DFT Kiril Tsemekhman (a), Eric Bylaska (b), Hannes Jonsson (a,c) (a) Department of Chemistry,