Configuraciones de TPU compatibles

A fin de determinar la configuración de TPU más eficaz para entrenar tu modelo, consulta la sección de versiones de hardware de TPU en el documento Arquitectura del sistema. Consulta la página de los modelos de referencia admitidos para ver una lista de los modelos compatibles con Cloud TPU.

En esta página, se proporcionan listas de los tipos de TPU y las versiones del entorno de ejecución de TPU. Cuando programas con marcos de trabajo de TensorFlow, el entorno de ejecución que usas depende de si usas la arquitectura de VM de TPU o de nodo TPU." Consulta la página Arquitectura del sistema para ver las diferencias entre las dos arquitecturas de VM.

En las siguientes tablas, se muestran los tipos de hardware de TPU compatibles.

Configuraciones de TPU v4

Las configuraciones de TPU v4 constan de dos grupos: aquellas con topologías menores a 64 chips (topologías pequeñas) y aquellas con topologías superiores a 64 chips (topologías grandes). Las configuraciones admitidas para cada una de estas topologías se describen en las siguientes secciones.

Topologías de v4 pequeñas

Cloud TPU es compatible con las porciones de TPU v4 más pequeñas que 64 chips, un cubo 4 x 4 x 4. Puedes crear estas topologías pequeñas v4 mediante su nombre basado en TensorCore (por ejemplo, v4-32) o su topología (por ejemplo, 2x2x4):

Nombre (según el recuento de TensorCore)	Cantidad de chips	Topología
v4-8	4	2 × 2 × 1
v4-16	8	2 × 2 × 2
v4-32	16	2 × 2 × 4
v4-64	32	2 × 4 × 4

Topologías de v4 grandes

Los fragmentos de TPU v4 están disponibles en incrementos de 64 chips, con formas que son múltiplos de 4 en las tres dimensiones. Las dimensiones también deben estar en orden ascendente. En la siguiente tabla, se muestran varios ejemplos. Algunas de estas topologías son topologías “personalizadas” que solo se pueden iniciar mediante la API de topología porque tienen la misma cantidad de chips que una topología con nombre más utilizada.

Nombre (según el recuento de TensorCore)	Cantidad de chips	Topología
v4-128	64	4 × 4 × 4
Versión 4-256	128	4x4x8
V4-512	256	4x8x8
Consulta la API de Topology	256	4 × 4 × 16
v4-1,024	512	8 × 8 × 8
v4-1,536	768	8 × 8 × 12
v4-2,048	1,024	8 × 8 × 16
Consulta la API de Topology	1,024	4x16x16
v4-4,096	2,048	8 × 16 × 16
…	…	…

Nota: Todas las porciones grandes de pod v4 tienen vínculos envolventes en las tres dimensiones, lo que proporciona un ancho de banda mayor para operaciones de comunicación colectivas. Consulta la página Tipos y topologías para obtener información sobre los beneficios de las topologías de pods grandes v4.

API de Topology

A fin de crear fragmentos de pod de Cloud TPU con topología personalizada, la API de TPU gcloud se puede usar de la siguiente manera:

     $ gcloud alpha compute tpus tpu-vm create tpu-name \
        --zone=us-central2-b \
        --subnetwork=tpusubnet \
        --type=v4 \
        --topology=4x4x16 \
        --version=runtime-version

Configuraciones de TPU v2 y v3

En la siguiente tabla, se enumeran los tipos de TPU v2 y v3 admitidos

Versión de TPU	Finalización de compatibilidad
v2-8	(Fecha de finalización aún no establecida)
v2-32	(Fecha de finalización aún no establecida)
v2-128	(Fecha de finalización aún no establecida)
v2-256	(Fecha de finalización aún no establecida)
v2-512	(Fecha de finalización aún no establecida)
v3-8	(Fecha de finalización aún no establecida)
v3-32	(Fecha de finalización aún no establecida)
v3-128	(Fecha de finalización aún no establecida)
v3-256	(Fecha de finalización aún no establecida)
v3-512	(Fecha de finalización aún no establecida)
v3-1024	(Fecha de finalización aún no establecida)
v3-2048	(Fecha de finalización aún no establecida)

Versiones de software de TPU

La versión del software de TPU que debes usar depende de la arquitectura de TPU, VM de TPU o Nodo TPU y el framework de AA que uses (TensorFlow, PyTorch o JAX).

VM de TPU

Cuando creas una VM de TPU, la versión más reciente de TensorFlow está preinstalada en la VM de TPU.

TensorFlow

