Implementa cargas de trabajo de TPU en GKE Standard

Estándar

En esta página, se muestra cómo solicitar e implementar cargas de trabajo que usan aceleradores de Cloud TPU (TPU) en Google Kubernetes Engine (GKE).

Antes de configurar y, además, implementar las cargas de trabajo de TPU en GKE, debes familiarizarte con los siguientes conceptos:

Antes de comenzar

Antes de comenzar, asegúrate de haber realizado las siguientes tareas:

Habilita la API de Google Kubernetes Engine.

Habilitar la API de Google Kubernetes Engine

Si deseas usar Google Cloud CLI para esta tarea, instala y, luego, inicializa gcloud CLI. Si ya instalaste gcloud CLI, ejecuta gcloud components update para obtener la versión más reciente.
Nota: Para las instalaciones de gcloud CLI existentes, asegúrate de configurar las propiedades compute/region y compute/zone. Cuando configuras las ubicaciones predeterminadas, puedes evitar errores en gcloud CLI como el siguiente: One of [--zone, --region] must be supplied: Please specify location.

Disponibilidad de TPU en GKE

Usa GKE para crear y administrar grupos de nodos con TPU. Puedes usar estos aceleradores para propósitos específicos a fin de realizar entrenamiento, inferencias y entrenamiento de modelos de IA a gran escala.

Consulta una lista de versiones de TPU compatibles en GKE.

Planifica tu configuración de TPU

Planifica tu configuración de TPU según el modelo de aprendizaje automático y la cantidad de memoria que requiere. Los siguientes son los pasos relevantes para planificar tu configuración de TPU:

Asegúrate de que la cuota sea suficiente para las VMs según demanda o Spot

Si creas un grupo de nodos TPU con VMs según demanda o Spot, debes tener suficiente cuota de TPU disponible en la región que deseas usar.

Crear un grupo de nodos TPU que consuma una reserva de TPU no requiere ninguna cuota de TPU.¹ Puedes omitir este paso de forma segura para TPU reservadas.

La creación de un grupo de nodos TPU según demanda o Spot en GKE requiere cuota de la API de Compute Engine. La cuota de la API de Compute Engine (compute.googleapis.com) no es lo mismo que la cuota de la API de Cloud TPU (tpu.googleapis.com), que es necesaria cuando se crean TPU con la API de Cloud TPU.

Si deseas verificar el límite y el uso actual de la cuota de la API de Compute Engine para las TPU, sigue estos pasos:

Dirígete a la página Cuotas en la consola de Google Cloud.

Ir a Cuotas

En la casilla Filtro, haz lo siguiente:

Selecciona la propiedad Servicio, ingresa API de Compute Engine y presiona Intro.
Selecciona la propiedad Tipo y elige Cuota.

Selecciona la propiedad Nombre y, luego, ingresa el nombre de la cuota según la versión de TPU y el tipo de máquina. Por ejemplo, si planeas crear nodos TPU v5e a pedido cuyo tipo de máquina comience con ct5lp-, ingresa TPU v5 Lite PodSlice chips.

Versión de TPU	El tipo de máquina comienza con	Nombre de la cuota para instancias bajo demanda	Nombre de la cuota para las instancias Spot²
TPU v4	`ct4p-`	`TPU v4 PodSlice chips`	`Preemptible TPU v4 PodSlice chips`
TPU v5e	`ct5l-`	`TPU v5 Lite Device chips`	`Preemptible TPU v5 Lite Device chips`
TPU v5e	`ct5lp-`	`TPU v5 Lite PodSlice chips`	`Preemptible TPU v5 Lite PodSlice chips`
TPU v5p	`ct5p-`	`TPU v5p chips`	`Preemptible TPU v5p chips`

Selecciona la propiedad Dimensiones (p. ej., ubicaciones) y, luego, ingresa region: seguido del nombre de la región en la que planeas crear TPUs en GKE. Por ejemplo, ingresa region:us-west4 si planeas crear nodos TPU en la zona us-west4-a. La cuota de TPU es regional, por lo que todas las zonas dentro de la misma región consumen la misma cuota de TPU.

Si ninguna cuota coincide con el filtro que ingresaste, no se le otorgó ninguna de las cuotas especificadas para la región deseada al proyecto y debes solicitar un aumento de la cuota de TPU.

Cuando crees un grupo de nodos TPU, usa las marcas --reservation y --reservation-affinity=specific para crear una instancia reservada. Las reservas de TPU están disponibles cuando se compra un compromiso. ↩
Cuando crees un grupo de nodos TPU, usa la marca --spot para crear una instancia Spot. ↩

Garantiza la disponibilidad de la reserva

La creación de un grupo de nodos TPU reservado, es decir un grupo de nodos TPU que consume una reserva, no requiere ninguna cuota de TPU. Sin embargo, la reserva debe tener suficientes chips disponibles o sin usar en el momento en que se crea el grupo de nodos.

Para ver qué reservas existen dentro de un proyecto, consulta una lista de tus reservas.

Para ver cuántos chips dentro de una reserva de TPU están disponibles, consulta los detalles de una reserva.

