Versiones de software de TPU

En este documento se ofrecen directrices para seleccionar la versión de software de TPU adecuada al crear TPUs de Cloud.

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

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

PyTorch y JAX

Usa las siguientes versiones de software comunes de TPU para PyTorch y JAX. A continuación, instala el framework que quieras 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
Versión 4 y anteriores tpu-ubuntu2204-base

Para obtener más información sobre la instalación y cómo empezar a usar PyTorch o JAX, consulta los artículos Ejecutar un cálculo en una VM de TPU de Cloud con PyTorch y Ejecutar un cálculo en una VM de TPU de Cloud con JAX.

TensorFlow

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

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

  • x.y.z: representa las versiones principal, secundaria y de parche de TensorFlow.
  • -pod (opcional): indica que estás usando un slice de TPU de varios hosts.
  • -{device_api} (opcional): especifica la API del dispositivo. Por ejemplo, -pjrt (si usas la API 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 admitidas 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 parche de TensorFlow, consulta Versiones de parche de TensorFlow compatibles.

TPU v6e, v5p y v5e

Las versiones v6e, v5e y v5p de TPU son compatibles con TensorFlow 2.15.0 y versiones posteriores. Para especificar la versión del software de TPU, usa el siguiente formato: tpu-vm-tf-x.y.z-{pod}-pjrt, donde x es la versión principal de TensorFlow, y es la versión secundaria y z es la versión del parche de TensorFlow. Añade 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 TPU tpu-vm-tf-2.16.0-pod-pjrt. En otras versiones de TensorFlow, sustituye 2.16.0 por las versiones principal y de parche de TensorFlow que estés usando. Si usas una TPU de un solo host, omite pod.

TPU v4

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

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 la versión 2 o 3 de las TPU, utiliza la versión del software de TPU que coincida con la versión de TensorFlow que estés usando. Por ejemplo, si usas TensorFlow 2.14.1, utiliza la versión de software de TPU tpu-vm-tf-2.14.1. Para otras versiones de TensorFlow, sustituye 2.14.1 por la versión de TensorFlow que estés usando. Si usas una TPU de varios hosts, añade "pod" al final de la versión del 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 del software. Por ejemplo, si usas TensorFlow 2.16.1 con la API de PJRT, usa la versión de software de TPU tpu-vm-tf-2.16.1-pjrt. Si usas la API del ejecutor de flujo 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 la 2.15.0 solo admiten stream_executor.

Compatibilidad con PJRT de TensorFlow

A partir de TensorFlow 2.15.0, puedes usar la interfaz PJRT para TensorFlow en TPU. PJRT incluye la desfragmentación automática de la memoria del dispositivo y simplifica la integración del hardware con los frameworks. Para obtener más información sobre PJRT, consulta PJRT: Simplifying ML Hardware and Framework Integration (PJRT: simplificación de la integración de hardware y frameworks de aprendizaje automático).

Accelerator Función Compatibilidad con PJRT Compatibilidad con el ejecutor de la secuencia
TPU v2 - v4 Computación densa (sin API de inserción de TPU) Sí Sí
TPU v2 - v4 API de computación densa + API de inserción de TPU No Sí
TPU v2 - v4 tf.summary/tf.print con colocación de dispositivos flexible No Sí
TPU v5e Computación densa (sin API de inserción de TPU) Sí No
TPU v5e API de inserción de TPU N/A No
TPU v5p Computación densa (sin API de inserción de TPU) Sí No
TPU v5p API de inserción de TPU Sí No

Siguientes pasos