PyTorch Lightning TPU support

[williamFalcon]

❓ Questions and Help

Hi! I spoke with @dlibenzi at NeurIPS about adding TPU support to Lightning. He had mentioned 3-4 steps we needed to do to support it. Any chance someone can list those out again?

The API needs to be implemented in this file:
https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/distrib_data_parallel.py

[mruberry]

Hey @williamFalcon,

Are you not on Workplace chat anymore? I was looking for you there. How are we supposed to organize our foodie tour of NY (and possibly other places)?

As for PyTorch Lightning + TPUs it sounds cool but the PyTorch/XLA team is already pretty committed for H1. If it's going to happen I think you'll have to do it. Is that cool? Also, while we can respond to specific questions about using the PyTorch/XLA API, we don't typically work out how to use it from a particular system.

Have you check out our docs yet? The code to run on TPUs is very similar to the code to run on other device types.

[williamFalcon]

@mruberry haha. I'm back at NYU for a bit :)
I am on the pytorch slack haha.

I can do it, I just remember there were 3 minor changes I had to make but I forgot.
Let me give it a shot and ping you for a PR review?

[williamFalcon]

Integrated!
Check out the demo.

https://colab.research.google.com/drive/1-_LKx4HwAxl5M6xPJmqAAu444LTDQoa3#scrollTo=kr8cql-aaKnC

[dlibenzi]

Did you add the ParallelLoader already?

[williamFalcon]

yup!

would love a sanity check here if you have time.

Lightning-AI/pytorch-lightning#868

[dlibenzi]

Posted a comment:

https://github.com/PyTorchLightning/pytorch-lightning/pull/868/files#r382832281