Versiones de software de TPU
En este documento, se proporciona orientación para seleccionar la versión de software de TPU adecuada cuando creas Cloud TPU.
Cuando creas recursos TPU, especificas la versión de software, también llamada versión de 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 las TPU.
Si usas Google Cloud CLI, especifica la versión de software
de la TPU con el parámetro --version o --runtime-version. Si usas
la consola de Google Cloud , selecciona una versión de software de TPU en la lista de versión de software
de TPU.
PyTorch y JAX
Usa las siguientes versiones de software comunes de TPU para PyTorch y JAX; luego, instala el framework que desees 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 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 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.
Existen versiones de software de TPU específicas para cada versión de TensorFlow. En la siguiente tabla, se muestran las versiones compatibles de TensorFlow 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 TPU v6e, v5e y v5p son compatibles con TensorFlow 2.15.0 y versiones posteriores. Especifica
la versión de software de la TPU con el siguiente formato: tpu-vm-tf-x.y.z-{pod}-pjrt, en el cual 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
TPU tpu-vm-tf-2.16.0-pod-pjrt. Para otras versiones de
TensorFlow, reemplaza 2.16.0 por la versión principal y de parche de
TensorFlow que estés usando. 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 para 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 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 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 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 el ejecutor de
transmisión.
Compatibilidad con PJRT de TensorFlow
A partir de TensorFlow 2.15.0, puedes usar la interfaz de 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: Simplifica la integración de hardware y frameworks de AA.
| Acelerador | Función | Compatibilidad con PJRT | Compatibilidad con el ejecutor de transmisión |
|---|---|---|---|
| TPU v2-v4 | Procesamiento denso (sin la API de embedding de TPU) | Sí | Sí |
| TPU v2-v4 | API de procesamiento denso y API de embedding de TPU | No | Sí |
| TPU v2-v4 | tf.summary/tf.print con colocación flexible del dispositivo | No | Sí |
| TPU v5e | Procesamiento denso (sin la API de embedding de TPU) | Sí | No |
| TPU v5e | API de embedding de TPU | N/A | No |
| TPU v5p | Procesamiento denso (sin la API de embedding de TPU) | Sí | No |
| TPU v5p | API de embedding de TPU | Sí | No |
¿Qué sigue?
- Consulta Arquitectura de TPU para obtener más información sobre la arquitectura de TPU.
- Consulta Cuándo usar las TPU para obtener información sobre los tipos de modelos que funcionan con Cloud TPU.