在 GCP 上的 Cloud TPU VM 上运行 Pytorch 会给出 INVALID_ARGUMENT: No matching devices found for \'/job:localservice/replica:0/task:0/device:TPU_SYSTEM:0\'

Question

我创建了一个TPU VM on GCP。

我正在关注有关如何 run a calculation on a Cloud TPU VM by using PyTorch 的文档页面

我已经在 VM 中设置了 XRT TPU 设备配置

export XRT_TPU_CONFIG="localservice;0;localhost:51011"

我创建了一个包含以下内容的 Python 文件：

import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

但是当我运行文件 python3 tpu-test.py 时，出现以下错误：

$ python3 tpu_test.py 
Traceback (most recent call last):
  File "tpu_test.py", line 6, in <module>
    dev = xm.xla_device()
  File "/usr/local/lib/python3.8/dist-packages/torch_xla/core/xla_model.py", line 244, in xla_device
    devices = get_xla_supported_devices(devkind=devkind)
  File "/usr/local/lib/python3.8/dist-packages/torch_xla/core/xla_model.py", line 138, in get_xla_supported_devices
    xla_devices = _DEVICES.value
  File "/usr/local/lib/python3.8/dist-packages/torch_xla/utils/utils.py", line 32, in value
    self._value = self._gen_fn()
  File "/usr/local/lib/python3.8/dist-packages/torch_xla/core/xla_model.py", line 20, in <lambda>
    _DEVICES = xu.LazyProperty(lambda: torch_xla._XLAC._xla_get_devices())
RuntimeError: tensorflow/compiler/xla/xla_client/xrt_computation_client.cc:1374 : Check failed: session.Run({tensorflow::Output(result, 0)}, &outputs) == ::tensorflow::Status::OK() (INVALID_ARGUMENT: No matching devices found for '/job:localservice/replica:0/task:0/device:TPU_SYSTEM:0' vs. OK)
*** Begin stack trace ***
    tensorflow::CurrentStackTrace[abi:cxx11]()
    xla::XrtComputationClient::InitializeAndFetchTopology(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, int, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, tensorflow::ConfigProto const&)
    xla::XrtComputationClient::InitializeDevices(std::unique_ptr<tensorflow::tpu::TopologyProto, std::default_delete<tensorflow::tpu::TopologyProto> >)
    xla::XrtComputationClient::XrtComputationClient(xla::XrtComputationClient::Options, std::unique_ptr<tensorflow::tpu::TopologyProto, std::default_delete<tensorflow::tpu::TopologyProto> >)
    xla::ComputationClient::Create()
    
    
    xla::ComputationClient::Get()
    
    
    PyCFunction_Call
    _PyObject_MakeTpCall
    _PyEval_EvalFrameDefault
    _PyFunction_Vectorcall
    _PyEval_EvalFrameDefault
    
    PyObject_GetAttr
    _PyEval_EvalFrameDefault
    _PyEval_EvalCodeWithName
    _PyFunction_Vectorcall
    _PyEval_EvalFrameDefault
    
    _PyEval_EvalFrameDefault
    _PyEval_EvalCodeWithName
    PyEval_EvalCode
    
    
    
    PyRun_SimpleFileExFlags
    Py_RunMain
    Py_BytesMain
    __libc_start_main
    _start
*** End stack trace ***

Answer 1

尝试使用以下标志

export TPU_NUM_DEVICES=[逻辑核心数]

例如：在 v3-8 上，

导出 TPU_NUM_DEVICES=8