V6e TPU-VM搭配TensorFlow无法初始化及识别设备求助

问题详情 #

尝试使用TPU训练模型时，TPU-VM始终无法找到或识别V6e TPU单元。

创建TPU-VM的基础命令为：

gcloud compute tpus tpu-vm create test123 --zone=us-central1-b --version=tpu-vm-tf-2.{x}.{y}-pjrt --accelerator-type=v6e-1 --spot

已测试过以下多种配置组合：

--version=tpu-vm-tf-2.16.1-pjrt --accelerator-type=v6e-1     
--version=tpu-vm-tf-2.16.2-pjrt --accelerator-type=v6e-1    
--version=tpu-vm-tf-2.16.1-pod-pjrt --accelerator-type=v6e-16    
--version=tpu-vm-tf-2.19.0-pod-pjrt --accelerator-type=v6e-16    

这些配置在v5litepod-1和v5litepod-8上可正常工作，但在V6e上完全失效，具体表现为TPU-VM初始化异常，找不到libtpu、libtpu.so文件，也无法识别任何TPU设备。

此外，tf-2.17至tf-2.18版本存在未解决的问题，只要导入tensorflow就会导致虚拟机崩溃。

内容的提问来源于stack exchange，提问作者Prz00124