Benchmarking TPU, GPU, and CPU Platforms for Deep Learning

In plain words

Picking hardware for training a neural network is hard, because the answer depends on the network. Most benchmarks run a few popular models, and one popular model can hide a weakness that another model would expose.

This paper builds a generator of many small-to-medium training models of three kinds: stacks of plain layers, image networks with residual blocks, and recurrent networks. Each can be sized by a few knobs, such as depth, width, batch size and input size. The authors run these generated models, plus six well-known real models, on a Google TPU, an NVIDIA GPU and a CPU. They then ask which machine is fastest for which kind of model.

The answer is that no single machine wins everywhere. The TPU is strongest with big batches and image networks. The GPU handles small batches and unusual operations better. The CPU has the most memory, so it can hold the largest models. The paper also looks inside the TPU to find where it stalls, and checks how much faster the software got over a few months.

TL;DR

• ParaDnn generates parameterized FC, CNN and RNN training models, and the paper runs them with six real models on Cloud TPU v2/v3, a V100 and a Skylake CPU §1 §2.
• The TPU has the highest training throughput for large batches and CNNs. The GPU is more flexible for small batches and non-matrix work. The CPU has the largest memory §5.
• On real models the TPU is 3x to 6.8x faster than the V100 (DenseNet to SqueezeNet). These are author-reported single-device numbers §5.3.
• The TPU stalls on memory-bound operations and on data infeed. Over 7 months of TensorFlow releases, the 90th-percentile speedup for parameterized RNNs was 9.7x §4.2 §4.4 §6.1.

Problem

Benchmarks built from a handful of popular models give a narrow picture. The authors give an example: Transformer trains about 3.5 times faster on the TPU than on the GPU, yet that one model would miss a TPU memory-bandwidth bottleneck that shows up in fully connected layers wider than 4k nodes §1.

Synthetic micro-benchmarks, such as single matrix multiplies, do not capture dependencies between operations in a full model §2. The goal is a benchmark that covers many model sizes, with end-to-end models §2.

The study is limited to training, on one GPU and one TPU board. It does not cover inference, multi-GPU or 256-TPU systems §1 §3 §7.

Method

Overview. The paper contributes ParaDnn, a sweep of model designs, and uses it on each platform with the same training setup. The main metric is examples per second, with speedup as the ratio of two platforms §5. The analysis then has four parts: FLOPS utilization on the TPU §4.1, a roofline view of memory versus compute §4.2, the cost of cross-core communication §4.3 and of host-to-device data feeding §4.4. Section 6 measures software versions.

Step 1. Model families. All three families are end-to-end and trainable §2.1.

• FC: an input, a stack of fully connected layers, and an output §2.1.
• CNN: four groups of residual or bottleneck blocks, followed by a fully connected layer. Each group has the same number of blocks §2.1.
• RNN: a stack of basic RNN, LSTM or GRU cells over an embedded token sequence §2.1.

Step 2. Sweep. Each family has a range for each hyperparameter, and for batch size §2.1 tab. 2. The paper says the ranges were chosen to cover the real models and to stay reasonable §2.1.

Step 3. Real models. Six real models are added: Transformer and ResNet-50 (TensorFlow, MLPerf), RetinaNet, DenseNet, MobileNet and SqueezeNet §2.2. Batch sizes are the largest the platform supports. On TPU with bfloat16 they are 64 for RetinaNet, 4k for Transformer and 1024 for the rest §2.2.

Step 4. Platform setup. The TPU runs data parallelism: one batch is split evenly over 8 cores, and each core keeps a full copy of the model §3. The GPU is one V100 package in a DGX-1 §3. Each platform's numeric format is listed in Hyperparameters.

Step 5. Throughput and speedup. Examples per second is the throughput §5. Speedup is the ratio of one platform's examples per second to another's §5.

Step 6. FLOPS utilization. Average FLOPS divided by peak FLOPS §4.1.

Step 7. Hyperparameter effects. A linear regression of FLOPS on the normalized hyperparameters, used for the signs and sizes of the weights, not for prediction §4.1. The model for FC is:

\mathrm{FLOPS} = w_0\,\mathrm{layer} + w_1\,\mathrm{node} + w_2\,\mathrm{input} + w_3\,\mathrm{output} + w_4\,\mathrm{batch} §4.1

Step 8. Roofline. Arithmetic intensity is FLOPs per byte moved from memory. The roofline has a slanted part (memory-bound) and a horizontal part at peak FLOPS (compute-bound) §4.2. The paper uses it for FC and CNN, and omits RNN because the TPU profiler reports incorrect memory-bandwidth numbers for RNN models §4.2.

# Model families and sweeps  [§2.1] [tab. 2]
# A ParaDnn model is built from these knobs; each sweep is a list of values.
def fc_model_space():
    # Layers, nodes per layer, input units, output units, batch  [tab. 2]
    return {
        "layers": [4, 8, 16, 32, 64, 128],
        "nodes": [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192],
        "input_units": [2000, 4000, 6000, 8000],
        "output_units": [200, 400, 600, 800, 1000],
        "batch": [64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384],
    }