Crea un clúster

Crea un clúster de GKE en modo estándar en una región con TPUs disponibles. Recomendamos que uses clústeres regionales, que proporcionan alta disponibilidad del plano de control de Kubernetes. Puedes usar Google Cloud CLI o la consola de Google Cloud.

gcloud container clusters create CLUSTER_NAME \
  --location LOCATION \
  --cluster-version VERSION

Reemplaza lo siguiente:

CLUSTER_NAME el nombre del clúster nuevo.
LOCATION: la región con la capacidad de TPU disponible.
VERSION: la versión de GKE, que debe ser compatible con el tipo de máquina que deseas usar. Ten en cuenta que es posible que la versión predeterminada de GKE no tenga disponibilidad para la TPU de destino. Para obtener información sobre cuáles son las versiones mínimas de GKE disponibles por tipo de máquina de TPU, consulta la disponibilidad de TPU en GKE

Crear un grupo de nodos

Porción de TPU de host único

Puedes crear un grupo de nodos de porción de TPU de host único con Google Cloud CLI, Terraform o la consola de Google Cloud.

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    [--num-nodes=NUM_NODES \]
    [--spot \]
    [--enable-autoscaling \]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME \]
    [--total-min-nodes TOTAL_MIN_NODES \]
    [--total-max-nodes TOTAL_MAX_NODES \]
    [--location-policy=ANY]

Reemplaza lo siguiente:

POOL_NAME: el nombre del grupo de nodos nuevo.
LOCATION: El nombre de la zona en función de la versión de TPU que deseas usar:
- Para TPU v4, usa us-central2-b.
- Para los tipos de máquina de TPU v5e que comienzan con ct5l-, usa us-central1-a o europe-west4-b.
- Para los tipos de máquinas de TPU v5e que comienzan con ct5lp-, usa us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b o europe-west4-a.
- Para TPU v5p, usa us-east1-d, us-east5-a o us-east5-c.
Para obtener más información, consulta la disponibilidad de TPU en GKE.
CLUSTER_NAME: Es el nombre del clúster.
NODE_ZONE: la lista separada por comas de una o más zonas en las que GKE crea el grupo de nodos.
MACHINE_TYPE: el tipo de máquina que se usará para los nodos. Para obtener más información sobre los tipos de máquinas compatibles con TPU, usa la tabla en Asignación de configuraciones de TPU.

De manera opcional, también puedes usar las siguientes marcas:

NUM_NODES: la cantidad inicial de nodos en el grupo de nodos en cada zona. Si omites esta marca, el valor predeterminado es 3. Si el ajuste de escala automático está habilitado para el grupo de nodos mediante la marca --enable-autoscaling, te recomendamos establecer NUM_NODES en 0, ya que el escalador automático aprovisiona nodos adicionales en cuanto tus cargas de trabajo los exigen.
RESERVATION_NAME: el nombre de la reserva que GKE usa cuando crea el grupo de nodos. Si omites esta marca, GKE usa las TPU disponibles. Para obtener más información sobre las reservas de TPU, consulta Reserva de TPU.
--enable-autoscaling: crea un grupo de nodos con el ajuste de escala automático habilitado.
- TOTAL_MIN_NODES: la cantidad mínima de todos los nodos en el grupo de nodos. Omite este campo, a menos que también se especifique el ajuste de escala automático.
- TOTAL_MAX_NODES: la cantidad máxima de todos los nodos en el grupo de nodos. Omite este campo, a menos que también se especifique el ajuste de escala automático.
--spot: Configura el grupo de nodos a fin de usar VMs Spot para los nodos del grupo de nodos. Esto no se puede cambiar después de la creación del grupo de nodos.

Terraform

Asegúrate de usar la versión 4.84.0 o una posterior del proveedor google.
Agrega el siguiente bloque a la configuración de Terraform:

resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]
  initial_node_count = NUM_NODES
  autoscaling {
    total_min_node_count = TOTAL_MIN_NODES
    total_max_node_count = TOTAL_MAX_NODES
    location_policy      = "ANY"
  }

  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    spot = true
  }
}

Reemplaza lo siguiente:

NODE_POOL_RESOURCE_NAME: el nombre del recurso del grupo de nodos en la plantilla de Terraform.
PROJECT_ID: ID del proyecto
CLUSTER_NAME: el nombre del clúster existente.
POOL_NAME: el nombre del grupo de nodos que se creará.
CLUSTER_LOCATION: las zonas de procesamiento del clúster. Especifica la región en la que está disponible la versión de TPU. Para obtener más información, consulta Selecciona una versión y una topología de TPU.
NODE_ZONES: la lista separada por comas de una o más zonas en las que GKE crea el grupo de nodos.
NUM_NODES: la cantidad inicial de nodos en el grupo de nodos en cada una de las zonas del grupo de nodos. Si se omite, el valor predeterminado es 3. Si el ajuste de escala automático está habilitado para el grupo de nodos mediante la plantilla de ajuste de escala automático, te recomendamos que configures NUM_NODES en 0, ya que GKE aprovisiona nodos TPU adicionales apenas los requiera tu carga de trabajo.
MACHINE_TYPE: el tipo de máquina de TPU que se usará. Para ver los tipos de máquinas compatibles con TPU, usa la tabla en Asignación de configuración de TPU.

