train with multiple tpus

[tweetyone]

I met ran out of memory problem.

Do you have any code to train on multiple TPUs? Thanks!

[jysohn23]

When you face device HBM OOM, try reducing the batch size first.

Training with multiple TPUs would not help in this case as long as we have the same per-core batch size, unless we use model parallelism, which we currently don't support. But we working to support distributed training on TPU Pods later this year.

[tweetyone]

Got, just my model is too big, to avoid the OOM error, the batch size can only be 16. Really hope you can support distributed training!
I also have another question: like the following image shows, why one step need to run on each xla device? why not just run one step on one core?

[dlibenzi]

This is not Intel x86 architecture ;)
Each core has its own memory (16GB in TPU v3), and the other cores cannot access it (or they could, at HUGE costs).
So you cannot run on a single core using the whole TPU combined memory.

[jysohn23]

I also have another question: like the following image shows, why one step need to run on each xla device? why not just run one step on one core?

I'm not sure I understand this question, but each TPU core (v3-8 has 8 TPU cores) is considered a single xla device so v3-8 would have xla:1 ... xla:8 devices.

Edit: oops looks like comments crossed, but basically what @dlibenzi said!

[tweetyone]

Sorry I didn't make my question clear enough. :)
My question is:
The 1～8 lines is the result of the same training step, right? So why we use 8 devices to run a same training step 8 times and get different loss results? Does this make the training process faster or you have another reason to do this?
Thanks!

[jysohn23]

The 1～8 lines is the result of the same training step, right? So why we use 8 devices to run a same training step 8 times and get different loss results? Does this make the training process faster or you have another reason to do this?

So each of the devices gets a different mini batch. So even though each of the devices has the exact same model weights (synchronous gradient descent) since the batches that each of the devices gets is different the losses are also different. So essentially during a single global step (where all 8 devices take a synchronized step) the global batch size is the --batch_size flag times 8. Is that a little more clear?

[tweetyone]

clear enough! thank you :)