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.
Что я пробовал, но не сработало:
- Настройка резервирования ресурсов согласно Docker-Compose’s Spec
- Установка цели бэкэнда 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