Phaseless Auxiliary-Field Quantum Monte Carlo ... - ACS Publications

Subscriber access provided by McMaster University Library

Quantum Electronic Structure

Phaseless Auxiliary-Field Quantum Monte Carlo on Graphical Processing Units James Shee, Evan J. Arthur, Shiwei Zhang, David R. Reichman, and Richard A Friesner J. Chem. Theory Comput., Just Accepted Manuscript • DOI: 10.1021/acs.jctc.8b00342 • Publication Date (Web): 13 Jun 2018 Downloaded from http://pubs.acs.org on June 14, 2018

Just Accepted “Just Accepted” manuscripts have been peer-reviewed and accepted for publication. They are posted online prior to technical editing, formatting for publication and author proofing. The American Chemical Society provides “Just Accepted” as a service to the research community to expedite the dissemination of scientific material as soon as possible after acceptance. “Just Accepted” manuscripts appear in full in PDF format accompanied by an HTML abstract. “Just Accepted” manuscripts have been fully peer reviewed, but should not be considered the official version of record. They are citable by the Digital Object Identifier (DOI®). “Just Accepted” is an optional service offered to authors. Therefore, the “Just Accepted” Web site may not include all articles that will be published in the journal. After a manuscript is technically edited and formatted, it will be removed from the “Just Accepted” Web site and published as an ASAP article. Note that technical editing may introduce minor changes to the manuscript text and/or graphics which could affect content, and all legal disclaimers and ethical guidelines that apply to the journal pertain. ACS cannot be held responsible for errors or consequences arising from the use of information contained in these “Just Accepted” manuscripts.

is published by the American Chemical Society. 1155 Sixteenth Street N.W., Washington, DC 20036 Published by American Chemical Society. Copyright © American Chemical Society. However, no copyright claim is made to original U.S. Government works, or works produced by employees of any Commonwealth realm Crown government in the course of their duties.

Page 1 of 40 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60

Journal of Chemical Theory and Computation

Phaseless Auxiliary-Field Quantum Monte Carlo on Graphical Processing Units James Shee,∗,† Evan J. Arthur,‡ Shiwei Zhang,¶ David R. Reichman,† and Richard A. Friesner† †Department of Chemistry, Columbia University, 3000 Broadway, New York, NY, 10027 ‡Schr¨odinger Inc., 120 West 45th Street, New York, New York 10036, United States ¶Department of Physics, College of William and Mary, Williamsburg, Virginia 23187-8795 E-mail: [email protected]

Abstract We present an implementation of phaseless Auxiliary-Field Quantum Monte Carlo (ph-AFQMC) utilizing graphical processing units (GPUs). The AFQMC method is recast in terms of matrix operations which are spread across thousands of processing cores and are executed in batches using custom Compute Unified Device Architecture kernels and the GPU-optimized cuBLAS matrix library. Algorithmic advances include a batched Sherman-Morrison-Woodbury algorithm to quickly update matrix determinants and inverses, density-fitting of the two-electron integrals, an energy algorithm involving a high-dimensional precomputed tensor, and the use of single-precision floating point arithmetic. These strategies accelerate ph-AFQMC calculations with both single- and multi-determinant trial wavefunctions, though particularly dramatic walltime reductions are achieved for the latter. For typical calculations we find speed-ups of roughly two orders of magnitude using just a single GPU card compared to a single modern CPU core. Furthermore, we achieve near-unity parallel efficiency using 8 GPU

1

ACS Paragon Plus Environment

Journal of Chemical Theory and Computation 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60

cards on a single node, and can reach moderate system sizes via a local memory-slicing approach. We illustrate the robustness of our implementation on hydrogen chains of increasing length, and through the calculation of all-electron ionization potentials of the first-row transition metal atoms. We compare long imaginary-time calculations utilizing a population control algorithm with our previously published correlated sampling approach, and show that the latter improves not only the efficiency but also the accuracy of the computed ionization potentials. Taken together, the GPU implementation combined with correlated sampling provides a compelling computational method that will broaden the application of ph-AFQMC to the description of realistic correlated electronic systems.