De manera opcional, también puedes usar las siguientes variables:

autoscaling: Crea un grupo de nodos con el ajuste de escala automático habilitado. Para una porción de TPU de host único, GKE escala entre los valores TOTAL_MIN_NODES y TOTAL_MAX_NODES.
- TOTAL_MIN_NODES: la cantidad mínima de todos los nodos en el grupo de nodos. Este campo es opcional, a menos que también se especifique el ajuste de escala automático.
- TOTAL_MAX_NODES: la cantidad máxima de todos los nodos en el grupo de nodos. Este campo es opcional, a menos que también se especifique el ajuste de escala automático.
RESERVATION_NAME: si usas la reserva de TPU, esta es la lista de etiquetas de los recursos de reserva que se usarán cuando se cree el grupo de nodos. Para obtener más información sobre cómo completar RESERVATION_LABEL_VALUES en el campo reservation_affinity, consulta Proveedor de Terraform Beta.
spot: configura el grupo de nodos a fin de usar VMs Spot para los nodos TPU. Esto no se puede cambiar después de la creación del grupo de nodos. Para obtener más información, consulta VMs Spot.

Consola

Para crear un grupo de nodos con TPUs, haz lo siguiente:

Ve a la página de Google Kubernetes Engine en la consola de Google Cloud.

Ir a Google Kubernetes Engine
En la lista de clústeres, haz clic en el nombre del clúster que deseas modificar.
Haz clic en Agregar grupo de nodos.
En la sección Detalles del grupo de nodos, marca la casilla Especificar las ubicaciones de los nodos.
Elige la zona según la versión de TPU que deseas usar:
- Para TPU v4, usa us-central2-b.
- Para los tipos de máquina de TPU v5e que comienzan con ct5l-, usa us-central1-a o europe-west4-b.
- Para los tipos de máquinas de TPU v5e que comienzan con ct5lp-, usa us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b o europe-west4-a.
- Para TPU v5p, usa us-east1-d, us-east5-a o us-east5-c.
Desde el panel de navegación, haz clic en Nodos.
En la sección Configuración de la máquina, selecciona TPUs.
En el menú desplegable Serie, selecciona una de las siguientes opciones:
- CT4P: TPU v4
- CT5LP: TPU v5e
- CT5P: TPU v5p
En el menú desplegable Tipo de máquina, elige el nombre de la máquina que se usará para los nodos. Usa la tabla Asignación de configuración de TPU para saber cómo definir el tipo de máquina y la topología de TPU que crean un grupo de nodos TPU de host único.
En el menú desplegable Topología de TPU, elige la topología física para la porción de TPU.
En el diálogo Cambios necesarios, haz clic en Hacer cambios.
Asegúrate de que el Tipo de disco de arranque sea Disco persistente estándar o Disco persistente SSD.
De manera opcional, selecciona la casilla de verificación Habilitar nodos en VMs Spot a fin de usar VMs Spot para los nodos en el grupo de nodos.
Haz clic en Crear.

Porción de TPU de varios hosts

Puedes crear un grupo de nodos de porción de TPU de varios hosts con Google Cloud CLI, Terraform o la consola de Google Cloud.

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONE \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    --num-nodes=NUM_NODES \
    [--spot \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME]

Reemplaza lo siguiente:

POOL_NAME: el nombre del grupo de nodos nuevo.
LOCATION: El nombre de la zona en función de la versión de TPU que deseas usar:
- Para TPU v4, usa us-central2-b.
- Los tipos de máquina de TPU v5e que comienzan con ct5l- nunca tienen hosts múltiples.
- Para los tipos de máquinas de TPU v5e que comienzan con ct5lp-, usa us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b o europe-west4-a.
- Para los tipos de máquinas TPU v5p que comienzan con ct5p-, usa us-east1-d, us-east5-a o us-east5-c.
Para obtener más información, consulta la disponibilidad de TPU en GKE.
CLUSTER_NAME: Es el nombre del clúster.
NODE_ZONE: la lista separada por comas de una o más zonas en las que GKE crea el grupo de nodos.
MACHINE_TYPE: el tipo de máquina que se usará para los nodos. Para obtener más información sobre los tipos de máquinas disponibles, consulta Asignación de configuración de TPU.
TPU_TOPOLOGY: la topología física para la porción de TPU. El formato de la topología depende de la versión de la TPU, de la siguiente manera:
- TPU v4 o v5p: Define la topología en 3 tuplas ({A}x{B}x{C}), por ejemplo, 4x4x4.
- TPU v5e: Define la topología en 2 tuplas ({A}x{B}), por ejemplo, 2x2.
Para obtener más información, consulta Topología.
NUM_NODES: la cantidad de nodos en el grupo de nodos. Debe ser cero o el producto de los valores definidos en TPU_TOPOLOGY ({A}x{B}x{C}) divididos por la cantidad de chips en cada VM. Para TPU v4 y TPU v5e de varios hosts, el número de chips en cada VM es cuatro. Por lo tanto, si tu TPU_TOPOLOGY es 2x4x4 (TPU v4 con cuatro chips en cada VM), la NUM_NODES es 32/4, que es igual a 8.

