3 ms·
A bit off topic, but I am looking at TPUs at the moment. Can I ask for clarity if you mean TPUs are easier to use than GPUs out vice versa? I thought TPUs are
by prennert 7y ago
A bit off topic, but I am looking at TPUs at the moment. Can I ask for clarity if you mean TPUs are easier to use than GPUs out vice versa?
I thought TPUs are harder to work with because they only support Tensorflow rather than Tensorflow and other high-level frameworks as well as low-level CUDA that are supported by GPUs
- sillysaurusx 7y agoSure! TPUs aren't necessarily easier to use – it's about the same – but they're powerful. I've documented some benchmarks in this tweet chain, where I trained GPT-2 1.5B to play chess using a technique called swarm training: https://twitter.com/theshawwn/status/1214013710173425665 https://twitter.com/theshawwn/status/1214013710173425665 The power turned out to be from the fact that every TPU gives you 8 cores at your disposal. I never use the Estimator API. I just scope Tensorflow operations to specific TPU cores. Works great. In terms of actual performance, I was delighted to discover that TPUs can be faster than GPUs when you use all 8 cores: https://twitter.com/theshawwn/status/1196593451174891520 https://twitter.com/theshawwn/status/1196593451174891520 (solution notebook: https://twitter.com/theshawwn/status/1205914446918492170 https://twitter.com/theshawwn/status/1205914446918492170) It also gives you flexibility. TPUv2-8 can apparently allocate up to 300GB (!) if you don't scope any operations to any cores. Meaning, you run it in a mode where you only get 1 core of performance, but you get 300GB of flexibility. And then you can connect multiple TPUs together as described in the tweet chain, which quickly makes up the difference. There is also the question of cost savings. A TPUv3-8 seems about as expensive as a V100. Which one is worth it? Well, it depends. In my experience a GPU is easier to use and quicker to set up if you only need one GPU of horsepower. But suppose you wanted to train a massive model in 24 hours. What's your best option? For us, it was TPUs. The reason is subtle: It's hard to find any single VM that can talk to 140 GPUs simultaneously. But you can talk to 140 TPUs from a single VM no problem. And since you get 800MB/s to and from the VM, you can average the parameters across all TPUs very quickly. This is similar to what TPU pods do internally. And while TPU pods are impressive, they are also impressively expensive. A TPUv3 pod will run you $192/hr at evaluation prices. Whereas you can play with a TPUv3-8 for $2.50/hr. You can also play with a TPUv2-8 for free using Colab: https://github.com/shawwn/colab-tricks https://github.com/shawwn/colab-tricks Yesterday I used that notebook to port forward Colab's free TPUv2-8 using ngrok, then trained using the new StyleGAN 2 codebase: https://twitter.com/theshawwn/status/1214245145664802817 https://twitter.com/theshawwn/status/1214245145664802817 I think a swarm of TPUs can cost significantly less than a cluster of V100s with less engineering effort. That said, right now most codebases are designed to work with V100's. It will take time before TPUs widely proliferate. But speaking as someone who was once skeptical of TPUs and who has spent several months trying to discover their secrets, I feel that TPUs can get the job done quicker and easier than a GPU cluster. The hardware is also more accessible, since you can more easily spin up 100 TPUs than 100 V100s. But mainly I like that it's all coordinated from a single machine. It's conceptually simpler to debug and to implement. If you run into any issues or have any trouble with TPUs, please feel free to ask here or DM me. I love talking about this stuff. EDIT: In regards to usability, the new Jax library works with TPUs out of the box. Google seems to be heading in the direction of Jax. My initial reaction was "Not another library..." but first impressions were positive. It's not quite the React of ML – an idea which I hope to see soon – but it does seem easier for certain research purposes. PyTorch also recently gained TPU support, and as far as I know they've put in some serious efforts to make sure things run quickly. As for how you use all 8 cores of a TPU using PyTorch, I haven't looked into it yet. But I'd be surprised if you couldn't. It seems unlikely that they would design an API that would hamstring you to just 1 out of 8 cores.
- prennert 7y agoThanks a lot. This is a lot of information to digest. I will check out your Twitter feed and I think I need to start to play around with TPUs then. Scaling seems to work fantastic for you. I am working more in computer vision and we are running sometimes into weird bottlenecks with our GPUs where neither GPUs not CPUs are under full load. Unfortunately, drilling down on where the bottlenecks come from is not easy at all. I am assuming the profiler from Tensorflow works with TPUs in the same way it does with GPUs?
- sillysaurusx 7y agoOooh, you're so lucky you get to work on those kinds of problems. I know sometimes it feels frustrating to hunt for bottlenecks, but man is it satisfying to find it. We had a similar situation at one point. The problem turned out to be that our CPU wasn't generating input data fast enough. So the first step is to confirm that your input pipeline isn't the issue. The next step would be to break down the problem: Can you extract the smallest part of the codebase into a separate program, and try to make that run under full load? That's not the technique I used, though. To figure out the multicore stuff, the trick for me was to comment out almost all of the code, until you're left with only a small part that actually runs on the device. Ideally the smallest part. Basically, change your code so that the model file returns tf.no_op() (or as close to that as possible while still letting your input pipeline run). You want to be in a situation where your training loop is doing an equivalent of while(true) { read_input(); } so that you can verify that your pipeline is able to peg your GPUs to 100% usage. If you get 100% usage, fantastic! That means you're left with an easy problem: start turning parts of the code back on until you find which part is reducing your performance. Then study that part to figure out why. If you're not at 100% usage, you're either running into a fundamental limitation (which sometimes happens) or the pipeline isn't designed correctly in some way. I would compare it against other popular codebases such as StyleGAN 2 https://github.com/NVlabs/stylegan2 https://github.com/NVlabs/stylegan2 which is designed to use 8 V100s. The optimizer.py file is pretty insightful: https://github.com/NVlabs/stylegan2/blob/eecd09cc8a067e09e127d2435b71232a658a024f/dnnlib/tflib/optimizer.py https://github.com/NVlabs/stylegan2/blob/eecd09cc8a067e09e12... Finally, my biggest tip would be to step back from the problem and think: is there something simple you can do to reframe the problem? When I find myself in a situation where I'm spending a lot of time and energy trying to get a certain thing to work, I can sometimes do X instead for 80% of the benefit. Try to find something like that in this case. FWIW the TPU profiler was the first tool I reached for. I never got it working. The bag of tricks above ended up giving me effective results on a variety of codebases with no profiler. (A usage graph is pretty crucial, though, which Colab TPUs don't provide.) So there are a bunch of general tips for solving weird bottlenecks blindfolded. To answer your question directly: I am assuming the profiler from Tensorflow works with TPUs in the same way it does with GPUs? Not really. You're supposed to use cloud_tpu_profiler: https://cloud.google.com/tpu/docs/cloud-tpu-tools https://cloud.google.com/tpu/docs/cloud-tpu-tools But yeah, if you give specifics (ideally a link to a codebase + dataset + script that runs it) then I can try to look for candidates of what might be the bottleneck.