Introduction Auxiliary-Field Quantum Monte Carlo (AFQMC) is a computational method capable of predicting ground-state observables of chemical systems with very high accuracy. We refer the reader to Refs. 1–4 for core papers presenting the constrained-path and phaseless variants of AFQMC and review publications, and to Refs. 5–24 for methodological advances and illustrative applications. For finite-sized systems such as molecules in Gaussian basis sets, AFQMC calculations scale with the fourth power of the system size, which compares favorably with traditional wavefunction methods such as second-order Møller-Plesset Perturbation Theory (MP2), 25 Coupled Cluster (CC) approaches, 26 and Complete Active Space methods such as CASSCF 27 and CASPT2. 28 However, the prefactor of a typical AFQMC calculation is relatively large. Recently we have introduced a correlated sampling (CS) approach for quantities involving energy differences which is capable of reducing computational prefactors by approximately an order of magnitude. 21 In this work we present a different but complementary strategy involving hardware optimization on graphical processing units (GPUs) which can drastically reduce the prefactors in calculations of general ground-state properties. GPUs have several distinct advantages over traditional Central Processing Units (CPUs), 2


Page 2 of 40

Page 3 of 40 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60


including the ability to perform efficiently parallelized matrix operations both in serial and in “batches,” and the use of single-precision (sp) floating-point arithmetic with significant gains in computational speed. We refer the reader to Ref. 29 for a lucid exposition of many general properties of GPU hardware. In recent years the use of GPUs has been extended well beyond traditional image visualization tasks into many fields such as machine learning 30 and molecular mechanics. 31 Of particular relevance to our work presented here is the progress in performing electronic structure calculations on GPUs. This hardware has been utilized to efficiently evaluate the integrals required in ab-initio calculations, 32–35 to perform HartreeFock 36,37 (HF) and Density Functional Theory (DFT) calculations, 38–40 and to study model systems such as the Hubbard Model within the dynamical cluster approximation 41 and the Ising model. 42,43 In addition there have been recent GPU implementations of MP2, 44–47 CC methods, 48 TDDFT, 49 Configuration Interaction (CI), 50,51 and CASSCF approaches. 52,53 Efficient algorithms to compute energy gradients 54,55 and tensor contractions 56 have also been developed. With respect to Quantum Monte Carlo (QMC) methods, GPU implementations have been formulated primarily for real-space approaches. For example, Diffusion Monte Carlo (DMC) with sp arithmetic has been accelerated by a factor of ∼6x on a GPU versus a quadcore CPU. 57 A recent study employing a multi-GPU implementation has reported speed-ups of a factor of 10-15x relative to a quad-core CPU for Variational and Diffusion MC for real materials. 58 Very recently an open-source QMC suite, QMCPACK, 59 has released scalable implementations of real-space methods including Variational, Diffusion, and Reptation MC. An implementation of auxiliary-field QMC is mentioned in Ref. 59, although data illustrating its efficiency and accuracy is not yet available. In this paper we detail our GPU implementation of the phaseless variant of auxiliary-field QMC (ph-AFQMC), and illustrate its performance and accuracy via calculations of the total energies of linear chains of hydrogen atoms and the all-electron ionization potentials (IPs) of the first-row transition metal (TM) atoms. We explicitly compare our GPU wall-times

3



Page 4 of 40

with CPU timings from a code of equivalent algorithmic sophistication. Speed-ups from the GPU port of two orders of magnitude are seen in large systems, with the potential for even greater reductions of the scaling prefactor depending on the system-size. The robustness and accuracy of our implementation are shown by comparing our calculated values to either exact numerical techniques or experiment. Our paper is organized as follows: In Sec. II we provide a concise review of AFQMC and the phaseless constraint. In Sec. III we detail our GPU implementation and highlight significant algorithmic additions. In Sec. IV and V we present timing and accuracy results for the hydrogen chains and TM IPs, respectively, and comment on the advantages of the correlated sampling approach. In Sec. VI we conclude with a summary of our results and a discussion of future work.

