Versiones de software de TPU

En este documento, se proporciona orientación para seleccionar la versión correcta del software de TPU cuando creas Cloud TPU.

Cuando creas recursos de TPU, especificas la versión de software, también llamada versión del entorno de ejecución, que hace referencia al entorno de software preinstalado en tu VM de TPU. Esto incluye el sistema operativo Ubuntu, Docker y otro software necesario para ejecutar tu código en TPU.

Si usas Google Cloud CLI, especificas la versión del software de la TPU con el parámetro --version o --runtime-version. Si usas la consola de Google Cloud , selecciona una versión del software de TPU de la lista Versión del software de TPU.

PyTorch y JAX

Usa las siguientes versiones comunes de software de TPU para PyTorch y JAX y, luego, instala el framework que deseas usar.

Versión de TPU Versión de software de TPU
Trillium (v6e) v2-alpha-tpuv6e
v5p v2-alpha-tpuv5
v5e v2-alpha-tpuv5-lite
v4 y versiones anteriores tpu-ubuntu2204-base

Para obtener más información sobre la instalación y cómo comenzar a usar PyTorch o JAX, consulta Cómo ejecutar un cálculo en una VM de Cloud TPU con PyTorch y Cómo ejecutar un cálculo en una VM de Cloud TPU con JAX.

TensorFlow

Las versiones de software de TPU para TensorFlow siguen una convención de nombres específica:

tpu-vm-tf-x.y.z[-{pod}][-{device_api}]

  • x.y.z: Representa las versiones principales, secundarias y de parche de TensorFlow.
  • -pod (opcional): Indica que usas una porción de TPU de varios hosts.
  • -{device_api} (opcional): Especifica la API del dispositivo, por ejemplo, -pjrt (si usas la API de PJRT).

Consulta las siguientes secciones para obtener más información sobre cómo especificar una versión de software de TPU.

Hay versiones de software de TPU específicas para cada versión de TensorFlow. En la siguiente tabla, se muestran las versiones de TensorFlow compatibles y las versiones de libtpu asociadas:

Versión de TensorFlow Versión de libtpu.so
2.18.0 1.12.0
2.17.1 1.11.1
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

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

TPU v6e, v5p y v5e

TPU v6e, v5e y v5p admiten TensorFlow 2.15.0 y versiones posteriores. Especifica la versión de software de la 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 versión de software de la 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, omite pod.

TPU v4

Si usas TPU v4 y TensorFlow 2.10.1 o una versión posterior, sigue las instrucciones para TPU v2 y v3. Si usas TensorFlow 2.10.0 o una versión anterior, usa una versión de software de TPU específica de la v4:

Versión de TensorFlow Versión de software 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 versión de software de TPU que coincida con la versión de TensorFlow que usas. Por ejemplo, si usas TensorFlow 2.14.1, usa la versión de software de la 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 una TPU de varios hosts, agrega pod al final de la versión de software de la 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 versión de software. Por ejemplo, si usas TensorFlow 2.16.1 con la API de PJRT, usa la versión de software de la TPU tpu-vm-tf-2.16.1-pjrt. Si usas la API del ejecutor de flujos con la misma versión de TensorFlow, usa la versión de software de TPU tpu-vm-tf-2.16.1-se. Las versiones de TensorFlow anteriores a 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 IA.

Accelerator Función Compatibilidad con PJRT Compatibilidad con el ejecutor de transmisión
TPU v2 a v4 Procesamiento intensivo (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 intensivo (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

¿Qué sigue?