Xla: Exepected performance relative to V100?

[myleott]

I'm doing some benchmarking to better understand the performance of TPUs vs GPUs and wanted to know if my results are expected.

I wrote a simple benchmark script that somewhat mirrors a 24-layer transformer but just uses MLPs and LayerNorm (i.e., no multihead attention), with 1024 dimension embed for a total of 480M parameters. Batch size 8, sequence length 512.

The script uses float16 on GPU and bfloat16 on TPU. There's a warmup of 50 steps, then I report the time for the subsequent 50 steps.

Benchmark script: https://gist.github.com/myleott/9017f443ab9d86ecf779d57fde58e1a2

On a single TPU v3 core this script reports: 18.5 seconds
On a single V100 GPU this script reports: 9.6 seconds

Is this expected?

More details:

• I'm on PyTorch 1.4 for the V100 and the latest pytorch-nightly/XLA for the TPU v3.
• I get the same performance regardless of whether I use XLA_USE_BF16=1, because I manually cast the model to bfloat16 in the script.
• I looked the metrics report and don't see any aten:: ops. Here's a copy of the metrics report from the end of the benchmark script: https://gist.github.com/myleott/fa49c10039c89b9472e6b0c59590b10b

[taylanbil]

Based on the equivalency "1 v100 GPU = 2 v3 TPU core", it sounds like these are performing similarly. Is 9.6 seconds on 1 chip (2 cores ?)

I get the same performance regardless of whether I use XLA_USE_BF16=1, because I manually cast the model to bfloat16 in the script.

That's good to know.

[dlibenzi]

The compare @taylanbil mentions is based on TCO.
As far as the metrics I don't see anything wrong.
The step time is about 350ms, which for 50 measure steps is about matching your measured 18s.
One thing I noticed is that it uses embedding, and we will have to check whether that lowering to XLA is right from a performance POV.
AFAICT we just copied the pytorch one, which might lead to unoptimized performance.

[myleott]

Ah, okay. So then I guess a more fair comparison is 4*V100 to a v3-8.

I've adjusted the script a bit so that the total batch size across all ranks is 128. Thus each V100 handles a batch size of 32 and each TPU core handles a batch size of 16. This is approximately the max batch size each of them can support, since each V100 has 32GB of memory and the v3 cores have 16GB. Both hit OOM if I increase the total batch size from 128 -> 256.

The updated script is here: https://gist.github.com/myleott/a66ba69601cbd21a5a2218a33b6363f8

On a TPU v3-8 the script reports: 36.8 seconds
On 4*V100 the script reports: 36.2 seconds

This is much closer, so closing this out. Thanks for the help :)

[taylanbil]

@dlibenzi any optimizations we can make to Embedding will have profound impact on all nlp usecases including fairseq imho.

[dlibenzi]

We can look into it.

[dlibenzi]

I have created a simple test to generate the HLO graph:

import torch
import torch.nn as nn
import torch_xla
import torch_xla.core.xla_model as xm


BATCH_SIZE = 8
SEQLEN = 512
NUM_EMBED = 50000
EMBED_DIM = 1024

device = xm.xla_device()

item = torch.arange(1, SEQLEN + 1, dtype=torch.long).to(device)
embed = nn.Embedding(
    num_embeddings=NUM_EMBED, embedding_dim=EMBED_DIM, padding_idx=0
).to(device)
x = embed(item)
x = x.sum()
x.backward()

g = [x]
for p in embed.parameters():
  if p.grad is not None:
    g.append(p.grad)

print(torch_xla._XLAC._get_xla_tensors_hlo(g))

Attached is the HLO graph.

xemb.hlo.txt

[dlibenzi]

Embedding does not seem to be a problem here.