Large Scale Distributed Linear Algebra With Tensor Processing Units

Abstract

We have repurposed Google Tensor Processing Units (TPUs), application-specific chips developed for machine learning, into large-scale dense linear algebra supercomputers. The TPUs' fast inter-core interconnects (ICI)s, physically two-dimensional network topology, and high-bandwidth memory (HBM) permit distributed matrix multiplication algorithms to rapidly become computationally bound. In this regime, the matrix-multiply units (MXU)s dominate the runtime, yielding impressive scaling, performance, and raw size: operating in float32 precision, a full 2048-core pod of third generation TPUs can multiply two matrices with linear size $N=220=1048576$ in about 2 minutes. Via curated algorithms emphasizing large, single-core matrix multiplications, other tasks in dense linear algebra can similarly scale. As examples, we present (i) QR decomposition; (ii) resolution of linear systems; and (iii) the computation of matrix functions by polynomial iteration, demonstrated by the matrix polar factorization.

arXiv comment12 pages, 8 figures

References in corpus (6)

• Machine learning accelerated computational fluid dynamics
• Learning data driven discretizations for partial differential equations
• Kohn-Sham equations as regularizer: building prior knowledge into machine-learned physics
• Simulation of quantum many-body dynamics with Tensor Processing Units: Floquet prethermalization
• Large-Scale Discrete Fourier Transform on TPUs
• High-resolution imaging on TPUs

Cited by in corpus (5)

• Dynamics of Transmon Ionization
• Density Matrix Renormalization Group with Tensor Processing Units
• Large scale quantum chemistry with Tensor Processing Units
• Solving the Discretised Neutron Diffusion Equations using Neural Networks
• Fast Time-Evolution of Matrix-Product States using the QR decomposition