De manera opcional, también puedes usar las siguientes marcas:

RESERVATION_NAME: el nombre de la reserva que GKE usa cuando crea el grupo de nodos. Si omites esta marca, GKE usa los grupos de nodos TPU disponibles. Para obtener más información sobre las reservas de TPU, consulta Reserva de TPU.
--spot: configura el grupo de nodos a fin de usar VMs Spot para los nodos TPU. Esto no se puede cambiar después de la creación del grupo de nodos. Para obtener más información, consulta VMs Spot.
--enable-autoscaling: Crea un grupo de nodos con el ajuste de escala automático habilitado. Cuando GKE escala un grupo de nodos de porción de TPU de varios hosts, escala de forma atómica desde el cero hasta el tamaño máximo.
- MAX_NODES: el tamaño máximo del grupo de nodos. La marca --max-nodes es obligatoria si se proporciona --enable-autoscaling y debe ser igual al producto de los valores definidos en TPU_TOPOLOGY ({A}x{B}x{C}) divididos por la cantidad de chips en cada VM.

Terraform

Asegúrate de usar la versión 4.84.0 o una posterior del proveedor google.
Agrega el siguiente bloque a la configuración de Terraform:
```
resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]
  initial_node_count = NUM_NODES

  autoscaling {
    max_node_count = MAX_NODES
    location_policy      = "ANY"
  }
  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    spot = true
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = TPU_TOPOLOGY
  }
}
```
Reemplaza lo siguiente:
- NODE_POOL_RESOURCE_NAME: el nombre del recurso del grupo de nodos en la plantilla de Terraform.
- PROJECT_ID: ID del proyecto
- CLUSTER_NAME: el nombre del clúster existente al que se agregará el grupo de nodos.
- POOL_NAME: el nombre del grupo de nodos que se creará.
- CLUSTER_LOCATION: la ubicación de procesamiento del clúster. Recomendamos tener un clúster regional para obtener una mayor confiabilidad del plano de control de Kubernetes. También puedes usar un clúster zonal. Para obtener más información, consulta Selecciona una versión y una topología de TPU.
- NODE_ZONES: la lista separada por comas de una o más zonas en las que GKE crea el grupo de nodos.
- NUM_NODES: la cantidad de nodos en el grupo de nodos. Debe ser cero o el producto del número de chips TPU dividido por cuatro, ya que en las porciones de TPU multihost cada nodo TPU tiene 4 chips. Por ejemplo, si TPU_TOPOLOGY es 4x8, entonces hay 32 chips, lo que significa que NUM_NODES debe ser 8. Para obtener más información sobre las topologías de TPU, usa la tabla de Asignación de configuración de TPU.
- TPU_TOPOLOGY: indica la topología física deseada para la porción de TPU. El formato de la topología depende de la versión de la TPU que usas:
  - Para TPU v4: define la topología en 3 tuplas ({A}x{B}x{C}), por ejemplo 4x4x4.
  - Para TPU v5e: define la topología en 2 tuplas ({A}x{B}), por ejemplo, 2x2.
De manera opcional, también puedes usar las siguientes variables:
- RESERVATION_NAME: si usas la reserva de TPU, esta es la lista de etiquetas de los recursos de reserva que se usarán cuando se cree el grupo de nodos. Para obtener más información sobre cómo completar RESERVATION_LABEL_VALUES en el campo reservation_affinity, consulta Proveedor de Terraform Beta.
- autoscaling: Crea un grupo de nodos con el ajuste de escala automático habilitado. Cuando GKE escala un grupo de nodos de porción de TPU de varios hosts, escala de forma atómica desde el cero hasta el tamaño máximo.
  - MAX_NODES: Es el tamaño máximo del grupo de nodos. Debe ser igual al producto de los valores definidos en TPU_TOPOLOGY ({A}x{B}x{C}) divididos por la cantidad de chips en cada VM.
- spot: permite que el grupo de nodos use VMs Spot para los nodos TPU. Esto no se puede cambiar después de la creación del grupo de nodos. Para obtener más información, consulta VMs Spot.

Consola

Para crear un grupo de nodos con TPUs, haz lo siguiente:

Ve a la página de Google Kubernetes Engine en la consola de Google Cloud.