# Throughput, speedup and utilization  [§5] [§4.1]
def examples_per_second(batch_size, steps, seconds):
    # Examples trained per second, used as an end-to-end proxy  [§5]
    return batch_size * steps / seconds

def speedup(throughput_a, throughput_b):
    # Ratio of one platform's examples/s to the other's  [§5]
    return throughput_a / throughput_b

def flops_utilization(average_flops, peak_flops):
    # Average achieved FLOPS over the platform's peak  [§4.1]
    return average_flops / peak_flops

# Roofline bound for one operation  [§4.2] [§4.5]
# intensity in FLOP/byte; bandwidth in bytes/s; FLOPS in FLOP/s
def attainable_flops(intensity, peak_flops, bandwidth):
    # Slanted part: limited by memory traffic  [§4.2]
    # Horizontal part: capped at the peak FLOPS  [§4.2]
    return min(peak_flops, intensity * bandwidth)

# Sign and size of each hyperparameter's effect  [§4.1]
# The paper uses these weights for ranking only, not for prediction.
def effect_weights(hparams, flops):
    # Put every column on one scale, so the weights compare  [§4.1]
    scaled = normalize_columns(hparams)     # (N, k)
    # Ordinary least squares; the paper does not state an intercept  [§4.1]
    weights = least_squares(scaled, flops)  # (k,) signed effect
    return weights

Hyperparameters

Values are the paper's sweep ranges and platform settings. Where the repository differs, the difference is in Paper vs code.

Results

Author-reported. Setup: ParaDnn plus six real models on Cloud TPU v2 (v3 in one section), one V100 package, and one Skylake CPU instance. Batch sizes are the largest supported on each platform §2.2 §3 tab. 3.

Platform comparison (TPU over GPU and CPU).

TPU architecture.

Software stack.

The paper states that the CUDA improvement figures are measured across CUDA versions but does not restate the starting version for each number §6.2.

Paper vs code

Official repository: https://github.com/Emma926/paradnn (Apache-2.0), read at code_commit 60a25ff. The owner account is Emma926. The README is signed "Yu (Emma) Wang", the paper's first author README.md. The README asks users to cite a 2020 MLSys paper with the same authors and a similar title, not this arXiv version README.md. The paper says ParaDnn will be open-sourced §2 footnote. I found the repository through a GitHub repository search for "ParaDnn". It is the tool release, not necessarily the exact artifact behind every number in the paper. Those differences are listed below.

• Software pinning. The paper's TensorFlow (1.7 to 1.12) and CUDA (9.0 to 10) versions §6.1 §6.2 are not pinned in the repository. The README asks for TensorFlow 1.6 to 1.13 README.md.
• FC optimizer constants. RMSProp decay 0.9, momentum 0.9, epsilon 1.0, and a fixed learning rate of 0.1 fc.py fc.py. The paper names RMSProp but gives no constants §4.2 §7.
• CNN learning rate. Starts at 0.1 and decays by a factor of 0.97 every 25,000 steps cnn.py. Not stated in the paper.
• RNN learning rate. Starts at 0.1 and decays by a factor of 0.9 every 10,000 steps rnn.py. Not stated in the paper.
• CNN stem and head. The code adds a 7x7 stride-2 convolution, batch norm, ReLU and 3x3 stride-2 max pool before the four groups, and a 7x7 average pool and dense layer after them resnet_model.py resnet_model.py. The paper describes four groups followed by a fully connected layer §2.1.
• CNN filters. The code starts at filters (default 16) and doubles it in each group, giving 16, 32, 64 and 128 cnn.py resnet_model.py. The paper's text says the minimum is 64, doubling to 512 §2.1. Its Table 2 says 16 to 32 tab. 2.
• CNN defaults. The code's default batch size is 8, input is 224, output is 1000 and block layers are 1,1,1,1 cnn.py cnn.py. The paper's sweeps are 64 to 1024 batch, 200 to 300 input and 500 to 1500 output tab. 2.
• CNN run script. The bottleneck sweep uses input 100 and output 500 cnn_tpu.sh cnn_tpu.sh. The input 100 is below the paper's range of 200 to 300 tab. 2. Tracing the shapes through the ResNet code, a 100-pixel input reaches a 4x4 map before the 7x7 final average pool resnet_model.py. This looks like it would fail to build. I did not run it.
• FC run script. It sweeps only nodes 32 and 64, with the paper's larger values commented out fc_tpu.sh. It sets both input and output units to the node count, ignoring the loop variables fc_tpu.sh. The paper's FC sweep uses input 2000 to 8000 and output 200 to 1000 tab. 2.
• RNN cell type. The cell is fixed to a basic RNN cell rnn.py rnn.py. The LSTM and GRU options exist only as a module constant that needs editing rnn.py. The paper says the family includes all three §2.1.
• RNN vocabulary in the run script. The script's only vocabulary value is 1 rnn_tpu.sh. The paper's minimum is 2 tab. 2. With one class, the softmax cross-entropy is zero by construction. I reasoned this from the code and did not run it.
• RNN output and lengths. The code projects outputs to the vocabulary with a linear layer rnn.py. It samples lengths uniformly from zero up to the maximum length rnn.py, so zero-length rows are possible. The paper mentions neither.
• Synthetic inputs. FC inputs are uniform in [-0.5, 0.5] and labels are uniform integers fc.py fc.py. CNN inputs come from NumPy's uniform [0, 1) cnn.py. RNN token ids are uniform over the vocabulary rnn.py. The paper says data is synthesized on the CPU host §4.4 but gives no distributions.
• Timing with warm-up. The FC script first trains to step 300, then calls train again with the absolute limit of 100 steps and times that call fc.py. Because the step limit is absolute in the TensorFlow Estimator, the timed call may run no steps. I did not run it, so the FC timings should be checked before use. The paper states no warm-up.
• Analytic FLOPS for FC. The script computes TFLOPS from the layer sizes, times batch, times 6 per training step, counting only dense layers fc.py. The paper gives no formula for the parameterized models' FLOPS utilization §4.1.
• TPU iterations. The code runs 100 TPU iterations per loop by default fc.py. The paper does not state this value.
• bfloat16 mechanics. The code keeps float32 variables and casts them on read through a custom getter fc.py. It casts the logits back to float32 before the loss fc.py. This matches the paper's description of bfloat16 compute with float32 accumulation §3.
• Data parallelism. The code shards over 8 cores by default fc.py, matching the paper §3.
• Real models. Transformer, ResNet-50, RetinaNet, DenseNet, MobileNet and SqueezeNet come from the TensorFlow TPU repository §2.2. They are not in this repository, so their TPU numbers cannot be checked here.
• Analysis code. The repository has scripts/plotting/ with a notebook and helpers. I did not review them, so I cannot say whether they reproduce the regression, roofline or Amdahl analyses.

