Acerca de las TPUs en GKE

En esta página, se describe cómo funciona Cloud TPU con Google Kubernetes Engine (GKE), incluida la terminología, los beneficios de las unidades de procesamiento de tensores (TPU) y las consideraciones sobre la programación de cargas de trabajo. Las TPU son circuitos integrados específicos de aplicaciones (ASIC) desarrollados por Google para acelerar las cargas de trabajo de AA que usan frameworks como TensorFlow, PyTorch y JAX.

Esta página está dirigida a los administradores y operadores de plataformas, y a los especialistas en IA y datos que ejecutan modelos de aprendizaje automático (AA) que tienen características como ser a gran escala, de larga duración o dominados por cálculos de matrices. Para obtener más información sobre los roles comunes y las tareas de ejemplo a las que hacemos referencia en el contenido de Google Cloud, consulta Roles y tareas comunes de los usuarios de GKE.

Antes de leer esta página, asegúrate de saber cómo funcionan los aceleradores de AA. Para obtener más detalles, consulta Introducción a Cloud TPU.

Beneficios de usar las TPUs en GKE

GKE proporciona compatibilidad total para la administración del ciclo de vida de los nodos y grupos de nodos de TPU, incluida la creación, configuración y eliminación de VMs de TPU. GKE también admite VM Spot y el uso de Cloud TPU reservada. Para obtener más información, consulta Opciones de consumo de Cloud TPU.

Estos son algunos de los beneficios de usar las TPU en GKE:

  • Entorno operativo coherente: Puedes usar una sola plataforma para todas las cargas de trabajo de aprendizaje automático y otras.
  • Actualizaciones automáticas: GKE automatiza las actualizaciones de versiones, lo que reduce la sobrecarga operativa.
  • Balanceo de cargas: GKE distribuye la carga, lo que reduce la latencia y mejora la confiabilidad.
  • Ajuste de escala adaptable: GKE ajusta automáticamente la escala de los recursos de TPU para satisfacer las necesidades de tus cargas de trabajo.
  • Administración de recursos: Con Kueue, un sistema de cola de trabajos nativo de Kubernetes, puedes administrar recursos en varios arrendatarios dentro de tu organización con la puesta en cola, la interrupción, la priorización y el uso compartido equitativo.
  • Opciones de zona de pruebas: GKE Sandbox ayuda a proteger tus cargas de trabajo con gVisor. Para obtener más información, consulta GKE Sandbox.

Beneficios de usar la TPU Trillium

Trillium es la TPU de sexta generación de Google. Trillium tiene los siguientes beneficios:

  • Trillium aumenta el rendimiento de procesamiento por chip en comparación con la TPU v5e.
  • Trillium aumenta la capacidad y el ancho de banda de la memoria de gran ancho de banda (HBM), y también incrementa el ancho de banda de la interconexión entre chips (ICI) en comparación con la TPU v5e.
  • Trillium está equipado con SparseCore de tercera generación, un acelerador especializado para procesar incorporaciones ultragrandes que son comunes en las cargas de trabajo avanzadas de clasificación y recomendación.
  • Trillium es más de un 67% más eficiente en términos de consumo energético que la TPU v5e.
  • Trillium puede escalar hasta 256 TPU en una sola porción de TPU con un alto ancho de banda y una latencia baja.
  • Trillium admite la programación de recopilación. La programación de recopilación te permite declarar un grupo de TPUs (grupos de nodos de porción de TPU de host único y de varios hosts) para garantizar una alta disponibilidad para las demandas de tus cargas de trabajo de inferencia.

En todas las plataformas técnicas, como las APIs y los registros, y en partes específicas de la documentación de GKE, usamos v6e o TPU Trillium (v6e) para hacer referencia a las TPU Trillium. Para obtener más información sobre los beneficios de Trillium, lee la entrada de blog del anuncio de Trillium. Para comenzar la configuración de la TPU, consulta Planifica las TPU en GKE.

Terminología relacionada con las TPU en GKE

En esta página, se usa la siguiente terminología relacionada con las TPU:

  • Tipo de TPU: El tipo de Cloud TPU, como v5e.
  • Nodo de porción de TPU: Es un nodo de Kubernetes representado por una sola VM que tiene uno o más chips de TPU interconectados.
  • Grupo de nodos de porción de TPU: Es un grupo de nodos de Kubernetes dentro de un clúster que tienen la misma configuración de TPU.
  • Topología de TPU: Es la cantidad y la disposición física de los chips TPU en una porción de TPU.
  • Atómico: GKE trata todos los nodos interconectados como una sola unidad. Durante las operaciones de ajuste de escala, GKE ajusta la escala de todo el conjunto de nodos a 0 y crea nodos nuevos. Si falla o se cierra una máquina del grupo, GKE vuelve a crear todo el conjunto de nodos como una nueva unidad.
  • Inmutables: No puedes agregar nodos nuevos de forma manual al conjunto de nodos interconectados. Sin embargo, puedes crear un grupo de nodos nuevo que tenga la topología de TPU que desees y programar cargas de trabajo en el grupo de nodos nuevo.

