Acerca de las TPU en GKE

Esta página presenta Cloud TPU y te muestra dónde encontrar información sobre el uso de Cloud TPU con Google Kubernetes Engine (GKE). Las unidades de procesamiento tensorial (TPUs) son los circuitos integrados específicos de aplicaciones (ASIC) desarrolladas por Google que se usan para acelerar las cargas de trabajo de aprendizaje automático que usan marcos de trabajo como TensorFlow,{101) }PyTorch y JAX.

Antes de usar las TPU en GKE, te recomendamos que completes la siguiente ruta de aprendizaje:

  1. Obtén información sobre cómo funcionan los aceleradores de aprendizaje automático con la Introducción a Cloud TPU.
  2. Obtén información sobre la disponibilidad actual de la versión de TPU con la arquitectura del sistema de Cloud TPU.

Para obtener información sobre cómo configurar Cloud TPU en GKE, consulta Implementa cargas de trabajo de TPU en GKE.

Beneficios de usar TPU en GKE

GKE proporciona compatibilidad total con la administración del ciclo de vida de las VM de TPU, que incluye la creación, configuración y eliminación de las VM de TPU. GKE también admite VM Spot y el uso de Cloud TPU reservada. Los beneficios de usar TPU en GKE son los siguientes:

  • Entorno operativo coherente: una sola plataforma para todo el aprendizaje automático y otras cargas de trabajo
  • Actualizaciones automáticas: GKE automatiza las actualizaciones de versión, lo que reduce la sobrecarga operativa.
  • Balanceo de cargas: GKE distribuye la carga que reduce la latencia y mejora la confiabilidad.
  • Escalamiento responsivo: GKE escala automáticamente los recursos TPU para satisfacer las necesidades de tus cargas de trabajo.
  • Administración de recursos: ConKueue , un sistema de colas de trabajo nativo de Kubernetes, puedes administrar recursos en varios usuarios de tu organización mediante la puesta en cola, la interrupción, la priorización y el uso compartido equitativo.

Terminología relacionada con TPU en GKE

TPU y Kubernetes usan algunos términos similares que debes considerar y diferenciar a medida que lees este documento:

  • VM de TPU: Una VM de Compute Engine que se ejecuta en una máquina física con hardware de TPU.
  • Pod de TPU: Una colección de chips TPU TPU interconectados La cantidad de chips TPU en un pod de TPU varía según la versión de la TPU.
  • Porción de TPU: Es un subconjunto de un pod de TPU completo.
  • Topología de TPU: La cantidad y la disposición física de los chips TPU en una porción de TPU.
  • Atomicidad: La propiedad de un grupo de nodos TPU de varios hosts en el que el grupo de nodos se trata como una sola unidad. No puedes cambiar el tamaño de un grupo de nodos de hosts múltiples. GKE escala un grupo de nodos de varios hosts mediante la creación o eliminación de todos los nodos en el grupo de nodos en un solo paso.

Cómo funcionan las TPU en GKE

La administración y prioridad de recursos de Kubernetes trata a las VM de TPU de la misma manera que a otros tipos de VM. Solicita los chips TPU con el nombre del recurso google.com/tpu:

    resources:
        requests:
          google.com/tpu: 4
        limits:
          google.com/tpu: 4

Cuando usas las TPU en GKE, debes tener en cuenta las siguientes características de TPU:

  • Una VM de TPU puede acceder hasta a 8 chips TPU.
  • Una porción de TPU contiene una cantidad fija de chips TPU que depende del tipo de máquina de TPU que elijas.
  • La cantidad de solicitudes google.com/tpu debe ser igual a la cantidad total de chips disponibles en el nodo TPU. Cualquier contenedor en un pod de GKE que solicite TPU debe consumir todos los chips TPU en el nodo. De lo contrario, la implementación fallará, ya que GKE no puede consumir recursos de forma parcial.
  • Se pueden programar varios Pods de Kubernetes en una VM de TPU, pero solo un contenedor en cada Pod puede acceder a los chips TPU.
  • Cada clúster debe tener al menos un grupo de nodos que no sea de TPU para crear pods del sistema de Kubernetes, como kube-dns.
  • De forma predeterminada, los nodos TPU tienen el google.com/tputaint que evita que se programen Pods que no sean de TPU en ellos. Los taints no garantizan que los recursos TPU se usen por completo. Te permite ejecutar cargas de trabajo que no usan TPU en nodos que no son TPU, lo que libera el procesamiento en nodos TPU para el código que usa TPU.
  • GKE recopila los registros emitidos por los contenedores que se ejecutan en las VMs de TPU. Para obtener más información, consulta Logging.
  • Las métricas de utilización de TPU, como el rendimiento del entorno de ejecución, están disponibles en Cloud Monitoring. Para obtener más información, consulta Observabilidad y métricas.

