Compatibilidad de Dataflow con las TPUs

Las TPUs de Google Cloud Platform son aceleradores de IA diseñados a medida por Google y optimizados para entrenar y usar modelos de IA grandes. Están diseñadas para escalar de forma rentable una amplia gama de cargas de trabajo de IA y ofrecen versatilidad para acelerar las cargas de trabajo de inferencia en frameworks de IA, como PyTorch, JAX y TensorFlow. Para obtener más información sobre las TPUs, consulta la introducción a las TPUs de Google Cloud Platform.

Requisitos para usar las TPU en Dataflow

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

Limitaciones

Esta oferta está sujeta a las siguientes limitaciones:

  • Solo se admiten aceleradores de TPU de un solo host: la oferta de TPU de Dataflow solo admite configuraciones de TPU de un solo host en las que cada trabajador de Dataflow gestiona uno o varios dispositivos de TPU que no están interconectados con las TPUs gestionadas por otros trabajadores.
  • Solo se admiten grupos de trabajadores de 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 TPUs se facturan por las horas de chip de TPU de los trabajadores consumidas y no 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 mediante la tupla (tipo de TPU, topología de TPU).

  • Tipo de TPU hace referencia al modelo del dispositivo de TPU.
  • La topología de TPU se refiere al número y la disposición física de los chips de TPU de un segmento.

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

Dataflow admite las siguientes configuraciones de TPU:

Tipo de TPU Topología Obligatorio 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 TPUs, consulta Regiones y zonas de TPU en la documentación de TPU de Cloud.

Siguientes pasos