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.
Compare that with a GPU. It took decades, but (at least for nvidia GPUs) they have finally reached the point where you can be at least reasonably certain that high-performance is the default. Sure, you can ruin your performance by doing some stupid things, but it's at least harder to do those stupid things. The defaults are high-performance.
It's exactly the opposite for Cloud TPUs. By default, it only uses one core. You have to use their TPUEstimator API, which is a byzantine mess. Since two weeks ago, I've spent roughly 3 solid working days of effort solely trying to read and understand (a) what is the TPUEstimator actually doing? (b) why is it supposedly so much faster?
I have some half-hearted justifications -- the answer seems to be "you have to colocate your gradients with the device; you have to scope your tensorflow graph to a specific device; and each TPU core is a separate device." But there are unanswered questions. For example, you're supposed to pass your tensorflow computation to tpu.rewrite(). Yet I've never done that. The defaults seem to just work. So does that mean it does the rewriting for me automatically? Is it being emulated in software, and I'm ruining my performance? Tensorflow, why don't you just crash instead of being so damn slow on TPUs? That would at least let me aim my optimizations!
If only I had low level access to the actual TPU operations, I could just write a compiler that specifically emits instructions to give me the performance I need. But this tensorflow graph abstraction makes everything "easy" yet exponentially more complicated.
Anyway. Yes. More of your mindset, please. Simplicity is such a lovely metric seldom optimized for.
(Anyone who's curious can see a dramatic 11x difference in performance on GPUs vs TPUs in Colab: https://twitter.com/theshawwn/status/1196593451174891520 ... this was very surprising, since Google's marketing would have you believe that TPUs are the bee's knees. Yet by default a TPUv2 is 11x slower than a K80 GPU for some basic matrix multiplications.)