Disponibilidad de TPU en GKE

En la siguiente tabla, se muestra la disponibilidad de TPU según el tipo de máquina y la versión:

Versión de TPU Tipo de máquina que comienza con Versión mínima de GKE Disponibilidad Zona
TPU v4 ct4p- 1.26.1-gke.1500 Disponible de manera general us-central2-b
TPU v5e ct5l- 1.27.2-gke.2100 Disponible de manera general us-central1-a
europe-west4-b
TPU v5e ct5lp- 1.26.1-gke.1500 Disponible de manera general us-west4-a
us-east5-b
us-east1-c

En las siguientes secciones, se describen las características de TPU que evalúas cuando planificas y configuras tus cargas de trabajo de TPU en GKE. Para obtener detalles sobre las versiones disponibles, los tipos de máquinas, las topologías válidas y su cantidad de chips, consulta Asignación de configuraciones de TPU en este documento.

Tipo de máquina

Los tipos de máquina que admiten recursos de TPU siguen una convención de nombres que incluye la versión de TPU y la cantidad de chips por nodo, como ct<version>-hightpu-<node-chip-count>t. Por ejemplo, el tipo de máquina ct5lp-hightpu-1t admite TPU v5e y contiene un chip TPU en total.

Tipos de grupos de nodos de TPU

Según la topología que definas, GKE coloca las cargas de trabajo de TPU en uno de los siguientes tipos de grupo de nodos:

  • Grupo de nodos de porción de TPU de host único: Un grupo de nodos que contiene una o más VM de TPU independientes. En los grupos de nodos de porción de TPU de host único, las TPU adjuntas a las VM no están interconectadas mediante interconexiones de alta velocidad.
  • Grupo de nodos de porción de TPU de varios hosts: Un grupo de nodos que contiene dos o más VM de TPU interconectadas. Este tipo de grupo de nodos es inmutable y atómico, lo que significa que no puedes agregar nodos al grupo de forma manual. En caso de falla o cierre de la máquina, GKE vuelve a crear todo el grupo de nodos como una unidad atómica nueva.

Topología

La topología define la disposición física de las TPU dentro de una porción de TPU. GKE aprovisiona una porción de TPU en topologías de dos o tres dimensiones según la versión de TPU. Debes especificar una topología como la cantidad de chips TPU en cada dimensión:

  • En TPU v4 programado en grupos de nodos de porción de TPU de varios hosts, debes definir la topología en 3 tuplas ({A}x{B}x{C}), por ejemplo, 4x4x4. El producto de {A}x{B}x{C} define la cantidad de chips en el grupo de nodos. Por ejemplo, puedes definir topologías pequeñas de menos de 64 chips con formas de topología como 2x2x2, 2x2x4 o 2x4x4. Si usas topologías de más de 64 chips, los valores que asignes a {A}, {B} y {C} deben cumplir las siguientes condiciones:

    • {A}, {B} y {C} son múltiplos de cuatro.
    • La topología más grande admitida es 12x16x16.
    • Los valores asignados mantienen el patrón A ≤ B ≤ C. Por ejemplo, 4x4x8 o 8x8x8.

  • Para TPU v5e, las topologías siguen un formato de 2 tuplas ({A}x{B}), por ejemplo 2x2. Para definir la configuración de tu TPU, consulta la tabla en Asignación de la configuración de TPU.

Asignación de la configuración de TPU

Usa la siguiente tabla para definir el tipo de máquina de TPU y la topología que se usará en función de tu caso de uso:

  • Para el entrenamiento o la inferencia de modelos a pequeña escala, usa TPU v4 o TPU v5e con grupos de nodos de porción de TPU de host único.
  • Para el entrenamiento o la inferencia de modelos a gran escala, usa TPU v4 o TPU v5e con grupos de nodos de porción de TPU de varios hosts.
Versión de TPU Tipo de máquina Topología Cantidad de chips TPU Cantidad de VM Tipo de grupo de nodos
TPU v4 ct4p-hightpu-4t 2x2x1 4 1 Host único
2x2x2 8 2 Varios hosts
2x2x4 16 4 Varios hosts
2x4x4 32 8 Varios hosts
{A}x{B}x{C} A*B*C (A*B*C/4)1 Varios hosts
TPU v5e ct5l-hightpu-1t 1x1 1 1 Host único
ct5l-hightpu-4t 2x2 4 1 Host único
ct5l-hightpu-8t 2x4 8 1 Host único
ct5lp-hightpu-1t 1x1 1 1 Host único
ct5lp-hightpu-4t 2x2 4 1 Host único
ct5lp-hightpu-8t 2x4 8 1 Host único
ct5lp-hightpu-4t 2x4 8 2 Varios hosts
4x4 16 4 Varios hosts
4x8 32 8 Varios hosts
8x8 64 16 Varios hosts
8x16 128 32 Varios hosts
16x16 256 64 Varios hosts

  1. Calculado por el producto de topología dividido por cuatro. ↩

