Imágenes de VM de TPU

Cuando creas recursos TPU, debes pasar el parámetro --version o --runtime-version, que especifica una imagen de VM de TPU. Las imágenes de VM de TPU contienen el sistema operativo (Ubuntu) y, opcionalmente, otro software necesario para ejecutar tu código en TPU. En este documento, se proporciona orientación para seleccionar la imagen de VM de TPU adecuada cuando creas Cloud TPU.

PyTorch y JAX

Usa las siguientes imágenes base de VM de TPU comunes para PyTorch y JAX y, luego, instala el framework que deseas usar.

  • tpu-ubuntu2204-base (predeterminado)
  • v2-alpha-tpuv5 (TPU v5p)

Consulta los documentos de inicio rápido de PyTorch/XLA y JAX para obtener las instrucciones de instalación.

TensorFlow

Hay imágenes de VM de TPU específicas para cada versión de TensorFlow. Las siguientes versiones de TensorFlow son compatibles con las Cloud TPU:

  • 2.16.1
  • 2.15.1
  • 2.15.0
  • 2.14.1
  • 2.14.0
  • 2.13.1
  • 2.13.0
  • 2.12.1
  • 2.12.0
  • 2.11.1
  • 2.11.0
  • 2.10.1
  • 2.10.0
  • 2.9.3
  • 2.9.1
  • 2.8.4
  • 2.8.3
  • 2.8.0
  • 2.7.4
  • 2.7.3

Para obtener más información sobre las versiones de parches de TensorFlow, consulta Versiones de parches compatibles de TensorFlow.

Para las versiones de TensorFlow 2.15.0 y posteriores, hay variantes de imágenes de VM de TPU basadas en la API del dispositivo (PJRT o ejecutor de transmisión) que usas.

Entrenamiento en v5p y v5e

TPU v5e y v5p son compatibles con TensorFlow 2.15.0 y versiones posteriores. Especifica la imagen de VM de TPU con el formato: tpu-vm-tf-x.y.z-{pod}-pjrt, en el que x es la versión principal de TensorFlow, y es la versión secundaria y z es la versión del parche de TensorFlow. Agrega pod después de la versión de TensorFlow si usas una TPU de hosts múltiples. Por ejemplo, si usas TensorFlow 2.16.0 en una TPU de varios hosts, usa la imagen de VM de TPU tpu-vm-tf-2.16.0-pod-pjrt. Para otras versiones de TensorFlow, reemplaza 2.16.0 por las versiones principal y de parche de TensorFlow que usas. Si usas una TPU de host único, omite pod.

Publicación en v5e

Hay imágenes de Docker de entrega que contienen todos los requisitos de software necesarios para entregar con TensorFlow, PyTorch y JAX. Para obtener más información, consulta Introducción a la inferencia v5e de Cloud TPU.

TPU v4

Si usas TPU v4 y TensorFlow 2.15.0 o una versión más reciente, sigue las instrucciones para entrenar con v5p y v5e. Si usas TensorFlow 2.10.0 o versiones anteriores, utiliza una imagen de VM de TPU específica de v4:

Versión de TensorFlow Versión de imagen de VM de TPU
2.10.0 tpu-vm-tf-2.10.0-v4
tpu-vm-tf-2.10.0-pod-v4
2.9.3 tpu-vm-tf-2.9.3-v4
tpu-vm-tf-2.9.3-pod-v4
2.9.2 tpu-vm-tf-2.9.2-v4
tpu-vm-tf-2.9.2-pod-v4
2.9.1 tpu-vm-tf-2.9.1-v4
tpu-vm-tf-2.9.1-pod-v4

TPU v2 y v3

Si usas TPU v2 o v3, usa la imagen de VM de TPU que coincida con la versión de TensorFlow que usas. Por ejemplo, si usas TensorFlow 2.14.1, utiliza la imagen de TPU tpu-vm-tf-2.14.1. Para otras versiones de TensorFlow, reemplaza 2.14.1 por la versión de TensorFlow que usas. Si usas un pod de TPU para varios hosts al final de la imagen de TPU, por ejemplo, tpu-vm-tf-2.14.1-pod.

A partir de TensorFlow 2.15.0, también debes especificar una API de dispositivo. Por ejemplo, si usas TensorFlow 2.16.1 con la API de PJRT, usa la imagen de TPU tpu-vm-tf-2.16.1-pjrt. Si usas la API del ejecutor de transmisiones con la misma versión de TensorFlow, usa la imagen de TPU tpu-vm-tf-2.16.1-se. Las versiones de TensorFlow anteriores a la 2.15.0 solo admiten el ejecutor de transmisiones.

Compatibilidad con PJRT de TensorFlow

A partir de TensorFlow 2.15.0, puedes usar la interfaz de PJRT para TensorFlow en TPU. PJRT cuenta con desfragmentación automática de la memoria del dispositivo y simplifica la integración del hardware con frameworks. Para obtener más información sobre PJRT, consulta PJRT: Simplifica la integración del framework y del hardware del AA.

Acelerador Atributo Compatibilidad con PJRT Compatibilidad con el ejecutor de transmisiones
TPU v2 - v4 Procesamiento denso (sin API de incorporación de TPU) Sí Sí
TPU v2 - v4 API de Dense Compute + TPU Embedding No Sí
TPU v2 - v4 tf.summary/tf.print con posición de dispositivo táctil No Sí
TPU v5e Procesamiento denso (sin API de incorporación de TPU) Sí No
TPU v5e API de TPU Embedding N/A No
TPU v5p Procesamiento denso (sin API de incorporación de TPU) Sí No
TPU v5p API de TPU Embedding Sí No

Versiones de Libtpu

Las imágenes de TensorFlow de VMs de TPU contienen una versión específica de TensorFlow y la biblioteca libtpu correspondiente. Si creas tu propia imagen de VM, usa las siguientes versiones de software de TensorFlow TPU y las versiones de libtpu correspondientes:

Versión de TensorFlow Versión de libtpu.so
2.16.1 1.10.1
2.15.1 1.9.0
2.15.0 1.9.0
2.14.1 1.8.1
2.14.0 1.8.0
2.13.1 1.7.1
2.13.0 1.7.0
2.12.1 1.6.1
2.12.0 1.6.0
2.11.1 1.5.1
2.11.0 1.5.0
2.10.1 1.4.1
2.10.0 1.4.0
2.9.3 1.3.2
2.9.1 1.3.0
2.8.3 1.2.3
2.8.0 1.2.0
2.7.3 1.1.2

¿Qué sigue?