Ir a Google Kubernetes Engine
En la lista de clústeres, haz clic en el nombre del clúster que deseas modificar.
Haz clic en Agregar grupo de nodos.
En la sección Detalles del grupo de nodos, marca la casilla Especificar las ubicaciones de los nodos.
Elige la zona según la versión de TPU que deseas usar:
- Para TPU v4, usa us-central2-b.
- Los tipos de máquina de TPU v5e que comienzan con ct5l- nunca tienen hosts múltiples.
- Para los tipos de máquinas de TPU v5e que comienzan con ct5lp-, usa us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b o europe-west4-a.
- Para los tipos de máquinas TPU v5p que comienzan con ct5p-, usa us-east1-d, us-east5-a o us-east5-c.
Desde el panel de navegación, haz clic en Nodos.
En la sección Configuración de la máquina, selecciona TPUs.
En el menú desplegable Serie, selecciona una de las siguientes opciones:
- CT4P: para TPU v4.
- CT5LP: para TPU v5e
En el menú desplegable Tipo de máquina, elige el nombre de la máquina que se usará para los nodos. Usa la tabla Asignación de configuración de TPU para saber cómo definir el tipo de máquina y la topología de TPU que crean un grupo de nodos TPU de varios hosts.
En el menú desplegable Topología de TPU, elige la topología física para la porción de TPU.
En el diálogo Cambios necesarios, haz clic en Hacer cambios.
Asegúrate de que el Tipo de disco de arranque sea Disco persistente estándar o Disco persistente SSD.
De manera opcional, selecciona la casilla de verificación Habilitar nodos en VMs Spot a fin de usar VMs Spot para los nodos en el grupo de nodos.
Haz clic en Crear.

Estado de aprovisionamiento

Si GKE no puede crear tu grupo de nodos de porción de TPU debido a la capacidad de TPU insuficiente, GKE devuelve un mensaje de error que indica que los nodos TPU no se pueden crear debido a la falta de capacidad.

Si creas un grupo de nodos de porción de TPU de host único, el mensaje de error será similar al siguiente:

2 nodes cannot be created due to lack of capacity. The missing nodes will be
created asynchronously once capacity is available. You can either wait for the
nodes to be up, or delete the node pool and try re-creating it again later.

Si creas un grupo de nodos de porción de TPU de varios hosts, el mensaje de error será similar al siguiente:

The nodes (managed by ...) cannot be created now due to lack of capacity. They
will be created asynchronously once capacity is available. You can either wait
for the nodes to be up, or delete the node pool and try re-creating it again
later.

Tu solicitud de aprovisionamiento de TPU puede permanecer en la cola durante mucho tiempo y permanecerá en el estado "Aprovisionamiento" mientras está en la fila.

Una vez que la capacidad está disponible, GKE crea los nodos restantes que no se crearon.

Si necesitas capacidad antes, considera probar las VMs Spot, aunque ten en cuenta que las VMs Spot consumen una cuota diferente que las instancias bajo demanda.

Puedes borrar la solicitud de TPU en cola con la eliminación del grupo de nodos de porción de TPU.

Ejecuta la carga de trabajo en nodos TPU

Preparación de la carga de trabajo

Las cargas de trabajo de TPU tienen los siguientes requisitos de preparación.

Los frameworks como JAX, PyTorch y TensorFlow acceden a las VMs de TPU con la biblioteca compartida libtpu. libtpu incluye el compilador XLA, el software del entorno de ejecución de TPU y el controlador de TPU. Cada versión de PyTorch y JAX requiere una versión determinada de libtpu.so. Para usar las TPUs en GKE, asegúrate de usar las siguientes versiones:

Tipo de TPU	Versión de `libtpu.so`
TPU v5e `tpu-v5-lite-podslice` `tpu-v5-lite-device`	Versión recomendada de jax[tpu]: v0.4.9 o posterior Versión recomendada de torchxla[tpuvm]: v2.1.0 o posterior
TPU v5p	Versión recomendada de jax[tpu]: 0.4.19 o posterior. Versión recomendada de torchxla[tpuvm]: Se sugiere usar una versión de compilación nocturna del 23 de octubre de 2023.
TPU v4 `tpu-v4-podslice`	jax[tpu]: v0.4.4 o posterior torchxla[tpuvm]: v2.0.0 o posterior

Configura las siguientes variables de entorno para el contenedor que solicita los recursos TPU:
- TPU_WORKER_ID: un número entero único para cada Pod. Este ID denota un ID de trabajador único en la porción de TPU. Los valores admitidos para este campo varían de cero a la cantidad de pods menos uno.
- TPU_WORKER_HOSTNAMES: una lista separada por comas de nombres de host de VM de TPU o direcciones IP que deben comunicarse entre sí dentro de la porción. Debe haber un nombre de host o una dirección IP para cada VM de TPU en la porción. El TPU_WORKER_ID ordena y cero indexa la lista de direcciones IP o nombres de host.
- En el manifiesto de la carga de trabajo, agrega un selector de nodos de Kubernetes para asegurarte de que GKE programe la carga de trabajo de TPU en el tipo de máquina de TPU y la topología de TPU que definiste:
```
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
    cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
  
```
  Reemplaza lo siguiente:
  - TPU_ACCELERATOR: Es el nombre del acelerador de TPU:
    - Para TPU v4, usa tpu-v4-podslice
    - Para los tipos de máquina de TPU v5e que comienzan con ct5l-, usa tpu-v5-lite-device
    - Para los tipos de máquina de TPU v5e que comienzan con ct5lp-, usa tpu-v5-lite-podslice
    - Para TPU v5p, usa tpu-v5p-slice
  - TPU_TOPOLOGY: la topología física para la porción de TPU. El formato de la topología depende de la versión de la TPU, de la siguiente manera:
    - TPU v4: Define la topología en 3 tuplas ({A}x{B}x{C}), por ejemplo, 4x4x4.
    - TPU v5e: Define la topología en 2 tuplas ({A}x{B}), por ejemplo, 2x2.
    - TPU v5p: Define la topología en 3 tuplas ({A}x{B}x{C}), por ejemplo, 4x4x4.