Brief Overview of AFQMC Detailed expositions of AFQMC and the phaseless constraint can be found elsewhere. 2,4,21 In this section we highlight only aspects that are directly relevant to this work. The ground state of a many-body wavefunction |Φi can be obtained via imaginary-time ˆ

propagation, i.e. limN →∞ (e∆τ (E0 −H) )N |Φi = |Φ0 i, if hΦ0 |Φi 6= 0. The general electronic Hamiltonian is ˆ = H

M X ij

M

Tij

X σ

c†iσ cjσ

X † † 1X Vijkl ciσ cjτ clτ ckσ , + 2 ijkl σ,τ

(1)

where M is the size of an orthonormal one-particle basis, σ and τ denote electron spin, and c†iσ and ciσ are the second-quantized fermionic creation and annihilation operators. The matrix elements, Vijkl , can be expressed as a sum over products of three-index quantities via Cholesky decomposition 13 (CD) or Density-Fitting 60 (DF) procedures, allowing Eq. (1) to be written as the sum of all one-body operators plus a two-body operator of the form ˆ 2 = − 1 P vˆα2 . H α 2 After utilizing a Trotter-Suzuki decomposition 61 and the Hubbard-Stratonovich trans4


Page 5 of 40 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60


formation written in the form 62,63

e

1 2 ∆τ vˆα 2

Z

∞

=

dxα

−∞

1 2 e− 2 xα √∆τ xα vˆα √ , e 2π

(2)

the imaginary-time propagator can be expressed as an exponential of a one-body operator integrated over auxiliary-fields

e

ˆ −∆τ H

Z =

ˆ dxP (x)B(x).

(3)

This integral is evaluated using a Monte Carlo technique in which the wavefunction is represented as a weighted sum of Slater determinants. A theorem due to Thouless 64 enables the propagation of the Slater determinants along “paths” defined by a set of auxiliary-fields, which in practice can be accomplished via matrix-matrix multiplications. Each auxiliaryfield is shifted by a force-bias chosen to minimize the fluctuations in a walker’s weight. With this choice the weights are updated after each propagation step by a factor proportional to e−∆τ EL , where the local energy is defined as EL =

ˆ hψT |H|φi , hψT |φi

and |ψT i is a trial wavefunction,

typically taken as a single Slater determinant or a linear combination of determinants. For the Coulomb interaction, the propagator in Eq. (3) has an imaginary component which rotates the walker orbitals and weights in the complex plane, leading inevitably to divergences and an exponentially decaying signal to noise ratio. The phaseless approximation can be used to constrain the weights to be real and positive. This involves taking the realpart of the local energy, and multiplying the resulting weights by max{0, cos(∆θ)}, where τ +∆τ

