High performance Monte Carlo simulation of ising model on TPU clusters

Abstract

Large-scale deep learning benefits from an emerging class of AI accelerators. Some of these accelerators' designs are general enough for compute-intensive applications beyond AI and Cloud TPU is one such example. In this paper, we demonstrate a novel approach using TensorFlow on Cloud TPU to simulate the two-dimensional Ising Model. TensorFlow and Cloud TPU framework enable the simple and readable code to express the complicated distributed algorithm without compromising the performance. Our code implementation fits into a small Jupyter Notebook and fully utilizes Cloud TPU's efficient matrix operation and dedicated high speed inter-chip connection. The performance is highly competitive: it outperforms the best published benchmarks to our knowledge by 60% in single-core and 250% in multi-core with good linear scaling. When compared to Tesla V100 GPU, the singlecore performance maintains a ∼10% gain. We also demonstrate that using low precision arithmetic-bfloat16-does not compromise the correctness of the simulation results.or heterogeneous nodes commonly seen in private or commercial clouds. Benefiting from the explosion of machine learning, especially deep learning, commercial clouds provide not only CPUs and GPUs, but also specialized chips such as FPGAs and other in-house processors. The Tensor Processing Unit ("Cloud TPU" or "TPU" for short)-an AI application-specific integrated circuit (ASIC) developed by Google for neural network machine learning-has received much attention in the machine learning community [18,17]. Its latest release, Cloud TPU v3, offers 420 × 10 12 floating-point operations per second (FLOPS) and 128GB of high bandwidth memory (HBM) 1 . Multiple units are connected to form a "POD" (Cloud TPU v3 Pod) through a dedicated high speed 2-D toroidal mesh network, allowing up to 100+ peta-FLOPS and 32TB of HBM 1 to be accessed by the application with very low latency and in lockstep. TPU is programmable via software frontends such as TensorFlow [1] or PyTorch [21], and can be deployed both for training huge deep neural networks and for performing low-latency online prediction. [14] reports impressive acceleration of training and online prediction.With the tremendous amount of computation resources that TPU offers, it is compelling to also consider the opportunities TPU brings for applications beyond machine learning. The programming frontends that are used for TPU, such as TensorFlow, also offer a rich set of functionalities that are highly relevant for scientific computations. The TensorFlow TPU programming stack also provides the additional benefits of allowing distributed algorithms to be expressed with simple and easy-to-understand code without sacrificing performance. In addition, the ability to program conventional scientific simulations in TensorFlow framework makes it easier to explore the hybrid approaches employing both conventional scientific computation methods and modern machine learning techniques on the same framework.Motivated by these observations, we developed a Single Instruc...

Machine learning–accelerated computational fluid dynamics

“…The ability to match DNS with a × coarser grid makes the learned interpolation solver much faster. We benchmark our solver on a single core of Google’s Cloud TPU v4, a hardware accelerator designed for accelerating ML models that is also suitable for many scientific computing use cases ( 45 – 47 ). The TPU is designed for high-throughput vectorized operations, with extremely high throughput matrix–matrix multiplication in low precision (bfloat16).…”

Dynamic critical exponent z of the three-dimensional Ising universality class: Monte Carlo simulations of the improved Blume-Capel model

“…As an example for recent work, where the two dimensional Ising model has been simulated on a GPU, see ref. [36].…”

Large-Scale Discrete Fourier Transform on TPUs

“…It is plausible with the following four reasons: (1) TPU is an ML application-specific integrated circuit (ASIC), devised for neural networks (NNs); NNs require massive amounts of multiplications and additions between the data and parameters and TPU can handle these computations in terms of matrix multiplications in a very efficient manner [29]; similarly, DFT can also be formulated as matrix multiplications between the input data and the Vandemonde matrix; (2) TPU chips are connected directly to each other with dedicated, high-speed, and lowlatency interconnects, bypassing host CPU or any networking resources; therefore, the large-scale DFT computation can be distributed among multiple TPUs with minimal communication time and hence very high parallel efficiency; (3) the large capacity of the in-package memory of TPU makes it possible to handle large-scale DFT efficiently; and (4) TPU is programmable with software front ends such as TensorFlow [30] and PyTorch [31], both of which make it straightforward to implement the parallel algorithms of DFT on TPUs. In fact, all the aforementioned four reasons have been verified in the high-performance Monte Carlo simulations on TPUs [32], [33].…”