Pitfalls

• Scope. Training only, one GPU and one TPU board §1 §3. Inference, multi-GPU and multi-node results are not covered §7. Do not extend the speedups to clusters.
• Real-model comparisons mix batch size with hardware. Real workloads use larger batches on the TPU than on the GPU §5.3. The batch sizes are the largest each platform supports §2.2. Part of any speedup is therefore a batch effect, not a hardware effect.
• Infeed excluded from parameterized models. ParaDnn synthesizes its data on the host to avoid infeed overhead §4.4. The paper shows that infeed can dominate real models: RetinaNet and SqueezeNet are infeed-bound §4.4. TPU throughput for ParaDnn models is therefore best read as a device number.
• Transformer infeed is missing. The TensorFlow implementation of Transformer does not support synthetic data, so its infeed bars are omitted §4.4.
• TPU v3 bandwidth is inferred. Memory bandwidth is not disclosed. The paper derives 1.5x from the 3x speedup of memory-bound operations, which it attributes to a 2x larger batch §4.5. The 420 TFLOPS and 3600 GB/s board figures are the paper's estimates tab. 3.
• Ridge point. The inflection point on the TPU v2 roofline is given as 75 operations per byte, from 180 TFLOPS over 2.4 TB/s §4.5. The paper gives no roofline for the GPU or CPU, so ridge points cannot be compared across platforms §4.2.
• Amdahl fractions have no formula. The paper says the non-parallel fraction can reach 60% for FC and 40% for CNN, "based on Amdahl's law" §4.3. It gives no equation or inputs, so the numbers cannot be recomputed.
• Variance is thin. The paper reports no repeat count. The only spread is one standard deviation of one-minute profiler samples, shown for the infeed case study §4.4.
• Memory limits decide which platform can run a model. Each TPU core holds a full model copy, so model size is limited per core, while batch size is limited by total board memory §3. The CPU is the only platform that runs the largest FC models §5.1. A GPU out-of-memory result appears as a missing cell in the FC grid fig. 7.
• Software versions move the results. The 90th-percentile TPU speedup for RNNs was 9.7x across TensorFlow 1.7 to 1.12 §6.1. Any comparison must name the software stack.
• Parameterized and real speedups disagree. Parameterized speedups range from below 1x to 10x, while the real models sit between 3x and 6.8x §5.3. The paper's real-model set is small, so a single model says little about the rest of the space.
• NVIDIA's ResNet-50 number was not reproduced. The paper quotes 4.2x for NVIDIA's version and notes it used unreleased libraries §5.3.
• Quantization is mixed in. The paper's bfloat16 and float16 results change memory traffic and batch size together §6.1 §4.4. Accuracy effects are out of scope §6.1 footnote.

Lineage

• Builds on: the roofline model of memory versus compute §4.2. Residual networks §2.1, Transformer §2.2 and the MLPerf reference models §1 are cited as the real workloads. The paper cites them by bibliography key, not by arXiv id.
• Compared against: MLPerf, Fathom, BenchNN and BenchIP as real-world suites, and DeepBench as a micro-benchmark suite §1 §2.
• Superseded by: not stated in the wiki. The repository README cites a 2020 MLSys version of the methodology README.md.