tl;dr — нажмите здесь, чтобы просмотреть репозиторий Github.

1 июня 2021 года Google, наконец, выпустила виртуальные машины TPU для общего доступа после того, как в конце прошлого года незаметно объявила о частной альфа-версии.

Основное ключевое различие между виртуальными машинами TPU и предыдущим способом использования TPU (то есть узлов TPU) заключалось в том, что у вас никогда не было прямого доступа к самому хосту TPU, и связь должна была осуществляться через grpc.

Это означало, что вам нужно было настроить отдельную виртуальную машину, которая затем взаимодействовала с хостом TPU, что часто приводило к некоторым проблемам, которые в противном случае могли бы возникнуть при обучении на графических процессорах.

Некоторые примеры:

  • Невозможно получить доступ к локальным файлам на диске
  • Узкое место, вызванное задержкой сети TPU ‹-› VM
  • Потенциально непредвиденные затраты на облачное хранилище из-за необходимости хранить все в GCS и если корзина не находится в том же регионе, что и TPU/VM.
  • Неожиданные ошибки TPU, за которыми часто следовали загадочные отладочные сообщения.

В то время как виртуальные машины TPU решают большую часть этих проблем, документация о конкретных вариантах использования скудна, если вообще существует.

Например, в Growth Engine мы используем TPU для экспериментов с крупномасштабными моделями NLP, и мы разработали систему, в которой мы контейнеризуем обучение модели на одном хосте виртуальной машины, подключенном к нескольким узлам TPU, выполняющим обучение отдельных заданий.

Поэтому, конечно, я хотел увидеть, как получить доступ к TPU напрямую через контейнер докеров. В отличие от узлов TPU, подключение к TPU в виртуальной машине TPU намного сложнее, поскольку вы не можете полагаться на простую настройку

os.environ[‘TPU_NAME’] = ‘ваш-тпу’

поверьте мне, я тоже так пробовал.

Если вам нужна короткая версия, просто нажмите на ссылку на репозиторий Github и сами посмотрите файл Docker Compose.

Что я пробовал, но не сработало:

  • Установка цели бэкэнда Jax в качестве внутреннего IP-адреса виртуальной машины TPU.

  • Попытка подключить необычные устройства, которые я нашел в /dev, которые выглядели так, как будто они могут быть TPU… становится теплее

  • Инициализация TF TPUClusterResolver до инициализации Jax. Он наполовину работал с подключенными выше устройствами, но имел другие ошибки.

… и множество других проб и ошибок, которые я не задокументировал после долгих часов тестирования различных конфигураций.

Изучив документацию TPU, исходный код Tensorflow и т. д., я наконец нашел ответ в исходном коде Jax. Конкретно этот крошечный кусочек.

Итак, что в конечном итоге сработало, так это комбинация:

  • Прикрепление «/dev:/dev» в качестве устройства
  • привилегированный: правда
  • Монтирование:
    — /var/run/docker.sock:/var/run/docker.sock (не уверен, что нужно)
    — / usr/share/tpu/:/usr/share/tpu/
    — /lib/libtpu.so:/lib/libtpu.so
  • Среды:
    — TPU_NAME=tpu_name
    — TF_CPP_MIN_LOG_LEVEL=0
    — XRT_TPU_CONFIG="localservice;0;localhost:51011' (вероятно, необходимо для Pytorch)
    — TF_XLA_FLAGS= — tf_xla_enable_xla_devices (Необходимо для TF/Jax)

Все вместе это выглядит примерно так:

И, наконец, вы можете получить доступ к TPU из виртуальной машины TPU в вашем контейнере Docker:

Надеюсь, что мои часы головной боли и отладки помогут нескольким людям, которые работают с виртуальными машинами TPU!

Последнее примечание: вам, вероятно, нужен VPC и прокси

В отличие от обычных виртуальных машин в GCP, виртуальные машины TPU не позволяют изменять настройки брандмауэра, а это означает, что общедоступные порты недоступны за пределами сети GCP.

Так что, к сожалению, вы не можете просто добавить внешний IP-адрес виртуальной машины TPU в свой DNS и покончить с этим.

В конце концов я запустил другую виртуальную машину в качестве прокси-сервера, используя контейнер nginx, который указывает на внутренний IP-адрес виртуальной машины TPU.

И виртуальная машина, и виртуальная машина TPU находились в одной сети VPC, поэтому я не уверен, как это могло бы работать, если бы это было не так.

Особая благодарность TRC за предоставление доступа к TPU и виртуальным машинам TPU