Imágenes de VM de TPU

Cuando creas recursos de TPU, pasas 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, Docker y 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 comunes de VM de TPU para PyTorch y JAX y, luego, instala el framework que quieras usar.

  • tpu-ubuntu2204-base (predeterminada)
  • v2-alpha-tpuv6e (TPU v6e)
  • v2-alpha-tpuv5 (TPU v5p)
  • v2-alpha-tpuv5-lite (TPU v5e)

Consulta los documentos de la guía de inicio rápido de PyTorch/XLA y JAX para obtener 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 Cloud TPU:

  • 2.18.0
  • 2.17.0
  • 2.16.2
  • 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 de TensorFlow compatibles.

En el caso de las versiones 2.15.0 y posteriores de TensorFlow, existen variantes de imágenes de VM de TPU según la API del dispositivo (PJRT o ejecutor de flujo) que usas.

Entrenamiento en v6e, v5p y v5e

TPU v6e, v5e y v5p admiten TensorFlow 2.15.0 y versiones posteriores. Especifica la imagen de la 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 varios hosts. 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 principales y de parches de TensorFlow que usas. Si usas una TPU de host único, omítelo.pod

Publicación en v6e y v5e

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

TPU v4

Si usas TPU v4 y TensorFlow 2.15.0 o versiones posteriores, sigue las instrucciones para el entrenamiento en v6e, v5p y v5e. Si usas TensorFlow 2.10.0 o una versión anterior, usa una imagen de VM de TPU específica para v4:

Versión de TensorFlow Versión de la imagen de la 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, usa 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 de varios hosts, agrégalo 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 como parte del nombre de la imagen. 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 transmisión 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 flujo.

Compatibilidad con PJRT de TensorFlow

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

Accelerator Atributo Compatibilidad con PJRT Compatibilidad con el ejecutor de transmisión
TPU v2 a v4 Procesamiento denso (sin API de incorporación de TPU) Sí Sí
TPU v2 a v4 API de procesamiento intensivo + API de incorporación de TPU No Sí
TPU v2 a v4 tf.summary/tf.print con ubicación de dispositivos flexible No Sí
TPU v5e Procesamiento denso (sin API de incorporación de TPU) Sí No
TPU v5e API de incorporación de TPU N/A No
TPU v5p Procesamiento intensivo (sin API de incorporación de TPU) Sí No
TPU v5p API de incorporación de TPU Sí No

Versiones de libtpu

Las imágenes de TensorFlow de las 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 correspondientes de libtpu:

Versión de TensorFlow Versión de libtpu.so
2.18.0 1.12.0
2.17.0 1.11.0
2.16.2 1.10.1
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?