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) personalizados de Google que se usan para acelerar las cargas de trabajo de AA que utilizan 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.
  • Escalado adaptable: 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 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.

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 Ironwood (TPU7x), consulta los siguientes recursos:

Cada una de las siguientes opciones para la creación de clústeres proporciona diferentes grados de facilidad y flexibilidad en la configuración del clúster y la programación de la carga de trabajo:

  • Usa el kit de procesamiento acelerado (XPK) para crear rápidamente clústeres de GKE y ejecutar cargas de trabajo para pruebas de concepto y pruebas. Para obtener más información, consulta la documentación de XPK.
  • Usa Google Cloud CLI para crear manualmente tu instancia de clúster de GKE y personalizar o expandir con precisión los entornos de producción de GKE existentes.

Terminología relacionada con las TPU en GKE

En esta página, 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: Una topología de 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 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 coordine 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 el desalojo 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 v5litepod-16 (v5e) de varios hosts. 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 TPUs adjuntas 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 del 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 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 el caso de las TPU v4, v5p y Ironwood (TPU7x) (vista previa) programadas 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 segmento de nodo, como ct<version>-hightpu-<node-chip-count>t. Por ejemplo, el tipo de máquina tpu7x-standard-4t admite Ironwood (TPU7x) y contiene cuatro chips 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 las TPU en GKE, ten en cuenta las siguientes características de las TPU:

  • Una VM puede acceder a hasta 8 chips de 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 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 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, 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 segmentación 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 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.

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

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.31.1-gke.1146000 y versiones posteriores
  • Ironwood (TPU7x) (Vista previa): 1.34.1-gke.2541000 o posterior

Otros tipos de Cloud TPU son compatibles con todas las versiones de GKE.

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

GKE escala automáticamente los grupos de nodo TPU de Cloud creados de forma manual o automática que usan el escalador automático de clústeres 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 ajuste de escala 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 vertical 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.

Política de cargas de trabajo

La política de carga de trabajo solo es compatible con Ironwood (TPU7x).

Una política de carga de trabajo es un tipo de política de recursos que te permite configurar la colocación física de la infraestructura subyacente. Una política de carga de trabajo proporciona opciones de configuración jerárquica para un control más detallado.

Ironwood (TPU7x) admite la política de carga de trabajo de alto rendimiento. Esta política te permite hacer lo siguiente:

  • Ejecuta cargas de trabajo distribuidas que requieren un alto rendimiento, como la computación de alto rendimiento (HPC) o el entrenamiento de aprendizaje automático (AA). La infraestructura subyacente está configurada para colocar las VMs de TPU en la misma ubicación y, así, reducir la latencia de red entre ellas.
  • Define la estrategia de mantenimiento para las VMs de TPU y, así, minimizar las interrupciones en las cargas de trabajo.

Programación de la recopilación

La programación de la recopilación solo se admite en TPU Trillium.

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 de 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: