Ошибка выполнения TPU V4-64: не удалось инициализировать TPU: не удалось установить канал grpc SliceBuilder ⇐ Python

[Anonymous]

Во время исследовательской программы TPU я пытался использовать TPU V4-64, поскольку у меня есть 32 бесплатных чипа TPU V4, которые можно заказать по запросу.
Однако, в отличие от TPU V4-8, тестовые коды, представленные в руководстве, не работали каждый раз, когда я использовал TPU V4-16 ~ TPU V4-64. Не только тестовые коды, но и мои примеры кодов тоже не работали.
Руководство по GCP TPU: https://cloud.google.com/tpu/docs/tutor ... et-pytorch< /p>
Я следую инструкциям по настройке torch_xla, приведенным выше.
[TRIAL 1]

• Создать виртуальную машину TPU

gcloud compute tpus queued-resources create RESOURCE_NAME \
--node-id NODE_NAME \
--project PROJECT_NAME \
--zone us-central2-b \
--accelerator-type v4-64 \
--runtime-version tpu-ubuntu2204-base

• Подключитесь к виртуальной машине TPU

gcloud compute tpus tpu-vm ssh NODE_NAME --zone=us-central2-b

• Установить фонарик/ torch_xla

pip install torch~=2.3.0 torch_xla[tpu]~=2.3.0 torchvision -f https://storage.googleapis.com/libtpu-r ... index.html

• Клонировать официальный репозиторий torch_xla на github

git clone --depth=1 --branch r2.3 https://github.com/pytorch/xla.git

• Выполните тестовый код № обучения ResNet

PJRT_DEVICE=TPU python3 xla/test/test_train_mp_imagenet.py --fake_data --batch_size=256 --num_epochs=1

[РЕЗУЛЬТАТ 1] Затем я получил приведенные ниже сообщения журнала без какого-либо прогресса в течение длительного времени. Когда я использовал TPU V4-8, обучающие коды работали правильно (= использование полного xla::0 ~ xla::3 и потеря обучения/теста при печати)
WARNING:root:PJRT is now the default runtime. For more information, see https://github.com/pytorch/xla/blob/master/docs/pjrt.md
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1720161247.044048 21553 pjrt_api.cc:100] GetPjrtApi was found for tpu at /home/hungwon3626/.local/lib/python3.10/site-packages/libtpu/libtpu.so
I0000 00:00:1720161247.044129 21553 pjrt_api.cc:79] PJRT_Api is set for device type tpu
I0000 00:00:1720161247.044136 21553 pjrt_api.cc:146] The PJRT plugin has PJRT API version 0.46. The framework PJRT API version is 0.46.

...

[TRIAL 2] Я приостановил выполнение, а затем запустил приведенные ниже коды:

>>> import os
>>> os.environ['TPU_NUM_DEVICES'] = '32'
# I also retried the same codes without the above lines, and got the same result.
>>> import torch_xla as xla
>>> xla.device()

[РЕЗУЛЬТАТ 2]: Затем я получил следующие сообщения об ошибках:
WARNING:root:PJRT is now the default runtime. For more information, see https://github.com/pytorch/xla/blob/master/docs/pjrt.md
WARNING:root:libtpu.so and TPU device found. Setting PJRT_DEVICE=TPU.
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1720161247.044048 21553 pjrt_api.cc:100] GetPjrtApi was found for tpu at /home/hungwon3626/.local/lib/python3.10/site-packages/libtpu/libtpu.so
I0000 00:00:1720161247.044129 21553 pjrt_api.cc:79] PJRT_Api is set for device type tpu
I0000 00:00:1720161247.044136 21553 pjrt_api.cc:146] The PJRT plugin has PJRT API version 0.46. The framework PJRT API version is 0.46.
Traceback (most recent call last):
File "", line 1, in
File "/home/hungwon3626/.local/lib/python3.10/site-packages/torch_xla/torch_xla.py", line 21, in device
return xm.xla_device(index)
File "/home/hungwon3626/.local/lib/python3.10/site-packages/torch_xla/core/xla_model.py", line 212, in xla_device
return runtime.xla_device(n, devkind)
File "/home/hungwon3626/.local/lib/python3.10/site-packages/torch_xla/runtime.py", line 95, in wrapper
return fn(*args, **kwargs)
File "/home/hungwon3626/.local/lib/python3.10/site-packages/torch_xla/runtime.py", line 124, in xla_device
return torch.device(torch_xla._XLAC._xla_get_default_device())
RuntimeError: Bad StatusOr access: UNKNOWN: TPU initialization failed: Failed to establish SliceBuilder grpc channel to 10.130.15.203:8471.

[ОЖИДАНИЕ] Я хочу использовать полные 32 чипа для обучения одной большой модели на устройствах xla.
[версия пакета]
torch 2.3.1
torch-xla 2.3.0
torchvision 0.18.1
Python 3.10.6


Подробнее здесь: https://stackoverflow.com/questions/787 ... h-slicebui