i T |φ }. The latter breaks the rotational symmetry of the random walk ∆θ = Im{ln hψhψ τ T |φ i

in the complex plane by choosing a unique gauge for |Φ0 i (and eliminates walkers whose weights have undergone phase rotations in excess of ± π2 ). We note that the severity of the phaseless constraint can be systematically reduced with the use of trial wavefunctions which even more closely represent the true ground-state.

5



GPU Implementation In contrast to traditional computing paradigms which utilize CPUs to execute all computing tasks, we employ a strategy in which CPUs offload a majority of the computational effort to one or more GPU cards. A typical GPU device has 4-12 GB of memory which is separate from that accessible by the CPU; therefore, data is usually allocated on both the host and device, and inter-communication between these different memory spaces requires the explicit copying of data back and forth. In this work great care is taken to minimize such transfers, and we create custom memory structures that organize memory addresses and facilitate switching between sp and double-precision (dp). A flowchart outlining the ph-AFQMC algorithm with 3 CPU/GPU pairs is presented in ˆ overlap Fig. 1. First the root CPU reads in relevant quantities such as matrix elements of H, integrals, and the initial and trial wavefunctions, and then completes a preliminary setup which includes transformations to an orthonormal basis, walker and operator initializations. These quantities are sent to all devices, after which the tensor (or slices of it) used in the energy evaluation is precomputed directly on the GPUs. Throughout the entire sequence of functions involved in propagating a walker, all operations are performed on the devices, i.e. without any data transfers or operations involving the CPUs. We utilize NVIDIA’s Compute Unified Device Architecture (CUDA) Basic Linear Algebra Subprograms (cuBLAS) library to execute, e.g., the matrix multiplications that propagate walker determinants by a one-body operator, and have supplemented this library with custom C/CUDA functions which can be classified roughly into two types. Those in the first constitute a matrix library of kernels which carry out, most notably, element-wise matrix additions and matrix sums (i.e. matrix → scalar). These are used frequently to compute the trace of a matrix product, which is utilized in the calculation of expectation values such as the force-bias and local energy. In addition to such library-type functions, we also wrote GPU kernels to sample auxiliary-fields, compute the force-bias, assemble and exponentiate one-body operators, carry out the Sherman-Morrison-Woodbury (SMW) updates, orthonormalize the orbitals of 6


Page 6 of 40

Page 7 of 40 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60


walker determinants, and measure the local energy. With the exception of the SMW and energy measurement functions, which we will subsequently detail, the GPU port of the above functions did not involve notable algorithmic improvements over our CPU implementation.

Figure 1: Flowchart of our AFQMC implementation, for 9 walkers (k = 0, . . . , 8) and 3 CPU-GPU pairs (CPU/GPU 0, 1, 2). If the size of the precomputed tensor used in the energy evaluation exceeds the memory capacity of each device, each GPU precomputes and stores only a slice of the tensor, and the energy of a walker is computed by circulating the walker to all other GPUs and tallying the partial energies obtained from the resident slices.

Once the code enters the loops shown in Fig. 1, data need only be transferred from the current device when the energy is measured, which for our typical time-step choice of ∆τ = 0.005 Ha−1 happens once in 20 propagation steps. An explanation of the slicing variant of the energy algorithm will be presented later in Sec. IV. For now we simply wish to emphasize that throughout the majority of an AFQMC calculation data does not need to leave the devices. This is in large part why the current implementation leads to 7



Page 8 of 40

such pronounced speed-ups compared to our initial attempts to simply offload the matrix multiplications. A relatively new addition to the cuBLAS library are so-called “batched” functions which perform many smaller operations simultaneously, e.g. a set of matrix-matrix multiplies or lower-upper (LU) decompositions. These batched functions are well-suited for operations that are individually too small to parallelize effectively across thousands of cores. We utilize this feature heavily in our implementation of SMW updates to quickly compute equal-time Green’s functions when multi-determinant trial functions are used. We note that previous Diffusion MC studies have utilized similar SMW updates, 58,65,66 and Ref. 67 has also presented fast updates and memory-saving techniques for multi-determinant CI trial wavefunctions in AFQMC. Given a reference matrix A, the following formulas are used to compute the determinants and inverses of matrices which differ from A by one or more row or column: det(A + Ui ViT ) = det(I + ViT A−1 Ui )det(A) (A +

Ui ViT )−1

=A

−1

−1

− A Ui (I +

(4)

ViT A−1 Ui )−1 ViT A−1 .

In the context of ph-AFQMC, suppose we use a multi-determinant trial wavefunction, P |ΨT i = i=0 ci |ψT,i i, where hψT,i |ψT,j i = δij . Then, for the k th walker determinant |φk i, A = [ψT,i=0 ]† [φk ], where the square brackets denote a matrix representation, Ui , Vi are of dimension Nσ x Ei , where Nσ is either the number of spin-up or spin-down electrons, and Ei is the number of excitations required to form the ith configuration of the multi-determinant expansion from the reference configuration. The determinant and inverse of the reference matrix corresponding to zero excitations (i = 0) is computed first for spin-up and spin-down configurations, followed by batched SMW updates for all i 6= 0. Sub-cubic scaling with respect to particle number is achieved since Ei

Phaseless Auxiliary-Field Quantum Monte Carlo ... - ACS Publications

Recommend Documents