Tipos de grupos de nodos de porción de TPU

GKE admite dos tipos de grupos de nodo TPU:

El tipo y la topología de TPU determinan si tu nodo de porción de TPU puede ser de host único o de varios hosts. Te recomendamos lo siguiente:

  • Para los modelos a gran escala, usa nodos de porción de TPU de varios hosts.
  • Para los modelos a pequeña escala, usa nodos de porción de TPU de host único.
  • Para el entrenamiento o la inferencia a gran escala, usa Pathways. Pathways simplifica los cálculos de aprendizaje automático a gran escala, ya que permite que un solo cliente de JAX organice cargas de trabajo en varias porciones de TPU grandes. Para obtener más información, consulta Pathways.

Grupos de nodos de porciones de TPU de varios hosts

Un grupo de nodos de porción de TPU de varios hosts es un grupo de nodos que contiene dos o más VMs de TPU interconectadas. Cada VM tiene un dispositivo TPU conectado. Las TPUs en una porción de TPU de varios hosts están conectadas a través de una interconexión de alta velocidad (ICI). Después de crear un grupo de nodos de porción de TPU de varios hosts, no puedes agregarle nodos. Por ejemplo, no puedes crear un grupo de nodos v4-32 y, luego, agregar un nodo de Kubernetes (VM de TPU) al grupo de nodos. Para agregar una porción de TPU a un clúster de GKE, debes crear un grupo de nodos nuevo.

Las VMs en un grupo de nodos de porción de TPU de varios hosts se tratan como una sola unidad atómica. Si GKE no puede implementar un nodo en la porción, no se implementará ningún nodo en el nodo de la porción de TPU.

Si un nodo dentro de una porción de TPU de varios hosts requiere reparación, GKE apaga todas las VMs en la porción de TPU, lo que fuerza la expulsión de todos los Pods de Kubernetes en la carga de trabajo. Una vez que todas las VMs de la porción de TPU estén en funcionamiento, los Pods de Kubernetes se pueden programar en las VMs de la nueva porción de TPU.

En el siguiente diagrama, se muestra una porción de TPU de hosts múltiples v5litepod-16 (v5e). Esta porción de TPU tiene cuatro VMs. Cada VM en la porción de TPU tiene cuatro chips TPU v5e conectados con interconexiones de alta velocidad (ICI), y cada chip TPU v5e tiene un TensorCore:

Diagrama de porción de TPU de varios hosts

En el siguiente diagrama, se muestra un clúster de GKE que contiene una porción de TPU v5litepod-16 (v5e) (topología: 4x4) y una porción de TPU v5litepod-8 (v5e) (topología: 2x4):

Diagrama del pod de TPU v5e

Grupos de nodos de porción de TPU de host único

Un grupo de nodos de porción de host único es un grupo de nodos que contiene una o más VMs de TPU independientes. Cada VM tiene un dispositivo TPU conectado. Si bien las VMs dentro de un grupo de nodos de porción de un solo host pueden comunicarse a través de la red del centro de datos (DCN), las TPU adjuntas a las VMs no están interconectadas.

En el siguiente diagrama, se muestra un ejemplo de una porción de TPU de un solo host que contiene siete máquinas v4-8:

Diagrama del grupo de nodos de la porción de host único

Características de las TPU en GKE

Las TPU tienen características únicas que requieren una planificación y configuración especiales.

Consumo de TPU

Para optimizar el uso de recursos y el costo, y, al mismo tiempo, equilibrar el rendimiento de la carga de trabajo, GKE admite las siguientes opciones de consumo de TPU:

Para elegir la opción de consumo que satisfaga los requisitos de tu carga de trabajo, consulta Acerca de las opciones de consumo de aceleradores para cargas de trabajo de IA/AA en GKE.

Antes de usar las TPU en GKE, elige la opción de consumo que mejor se adapte a los requisitos de tu carga de trabajo.

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 bidimensionales o tridimensionales, según la versión de TPU. Debes especificar una topología como la cantidad de chips TPU en cada dimensión de la siguiente manera:

En TPU v4 y v5p, que se programan en grupos de nodos de porciones 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 TPU en el grupo de nodos. Por ejemplo, puedes definir topologías pequeñas que tengan menos de 64 chips de TPU con formas de topología como 2x2x2, 2x2x4 o 2x4x4. Si usas topologías más grandes que tienen más de 64 chips de TPU, los valores que asignes a {A}, {B} y {C} deben cumplir las siguientes condiciones:

  • {A}, {B} y {C} deben ser múltiplos de cuatro.
  • La topología más grande compatible con v4 es 12x16x16 y v5p es 16x16x24.
  • Los valores asignados deben mantener el patrón A ≤ B ≤ C. Por ejemplo, 4x4x8 o 8x8x8.

Tipo de máquina

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

Modo con privilegios

Si usas versiones de GKE anteriores a la 1.28, debes configurar tus contenedores con capacidades especiales para acceder a las TPU. En los clústeres en modo Standard, puedes usar el modo con privilegios para otorgar este acceso. El modo privilegiado anula muchos de los otros parámetros de configuración de seguridad en securityContext. Para obtener más detalles, consulta Ejecuta contenedores sin modo privilegiado.