Después de completar la preparación de la carga de trabajo, puedes ejecutar un trabajo que use TPUs.

En las siguientes secciones, se muestran ejemplos sobre cómo ejecutar un trabajo que realiza procesamiento simple con las TPUs.

Ejemplo 1: Ejecuta una carga de trabajo que muestre la cantidad de chips TPU disponibles en un grupo de nodos TPU

La siguiente carga de trabajo devuelve la cantidad de chips TPU en todos los nodos de una porción de TPU de varios hosts. Para crear una porción de varios hosts, la carga de trabajo tiene los siguientes parámetros:

Versión de TPU: TPU v4
Topología: 2 x 2 x 4

Esta elección de versión y topología dan como resultado una porción de varios hosts.

Guarda el siguiente manifiesto como available-chips-multihost.yaml:

apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        securityContext:
          privileged: true
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4

Implementa el manifiesto:
```
kubectl create -f available-chips-multihost.yaml
```
GKE ejecuta una porción de TPU v4 con cuatro VMs de TPU (porción de TPU de varios hosts). La porción tiene 16 chips interconectados.

Verifica que Job haya creado cuatro pods:

kubectl get pods

El resultado es similar al siguiente:

NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s

Obtén los registros de uno de los Pods:
```
kubectl logs POD_NAME
```
Reemplaza POD_NAME por el nombre de uno de los Pods creados. Por ejemplo, tpu-job-podslice-0-5cd8r

El resultado es similar al siguiente:
```
TPU cores: 16
```

Ejemplo 2: Ejecuta una carga de trabajo que muestre la cantidad de chips TPU disponibles en la VM de TPU

La siguiente carga de trabajo es un Pod estático que muestra la cantidad de chips TPU adjuntos a un nodo específico. Para crear un nodo de host único, la carga de trabajo tiene los siguientes parámetros:

Versión de TPU: TPU v5e
Topología: 2 x 4

Esta elección de versión y topología da como resultado una porción de host único.

Guarda el siguiente manifiesto como available-chips-singlehost.yaml:

apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    securityContext:
      privileged: true
    command:
    - bash
    - -c
    - |
      pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8

Implementa el manifiesto:
```
kubectl create -f available-chips-singlehost.yaml
```
GKE aprovisiona nodos con ocho porciones de TPU de host único que usan TPU v5e. Cada VM de TPU tiene ocho chips (porción de TPU de host único).
Obtén los registros del pod:
```
kubectl logs tpu-job-jax-v5
```
El resultado es similar al siguiente:
```
Total TPU chips: 8
```

Actualiza los grupos de nodos mediante aceleradores (GPU y TPU)

GKE actualiza de forma automática los clústeres de Standard, incluidos los grupos de nodos. También puedes actualizar manualmente los grupos de nodos si deseas que los nodos se usen en una versión posterior antes. Para controlar cómo funcionan las actualizaciones de tu clúster, usa canales de versiones, períodos de mantenimiento y exclusiones, y secuenciación de lanzamientos.

También puedes configurar una estrategia de actualización de nodos para tu grupo de nodos, como actualizaciones de aumento o actualizaciones azul-verde. Si configuras estas estrategias, puedes asegurarte de que los grupos de nodos se actualicen de una manera que logre el equilibrio óptimo entre velocidad e interrupción de tu entorno. Para los grupos de nodos de porción de TPU de varios hosts, en lugar de usar la estrategia de actualización de nodos configurada, GKE vuelve a crear de forma atómica todo el grupo de nodos en un solo paso. Para obtener más información, consulta la definición de atomicidad en Terminología relacionada con la TPU en GKE.

El uso de una estrategia de actualización de nodos requerirá que GKE aprovisione de forma temporal recursos adicionales, según la configuración. Si Google Cloud tiene una capacidad limitada para los recursos de tu grupo de nodos, por ejemplo, si ves errores de disponibilidad de recursos cuando intentas crear más nodos con GPU o TPU, consulta Actualiza en un entorno con recursos restringidos.

Limpia

Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos usados en esta guía, considera borrar los grupos de nodos TPU que ya no tengan cargas de trabajo programadas. Si las cargas de trabajo que se ejecutan deben finalizarse de forma correcta, usa kubectl drain para limpiar las cargas de trabajo antes de borrar el nodo.

Borra un grupo de nodos TPU:
```
gcloud container node-pools delete POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME
```
Reemplaza lo siguiente:
- POOL_NAME: Es el nombre del grupo de nodos.
- CLUSTER_NAME: Es el nombre del clúster.
- LOCATION: Es la ubicación de procesamiento del clúster.

Parámetros de configuración adicionales

En las siguientes secciones, se describen las configuraciones adicionales que puedes aplicar a tus cargas de trabajo de TPU.

Porciones múltiples

Puedes hacer una agregación de porciones más pequeñas que formen una porción múltiple para manejar cargas de trabajo de entrenamiento más grandes. Para obtener más información, consulta TPUs de porciones múltiples en GKE.

Migra tu reserva de TPU

Si tienes reservas de TPU existentes, primero debes migrar tu reserva de TPU a un sistema de reserva nuevo basado en Compute Engine. También puedes crear un sistema de reserva basado en Compute Engine en el que no se necesite migración. Para obtener información sobre cómo migrar tus reservas de TPU, consulta Reserva de TPU.

Registros

Los registros que emiten los contenedores que se ejecutan en nodos de GKE, incluidas las VMs de TPU, son recopiladas por el agente de Logging de GKE, se envían a Logging y son visibles en Logging.

Usa el aprovisionamiento automático de nodos de GKE

Puedes configurar GKE a fin de crear y borrar grupos de nodos de forma automática para satisfacer las demandas de recursos de tus cargas de trabajo de TPU. Para obtener más información, consulta Configura Cloud VPN.

Reparación automática de nodos TPU

Si un nodo TPU en un grupo de nodos de porción de TPU de varios hosts está en mal estado, se vuelve a crear todo el grupo de nodos. Las condiciones que dan como resultado nodos TPU en mal estado incluyen:

Cualquier nodo TPU con condiciones de nodo comunes
Cualquier nodo TPU con un recuento de TPU no asignables superior a cero
Cualquier instancia de VM de TPU que se detenga (debido a una interrupción) o que se cierre
Mantenimiento de nodos: si algún nodo TPU (VM) dentro de un grupo de nodos de porción de TPU de varios hosts queda inactivo para el mantenimiento del host, GKE vuelve a crear toda la porción de TPU.

Puedes ver el estado de reparación (incluido el motivo de la falla) en el historial de operaciones. Si la falla se debe a una cuota insuficiente, comunícate con tu representante de cuenta de Google Cloud para aumentar la cuota correspondiente.

Configura la finalización correcta del nodo TPU

En los clústeres de GKE con el plano de control que ejecuta 1.29.1-gke.1425000 o una versión posterior, los nodos TPU admiten los indicadores SIGTERM que alertan al nodo de un cierre inminente. La notificación de cierre inminente se puede configurar hasta cinco minutos en nodos TPU.

Puedes configurar GKE para finalizar tus cargas de trabajo del AA de forma correcta en el plazo de este período de notificación. Durante la finalización ordenada, tus cargas de trabajo pueden llevar a cabo procesos de limpieza, como almacenar datos de la carga de trabajo para reducir la pérdida de datos. Para obtener el tiempo máximo de notificación, en el manifiesto de tu Pod, establece el campo spec.terminationGracePeriodSeconds en 300 segundos (cinco minutos) de la siguiente manera:

    spec:
      terminationGracePeriodSeconds: 300

GKE hace su mejor esfuerzo para finalizar estos Pods de forma correcta y ejecutar la acción de finalización que definiste, por ejemplo, guardar un estado de entrenamiento.

Ejecutar contenedores sin modo privilegiado

Si tu nodo TPU ejecuta versiones anteriores a la 1.28, lee la siguiente sección:

Un contenedor que se ejecuta en una VM de TPU necesita acceso a límites más altos en la memoria bloqueada para que el controlador pueda comunicarse con los chips TPU mediante el acceso directo a la memoria (DMA). Para habilitar esto, debes configurar un ulimit más alto. Si deseas reducir el permiso de permisos en tu contenedor, completa los siguientes pasos:

Edita securityContext para que incluya los siguientes campos:

securityContext:
  capabilities:
    add: ["SYS_RESOURCE"]

Aumenta ulimit mediante la ejecución del siguiente comando dentro del contenedor antes de configurar las cargas de trabajo para que usen recursos TPU:
```
ulimit -l 68719476736
```

Nota: En TPU v5e, la ejecución de contenedores sin modo privilegiado está disponible en los clústeres en la versión 1.27.4-gke.900 y posteriores.

Observabilidad y métricas

Panel

En la página Clústeres de Kubernetes de la consola de Google Cloud, la pestaña Observabilidad muestra las métricas de observabilidad de TPU. Para obtener más información, consulta Métricas de observabilidad de GKE.

El panel de TPU solo se propaga si tienes las métricas del sistema habilitadas en tu clúster de GKE.

Métricas del entorno de ejecución

En la versión 1.27.4-gke.900 o posterior de GKE, las cargas de trabajo de TPU que usan una versión de JAX 0.4.14 o posterior y especifican containerPort: 8431 la exportación de las métricas de utilización de TPU como las métricas del sistema de GKE. Las siguientes métricas están disponibles en Cloud Monitoring para supervisar el rendimiento del entorno de ejecución de tu carga de trabajo de TPU:

Ciclo de trabajo: Porcentaje de tiempo durante el último período de muestreo (60 segundos) durante el cual los TensorCores se procesaron de forma activa en un chip TPU. Un porcentaje mayor significa un mejor uso de TPU.
Uso de memoria: cantidad de memoria del acelerador asignada en bytes. Se tomaron muestras cada 60 segundos.
Memoria total: el total de memoria del acelerador en bytes. Se hace un muestreo cada 60 segundos.

Estas métricas se encuentran en el esquema de nodos de Kubernetes (k8s_node) y de contenedor de Kubernetes (k8s_container).

Contenedor de Kubernetes:

kubernetes.io/container/accelerator/duty_cycle
kubernetes.io/container/accelerator/memory_used
kubernetes.io/container/accelerator/memory_total

Nodo de Kubernetes:

kubernetes.io/node/accelerator/duty_cycle
kubernetes.io/node/accelerator/memory_used
kubernetes.io/node/accelerator/memory_total

Métricas de host

En la versión 1.28.1-gke.1066000 o posterior de GKE, la VM de TPU exporta las métricas de uso de TPU como métricas del sistema de GKE. Las siguientes métricas están disponibles en Cloud Monitoring para supervisar el rendimiento de tu host TPU:

Uso de TensorCore: porcentaje actual del TensorCore que se usa. El valor de TensorCore es igual a la suma de las unidades de multiplicación de matriz (MXUs) más la unidad vectorial. El valor de uso de TensorCore es la división de las operaciones de TensorCore que se llevaron a cabo durante el último período de muestra (60 segundos) por la cantidad compatible de operaciones de TensorCore. durante el mismo período. Un valor mayor significa un mejor uso.
Uso de ancho de banda de memoria: Porcentaje actual del ancho de banda de memoria del acelerador que se usa. Se calcula dividiendo el ancho de banda de memoria usado durante un período de muestra (60 s) por el ancho de banda máximo admitido durante el mismo período de muestra.

Estas métricas se encuentran en el esquema de nodos de Kubernetes (k8s_node) y de contenedor de Kubernetes (k8s_container).

Contenedor de Kubernetes:

kubernetes.io/container/accelerator/tensorcore_utilization
kubernetes.io/container/accelerator/memory_bandwidth_utilization

Nodo de Kubernetes:

kubernetes.io/container/node/tensorcore_utilization
kubernetes.io/container/node/memory_bandwidth_utilization

Para obtener más información, consulta Métricas de Kubernetes y Métricas del sistema de GKE.

Problemas conocidos

El escalador automático del clúster podría calcular de forma incorrecta la capacidad de los nodos TPU nuevos antes de que esos nodos informen las TPUs disponibles. Luego, el escalador automático del clúster puede realizar un escalamiento vertical adicional y, como resultado, crear más nodos de los necesarios. El escalador automático del clúster reducirá la escala verticalmente de los nodos adicionales, si no son necesarios, después de la operación de reducción de escala vertical regular.
El escalador automático de clústeres cancela el escalamiento vertical de los grupos de nodos TPU que permanecen en estado de espera durante más de 10 horas. El escalador automático del clúster volverá a intentar esas operaciones de escalamiento vertical más adelante. Este comportamiento puede reducir la capacidad de obtención de TPUs para los clientes que no usan reservas.
Las cargas de trabajo que no sean de TPU y tengan tolerancia al taint de TPU pueden impedir la reducción de la escala vertical del grupo de nodos si se vuelven a crear durante el vaciado del grupo de nodos TPU.

Implementa cargas de trabajo de TPU en GKE Standard

Antes de comenzar

Disponibilidad de TPU en GKE

Planifica tu configuración de TPU

Asegúrate de que la cuota sea suficiente para las VMs según demanda o Spot

Garantiza la disponibilidad de la reserva

Crea un clúster

Crear un grupo de nodos

Porción de TPU de host único

gcloud

Terraform

Consola

Porción de TPU de varios hosts

gcloud

Terraform

Consola

Estado de aprovisionamiento

Ejecuta la carga de trabajo en nodos TPU

Preparación de la carga de trabajo

Ejemplo 1: Ejecuta una carga de trabajo que muestre la cantidad de chips TPU disponibles en un grupo de nodos TPU

Ejemplo 2: Ejecuta una carga de trabajo que muestre la cantidad de chips TPU disponibles en la VM de TPU

Actualiza los grupos de nodos mediante aceleradores (GPU y TPU)

Limpia

Parámetros de configuración adicionales

Porciones múltiples

Migra tu reserva de TPU

Registros

Usa el aprovisionamiento automático de nodos de GKE

Reparación automática de nodos TPU

Configura la finalización correcta del nodo TPU

Ejecutar contenedores sin modo privilegiado

Observabilidad y métricas

Panel

Métricas del entorno de ejecución

Métricas de host

Problemas conocidos

¿Qué sigue?