在GCP上的云TPU虚拟机上运行Pytorch会产生INVALID_ARGUMENT:未找到与“/job:本地服务/副本:0/任务:0/设备:TPU_SYSTEM:0”匹配的设备

q3aa0525  于 2022-12-13  发布在  其他
关注(0)|答案(1)|浏览(132)

我创建了一个TPU VM on GCP
我正在查看有关如何run a calculation on a Cloud TPU VM by using PyTorch的文档页
我已经在虚拟机中设置了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 ***
shyt4zoc

shyt4zoc1#

尝试使用下方标志
导出TPU_NUM_DEVICES=[逻辑核心数]
例如:在V3-8上,
导出TPU设备数量=8

相关问题