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, las ventajas de las unidades de procesamiento de tensor (TPUs) y las consideraciones sobre la programación de cargas de trabajo. Las TPUs son circuitos integrados para aplicaciones específicas (ASIC) desarrollados a medida por Google que agilizan las cargas de trabajo de aprendizaje automático que usan frameworks como TensorFlow, PyTorch y JAX.

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

Antes de leer esta página, asegúrate de que sabes cómo funcionan los aceleradores de aprendizaje automático. Para obtener más información, consulta el artículo Introducción a Cloud TPU.

Ventajas de usar TPUs en GKE

GKE ofrece asistencia completa para la gestión del ciclo de vida de los nodos y los grupos de nodos de TPU, lo que incluye la creación, la configuración y la eliminación de máquinas virtuales de TPU. GKE también admite máquinas virtuales de Spot y el uso de TPUs de Cloud reservadas. Para obtener más información, consulta las opciones de consumo de TPU de Cloud.

Entre las ventajas de usar TPUs en GKE se incluyen las siguientes:

  • Entorno operativo coherente: puedes usar una sola plataforma para todas las cargas de trabajo de aprendizaje automático y de otro tipo.
  • Actualizaciones automáticas: GKE automatiza las actualizaciones de versiones, lo que reduce la sobrecarga operativa.
  • Balanceo de carga: GKE distribuye la carga, lo que reduce la latencia y mejora la fiabilidad.
  • Escalado adaptable: GKE escala automáticamente los recursos de TPU para satisfacer las necesidades de tus cargas de trabajo.
  • Gestión de recursos: con Kueue, un sistema de colas de trabajos nativo de Kubernetes, puedes gestionar recursos en varios arrendatarios de tu organización mediante colas, apropiación, priorización y uso compartido equitativo.
  • Opciones de aislamiento: GKE Sandbox te ayuda a proteger tus cargas de trabajo con gVisor. Para obtener más información, consulta GKE Sandbox.

Empezar 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 crear clústeres ofrece distintos grados de facilidad y flexibilidad en la configuración de clústeres y la programación de cargas 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 la CLI de Google Cloud para crear manualmente la instancia de clúster de GKE si quieres personalizarla con precisión o ampliar los entornos de producción de GKE que ya tengas.

Terminología relacionada con las TPUs en GKE

En esta página se utiliza la siguiente terminología relacionada con las TPUs:

  • Resistencia de ICI de Cloud TPU: una función que ayuda a mejorar la tolerancia a fallos de los enlaces ópticos y los conmutadores de circuitos ópticos (OCS) que conectan las TPUs entre cubos. Para obtener más información, consulta la arquitectura de las TPU.
  • Cubo de TPUs: una topología 4x4x4 de chips de TPU interconectados. Esto solo se aplica a las topologías en tuplas de 3 elementos ({A}x{B}x{C}).
  • Tipo de TPU: el tipo de TPU de Cloud, como v5e.
  • Slice de TPU: conjunto de chips ubicados en el mismo pod de TPU y conectados mediante interconexiones entre chips (ICI) de alta velocidad. Las porciones se describen en términos de chips o TensorCores, según la versión de la TPU.
  • Nodo de segmento de TPU: un nodo de Kubernetes representado por una sola VM que tiene uno o varios chips de TPU interconectados.
  • Grupo de nodos de segmento de TPU: un grupo de nodos de Kubernetes dentro de un clúster que tienen la misma configuración de TPU.
  • Topología de la TPU: número y disposición física de los chips de TPU en una porción de TPU.
  • Atómico: GKE trata todos los nodos interconectados como una sola unidad. Durante las operaciones de escalado, 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.
  • Inmutables: no puedes añadir manualmente nodos nuevos al conjunto de nodos interconectados. Sin embargo, puedes crear un nuevo pool de nodos que tenga la topología de TPU que quieras y programar cargas de trabajo en él.

Tipos de grupos de nodos de slices de TPU

GKE admite dos tipos de grupos de nodos de TPU:

El tipo y la topología de TPU determinan si el nodo de segmento de TPU puede ser de varios hosts o de un solo host. Te recomendamos que hagas lo siguiente:

  • En el caso de los modelos a gran escala, usa nodos de segmento de TPU de varios hosts.
  • En el caso de los modelos a pequeña escala, usa nodos de segmento de TPU de un solo host.
  • Para el entrenamiento o la inferencia a gran escala, usa Pathways. Pathways simplifica los cálculos de aprendizaje automático a gran escala al permitir que un solo cliente JAX coordine las cargas de trabajo en varios sectores de TPU grandes. Para obtener más información, consulta Pathways.

Grupos de nodos de slices de TPU de varios hosts

