Compatibilidad de Dataflow con las TPU

Las TPU de Google Cloud son aceleradores de IA con diseño personalizado creados por Google que están optimizados para el entrenamiento y el uso de grandes modelos de IA. Están diseñadas para escalar de manera rentable en una amplia variedad de cargas de trabajo de IA y proporcionan versatilidad para acelerar las cargas de trabajo de inferencia en frameworks de IA, como PyTorch, JAX y TensorFlow. Para obtener más detalles sobre las TPU, consulta Introducción a las TPU de Google Cloud.

Requisitos previos para usar TPUs en Dataflow

  • Tus proyectos Google Cloud deben estar aprobados para usar esta oferta de GA.

Limitaciones

Esta oferta está sujeta a las siguientes limitaciones:

  • Solo se admiten aceleradores de TPU de host único: La oferta de TPU de Dataflow solo admite configuraciones de TPU de host único en las que cada trabajador de Dataflow administra uno o varios dispositivos de TPU que no están interconectados con las TPU administradas por otros trabajadores.
  • Solo se admiten grupos de trabajador TPU homogéneos: Las funciones como el ajuste adecuado de Dataflow y Dataflow Prime no admiten cargas de trabajo de TPU.

Precios

Los trabajos de Dataflow que usan TPU se facturan por las horas de chip de TPU de los trabajadores consumidas y no se facturan por la CPU ni la memoria de los trabajadores. Para obtener más información, consulta la página de precios de Dataflow.

Disponibilidad

Están disponibles los siguientes aceleradores de TPU y regiones de procesamiento.

Aceleradores de TPU compatibles

Las combinaciones de aceleradores de TPU admitidas se identifican con la tupla (tipo de TPU, topología de TPU).

  • El tipo de TPU hace referencia al modelo del dispositivo de TPU.
  • La topología de TPU hace referencia a la cantidad y la disposición física de los chips TPU en una porción.

Para configurar el tipo y la topología de las TPU para los trabajadores de Dataflow, usa la opción de canalización worker_accelerator con el formato type:TPU_TYPE;topology:TPU_TOPOLOGY.

Las siguientes configuraciones de TPU son compatibles con Dataflow:

Tipo de TPU Topología Se requiere worker_machine_type
tpu-v5-lite-podslice 1x1 ct5lp-hightpu-1t
tpu-v5-lite-podslice 2x2 ct5lp-hightpu-4t
tpu-v5-lite-podslice 2x4 ct5lp-hightpu-8t
tpu-v6e-slice 1x1 ct6e-standard-1t
tpu-v6e-slice 2x2 ct6e-standard-4t
tpu-v6e-slice 2x4 ct6e-standard-8t
tpu-v5p-slice 2x2x1 ct5p-hightpu-4t

Regiones

Para obtener información sobre las regiones y zonas disponibles para las TPU, consulta Regiones y zonas de TPU en la documentación de Cloud TPU.

¿Qué sigue?