Ejecutar un flujo de procesamiento con TPUs

En esta página se explica cómo ejecutar una canalización de Apache Beam en Dataflow con TPUs. Las tareas que usan TPUs generan cargos según lo especificado en la página de precios de Dataflow.

Para obtener más información sobre el uso de las TPUs con Dataflow, consulta el artículo sobre la compatibilidad de Dataflow con las TPUs.

Opcional: Hacer una reserva específica para usar aceleradores

Aunque puedes usar las TPUs bajo demanda, te recomendamos que uses las TPUs de Dataflow con reservas de Google Cloud Platformespecíficas. De esta forma, te aseguras de tener acceso a los aceleradores disponibles y de que los trabajadores se inicien rápidamente. Las canalizaciones que consumen una reserva de TPU no requieren cuota de TPU adicional.

Si no haces una reserva y decides usar las TPUs bajo demanda, proporciona cuota de TPU antes de ejecutar tu canalización.

Opcional: Aprovisionar cuota de TPU

Puedes usar las TPUs con capacidad bajo demanda o mediante una reserva. Si quieres usar las TPUs bajo demanda, debes aprovisionar cuota de TPU antes de hacerlo. Si usas una reserva específicamente orientada, puedes saltarte esta sección.

Para usar las TPUs bajo demanda sin reserva, comprueba el límite y el uso actual de tu cuota de la API de Compute Engine para TPUs de la siguiente manera:

Consola

1. Ve a la página Cuotas de la Google Cloud consola:

   Ir a Cuotas

2. En el cuadro Filtrar, haz lo siguiente:

   1. Usa la siguiente tabla para seleccionar y copiar la propiedad de la cuota en función de la versión de TPU y del tipo de máquina. Por ejemplo, si tienes previsto crear nodos de TPU v5e bajo demanda cuyo tipo de máquina empiece por ct5lp-, introduce Name: TPU v5 Lite PodSlice chips.

      Versión de TPU, el tipo de máquina empieza por	Propiedad y nombre de la cuota de instancias bajo demanda
      TPU v5e ct5lp-	Name: TPU v5 Lite PodSlice chips
      TPU v5p ct5p-	Name: TPU v5p chips
      TPU v6e ct6e-	Dimensions (e.g. location): tpu_family:CT6E

   2. Seleccione la propiedad Dimensiones (por ejemplo, ubicación) e introduzca region: seguido del nombre de la región en la que quiera iniciar su canalización. Por ejemplo, introduzca region:us-west4 si tiene previsto usar la zona us-west4-a. La cuota de TPU es regional, por lo que todas las zonas de la misma región consumen la misma cuota de TPU.

Configurar una imagen de contenedor personalizada

Para interactuar con las TPUs en los flujos de procesamiento de Dataflow, debes proporcionar software que pueda funcionar en dispositivos XLA en tu entorno de tiempo de ejecución del flujo de procesamiento. Para ello, debes instalar bibliotecas de TPU en función de las necesidades de tu canalización y configurar variables de entorno en función del dispositivo de TPU que utilices.

Para personalizar la imagen de contenedor, instala Apache Beam en una imagen base estándar que tenga las bibliotecas de TPU necesarias. También puedes instalar el software de TPU en las imágenes publicadas con las versiones del SDK de Apache Beam.

Para proporcionar una imagen de contenedor personalizada, usa la opción sdk_container_image de la canalización. Para obtener más información, consulta Usar contenedores personalizados en Dataflow.

Cuando usas un acelerador de TPU, debes definir las siguientes variables de entorno en la imagen de contenedor.

ENV TPU_SKIP_MDS_QUERY=1 # Don't query metadata
ENV TPU_HOST_BOUNDS=1,1,1 # There's only one host
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0 # Always 0 for single-host TPUs

En función del acelerador que utilices, también tendrás que definir las variables de la siguiente tabla.

tipo	topología	Dataflow obligatorio worker_machine_type	Variables de entorno adicionales
tpu-v5-lite-podslice	1x1	ct5lp-hightpu-1t	TPU_ACCELERATOR_TYPE=v5litepod-1 TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v5-lite-podslice	2x2	ct5lp-hightpu-4t	TPU_ACCELERATOR_TYPE=v5litepod-4 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v5-lite-podslice	2x4	ct5lp-hightpu-8t	TPU_ACCELERATOR_TYPE=v5litepod-8 TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v6e-slice	1x1	ct6e-standard-1t	TPU_ACCELERATOR_TYPE=v6e-1 TPU_CHIPS_PER_HOST_BOUNDS=1,1,1
tpu-v6e-slice	2x2	ct6e-standard-4t	TPU_ACCELERATOR_TYPE=v6e-4 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1
tpu-v6e-slice	2x4	ct6e-standard-8t	TPU_ACCELERATOR_TYPE=v6e-8 TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
tpu-v5p-slice	2x2x1	ct5p-hightpu-4t	TPU_ACCELERATOR_TYPE=v5p-8 TPU_CHIPS_PER_HOST_BOUNDS=2,2,1

Un ejemplo de Dockerfile para la imagen de contenedor personalizada podría ser el siguiente:

FROM python:3.11-slim

COPY --from=apache/beam_python3.11_sdk:2.66.0 /opt/apache/beam /opt/apache/beam

# Configure the environment to access TPU device

ENV TPU_SKIP_MDS_QUERY=1
ENV TPU_HOST_BOUNDS=1,1,1
ENV TPU_WORKER_HOSTNAMES=localhost
ENV TPU_WORKER_ID=0