Un grupo de nodos de un sector de TPU multihost es un grupo de nodos que contiene dos o más VMs de TPU interconectadas. Cada máquina virtual tiene un dispositivo TPU conectado. Las TPUs de una porción de TPU multihost están conectadas a través de una interconexión de alta velocidad (ICI). Una vez que se ha creado un grupo de nodos de un segmento de TPU de varios hosts, no se pueden añadir nodos. Por ejemplo, no puedes crear un grupo de nodos v4-32 y, más adelante, añadir un nodo de Kubernetes (VM de TPU) al grupo de nodos. Para añadir una porción de TPU a un clúster de GKE, debes crear un grupo de nodos.

Las VMs de un grupo de nodos de un segmento de TPU multihost se tratan como una sola unidad atómica. Si GKE no puede desplegar un nodo en el segmento, no se desplegará ningún nodo en el nodo del segmento de TPU.

Si es necesario reparar un nodo de una porción de TPU de varios hosts, GKE cierra todas las VMs de la porción de TPU, lo que obliga a expulsar todos los pods de Kubernetes de 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 podrán programar en las VMs de la nueva porción de TPU.

En el siguiente diagrama se muestra un segmento de TPU de varios hosts (v5e) v5litepod-16. Esta slice de TPU tiene cuatro VMs. Cada máquina virtual de la porción de TPU tiene cuatro chips de TPU v5e conectados con interconexiones de alta velocidad (ICI), y cada chip de TPU v5e tiene un Tensor Core:

Diagrama de un segmento de TPU de varios hosts

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

Diagrama del pod de TPU v5e

Grupos de nodos de slices de TPU de un solo host

Un grupo de nodos de slice de un solo host es un grupo de nodos que contiene una o varias VMs de TPU independientes. Cada máquina virtual tiene un dispositivo TPU conectado. Aunque las VMs de un mismo nodo de segmento de un solo host pueden comunicarse a través de la red del centro de datos (DCN), las TPUs conectadas a las VMs no están interconectadas.

En el siguiente diagrama se muestra un ejemplo de un segmento de TPU de un solo host que contiene siete máquinas v4-8:

Diagrama de un grupo de nodos de un solo host

Características de las TPUs en GKE

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

Consumo de TPUs

Para optimizar el uso de recursos y los costes, al tiempo que se equilibra el rendimiento de las cargas de trabajo, GKE admite las siguientes opciones de consumo de TPU:

Para elegir la opción de consumo que se adapte a los requisitos de tu carga de trabajo, consulta Acerca de las opciones de consumo de aceleradores para cargas de trabajo de IA y aprendizaje automático 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 TPUs en un segmento de TPU. GKE aprovisiona un segmento de TPU en topologías bidimensionales o tridimensionales, según la versión de la TPU. Para especificar una topología, indica el número de chips de TPU en cada dimensión de la siguiente manera:

En el caso de las versiones 4 y 5p de TPU y de Ironwood (TPU7x) (versión preliminar) programadas en grupos de nodos de segmentos de TPU multihost, la topología se define en tuplas de 3 elementos ({A}x{B}x{C}), como 4x4x4. El producto de {A}x{B}x{C} define el número de chips de TPU del 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 tengan 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 admitida en la versión 4 es 12x16x16 y en la versión 5p es 16x16x24.
  • Los valores asignados deben seguir 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 la TPU y el número 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 de TPU.

Modo con privilegios

Si usas versiones de GKE anteriores a la 1.28, debes configurar tus contenedores con funciones especiales para acceder a las TPUs. En los clústeres en modo Estándar, puedes usar el modo Privilegiado para conceder este acceso. El modo con privilegios anula muchos de los otros ajustes de seguridad de securityContext. Para obtener más información, consulta Ejecutar contenedores sin modo privilegiado.

Las versiones 1.28 y posteriores no requieren el modo con privilegios ni funciones especiales.

Cómo funcionan las TPUs en GKE

