Imágenes de VM de TPU
Cuando creas recursos TPU, debes pasar --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
y, opcionalmente, otro software necesario para ejecutar tu código en TPU. Este documento
proporciona orientación sobre cómo seleccionar la imagen de VM de TPU adecuada cuando se crea
Cloud TPU
PyTorch y JAX
Usa las siguientes imágenes base de VM de TPU comunes para PyTorch y JAX y, luego, instálalas el framework que quieres usar.
- tpu-ubuntu2204-base (predeterminado)
- v2-alpha-tpuv5 (TPU v5p)
- v2-alpha-tpuv5-lite (TPU v5e)
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. El las siguientes versiones de TensorFlow son compatibles con Cloud TPU:
- 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.
Para las versiones 2.15.0 y posteriores de TensorFlow, hay variantes de imágenes de VM de TPU según la API del dispositivo (PJRT o ejecutor de transmisión) que uses.
Entrenamiento en v5p y v5e
TPU v5e y v5p son compatibles con TensorFlow 2.15.0 y versiones posteriores. Especifica la VM de TPU
con el formato tpu-vm-tf-x.y.z-{pod}-pjrt
, en el que x
es el valor principal
versión de TensorFlow, y
es la versión secundaria y z
es la versión
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 la principal y la parche
versiones de TensorFlow que uses. Si usas una TPU de host único,
omitir pod
Publicación en v5e
Hay imágenes de Docker entregadas que contienen todos los requisitos de software necesarios para y entrega 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 posterior, sigue las Instrucciones de entrenamiento para v5p y v5e. Si utilizas TensorFlow 2.10.0 o versiones anteriores usa 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 estás usando. 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
usan. Si usas un Pod de adjuntos de TPU de múltiples hosts al final 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.
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 TPU tpu-vm-tf-2.16.1-se
imagen. 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 la desfragmentación automática de la memoria del dispositivo y simplifica la integración del hardware con los frameworks. Más información sobre PJRT, consulta PJRT: Simplificación de la integración del framework y del hardware del AA.
Accelerator | 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:
Versión de TensorFlow | Versión de libtpu.so |
---|---|
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?
- Obtén más información sobre la arquitectura de TPU en la sección Sistema Arquitectura.
- Consulta Cuándo utilizar las TPU para obtener más información de modelos que funcionan bien con Cloud TPU.