Loss does not decrease on single TPU

Nobody has claimed this yet.

https://github.com/pytorch/xla/issues/2735#issue-786787271

> ## 🐛 Bug
> With pytorch-lightning, the trainer hangs up at the end of epoch. By keyboard interruption, it seems that it hangs up at fd_event_list = self._poll.poll(timeout)
>
> ```
> ---------------------------------------------------------------------------
> KeyboardInterrupt Traceback (most recent call last)
> in ()
> 4 autoencoder_1 = LitAutoEncoder()
> 5 trainer_1 = pl.Trainer(tpu_cores=8, max_epochs=1, progress_bar_refresh_rate=20)
> ----> 6 trainer_1.fit(autoencoder_1, DataLoader(dataset_mnist_train))
>
> 6 frames
> /usr/local/lib/python3.6/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloader, val_dataloaders, datamodule)
> 471 self.call_hook('on_fit_start')
> 472
> --> 473 results = self.accelerator_backend.train()
> 474 self.accelerator_backend.teardown()
> 475
>
> /usr/local/lib/python3.6/dist-packages/pytorch_lightning/accelerators/tpu_accelerator.py in train(self)
> 111 args=(model, self.trainer, self.mp_queue),
> 112 nprocs=self.trainer.tpu_cores,
> --> 113 start_method=self.start_method
> 114 )
> 115
>
> /usr/local/lib/python3.6/dist-packages/torch_xla/distributed/xla_multiprocessing.py in spawn(fn, args, nprocs, join, daemon, start_method)
> 393 join=join,
> 394 daemon=daemon,
> --> 395 start_method=start_method)
> 396
> 397
>
> /usr/local/lib/python3.6/dist-packages/torch/multiprocessing/spawn.py in start_processes(fn, args, nprocs, join, daemon, start_method)
> 155
> 156 # Loop on join until it returns True or raises an exception.
> --> 157 while not context.join():
> 158 pass
> 159
>
> /usr/local/lib/python3.6/dist-packages/torch/multiprocessing/spawn.py in join(self, timeout)
> 75 ready = multiprocessing.connection.wait(
> 76 self.sentinels.keys(),
> ---> 77 timeout=timeout,
> 78 )
> 79
>
> /usr/lib/python3.6/multiprocessing/connection.py in wait(object_list, timeout)
> 909
> 910 while True:
> --> 911 ready = selector.select(timeout)
> 912 if ready:
> 913 return [key.fileobj for (key, events) in ready]
>
> /usr/lib/python3.6/selectors.py in select(self, timeout)
> 374 ready = []
> 375 try:
> --> 376 fd_event_list = self._poll.poll(timeout)
> 377 except InterruptedError:
> 378 return ready
>
> KeyboardInterrupt:
> ```
>
> ## To Reproduce
> Steps to reproduce the behavior:
>
> 1. Run [this Colab](https://colab.research.google.com/drive/1qRnR8wUbB3HT2_AUo2pHxDBKe8lP0QQF?usp=sharing) (or [another Colab](https://colab.research.google.com/drive/1LJU5mtBAjNP3rUgMuKRMGOKGAqgtqIub?usp=sharing) for issue reporting) with TPU.
> 2. Wait until the first (and the last) epoch finishes. (all steps are completed but it keeps running)
> 3. Interrupt kernel
>
> ## Expected behavior
> The training should be completed immediately after all steps are completed.
>
> ## Environment
> * Reproducible on XLA backend [CPU/TPU]: TPU
> * torch_xla version: 1.7 and nightly

cc @carmocca @JackCaoG @Liyang90 @gkroiz @kaushikb11 @rohitgr7 @akihironitta

Start with the linked Colab reproducer and trace the TPU training path through tpu_accelerator.py and xla_multiprocessing.py, especially the process-join behavior shown in the traceback. Compare the end-of-epoch lifecycle with the expected completion behavior and verify any change by rerunning the one-epoch TPU reproduction.

Written by the indexing model from the issue text.