TPU Spawn Rank & root device Error (#7074)

* TPU Spawn Rank Error

* Update tpu spawn

* Fix root device property for tpu spawn

* Update changelog

CHANGELOG.md

@@ -285,6 +285,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed process rank not being available right away after `Trainer` instantiation ([#6941](https://github.com/PyTorchLightning/pytorch-lightning/pull/6941))
 
 
+- Fixed the order to call for world ranks & the `root_device` property in `TPUSpawnPlugin` ([#7074](https://github.com/PyTorchLightning/pytorch-lightning/pull/7074))
+
+
 ## [1.2.7] - 2021-04-06
 
 ### Fixed

pytorch_lightning/plugins/training_type/tpu_spawn.py

@@ -65,6 +65,10 @@ def local_rank(self) -> int:
     def world_size(self) -> int:
         return self.num_processes
 
+    @property
+    def root_device(self) -> torch.device:
+        return self.device
+
     @staticmethod
     def _validate_dataloader(dataloaders: Union[List['DataLoader'], 'DataLoader']):
         if not isinstance(dataloaders, list):
@@ -116,9 +120,7 @@ def is_distributed(self):
 
     def process_dataloader(self, dataloader: 'DataLoader') -> MpDeviceLoader:
         TPUSpawnPlugin._validate_dataloader(dataloader)
-        device = xm.xla_device()
-        dataloader = MpDeviceLoader(dataloader, device)
-        return dataloader
+        return MpDeviceLoader(dataloader, self.device)
 
     def configure_ddp(self) -> None:
         pass
@@ -127,8 +129,7 @@ def init_ddp_connection(self, global_rank: int, world_size: int) -> None:
         pass
 
     def set_world_ranks(self, process_idx: int = 0) -> None:
-        self.tpu_local_core_rank = xm.get_local_ordinal()
-        self.tpu_global_core_rank = xm.get_ordinal()
+        pass
 
     def new_process(self, process_idx: int, trainer, mp_queue) -> None:
         self.mp_queue = mp_queue
@@ -137,7 +138,8 @@ def new_process(self, process_idx: int, trainer, mp_queue) -> None:
         if seed is not None:
             seed_everything(int(seed))
 
-        self.set_world_ranks()
+        self.tpu_local_core_rank = xm.get_local_ordinal()
+        self.tpu_global_core_rank = xm.get_ordinal()
 
         # set warning rank
         rank_zero_only.rank = self.global_rank
@@ -163,7 +165,8 @@ def new_process(self, process_idx: int, trainer, mp_queue) -> None:
             time.sleep(2)
 
     def model_to_device(self) -> None:
-        self._model.to(xm.xla_device())
+        self.device = xm.xla_device()
+        self.model.to(self.device)
 
     def barrier(self, name: Optional[str] = None) -> None:
         rendezvous(name)