# Configure the environment for the chosen accelerator.
# Adjust according to the accelerator you use.
ENV TPU_ACCELERATOR_TYPE=v5litepod-1
ENV TPU_CHIPS_PER_HOST_BOUNDS=1,1,1

# Install TPU software stack.
RUN pip install jax[tpu] apache-beam[gcp]==2.66.0 -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

ENTRYPOINT ["/opt/apache/beam/boot"]

Ejecutar un trabajo con TPUs

Para ejecutar un trabajo de Dataflow con TPUs, debes tener en cuenta lo siguiente:

• Como los contenedores de TPU pueden ser grandes, para evitar quedarte sin espacio en el disco, aumenta el tamaño predeterminado del disco de arranque a 50 gigabytes o al tamaño adecuado que requiera tu imagen de contenedor mediante la opción de canalización --disk_size_gb.
• Limita el paralelismo entre trabajadores.

TPUs y paralelismo de trabajadores

En la configuración predeterminada, las canalizaciones de Python de Dataflow inician un proceso del SDK de Apache Beam por cada núcleo de VM. Los tipos de máquinas TPU tienen un gran número de núcleos de vCPU, pero solo un proceso puede realizar cálculos en un dispositivo TPU. Además, un proceso puede reservar un dispositivo TPU durante toda su duración. Por lo tanto, debes limitar el paralelismo entre trabajadores al ejecutar una canalización de TPU de Dataflow. Para limitar el paralelismo de los trabajadores, sigue estas instrucciones:

• Si tu caso práctico implica ejecutar inferencias en un modelo, usa la API Beam RunInference. Para obtener más información, consulta Inferencia de modelos de lenguaje de gran tamaño en Beam.
• Si no puedes usar la API Beam RunInference, usa los objetos compartidos entre varios procesos de Beam para restringir determinadas operaciones a un solo proceso.
• Si no puedes usar las recomendaciones anteriores y prefieres iniciar solo un proceso de Python por trabajador, define la opción de canalización --experiments=no_use_multiple_sdk_containers.
• En el caso de los trabajadores con más de 100 vCPUs, reduce el número de hilos mediante la opción de canalización --number_of_worker_harness_threads. Consulta la siguiente tabla para ver si tu tipo de TPU usa más de 100 vCPUs.

En la siguiente tabla se indican los recursos de computación totales por trabajador de cada configuración de TPU.

Tipo de TPU	topología	tipo de equipo	Chips de TPU	vCPU	RAM (GB)
tpu-v5-lite-podslice	1x1	ct5lp-hightpu-1t	1	24	48
tpu-v5-lite-podslice	2x2	ct5lp-hightpu-4t	4	112	192
tpu-v5-lite-podslice	2x4	ct5lp-hightpu-8t	8	224	384
tpu-v6e-slice	1x1	ct6e-standard-1t	1	44	176
tpu-v6e-slice	2x2	ct6e-standard-4t	4	180	720
tpu-v6e-slice	2x4	ct6e-standard-8t	8	360	1440
tpu-v5p-slice	2x2x1	ct5p-hightpu-4t	4	208	448

Ejecutar un flujo de procesamiento con TPUs

Para ejecutar un trabajo de Dataflow con TPUs, usa el siguiente comando.

python PIPELINE \
  --runner "DataflowRunner" \
  --project "PROJECT" \
  --temp_location "gs://BUCKET/tmp" \
  --region "REGION" \
  --dataflow_service_options "worker_accelerator=type:TPU_TYPE;topology:TPU_TOPOLOGY" \
  --worker_machine_type "MACHINE_TYPE" \
  --disk_size_gb "DISK_SIZE_GB" \
  --sdk_container_image "IMAGE" \
  --number_of_worker_harness_threads NUMBER_OF_THREADS

Haz los cambios siguientes:

• PIPELINE: el archivo de código fuente de tu canalización.
• PROJECT: el nombre del proyecto. Google Cloud
• BUCKET: el segmento de Cloud Storage.
• REGION: una región de Dataflow, como us-central1.
• TPU_TYPE: un tipo de TPU compatible, por ejemplo, tpu-v5-lite-podslice. Para ver una lista completa de los tipos y las topologías, consulta Aceleradores de TPU compatibles.
• TPU_TOPOLOGY: la topología de la TPU. Por ejemplo, 1x1.
• MACHINE_TYPE: el tipo de máquina correspondiente, por ejemplo, ct5lp-hightpu-1t.
• DISK_SIZE_GB: tamaño del disco de arranque de cada máquina virtual de trabajador. Por ejemplo, 100.
• IMAGE: la ruta de Artifact Registry de tu imagen de Docker.
• NUMBER_OF_THREADS: opcional. Número de hilos de controlador de trabajadores.

Verificar la tarea de Dataflow

Para confirmar que el trabajo usa máquinas virtuales de trabajador con TPUs, sigue estos pasos:

1. En la Google Cloud consola, ve a la página Dataflow > Trabajos.

   Ir a Tareas

2. Selecciona un trabajo.

3. Haz clic en la pestaña Métricas de los trabajos.

4. En la sección Escalado automático, comprueba que haya al menos una VM Trabajadores actuales.

5. En el panel lateral Información del trabajo, comprueba que machine_type empieza por ct. Por ejemplo, ct6e-standard-1t. Indica el uso de la TPU.

Solucionar problemas de una tarea de Dataflow

Si tienes problemas para ejecutar tu trabajo de Dataflow con TPUs, consulta el artículo Solucionar problemas de trabajos de Dataflow con TPUs.

Siguientes pasos

• Más información sobre la compatibilidad con TPU en Dataflow
• Consulta información sobre la inferencia de modelos grandes en Beam.