No TPU Devices found for TPU Pod Training, even after recognizing it

[kaushikb11]

🐛 Bug

No TPU Devices found for TPU Pod Training, even after recognizing it.

2021-04-23 18:14:43 10.164.0.29 [3]   File "/home/kaushikbokka/pytorch-lightning/pytorch_lightning/trainer/trainer.py", line 316, in __init__
2021-04-23 18:14:43 10.164.0.29 [3]     replace_sampler_ddp, deterministic, precision, amp_backend, amp_level, plugins
2021-04-23 18:14:43 10.164.0.29 [3]   File "/home/kaushikbokka/pytorch-lightning/pytorch_lightning/trainer/connectors/accelerator_connector.py", line 97, in __init__
2021-04-23 18:14:43 10.164.0.29 [3]     self.tpu_cores = device_parser.parse_tpu_cores(tpu_cores)
2021-04-23 18:14:43 10.164.0.29 [3]   File "/home/kaushikbokka/pytorch-lightning/pytorch_lightning/utilities/device_parser.py", line 104, in parse_tpu_cores
2021-04-23 18:14:43 10.164.0.29 [3]     raise MisconfigurationException('No TPU devices were found.')
2021-04-23 18:14:43 10.164.0.29 [3] pytorch_lightning.utilities.exceptions.MisconfigurationException: No TPU devices were found.
2021-04-23 18:14:46 10.164.0.27 [0] GPU available: False, used: False
2021-04-23 18:14:46 10.164.0.27 [0] TPU available: True, using: 8 TPU cores

TPU Pod v2-32, running bug_report_model.py on 4 VMs with 8 tpu cores each.

[Borda]

@kaushikb11 is this on master or on past v1.2.x release?

[kaushikb11]

@Borda It's everywhere. Single Cloud TPU (v2-8 or v3-8) + VM works well and is supported. I don't think Cloud TPU Pod has been supported well enough. Ref: pytorch-lightning.readthedocs.io/en/stable/advanced/tpu.html#tpu-pod, cloud.google.com/tpu/docs/tutorials/pytorch-pod.

[kaushikb11]

On non-master VMs

Traceback (most recent call last):
   File "/home/kaushikbokka/pytorch-lightning/pytorch_lightning/utilities/xla_device.py", line 31, in inner_f
     queue.put(func(*args, **kwargs))
   File "/home/kaushikbokka/pytorch-lightning/pytorch_lightning/utilities/xla_device.py", line 70, in _is_device_tpu
     return len(xm.get_xla_supported_devices("TPU")) > 0
   File "/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 136, in get_xla_supported_devices
     xla_devices = _DEVICES.value
   File "/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch_xla/utils/utils.py", line 32, in value
     self._value = self._gen_fn()
   File "/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch_xla/core/xla_model.py", line 18, in <lambda>
     _DEVICES = xu.LazyProperty(lambda: torch_xla._XLAC._xla_get_devices())
 RuntimeError: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:258 : Check failed: default_device_target != options_.global_device_map.end()

[kaushikb11]

Some start...