Can you be more specific? I have seen many such claims, and every time I try to reproduce the results, there always seems to be some catch.
We're trying to use TPUs to fine tune GPT-2 1.5B. The model takes up 5.8GB memory, which is well over half of a TPUv2 core (8GB). It always OOMs when I try to do a training step, due to the gradient calculations requiring memory. It even OOMs on a TPUv3, which has 16GB per core. I've tried using bfloat16 (which ought to cut memory usage in half) and using Ada optimizer (which should be no more expensive than plain old SGD). Yet if I colocate the gradients to the same core as the model, I always OOM. (Colocation just means "don't use any memory except the memory physically on this one core.") With colocation off, I don't OOM, and I do see some speed gains using all 8 cores. But it's no more than a factor of 2x, and in fact closer to like 1.15x (i.e. it's roughly equivalent to just using larger batch sizes on a single core). And I don't understand why I'd OOM in the TPUv3 case; even with float32, the model is only using 5.8GB out of 16GB. Are gradient computations really taking up more than 10GB for the optimizer? (That leads to https://github.com/cybertronai/gradient-checkpointing and such, but I haven't tried it yet.)
If I try the same experiment with a much smaller model (117M, or about 13x smaller), I can successfully colocate the gradients onto the same core as the model. And when I use all eight cores, I'm able to get 1225 tokens/sec (roughly 1 example per second, since 1 example = 1024 tokens for GPT-2), vs the standard case of around 400 tokens/sec when using only one core. But that's still "only" a 3x speedup.
So when I see "50-100x increases," alarm bells start going off. I'm missing something fundamental here. Either you are getting 100x speedups, or I am somehow missing something fundamental.
People have even started asking me for answers regarding the TPU case, and I'm forced to be like "Yeah! I expected TPUs to be so much faster too. Everyone says they're getting 100x speed gains. Yet we're 11x slower than the GPU case, and here's a notebook showing a 11x slowdown."
https://github.com/shawwn/gpt-2/issues/5
I'm suspecting that memory bandwidth is the bottleneck here for large models. This paper even pretty much says "GPUs are more flexible and faster when memory bandwidth is an issue": https://twitter.com/mosicr/status/1196749286815481856
The closest I've come to finding an actual example of a speedup to aim for is this: https://github.com/imcaspar/gpt2-ml
They used a TPUv3-512 pod to train a GPT-2 1.5B model to 99k steps in 50 hours. If you work out the math, that's about 1 example per second. We're getting about 0.08 examples per second on a single TPUv2 core. So yes, it's a big speedup (12.5x) but certainly not 50-100x. Yet it has 64x the cores as my TPUv2; why isn't it 64x faster? And we're only using 1 core; why not 512x faster?
I have also tried this on TPUv3-8, and we're getting about the same examples/sec, further increasing the plausibility of the theory that memory bandwidth is the bottleneck.