Acerca de las TPUs en GKE

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

Este documento está dirigido 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 del usuario de GKE.

Antes de leer este documento, asegúrate de conocer cómo funcionan los aceleradores de AA. Para obtener más información, consulta Introducción a Cloud TPU.

Beneficios de usar las TPUs en GKE

GKE proporciona compatibilidad total con la administración del ciclo de vida de los nodos y grupos de nodos TPU, lo que incluye crear, configurar y borrar 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.

Entre los beneficios de usar las TPU en GKE, se incluyen los siguientes:

  • Entorno operativo coherente: Puedes usar una sola plataforma para todo el aprendizaje automático y otras cargas de trabajo.
  • 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.
  • Escalamiento responsivo: GKE escala automáticamente los recursos de TPU para satisfacer las necesidades de tus cargas de trabajo.
  • Administración de recursos: Con Kueue, un sistema de colocación de trabajos en cola nativo de Kubernetes, puedes administrar recursos en varios usuarios dentro de tu organización con la colocación 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.

Comienza a usar Ironwood (TPU7x)

Ironwood (TPU7x) es la TPU de séptima generación de Google, diseñada para cargas de trabajo de IA a gran escala. Para obtener más información sobre los beneficios de Ironwood (TPU7x), consulta Acerca de Ironwood (TPU7x) en GKE.

Terminología relacionada con las TPU en GKE

En este documento, se usa la siguiente terminología relacionada con las TPUs:

  • Resiliencia de ICI de Cloud TPU: Es una función que ayuda a mejorar la tolerancia a errores de los vínculos ópticos y los conmutadores de circuitos ópticos (OCS) que conectan las TPU entre cubos. Para obtener más información, consulta Arquitectura de TPU.
  • Cubo de TPU: Es una topología 4x4x4 de chips TPU interconectados. Esto solo se aplica a las topologías en 3 tuplas ({A}x{B}x{C}).
  • Tipo de TPU: el tipo de Cloud TPU, como v5e.
  • Porción de TPU: Es una colección de chips ubicados dentro del mismo pod de TPU y conectados por interconexiones entre chips (ICI) de alta velocidad. Las porciones se describen en términos de chips o TensorCores, según la versión de TPU.
  • Nodo de porción de TPU: Es un nodo de Kubernetes representado por una sola VM que tiene uno o más chips 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 escalamiento, GKE escala todo el conjunto de nodos a 0 y crea nodos nuevos. Si una máquina del grupo falla o finaliza, GKE vuelve a crear todo el conjunto de nodos como una unidad nueva.
  • Inmutable: 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 el nodo de porción de TPU puede ser de varios hosts o de un solo host. 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 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 las 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 TPU 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 de 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 porción de TPU.

Si un nodo dentro de una porción de TPU de varios hosts requiere reparación, GKE cierra todas las VMs de la porción de TPU, lo que fuerza la expulsión de todos los Pods de Kubernetes en la carga de trabajo. Después de 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 varios hosts 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 de un 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 conectadas a las VMs no están interconectadas.

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

Diagrama de un grupo de nodos de 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 las cargas de trabajo, GKE admite las siguientes opciones de consumo de TPU:

  • Inicio flexible: Para aprovisionar VMs de inicio flexible por un máximo de siete días, GKE asigna automáticamente el hardware de la mejor manera posible según la disponibilidad. Para obtener más información, consulta Acerca del aprovisionamiento de GPU, TPU y H4D con el modo de aprovisionamiento de inicio flexible.
  • VMs Spot: Para aprovisionar VMs Spot, puedes obtener descuentos significativos, pero las VMs Spot se pueden interrumpir en cualquier momento, con una advertencia de 30 segundos. Para obtener más información, consulta VMs Spot.
  • Reserva futura de hasta 90 días (en modo de calendario): para aprovisionar recursos TPU por hasta 90 días, durante un período específico. Para obtener más información, consulta Solicita TPU con reserva futura en el modo de calendario.
  • Reservas de TPU: para solicitar una reserva futura por un año o más.
  • On-demand: para consumir TPUs sin organizar la capacidad con anticipación. Antes de solicitar recursos, debes tener suficiente cuota a pedido para el tipo y la cantidad específicos de VMs de TPU. La opción a pedido es la más flexible, pero no se garantiza que haya suficientes recursos a pedido disponibles para satisfacer tu solicitud.

A pedido es el modelo de consumo predeterminado para las TPU en GKE si no especificas otra opción. Para elegir la opción de consumo que cumpla con 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 TPUs 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 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 de la siguiente manera:

En TPU v4, v5p y Ironwood (TPU7x), 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 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 TPU, los valores que asignes a {A}, {B} y {C} deben cumplir con 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.

Nombres de los tipos de máquinas