Características de TPU v5e

Las máquinas de TPU v5e tienen las siguientes características técnicas:

Tipo de máquina Cantidad de CPU virtuales Memoria (GB) Cantidad de nodos de NUMA Probabilidad de ser adelantado
ct5l-hightpu-1t 24 48 1 ROI
ct5l-hightpu-4t 112 192 1 Media
ct5l-hightpu-8t 224 384 2 Inferior
ct5lp-hightpu-1t 24 48 1 ROI
ct5lp-hightpu-4t 112 192 1 Media
ct5lp-hightpu-8t 224 384 2 Inferior

Reserva de TPU

Para asegurarte de que los recursos TPU estén disponibles cuando los necesites, puedes usar reservas de TPU en las siguientes situaciones:

  • Si tienes reservas de TPU existentes, debes trabajar con tu equipo de cuentas de Google Cloud para migrar tu reserva de TPU a un sistema de reservas nuevo basado en Compute Engine.

  • Si no tienes una reserva de TPU existente, puedes crear una reserva de TPU y no se necesita migración.

Limitaciones

  • Cuando se usan TPU en GKE, SPECIFIC es el único valor admitido para la marca --reservation-affinity de gcloud container node-pools create.
  • Las TPU no están disponibles en los clústeres de GKE Autopilot.
  • La asignación de costos y la medición de uso de GKE no incluye datos sobre el uso o los costos de la TPU v4 reservada.

Consideraciones sobre la programación de cargas de trabajo

Las TPU tienen características únicas que requieren una programación y administración especiales de cargas de trabajo en Kubernetes. En las siguientes secciones, se describen las prácticas recomendadas de programación.

CPU

Para programar una carga de trabajo en la CPU integrada en una VM de TPU, asegúrate de que el Pod de GKE pueda tolerar el taint google.com/tpu. Si quieres que la carga de trabajo se implemente en nodos específicos, usa los selectores de nodos.

La administración y prioridad de recursos de Kubernetes trata a las VM de TPU de la misma manera que a otros tipos de VM. Para otorgar a los pods que requieren TPU una programación de prioridad sobre otros pods en los mismos nodos, solicita la CPU o la memoria máximas para esos pods de TPU. Los pods de prioridad baja deben hacer lo siguiente:

  1. Establece solicitudes de CPU y memoria bajas a fin de garantizar que el nodo tenga suficientes recursos asignables para las cargas de trabajo de TPU. Para obtener más información, consulta Cómo Kubernetes aplica solicitudes y límites de recursos.
  2. No establezcas un límite de CPU (ilimitado) para garantizar que los pods puedan tener aumentos de actividad a fin de usar todos los ciclos sin usar.
  3. Establece un límite de memoria alto para garantizar que los Pods puedan usar la mayor parte de la memoria sin usar y, al mismo tiempo, mantener la estabilidad de los nodos.

Si un pod de Kubernetes no solicita CPU y memoria (incluso si solicita TPU), entonces Kubernetes lo considera un pod de mejor esfuerzo y no hay garantía de que necesite CPU y memoria. Solo los Pods que solicitan CPU y memoria de forma explícita tienen esas garantías. Si deseas obtener más información, consulta Administración de recursos para Pods y contenedores.

Para obtener más información, consulta las prácticas recomendadas de Kubernetes sobre límites y solicitudes de recursos.

Maximiza el entorno de ejecución de cargas de trabajo sin interrupciones

Si usas TPU para entrenar un modelo de aprendizaje automático y se interrumpe la carga de trabajo, se pierde todo el trabajo realizado desde el último punto de control. Para disminuir la probabilidad de que se interrumpa la carga de trabajo, haz lo siguiente:

  • Establece una prioridad más alta para este trabajo que para todos los demás trabajos: Si los recursos son escasos, el programador de GKE interrumpe los trabajos de prioridad más baja para programar un trabajo de mayor prioridad. Esto también garantiza que tu carga de trabajo de mayor prioridad reciba todos los recursos que necesita (hasta el total de recursos disponibles en el clúster). Para obtener más información, consulta Interrupción y prioridad del pod.
  • Configura la exclusión de mantenimiento: una exclusión de mantenimiento es un período no recurrente durante el cual se prohíbe el mantenimiento automático. Para obtener más información, consulta Exclusiones de mantenimiento.

Maximiza el uso de TPU

Para maximizar tu inversión en las TPU, programa una combinación de prioridades del trabajo y agrégalas a la cola para maximizar la cantidad de tiempo que funcionan las TPU. Si quieres programar e interrumpir los niveles de un trabajo, debes usar un complemento de Kubernetes que organice los trabajos en colas. Te recomendamos usar Kueue para ese caso de uso.

¿Qué sigue?