Usa la versión de software de TPU que coincida con la versión de TensorFlow con la que se escribió tu modelo. Por ejemplo, si usas TensorFlow 2.11.0, usa la versión de software de TPU tpu-vm-tf-2.11.0. Si usas un pod de TPU, usa tpu-vm-tf-2.11.0-pod. Las versiones de software de VM de TensorFlow compatibles actuales con las TPU son las siguientes:

• tpu-vm-tf-2.11.0
• tpu-vm-tf-2.10.1
• tpu-vm-tf-2.10.0
• tpu-vm-tf-2.9.3
• tpu-vm-tf-2.9.1.
• tpu-vm-tf-2.8.4
• tpu-vm-tf-2.8.3
• tpu-vm-tf-2.8.0
• tpu-vm-tf-2.7.4
• tpu-vm-tf-2.7.3

Para especificar un pod de VM de TPU, agrega -pod a la versión del software de TPU que deseas usar, por ejemplo, tpu-vm-tf-2.11.0-pod.

VM de TPU con TPU v4

Si entrenas un modelo en la VM de TPU v4 con TensorFlow, usa una de las versiones de v4 que se muestran a continuación.

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
2.10.0 - tpu-vm-tf-2.10.0-v4, tpu-vm-tf-2.10.0-pod-v4

A partir de TensorFlow 2.10.1, solo hay una imagen y no necesitas especificar -v4 con el número de versión.

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

Las VM de TPU se crean con TensorFlow y la biblioteca Libtpu correspondiente preinstalada. Si creas tu propia imagen de VM, especifica las siguientes versiones de software de TensorFlow TPU y las versiones de Libtpu correspondientes:

Versión de TensorFlow	Versión de libtpu.so
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.*	1.2.0
2.7.3	1.1.2

PyTorch

Usa la versión de software de TPU que coincida con la versión de PyTorch con la que se escribió tu modelo. Por ejemplo, si usas PyTorch 1.13, usa la versión de software de TPU tpu-vm-pt-1.13 para v2 y v3, o la versión de software de TPU tpu-vm-v4-pt-1.13 para v4. Se usa la misma versión de software de TPU para los pods de TPU (p.ej., v2-32, v3-128 y v4-32). Las versiones de software de TPU compatibles son las siguientes:

TPU v2/v3:

• tpu-vm-pt-1.13 (pytorch-1.13).
• tpu-vm-pt-1.12 (pytorch-1.12)
• tpu-vm-pt-1.11 (pytorch-1.11)
• tpu-vm-pt-1.10 (pytorch-1.10).
• v2-alpha (pytorch-1.8.1)

TPU v4:

• tpu-vm-v4-pt-1.13 (pytorch-1.13)

Cuando creas una VM de TPU, la versión más reciente de PyTorch se instala de forma previa en la VM de TPU. La versión correcta de libtpu.so se instala automáticamente cuando instalas PyTorch.

Para cambiar la versión actual del software de PyTorch, consulta Cambia la versión de PyTorch.

JAX

Debes instalar JAX de forma manual en tu VM de TPU, ya que no hay una versión de software de TPU específica de JAX. Para las configuraciones de v2 y v3, usa la versión de software del TPU tpu-vm-base. Para las configuraciones de v4, usa tpu-vm-v4-base. La versión correcta de libtpu.so se instala automáticamente cuando instalas JAX.

Nodo TPU

TensorFlow

Usa la versión de software de TPU que coincida con la versión de TensorFlow con la que se escribió tu modelo. Por ejemplo, si usas TF 2.11.0, debes usar la versión de software de TPU 2.11.0. Las versiones de software de TPU específicas de TensorFlow son las siguientes:

• 2.11.0
• 2,10.1
• 2.10.0
• 2,9,3
• 2.9.1
• 2.8.4
• 2.8.2
• 2.7.3

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

Cuando creas un nodo TPU, la versión más reciente de TensorFlow está preinstalada en el nodo TPU.

PyTorch

Usa la versión de software de TPU que coincida con la versión de PyTorch con la que se escribió tu modelo. Por ejemplo, si usas PyTorch 1.9, usa la versión de software pytorch-1.9.

Las versiones de software de TPU específicas de PyTorch son las siguientes:

• pytorch-1.13
• pytorch-1.12.
• pytorch-1.11.
• pytorch-1.10.
• pytorch-1.9
• pytorch-1.8
• pytorch-1.7

• pytorch-1.6

• pytorch-nightly

Cuando creas un nodo TPU, la versión más reciente de PyTorch se instala previamente en el nodo TPU.

JAX

Debes instalar JAX de forma manual en tu VM de TPU a fin de que no haya una versión de software de TPU específica de JAX preinstalada. Puedes usar cualquiera de las versiones de software enumeradas para TensorFlow.