Add support for TPU devices

[vfdev-5]

Clear and concise description of the problem

It would be good to provide an option to select accelerator as TPU instead of GPU
We can also auto-select TPU accelerator if open with Colab + add torch_xla installation steps.

What to do:
0) Try a template with TPUs. Choose distributed training option with 8 processes and spawning option. "Open in colab" one template, for example, vision classification template, install manually torch_xla (see https://colab.research.google.com/drive/1E9zJrptnLJ_PKhmaP5Vhb6DTVRvyrKHx) and run the code with backend xla-tpu: python main.py --nproc_per_node 8 --backend nccl. If everything is correctly done, training should probably run

1. Update UI

• Add a drop-out menu for backend selection: "nccl" and "xla-tpu" in "Training Options"
• when user selects "xla-tpu", training should be only distributed with 8 processes and "Run the training with torch.multiprocessing.spawn".

2. Update content: README.md and other impacted files
3. if exported to Colab, we need to make sure that accelerator is "TPU"

Suggested solution

Alternative

Additional context

[afzal442]

hi there, trying to refactor it with a little digging into this following this way https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/getting-started.ipynb

We can allow it with this nature;
cc @vfdev-5

[afzal442]

something like this as follow;

[afzal442]

// @ydcjeff

[vfdev-5]

@sayantan1410 I updated issue description adding few initial steps on how I would tackle this issue.

[sayantan1410]

@vfdev-5 I will start to work as per the description, and will let you know if I face some problem.

[sayantan1410]

0. Try a template with TPUs. Choose distributed training option with 8 processes and spawning option. "Open in colab" one template, for example, vision classification template, install manually torch_xla (see https://colab.research.google.com/drive/1E9zJrptnLJ_PKhmaP5Vhb6DTVRvyrKHx) and run the code with backend xla-tpu: python main.py --nproc_per_node 8 --backend nccl. If everything is correctly done, training should probably run

@vfdev-5 Hello, should i do this from the code-generator website or by running it locally or both works ?

[vfdev-5]

0. Try a template with TPUs. Choose distributed training option with 8 processes and spawning option. "Open in colab" one template, for example, vision classification template, install manually torch_xla (see https://colab.research.google.com/drive/1E9zJrptnLJ_PKhmaP5Vhb6DTVRvyrKHx) and run the code with backend xla-tpu: python main.py --nproc_per_node 8 --backend nccl. If everything is correctly done, training should probably run

@vfdev-5 Hello, should i do this from the code-generator website or by running it locally or both works ?

From code-generator and exporting in colab as I explained. You can't test that locally as you should use TPUs.