arXiv:2202.01255 · Journal of Chemical Theory and Computation
Large Scale Quantum Chemistry with Tensor Processing Units
Ryan D. Pederson, John Kozlowski, Ruyi Song, Jackson Beall, Martin Ganahl, Markus Hauru, Adam G. M. Lewis, Yi Yao, Shrestha Basu Mallick, Volker Blüm, Guifré Vidal
We demonstrate the use of Googles cloud-based Tensor Processing Units (TPUs) to accelerate and scale up conventional (cubic-scaling) density functional theory (DFT) calculations. Utilizing 512 TPU cores, we accomplish the largest such DFT computation to date, with 247848 orbitals, corresponding to a cluster of 10327 water molecules with 103270 electrons, all treated explicitly. Our work thus paves the way toward accessible and systematic use of conventional DFT, free of any system-specific constraints, at unprecedented scales.
We demonstrate the use of Googles cloud-based Tensor Processing Units (TPUs) to accelerate and scale up conventional (cubic-scaling) density functional theory (DFT) calculations. Utilizing 512 TPU cores, we accomplish the largest such DFT computation to date, with 247848 orbitals, corresponding to a cluster of 10327 water molecules with 103270 electrons, all treated explicitly. Our work thus paves the way toward accessible and systematic use of conventional DFT, free of any system-specific constraints, at unprecedented scales.
Tensor processing units (TPUs) are specialized hardware accelerators developed by Google to support large-scale machine-learning tasks but they can also be leveraged to accelerate and scale other linearalgebra-intensive computations.In this paper, we demonstrate the usage of TPUs for massively parallel classical simulations of quantum many-body dynamics on long time scales.We apply our methods to study the phenomenon of Floquet prethermalization, i.e., exponentially slow heating in quantum spin chains subject to high-frequency periodic driving.We simulate the dynamics of L = 34 qubits for over 10 5 Floquet periods, corresponding to circuits with 4 × 10 6 nearest-neighbor two-qubit gates.The circuits simulated have no additional symmetries and represent a pure-state evolution in the full 2 L -dimensional Hilbert space.This is achieved by distributing the computation over 128 TPU cores.On that size TPU cluster, we find speed-ups in wall-clock run time of 230 times and 15 times when compared to reference CPU and single-graphics-processing-unit (GPU) simulations, respectively, for shorter-time 30-qubit simulations that can be handled by all three platforms.We study the computational cost of the simulations, as a function of both the number of qubits and the number of TPU cores used, up to our maximum capacity of L = 40 qubits, which requires a "full pod" of 2048 TPU cores with tens of terabytes of memory in total.For these simulations, an eight-TPU-core machine is comparable to a single A100 GPU and thus the full TPU pod is comparable to a machine with hundreds of top-of-the-line GPUs.However, the TPU pod is more energy and cost efficient and readily accessible (via Google Cloud), unlike such large many-GPU configurations.We also study the accumulation of numerical error as a function of circuit depth in very deep circuits.Our work demonstrates that TPUs can offer significant advantages for state-of-the-art simulations of quantum many-body dynamics.
Tensor Processing Units (TPUs) were developed by Google exclusively to support large-scale machine learning tasks. TPUs can, however, also be used to accelerate and scale up other computationally demanding tasks. In this paper we repurpose TPUs for the challenging problem of simulating quantum spin systems. Consider a lattice model made of N spin-21 quantum spins, or qubits, with a Hamiltonian H=∑ihi that is a sum of local terms hi and a wavefunction ∣Ψ⟩ consisting of 2N complex amplitudes. We demonstrate the usage of TPUs for both (i) computing the ground state ∣Ψgs⟩ of the Hamiltonian H, and (ii) simulating the time evolution ∣Ψ(t)⟩=e−itH∣Ψ(0)⟩ generated by this Hamiltonian starting from some initial state ∣Ψ(0)⟩. The bottleneck of the above tasks is computing the product H∣Ψ⟩, which can be implemented with remarkable efficiency utilising the native capabilities of TPUs. With a TPU v3 pod, with 2048 cores, we simulate wavefunctions ∣Ψ⟩ of up to N=38 qubits. The dedicated matrix multiplication units (MXUs), the high bandwidth memory (HBM) on each core, and the fast inter-core interconnects (ICIs) together provide performance far beyond the capabilities of general purpose processors.
Electronic structure calculations have been one of the most successful scientific fields in the past 50 years. Density functional theory (DFT) with its nice balance between accuracy and efficiency has now become the standard technique for many materials research. On the other hand, the size of systems that wavefunction-based methods can handle has also been improved thanks to the recent developments in both computation power and more advanced methodologies. However, these methods are still suffering from the curse of dimensionality. The cubic scaling of a typical DFT calculation restricts the handleable system size of DFT to only a few thousand electrons. To this end, tensor-structured techniques provide a route to constructing a reduced-order algorithm and hence are promising to improve the computational efficiency. In this dissertation work, we aim at exploring the applications of the tensor-structured techniques on different aspects of electronic structure calculations. In the first part of this work, we present a tensor-structured algorithm for efficient large-scale DFT calculations by constructing a Tucker tensor basis that is adapted to the Kohn–Sham Hamiltonian and localized in real-space. The proposed approach uses an additive separable approximation to the Kohn–Sham Hamiltonian and an L1 localization technique to generate the 1-D localized functions that constitute the Tucker tensor basis. Numerical results show that the resulting Tucker tensor basis exhibits exponential convergence in the ground-state energy with increasing Tucker rank. Further, the proposed tensor-structured algorithm demonstrated sub-quadratic scaling with system-size for both systems with and without a gap, and involving many thousands of atoms. This reduced-order scaling has also resulted in the proposed approach outperforming plane-wave DFT implementation for systems beyond 2000 electrons. In the second part of this work, we present an acceleration scheme for the proposed tensor-structured algorithm. In the proposed algorithm, the discrete Kohn-Sham problem is solved using the Chebyshev filtering subspace iteration method that relies on matrix-matrix multiplications of a sparse symmetric Hamiltonian matrix and a dense wavefunction matrix, expressed in the localized Tucker tensor basis. These matrix-matrix multiplication operations, which constitute the most computationally intensive step of the solution procedure, are GPU accelerated providing ~8-fold GPU-CPU speedup for these operations on the largest systems studied. The computational performance of the GPU accelerated code is presented using benchmark studies on aluminum nano-particles and silicon quantum dots with system sizes ranging up to ~7,000 atoms. Finally, we present a computation kernel using tensor-structured techniques to evaluate the one-electron and two-electron integrals, which are the central quantities to construct the Hamiltonian matrix of a full configuration interaction (FCI) calculation, for any molecular orbitals projected on a tensor-structured finite-element mesh. The proposed computation kernel is used to compute the integrals using Hartree-Fock molecular orbitals, which is subsequently used to perform an FCI calculation. In our numerical study, the FCI energy using integrals from the proposed computation kernel is confirmed to be consistent with a standard FCI calculation using the same Hartree-Fock orbitals. The proposed computation kernel provides a useful tool to investigate the behavior of Kohn-Sham orbitals as a basis for FCI calculations in the future, which could be a route to construct a reduced-order basis for an FCI calculation to reduce the computational cost and enable FCI calculation for the previously computationally inaccessible systems.
ABSTRACT Applying quantum chemistry algorithms to large‐scale systems requires substantial computational resources scaling with the system size and the desired accuracy. To address this, ByteQC, a fully functional and efficient package for large‐scale quantum chemistry simulations, has been open sourced at https://github.com/bytedance/byteqc , leveraging recent advances in computational power and many‐body algorithms. Regarding computational power, several standard algorithms are efficiently implemented on modern GPUs, ranging from mean‐field calculations (Hartree‐Fock and density functional theory) to post‐Hartree‐Fock methods such as Møller‐Plesset perturbation theory and coupled cluster methods. For the algorithmic approach, we also employ a quantum embedding method, which significantly expands the tractable system size while preserving high accuracy at the gold‐standard level. All these features have been systematically benchmarked. For standalone algorithms, the benchmark results demonstrate up to a 60× speedup when compared to 100‐core CPUs. Additionally, the tractable system sizes have been significantly expanded: 1610 orbitals for coupled cluster with single and double excitations (1380 orbitals with perturbative triple excitations), 11,040 orbitals for Møller‐Plesset perturbation theory of second order, 37,120 orbitals for mean‐field calculations under open boundary conditions, and over 100,000 orbitals for periodic boundary conditions. For the advanced quantum embedding feature, two representative examples are demonstrated: the water cluster problem (2752 orbitals) and a water monomer adsorbed on a boron nitride surface (3929 orbitals), achieving the gold‐standard accuracy. With these efforts, ByteQC is expected to significantly advance research in quantum chemistry, particularly in large‐scale, high‐accuracy calculations.
Quantum-mechanical calculations based on Kohn–Sham density functional theory (DFT) played a significant role in accurately predicting various aspects of materials behavior over the past decade. The Kohn–Sham approach to DFT reduces the many-body Schrodinger (eigen value) problem of interacting electrons into an equivalent problem of noninteracting electrons in an effective mean field that is governed by electron-density. Despite the reduced computational complexity of Kohn–Sham DFT, large-scale DFT calculations are still computationally very demanding with the resulting computational complexity scaling cubically with number of atoms in a given materials system. Numerical algorithms with reduced scaling behavior which are robust, computationally efficient and scalable on parallel computing architectures are always desirable to enable simulations at larger scales and on more complex systems. Following this line of thought, this study explores the use of tensor structured methods for ab-initio numerical solution of Kohn–Sham equations arising in DFT calculations. Earlier studies on tensor-structured methods have been quite successful in the accurate calculation of Hartree and the nonlocal exchange operators arising in the Hartree–Fock equations. A recent investigation of low-rank Tucker-type decomposition of the electron-density of large aluminum clusters (obtained from the finite-element discretization of orbital free DFT) shows the exponential decay of approximation error with respect to Tucker rank (number of tensor-basis functions in Tucker type representation). The results also indicate a smaller Tucker rank for the accurate representation of the electron density and is only weakly dependent on the system sizes studied. The promising success of tensor-structured techniques in resolving the electronic structure of material systems has enabled us to take a step further. In this study, we propose a systematic way of computing a globally adapted Tucker-type basis for solving the Kohn–Sham DFT problem by using a separable approximation of the Kohn–Sham Hamiltonian. Further, the resulting Kohn–Sham eigenvalue problem is projected into the aforementioned Tucker basis and is solved for ground-state energy using a self-consistent field iteration. The rank of the resulting Tucker representation and the computational complexity of the calculation are examined on representative benchmark examples involving metallic and insulating systems.
Abstract We present a tensor-structured algorithm for efficient large-scale density functional theory (DFT) calculations by constructing a Tucker tensor basis that is adapted to the Kohn–Sham Hamiltonian and localized in real-space. The proposed approach uses an additive separable approximation to the Kohn–Sham Hamiltonian and an L1 localization technique to generate the 1-D localized functions that constitute the Tucker tensor basis. Numerical results show that the resulting Tucker tensor basis exhibits exponential convergence in the ground-state energy with increasing Tucker rank. Further, the proposed tensor-structured algorithm demonstrated sub-quadratic scaling with system-size for both systems with and without a gap, and involving many thousands of atoms. This reduced-order scaling has also resulted in the proposed approach outperforming plane-wave DFT implementation for systems beyond 2000 electrons.