Los nombres de los tipos de máquinas para las TPU en GKE varían según el modo del clúster y la versión de la TPU:

  • GKE Standard: Seleccionas un tipo de máquina de Compute Engine específico, por ejemplo, ct6e-standard-1t para TPU Trillium (v6e).

  • GKE Autopilot: No seleccionas tipos de máquinas directamente. En su lugar, solicita TPUs con un tipo de acelerador en el manifiesto de tu carga de trabajo. Por ejemplo, usa tpu-v6e-slice para TPU Trillium (v6e) o tpu-v5-lite-podslice para TPU v5e. Luego, GKE Autopilot aprovisiona los nodos subyacentes con los tipos de máquinas adecuados para satisfacer la solicitud.

Para conocer los tipos de máquinas exactos disponibles para cada versión de TPU, consulta las tablas en Planifica las TPU en GKE.

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 el modo privilegiado ni capacidades especiales.

Cómo funcionan las TPUs en GKE

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

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

Cuando uses las 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 TPU, y la cantidad depende del tipo de máquina TPU que elijas.
  • La cantidad de google.com/tpu solicitadas debe ser igual a la cantidad total de chips TPU disponibles en el nodo de 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, tu 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 TPU cada uno, lo que da un total de ocho chips TPU. Con este tipo de máquina, puedes hacer lo siguiente:
    • No se puede implementar un Pod de GKE que requiera ocho chips TPU en los nodos de este grupo de nodos.
    • Se pueden implementar dos Pods que requieran cuatro chips TPU cada uno, cada Pod en uno de los dos nodos de este grupo de nodos.
    • 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 TPU en cada réplica, de una a cuatro réplicas.
  • En los clústeres de Standard, se pueden programar varios Pods de Kubernetes en una VM, pero solo un contenedor de cada Pod puede acceder a los chips TPU.
  • Para crear Pods de kube-system, como kube-dns, cada clúster de Standard debe tener al menos un grupo de nodos de porción que no sea de TPU.
  • De forma predeterminada, los nodos de porción de TPU tienen el taint google.com/tpu, que evita 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 procesamiento en los nodos de porción de TPU para el código que usa TPU. Ten en cuenta que la tacha no garantiza que los recursos de TPU se usen 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.

Creación automática de grupos de nodos con TPU

La creación automática de grupos de nodos admite las siguientes Cloud TPU solo en versiones específicas de GKE:

  • TPU v3: 1.31.0 y versiones posteriores.
  • TPU v5 y TPU v4: 1.29.0 y versiones posteriores.
  • TPU Trillium: 1.32.0 y versiones posteriores.
  • Ironwood (TPU7x): 1.34.1-gke.2541000 o posterior.

Otros tipos de Cloud TPU son compatibles con todas las versiones de GKE. Para obtener más información sobre las versiones de GKE disponibles para las TPU, consulta Valida la disponibilidad de TPU en GKE.

Ajuste de escala automático del grupo de nodos de Cloud TPU

GKE escala los grupos de nodos de Cloud TPU creados de forma automática o manual que usan el escalador automático del clúster de una de las siguientes maneras:

  • Grupo de nodos de porción de TPU de host único: GKE agrega o quita nodos TPU en el grupo de nodos existente. El grupo de nodos puede contener cualquier cantidad de nodos TPU entre cero y el tamaño máximo del grupo de nodos según lo determinado por las marcas de escalado automático --max-nodes y --total-max-nodes. Todos los nodos TPU del grupo de nodos tienen el mismo tipo de máquina y topología. Para obtener más información sobre cómo crear un grupo de nodos de porción de TPU de host único, consulta Crea un grupo de nodos de porción de TPU de host único.
  • Grupo de nodos de porción de TPU de varios hosts: GKE escala verticalmente el grupo de nodos de forma atómica desde cero hasta la cantidad de nodos necesarios para satisfacer la topología de TPU. Por ejemplo, con un grupo de nodo TPU que tiene el tipo de máquina ct5lp-hightpu-4t y una topología de 16x16, el grupo de nodos siempre tiene 64 nodos o cero nodos. GKE reduce la escala verticalmente del grupo de nodos si no hay cargas de trabajo de TPU en él. Para reducir la escala del grupo de nodos, GKE expulsa todos los Pods programados y quita todos los nodos del grupo de nodos. Para obtener más información sobre cómo crear un grupo de nodos de porción de TPU de varios hosts, consulta Crea un grupo de nodos de porción de TPU de varios hosts.

Programación de la recopilación

La programación de colecciones solo se admite para TPU Trillium.

En TPU Trillium, puedes usar la programación de recopilación para agrupar nodos de porción de TPU. Agrupar estos nodos de porciones 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 porciones dentro de la colección para entregar tráfico.

TPU Trillium admite la programación de colecciones para grupos de nodos de host único y de varios hosts que ejecutan cargas de trabajo de inferencia. A continuación, se describe cómo el comportamiento de la programación de la colección depende del tipo de porción de TPU que uses:

  • Porción de TPU de varios hosts: GKE agrupa 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 recopilación tiene las siguientes limitaciones:

  • Solo puedes programar colecciones para 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 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: