Implementa TPU Multislices en GKE

En esta página, se muestra cómo implementar cargas de trabajo en Google Kubernetes Engine (GKE) mediante la configuración de Cloud TPU Multislice para un entrenamiento rentable y a gran escala.

Antes de configurar Multislice en GKE, debes familiarizarte con los siguientes conceptos:

  1. Introducción a Cloud TPU
  2. Arquitectura de sistemas de Cloud TPU
  3. Acerca de las TPU en GKE

¿Qué es TPU Multislice?

TPU Multislice es la organización arquitectónica de las VMs en una porción de TPU en la que dos o más porciones de Cloud TPU se comunican a través de la red del centro de datos (DCN). Multislice permite el entrenamiento de pila completa, rentable y a gran escala con escalamiento casi lineal hasta decenas de miles de chips TPU. En una configuración Multislice, GKE implementa una carga de trabajo Multislice en varias porciones de TPU. La comunicación entre los chips TPU dentro de una porción se realiza a través de interconexiones entre chips (ICI). La comunicación entre porciones se realiza a través del DCN.

Te recomendamos que uses Multislice si tu trabajo es demasiado grande para caber en una porción de TPU única.

Disponibilidad de Multislice en GKE

  • El estándar es compatible con Multislice en la versión 1.27.4-gke.900 y versiones posteriores.
  • Autopilot admite Multislice en la versión 1.29.2-gke.1521000 y posteriores.
  • Multislice admite frameworks de JAX y PyTorch. La versión mínima admitida de JAX es 2.1.
  • Multislice solo admite grupos de nodos de porciones de TPU de varios hosts. Por ejemplo, no puedes usar Multislice con un ct4p-hightpu-4t con una topología 2x2x1 o un ct5lp-hightpu-4t con una topología 2x2, ya que estos son grupos de nodos de porción de TPU de host único.
  • Multislice solo admite el entrenamiento de varios controladores síncrono.
  • Las cargas de trabajo de Multislice solo pueden ejecutarse en porciones de TPU que comparten el mismo tipo, tamaño y topología de TPU.

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.

Ejecuta una carga de trabajo en una Multislice

En esta sección, se muestra cómo ejecutar una carga de trabajo en una Multislice. Si usas el modo Autopilot de GKE, ve a la sección Ejecuta una carga de trabajo de Multislice. Los clústeres de Autopilot que ejecutan la versión 1.29.2-gke.1521000 o posterior habilitan las TPUs de forma predeterminada.

Prepara un grupo de nodos de modo Standard

En esta sección, se incluyen los siguientes pasos:

  1. Crea tres grupos de nodos de porciones de TPU de varios hosts
  2. Verifica el estado del grupo de nodos

Crea el grupo de nodos de porción de TPU

Puedes crear más de un grupo de nodos de porción de TPU de varios hosts. Para los fines de esta guía, crea tres grupos de nodos de porción de TPU de varios hosts para ejecutar una carga de trabajo de Multislice. 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 de porción de 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 de porción de 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

  1. Asegúrate de usar la versión 4.84.0 o una posterior del proveedor google.

  2. 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 de porción de 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 de porción de 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.

Console

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

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

    Ir a Google Kubernetes Engine

  2. En la lista de clústeres, haz clic en el nombre del clúster que deseas modificar.

  3. Haz clic en Agregar grupo de nodos.

  4. En la sección Detalles del grupo de nodos, marca la casilla Especificar las ubicaciones de los nodos.

  5. 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.

  6. Desde el panel de navegación, haz clic en Nodos.

  7. En la sección Configuración de la máquina, selecciona TPUs.

  8. En el menú desplegable Serie, selecciona una de las siguientes opciones:

    • CT4P: para TPU v4.
    • CT5LP: para TPU v5e

  9. 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 de porción de TPU de varios hosts.

  10. En el menú desplegable Topología de TPU, elige la topología física para la porción de TPU.

  11. En el diálogo Cambios necesarios, haz clic en Hacer cambios.

  12. Asegúrate de que el Tipo de disco de arranque sea Disco persistente estándar o Disco persistente SSD.

  13. 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.

  14. Haz clic en Crear.

Verifica el estado del grupo de nodos

  1. Obtén credenciales de modo que puedas usar kubectl para acceder al clúster:

    gcloud container clusters get-credentials CLUSTER_NAME \
    --project=PROJECT_ID

    Reemplaza lo siguiente:

    • CLUSTER_NAME: Es el nombre del clúster.
    • PROJECT_ID: ID del proyecto

  2. Usa kubectl en Cloud Shell para ver tus nodos de porción de TPU:

    kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=TPU_ACCELERATOR \
   -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGY

    Reemplaza lo siguiente:

    • TPU_ACCELERATOR: El tipo de acelerador de TPU que usaste cuando creaste los grupos de nodos. Por ejemplo, tpu-v4-podslice, tpu-v5-lite-device o tpu-v5-lite-podslice.
    • TPU_TOPOLOGY: la topología física para la porción de TPU.

    El resultado es similar al siguiente:

     NAME                                    STATUS   ROLES    AGE    VERSION
 gke-tpu-20ee2cce-5tv6                   Ready    <none>   34h     v1.28.1-gke.1066000

Ejecuta una carga de trabajo Multislice

En esta sección, ejecutarás una carga de trabajo de JAX que muestra la cantidad global de chips TPU en la porción de TPU y, luego, se cierra.

Para ejecutar una carga de trabajo de JAX, haz lo siguiente:

  1. Crea el siguiente manifiesto tpu-multislice.yaml:

    Autopilot

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-job
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: NUM_SLICES
      template:
        spec:
          parallelism: NUM_NODES
          completions: NUM_NODES
          backoffLimit: 0
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                command:
                - bash
                - -c
                - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  python -c 'import jax; print("Global device count:", jax.device_count())'
                  sleep 60
                resources:
                 limits:
                    google.com/tpu: NUM_CHIPS

    Estándar

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-job
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: NUM_SLICES
      template:
        spec:
          parallelism: NUM_NODES
          completions: NUM_NODES
          backoffLimit: 0
          template:
            spec:
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                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("Global device count:", jax.device_count())'
                  sleep 60
                resources:
                  limits:
                   google.com/tpu: NUM_CHIPS

    Reemplaza lo siguiente:

    • NUM_SLICES: La cantidad de grupos de nodos de porción de TPU. En este caso, NUM_SLICES es igual a 3.
    • ACCELERATOR_TYPE: El tipo de acelerador de TPU que usaste cuando creaste los grupos de nodos. Por ejemplo, tpu-v4-podslice, tpu-v5-lite-device o tpu-v5-lite-podslice.
    • TPU_TOPOLOGY: la topología física para la porción de TPU. Por ejemplo, 4x4x4 o 2x2, según la versión de TPU.
    • 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 TPU en cada VM. Para TPU v4 de varios hosts, el número de chips TPU en cada VM es cuatro. Para TPU v5e de varios hosts, el número de chips TPU en cada VM es uno, cuatro u ocho. Por lo tanto, si tu TPU_TOPOLOGY es 2x4x4 (TPU v4 con cuatro chips TPU en cada VM), la NUM_NODES es 32/4, que es igual a 8.
    • NUM_CHIPS: Para TPU v4 de varios hosts, la cantidad de chips TPU en cada VM es cuatro. Para TPU v5e de varios hosts, el número de chips TPU en cada VM es uno, cuatro u ocho. Para obtener más información, consulta Chips TPU en la VM en una porción de TPU.

    En el manifiesto se muestra lo siguiente:

    • El JobSet es un Service sin interfaz gráfica con el mismo nombre que el nombre del JobSet; en este caso es multislice-job.
    • El maxRestarts: 4 indica la cantidad máxima de veces que GKE reinicia el JobSet cuando falla un trabajo secundario. Si el JobSet se reinicia alcanza el máximo definido, el JobSet se marca como con errores.
    • Los campos parallelism y completions son iguales a la cantidad de nodos en cada grupo de nodos.
    • backoff es 0 porque Multislice solo admite el entrenamiento de varios controladores síncrono. Se debe configurar como 0. Haz que falle el trabajo cuando falle cualquier Pod.
    • Los valores en la sección de afinidad garantizan que solo haya una carga de trabajo de TPU Multislice en ejecución en un grupo de Multislices.
    • containerPort: 8080 es el puerto para el coordinador de MXLA
    • containerPort: 8431 es el puerto para exportar las métricas de uso de TPU
    • El securityContext: privileged: true indica que los nodos tienen habilitado el modo privilegiado para acceder a las TPU. Los nodos en GKE versión 1.28 o posterior no necesitan tener habilitado el modo privilegiado para acceder a las TPU. Para obtener más información, consulta Ejecuta contenedores sin modo privilegiado.

  2. Aplica el manifiesto

    kubectl apply -f tpu-multislice.yaml

  3. Confirma que se admita la carga de trabajo:

    kubectl get jobsets

    El resultado es similar al siguiente:

    NAME            RESTARTS   COMPLETED   AGE
multislice-job                         3s

  4. Supervisa el estado de los Pods aprovisionados:

    kubectl get pods

    El resultado es similar al siguiente:

     NAME                                READY   STATUS      RESTARTS   AGE
 multislice-job-slice-0-0-wzq9t      0/1     Completed   0          2m31s
 multislice-job-slice-0-1-zf4dp      0/1     Completed   0          2m30s
 multislice-job-slice-1-0-hbfn5      0/1     Completed   0          2m31s
 multislice-job-slice-1-1-45fgl      0/1     Completed   0          2m30s
 multislice-job-slice-2-0-wjbp4      0/1     Completed   0          2m30s
 multislice-job-slice-2-1-lwnvs      0/1     Completed   0          2m30s

El JobSet multislice-job programa, crea y, luego, ejecuta los Pods hasta su finalización. Los nombres de los Pods tienen el formato <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>. El prefijo jobsetName determina el JobSet al que pertenece el Pod.

Parámetros de configuración adicionales

En las siguientes secciones, se describen las configuraciones adicionales que puedes aplicar a tu Multislice.

Habilita hostNetwork en tus Pods de GKE Standard

Para mejorar el rendimiento de la red entre porciones de TPU, te recomendamos que actives hostNetworking. Usa hostNetwork: true en las especificaciones de tu Pod para omitir toda la pila de herramientas de redes de Kubernetes y permitir que tus Pods de Kubernetes usen la red del host directamente para la comunicación de VM a VM.

Para activar hostNetworking, quita las siguientes dos líneas de tu especificación del Pod:

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

Si quieres seguir usando podHostnames para el descubrimiento de nodos trabajadores con hostNetwork, configura dnsPolicy: ClusterFirstWithHostNet. Esto es importante cuando ejecutas trabajos de entrenamiento de reanudación automática y necesitas tener los mismos nombres para volver a cargar los mismos puntos de control.

Logging

Los registros que emiten los contenedores que se ejecutan en nodos de GKE, incluidos los nodos de porción de TPU, se pueden ver en el Explorador de registros, si tienes el registro del sistema de GKE habilitado en tu clúster.

Puedes ver tus registros desde GKE mediante el Explorador de registros con el siguiente filtro a fin depara ver los registros del contenedor para tu carga de trabajo:

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME

Usa el siguiente filtro para la porción y los trabajadores de TPU:

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME
resource.labels.pod_name:<jobSetName>-<replicateJobName>-<job-index>-<worker-index>

Para obtener más información, consulta Visualiza los registros de TPU de GKE.

Observabilidad y métricas

Además de las métricas de TPU generales, hay 4 métricas adicionales del entorno de ejecución de TPU específicas de varias porciones. Estas métricas están disponibles en GKE versión 1.29.1-gke.1016000 o posterior. La carga de trabajo de TPU debe usar la versión de JAX 0.4.24

Las siguientes son las métricas de varias porciones disponibles:

  • Latencias de transferencia de DCN (red de centros de datos): distribución de las latencias de transferencia de red para el tráfico de varias porciones.
  • Latencias colectivas: distribución de la latencia colectiva de extremo a extremo para el tráfico de varias porciones.
  • Latencias de transferencia de host a dispositivo: distribución de la latencia de transferencia de host al dispositivo de cada fragmento de datos para el tráfico de varias porciones.
  • Latencias de transferencia de dispositivo a host: distribución de la latencia de transferencia d dispositivo al host de cada fragmento de datos para el tráfico de varias porciones.

Estas métricas se encuentran en el esquema del contenedor de Kubernetes (k8s_container):

  • kubernetes.io/container/multislice/network/dcn_transfer_latencies
  • kubernetes.io/container/multislice/network/collective_end_to_end_latencies
  • kubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencies
  • kubernetes.io/container/multislice/accelerator/device_to_host_transfer_latencies

Porción de TPU frente a Multislice

En la siguiente tabla, se diferencia la organización de arquitectura de una porción de TPU y una Multislice:

Porción de TPU Multislice
Interconectividad La carga de trabajo se ejecuta en una porción de TPU única. Todos los chips TPU en una porción están conectados con ICI. La carga de trabajo se ejecuta en varias porciones de TPU. La comunicación dentro de una porción se realiza a través de ICI. La comunicación entre porciones se produce a través de DCN.
Grupos de nodos compatibles Porción de TPU de host único y porción de TPU de varios hosts Grupos de porciones de TPU de varios hosts
Tipo de carga de trabajo recomendado IndexedJob o JobSet JobSet

¿Qué sigue?