La gestión y la priorización de recursos de Kubernetes tratan las VMs de las TPUs de la misma forma que otros tipos de VMs. 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 TPUs en GKE, ten en cuenta las siguientes características de las TPU:

  • Una VM puede acceder a un máximo de 8 chips de TPU.
  • Una porción de TPU contiene un número fijo de chips de TPU, que depende del tipo de máquina de TPU que elijas.
  • El número de google.com/tpu solicitados debe ser igual al número total de chips de TPU disponibles en el nodo de la porción de TPU. Cualquier contenedor de un pod de GKE que solicite TPUs debe consumir todos los chips de TPU del nodo. De lo contrario, la implementación 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 segmento de TPU con cuatro chips de TPU cada uno, lo que hace un total de ocho chips de TPU. Con este tipo de máquina, puedes hacer lo siguiente:
    • No se puede desplegar un pod de GKE que requiera ocho chips de TPU en los nodos de este grupo de nodos.
    • Puede desplegar dos pods que requieran cuatro chips de TPU cada uno, y cada pod en uno de los dos nodos de este grupo de nodos.
    • La versión 5e de TPU con topología 4x4 tiene 16 chips de TPU en cuatro nodos. La carga de trabajo de Autopilot de GKE que seleccione esta configuración debe solicitar cuatro chips de TPU en cada réplica, de una a cuatro réplicas.
  • En los clústeres estándar, se pueden programar varios pods de Kubernetes en una máquina virtual, pero solo un contenedor de 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 las slices de TPU tienen el taint google.com/tpu, lo que impide que las cargas de trabajo que no son de TPU se programen en los nodos de las slices de TPU. Las cargas de trabajo que no usan TPUs se ejecutan en nodos que no son de TPU, lo que libera recursos de computación en los nodos de los sectores de TPU para el código que usa TPUs. Ten en cuenta que la contaminación no garantiza que los recursos de TPU se utilicen por completo.
  • GKE recoge los registros emitidos por los contenedores que se ejecutan en los nodos de las slices de TPU. Para obtener más información, consulta Registro.
  • Las métricas de utilización de las TPU, como el rendimiento del tiempo de ejecución, están disponibles en Cloud Monitoring. Para obtener más información, consulta Observabilidad y métricas.
  • Puedes usar GKE Sandbox para crear un entorno aislado para tus cargas de trabajo de TPU. 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 solo admite las siguientes TPUs de Cloud 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 versiones posteriores.

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

Escalado automático de grupos de nodos de TPU de Cloud

GKE escala automáticamente los grupos de nodos de Cloud TPU creados automáticamente o manualmente que usan la herramienta de ajuste automático de clústeres de una de las siguientes formas:

  • Grupo de nodos de segmento de TPU de un solo host: GKE añade o elimina nodos de TPU en el grupo de nodos. El grupo de nodos puede contener cualquier número de nodos de TPU entre cero y el tamaño máximo del grupo de nodos, determinado por las marcas de autoescalado --max-nodes y --total-max-nodes. Todos los nodos de TPU del pool de nodos tienen el mismo tipo de máquina y la misma topología. Para obtener más información sobre cómo crear un pool de nodos de una porción de TPU de un solo host, consulta Crear un pool de nodos de una porción de TPU de un solo host.
  • Grupo de nodos de segmento de TPU de varios hosts: GKE aumenta de forma atómica el grupo de nodos de cero al número de nodos necesarios para satisfacer la topología de la TPU. Por ejemplo, en un grupo de nodos de TPU con el ct5lp-hightpu-4ttipo de máquina y una topología de 16x16, el grupo de nodos siempre tiene 64 nodos o ninguno. GKE reduce el tamaño del grupo de nodos si no hay cargas de trabajo de TPU en él. Para reducir el tamaño del grupo de nodos, GKE expulsa todos los pods programados y elimina todos los nodos del grupo de nodos. Para obtener más información sobre cómo crear un grupo de nodos de segmento de TPU multihost, consulta Crear un grupo de nodos de segmento de TPU multihost.

Política de cargas de trabajo

La política de cargas de trabajo solo se admite en Ironwood (TPU7x).

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

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

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

Programación de colecciones

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

En TPU Trillium, puedes usar la programación de colecciones para agrupar nodos de slices de TPU. Al agrupar estos nodos de segmento de TPU, es más fácil ajustar el número de réplicas para satisfacer la demanda de la carga de trabajo. Google Cloud controla las actualizaciones de software para asegurarse de que siempre haya suficientes segmentos disponibles en la colección para servir el tráfico.

TPU Trillium admite la programación de recolecciones para grupos de nodos de un solo host 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 recogida en función del tipo de segmento de TPU que utilices:

  • Slice de TPU de varios hosts: GKE agrupa los slices de TPU de varios hosts para formar una colección. Cada grupo de nodos de GKE es una réplica de esta colección. Para definir una colección, crea un segmento de TPU multihost y asigna un nombre único a la colección. Para añadir más porciones de TPU a la colección, crea otro grupo de nodos de porción de TPU multihost con el mismo nombre de colección y tipo de carga de trabajo.
  • Slice de TPU de un solo host: GKE considera todo el pool de nodos de slice de TPU de un solo host como una colección. Para añadir más porciones de TPU a la colección, puedes cambiar el tamaño del grupo de nodos de la porción de TPU de un solo host.

La programación de colecciones 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 de Spot.
  • Las colecciones que contengan grupos de nodos de slices de TPU multihost deben usar el mismo tipo de máquina, la misma topología y la misma versión para todos los grupos de nodos de la colección.

Puedes configurar la programación de la recogida en los siguientes casos:

Siguientes pasos

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