{'fraction_non_alphanumeric': 0.02864816472694718, 'fraction_numerical': 0.02018393423944006, 'mean_word_length': 4.616087751371115, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Benchmarking GPU and TPU Performance with Graph Neural Networks

Many artificial intelligence (AI) devices have been developed to accelerate the training and inference of neural networks models. The most common ones are the Graphics Processing Unit (GPU) and Tensor Processing Unit (TPU). They are highly optimized for dense data representations. However, sparse representations such as graphs are prevalent in many domains, including science. It is therefore important to characterize the performance of available AI accelerators on sparse data. This work analyzes and compares the GPU and TPU performance training a Graph Neural Network (GNN) developed to solve a real-life pattern recognition problem. Characterizing the new class of models acting on sparse data may prove helpful in optimizing the design of deep learning libraries and future AI accelerators.


Graph Neural Network Benchmark
The data are from the TrackML [6,7] dataset, which simulates top quark pair production from proton-proton collisions at the Large Hadron Collider. Graphs are constructed with the embedding learning and filtering method described in Ref [8]. In this application, the graph contains all data from a collision event. Each node represents a 3D spacepoint measurement (hit), and each edge represents a directional connection of one hit (sender) pointing to another (receiver). Hit positions in the cylindrical coordinate system are assigned as node attributes. The objective of graph neural networks is to assign a score to each edge so as to indicate the probability that the edge is true, i.e. that it connects two hits from the same particle. The graph size varies for each collision event, as shown in Figure 1. On average, there are about 50,000 nodes and 250,000 edges. Out of the 250,000 edges, the expected number of true edges is about 50,000.  The GNN architecture has three trainable components [9]: an encoder that transforms input graphs to their latent representations, an interaction network [10] and a decoder that computes the edge scores. The encoder has two independent networks: node network and edge network, which use two layers of fully connected network with sizes of [128,64]. The interaction network also has node and edge networks. It first updates node features and then updates the edge features. Node features are updated with a multilayer perceptron (MLP) to which the inputs are the concatenated node features and the aggregated neighboring edge features. Edge features are updated with another MLP to which the inputs are the concatenated edge features and the sender/receiver node features. In this way, messages are passed between nodes and edges via a message passing operation, one of which is the UnsortedSegmentSum in TensorFlow [11]. After performing the update eight times, edge features are fed to the decoder to predict edge scores. There are 132,291 trainable parameters in the GNN.

Hardware and Software
Our selection of hardware reflects the latest configurations available on the Google Cloud Platform (GCP) at the time of writing. The GPU is an NVIDIA V100 in a computing node at the National Energy Research Scientific Computing Center (NERSC). Each computing node contains eight V100 packages (PCIe) connected via 25 GB/s NVlink connection, and each V100 armed with 5120 CUDA cores has 16 GB of memory and 900 GB/s memory bandwidth. It reaches peak performance of 14 TFLOPS in single precision (float32) and seven teraflops in double precision (float64). In addition, A V100 has 640 tensor cores and can run mixed-precision training using half-precision (float16) to compute and single precision to accumulate, making its peak performance 125 TFLOPS. However, studies are needed to make sure the model converges to optimal results. Tensor cores are not used in this study. In May 2020, Nvidia announced the Nvidia A100, to which we were granted early beta access on GCP. The peak performance in single precision of the Nvidia GPU A100 is about 40% better than that of Nvidia GPU V100. Our experiments comparing the GNN training time between V100 and A100 are consistent with the expected peak performance improvement. Only the results with V100 are presented here.
The TPU is a Cloud TPU instance to which we were given academic access in Spring 2020. TPUs are designed to run whole inference models in the TPU to reduce communications with the host CPU. Details of TPU architectures can be found in Ref [2]. The heart of the TPU is the Matrix Unit (MXU). It can perform 16,000 multiply-accumulate operations in each cycle at reduced precision (bfloat16). It supports mixed-precision training, using bfloat16 to compute and float32 to accumulate. There are two versions of TPU cores: v2 and v3. A TPU v2 core has 8 GB of memory and one MXU, while a TPU v3 core doubles the memory size and MXU. However, the number of cores in each TPU can be configured [12]. We chose 32 v2 cores and eight v3 cores in the GCP us-central1-a region for this study.
The software stack is based on the TensorFlow v2.3 [11], which unifies the usage of CPUs, GPUs, and TPUs via the tf.distribute package. It allows users to perform data distributed training for all kinds of devices in a consistent way. In addition, we use tf.data to build TensorFlow input pipelines, and use tf.function to compile the model and generate a computation graph in gRPC format [13]. The gRPC file is sent via the cloud to the TPU host and compiled by XLA, and in the end, the binary files are executed in TPUs. As of writing 1 , the TPU does not support dynamic graph sizes. The input graph size has to be the same. Padding all graphs to the largest graph in the training data is expensive. Instead, we padded to the graph whose number of nodes and edges have a quantile of 99%, saving the training time by 30% compared to that padding to the largest graph size.

Comparison of TPU and GPU
In this section, we use the following key metrics to compare the computational performances of TPUs and GPUs: accuracy, latency, cost, power consumption.
Accuracy It measures the quality of the trained model and is evaluated by the Area Under the receiver operating characteristic curve (AUC). Figure 2 shows the comparison of accuracy achieved with TPU and GPU. The difference between the two training results lies in the batch size. Training with GPU uses a batch size of 1, while training with TPU v3 uses 8. GPU and TPU performances are very close, even though model hyperparameters, particularly the learning rate, were not tuned for TPUs.
Latency It measures the time it takes to finish one training epoch. It is one of the critical metrics for HEP online data processing. Figure 3 shows the comparison of the averaged time it takes to finish one training epoch. A TPU with 32 TPU v2 cores runs as fast as 4 GPU V100 does and A TPU with 8 TPU v3 cores runs as fast as 2 GPU V100 does.
Cost It measures the money it costs to train one epoch in US dollars, as shown in Figure 4. The prices for using the devices are taken from Google Cloud Platform, charged by US dollars per hour. It is found that using GPUs is more economical than using TPUs. Simultaneously using multiple GPUs for training will increase the cost a bit due to the imperfect strong scaling efficiency.
Power Consumption It measures the energy cost per epoch, a key metric to evaluate if the device is environment friendly. It is calculated as the thermal design watts times the latency, as shown in Figure 5. For simplicity, we assume the device is 100% busy during the training, which usually does not hold. It is found the GPU is more environmentally friendly.

Profiling and Roofline analysis
To understand the results shown in Section 4, we use Tensor Board v2.3 [11] to profile the training process for TPU v3 and GPU. This section describes the profiled results and the performance analysis with the Roofline model [14].
We break the time usage of each device into different TensorFlow operations. For TPU v3, 37.7% of total training time is spent in UnsortedSegmentSum, which aggregates the node or edge features, 13.6% in idling, and 12.7% in matrix multiplications. On the other hand, for GPU, 37.4% of total training time is spent in matrix multiplications, 10% in idling, and 6.4% in UnsortedSegmentSum. In both cases, about half of the time is spent in gradient calculations and back-propagation.
To determine which kernel takes most of the training time, we ranked all kernels in descending order of their total duration in each training step. The top three CUDA kernels account for 98.7% of total FLOPs but only for 40.7% of total training time. They are to perform matrix multiplications. The following two CUDA kernels account for 13.8% of total training time but zero FLOPs. Specifically, the two CUDA kernels are to concatenate and slice tensors. The top 5 to 20 ranked kernels are led by the message-passing operation

UnsortedSegmentSum.
FLOP utilization, defined as the fraction of the FLOPS used over the peak FLOPS offered by the device, is an essential metric in understanding how well an application utilizes the accelerator's total computation capacity. The GNN trained with TPU v3 has a FLOP utilization of 2.3%, but that with GPU V100 has a FLOP utilization of about 30% for single precision. The TensorCore in V100 was not used because it requires the benchmark to support mixed-precision training.
FLOPS utilization is only part of the problem when designing an accelerator. In particular, memory bandwidth is another important aspect that can have a significant impact on performance. We use the Roofline model to study the GNN training runs for GPU V100. Figure 6 shows the Roofline studies for different memory hierarchical levels. The Roofline has a slanted part and a horizontal part. It represents the highest achievable FLOPS at a given arithmetic intensity. Any data point (x, y) on the slanted part has y/x = memory bandwidth (GB/s). The horizontal part is the peak FLOPS on the device. A workload, which is a CUDA kernel in our case, close to the slanted roofline is memory-bound; one close to the horizontal part is compute-bound. Figure 6 (bottom left) shows that the kernels performing message-passing operations indicated in yellow squares are bounded by the high bandwidth memory. It is the interest of future device design to improve memory bandwidth and add sparse operation support. The Roofline model ignores the CUDA kernels with zero arithmetic intensity (zero-AI kernels). The zero-AI kernels, summing to 1/3 of total CUDA kernels, take 44.8% of total training time. They account for 20.6% of data transactions in L1, 38.0% in L2, and 54.0% in HBM. The total overhead in creating all zero-AI kernels is about 2.4% of the total training time.

Conclusions
Graph Neural Networks combine deep learning operations (e.g., matrix multiplication and convolution) with iterative message passing operations such as UnsortedSegementSum, which aggregates node (edge) features to edge (node) networks. GNN is a unique group of underrepresented neural networks in current popular benchmark suites such as MLPerf.
We use a GNN model that solves a scientific problem as a benchmark to compare the computational performances of GPU and TPU. In order to utilize more than one TPU core, the distributed training of GNN was implemented. Different software implementations result in different training time for one epoch and different strong scaling efficiencies. The GNN model's accuracy trained with TPU and GPU is very similar; however, they manifest different latencies. A TPU with 32 TPU v2 cores runs as fast as 4 GPU V100 does. TPUs spend most of the training time in running message passing operations but GPUs in running matrix multiplications. GPUs cost less money and consume less power compared with TPUs. The Roofline analysis reveals that the bottlenecks of training GNN on GPUs are the computing capability for the matrix operations and the memory bandwidth for the message passing operations. Besides GPUs and TPUs, Field Programmable Gate Arrays (FPGA) [15] and Intelligence Processing Units (IPU) [16] offer alternative approaches to AI acceleration that show promise for sparse data. It would be interesting for future studies to include IPUs and FPGAs in the comparison.