Las versiones 1.28 y posteriores no requieren modo privilegiado ni capacidades especiales.

Cómo funcionan las TPUs en GKE

La administración de recursos y la priorización de Kubernetes tratan a las VMs en las TPUs de la misma manera que a otros tipos de VM. Para solicitar chips de TPU, usa el nombre del recurso google.com/tpu:

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

Cuando uses TPU en GKE, ten en cuenta las siguientes características de las TPU:

  • Una VM puede acceder a hasta 8 chips TPU.
  • Una porción de TPU contiene una cantidad fija de chips de TPU, y la cantidad depende del tipo de máquina TPU que elijas.
  • La cantidad de google.com/tpu solicitados debe ser igual a la cantidad total de chips TPU disponibles en el nodo de la porción de TPU. Cualquier contenedor en un Pod de GKE que solicite TPU debe consumir todos los chips TPU en el nodo. De lo contrario, la Deployment fallará porque GKE no puede consumir parcialmente los recursos de TPU. Considera las siguientes situaciones:
    • El tipo de máquina ct5lp-hightpu-4t con una topología 2x4 contiene dos nodos de porción de TPU con cuatro chips de TPU cada uno, lo que da un total de ocho chips de TPU. Con este tipo de máquina, puedes hacer lo siguiente:
    • No se puede implementar un Pod de GKE que requiere ocho chips TPU en los nodos de este grupo de nodos.
    • Se pueden implementar dos Pods que requieran cuatro chips TPU cada uno, y cada Pod en uno de los dos nodos de este grupo de nodos.
    • La TPU v5e con topología 4x4 tiene 16 chips TPU en cuatro nodos. La carga de trabajo de GKE Autopilot que selecciona esta configuración debe solicitar cuatro chips de TPU en cada réplica, para una a cuatro réplicas.
  • En los clústeres estándar, se pueden programar varios Pods de Kubernetes en una VM, pero solo un contenedor en cada Pod puede acceder a los chips de TPU.
  • Para crear Pods de kube-system, como kube-dns, cada clúster estándar debe tener al menos un grupo de nodos de segmento que no sea de TPU.
  • De forma predeterminada, los nodos de porción de TPU tienen el taint google.com/tpu que impide que se programen cargas de trabajo que no sean de TPU en los nodos de porción de TPU. Las cargas de trabajo que no usan TPU se ejecutan en nodos que no son de TPU, lo que libera capacidad de procesamiento en los nodos de la porción de TPU para el código que usa TPU. Ten en cuenta que el aislamiento no garantiza que los recursos de TPU se utilicen por completo.
  • GKE recopila los registros emitidos por los contenedores que se ejecutan en los nodos de porción de TPU. Para obtener más información, consulta Logging.
  • Las métricas de uso 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.
  • Puedes ejecutar tus cargas de trabajo de TPU en una zona de pruebas con GKE Sandbox. GKE Sandbox funciona con los modelos de TPU v4 y versiones posteriores. Para obtener más información, consulta GKE Sandbox.

Cómo funciona la programación de la recopilación

En la TPU Trillium, puedes usar la programación de recopilación para agrupar los nodos de porción de TPU. Agrupar estos nodos de segmentación de TPU facilita el ajuste de la cantidad de réplicas para satisfacer la demanda de la carga de trabajo. Google Cloud controla las actualizaciones de software para garantizar que siempre haya suficientes segmentaciones disponibles en la colección para atender el tráfico.

La TPU Trillium admite la programación de recopilación para grupos de nodos de host único y varios hosts que ejecutan cargas de trabajo de inferencia. A continuación, se describe cómo se comporta la programación de la recopilación según el tipo de porción de TPU que uses:

  • Porción de TPU de varios hosts: GKE agrupa las porciones de TPU de varios hosts para formar una colección. Cada grupo de nodos de GKE es una réplica dentro de esta colección. Para definir una colección, crea una porción de TPU de varios hosts y asígnale un nombre único. Para agregar más porciones de TPU a la colección, crea otro grupo de nodos de porción de TPU de varios hosts con el mismo nombre de colección y tipo de carga de trabajo.
  • Porción de TPU de host único: GKE considera todo el grupo de nodos de porción de TPU de host único como una colección. Para agregar más porciones de TPU a la colección, puedes cambiar el tamaño del grupo de nodos de porción de TPU de host único.

La programación de la recopilación tiene las siguientes limitaciones:

  • Solo puedes programar recopilaciones para la TPU Trillium.
  • Solo puedes definir colecciones durante la creación del grupo de nodos.
  • No se admiten las VMs Spot.
  • Las colecciones que contienen grupos de nodos de porción de TPU de varios hosts deben usar el mismo tipo de máquina, topología y versión para todos los grupos de nodos dentro de la colección.

Puedes configurar la programación de la recopilación en las siguientes situaciones:

¿Qué sigue?

Para obtener información sobre cómo configurar Cloud TPU en GKE, consulta las siguientes páginas: