tl;dr — щракнете тук, за да видите Github Repo.

На 1 юни 2021 г. Google най-накрая пусна „TPU VM за обществеността“, след като тихо обяви частната Alpha в края на миналата година.

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

Това означаваше, че трябва да настроите отделна виртуална машина, която след това комуникира с TPU хоста, което често въвежда някои проблеми, които бихте могли да имате в противен случай с обучението на GPU.

Някои примери включват:

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

Докато TPU VM решават повечетоот тези проблеми, тъй като е толкова нов, документацията за конкретни случаи на употреба е оскъдна, ако изобщо съществува.

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

Така че, разбира се, исках да видя как да имам достъп до TPU директно през докер контейнер. За разлика от TPU възлите, свързването към TPU в TPU VM е много по-трудно, тъй като не можете да разчитате просто на настройка

os.environ[‘TPU_NAME’] = ‘вашият-tpu’

повярвайте ми, и аз опитах това.

Ако искате кратката версия, просто щракнете върху връзката Github repo и вижте сами файла Docker Compose.

Неща, които опитах и ​​които не проработиха:

  • Задаване на целта на задния край на Jax да бъде вътрешният IP адрес на TPU VM

  • Опитвайки се да прикача необичайните устройства, които намерих в /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 VM във вашия Docker контейнер:

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

Последна забележка: Вероятно искате VPC и прокси

За разлика от обикновените VM в рамките на GCP, TPU VM не ви позволяват да променяте настройките на защитната стена, което означава, че публичните портове не са достъпни извън GCP мрежата.

Така че, за съжаление, не можете просто просто да добавите външния IP на TPU VM към вашия DNS и да го наречете ден.

Това, което в крайна сметка направих, беше да пусна друга виртуална машина като прокси, използвайки nginx контейнер, който сочи към вътрешния IP адрес на TPU VM.

И VM, и TPU VM бяха в една и съща VPC мрежа, така че не съм сигурен как може да работи това, ако не са.

Специални благодарности на TRC за предоставяне на достъп до TPU и TPU VM