Desplegar Multislices de TPU en GKE


En esta página se explica cómo desplegar cargas de trabajo en Google Kubernetes Engine (GKE) mediante la configuración de Multislice de TPU de Cloud para llevar a cabo entrenamientos a gran escala y rentables.

Este tutorial está dirigido a ingenieros de aprendizaje automático y administradores y operadores de plataformas que quieran usar la orquestación de contenedores de Kubernetes para gestionar cargas de trabajo de entrenamiento, ajuste y inferencia de modelos a gran escala con TPUs. Para obtener más información sobre los roles habituales y las tareas de ejemplo que se mencionan en el contenido de Google Cloud , consulta Roles y tareas de usuario habituales de GKE.

Antes de configurar Multislice en GKE, familiarízate con los siguientes conceptos:

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

Qué es Multislice de TPU

Multisector de TPU es la organización arquitectónica de las VMs en un sector de TPU en la que dos o más sectores de TPU de Cloud se comunican a través de la red del centro de datos (DCN). Multislice permite realizar entrenamientos a gran escala, rentables y de pila completa con un escalado casi lineal de hasta decenas de miles de chips de TPU. En una configuración de varios sectores, GKE despliega una carga de trabajo de varios sectores en varios sectores de TPU. La comunicación entre los chips de TPU de un slice se produce a través de interconexiones entre chips (ICI). La comunicación entre las porciones se produce a través de la DCN.

Te recomendamos que uses Multislice si tu trabajo es demasiado grande para caber en un solo slice de TPU.

Disponibilidad de Multislice en GKE

  • Standard admite Multislice en la versión 1.27.4-gke.900 y posteriores.
  • Autopilot admite Multislice en la versión 1.29.2-gke.1521000 y posteriores.
  • Multislice es compatible con los frameworks JAX y PyTorch. La versión mínima compatible de JAX es la 2.1.
  • Multislice solo admite grupos de nodos de slices de TPU multihost. Por ejemplo, no puedes usar Multislice con una ct4p-hightpu-4t con una topología 2x2x1 o una ct5lp-hightpu-4t con una topología 2x2, porque se trata de grupos de nodos de segmentos de TPU de un solo host.
  • Multislice solo admite el entrenamiento multicontrolador síncrono.
  • Las cargas de trabajo de Multislice solo se pueden ejecutar en slices de TPU que compartan el mismo tipo, tamaño y topología de TPU.
  • Multislice no es compatible con la TPU v3.

Antes de empezar

Antes de empezar, asegúrate de que has realizado las siguientes tareas:

  • Habilita la API de Google Kubernetes Engine.
  • Habilitar la API de Google Kubernetes Engine
  • Si quieres usar Google Cloud CLI para esta tarea, instálala y, a continuación, inicialízala. Si ya has instalado la gcloud CLI, obtén la versión más reciente ejecutando gcloud components update.

Ejecutar una carga de trabajo en un Multislice

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

Preparar un grupo de nodos en modo Estándar

En esta sección se explican los siguientes pasos:

  1. Crea tres grupos de nodos de segmento de TPU multihost
  2. Verificar el estado del grupo de nodos

Crear el grupo de nodos del slice de TPU

Puedes crear más de un grupo de nodos de segmento de TPU de varios hosts. Para esta guía, crea tres grupos de nodos de slices de TPU de varios hosts para ejecutar una carga de trabajo Multislice. Puedes crear un nodo de slice de TPU multihost mediante la CLI de Google Cloud, Terraform o la Google Cloud consola.

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    [--num-nodes=NUM_NODES] \
    [--spot \]
    [--flex-start \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME] \
    [--node-labels cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME,cloud.google.com/gke-workload-type=HIGH_AVAILABILITY]
    [--placement-type=COMPACT]

Haz los cambios siguientes:

  • POOL_NAME: el nombre del nuevo grupo de nodos.
  • LOCATION: nombre de la zona en función de la versión de TPU que quieras usar. Para identificar una ubicación disponible, consulta Disponibilidad de las TPUs en GKE.
  • CLUSTER_NAME: el nombre del clúster.
  • NODE_ZONES: lista separada por comas de una o varias zonas en las que GKE crea el grupo de nodos.
  • MACHINE_TYPE: el tipo de máquina que se va a usar en los nodos. Para obtener más información sobre los tipos de máquinas disponibles, consulta Elegir la versión de TPU.
  • TPU_TOPOLOGY: la topología física del segmento de TPU. El formato de la topología depende de la versión de la TPU. Para obtener más información sobre las topologías de TPU, consulta la tabla de la sección Elegir una topología.

    Para obtener más información, consulta Topología.

También puedes usar los siguientes elementos opcionales:

  • NUM_NODES: número de nodos del grupo de nodos. Debe ser cero o el producto de los valores definidos en TPU_TOPOLOGY ({A}x{B}x{C}) dividido entre el número de chips de cada máquina virtual. En el caso de las TPU v4 y v5e multihost, el número de chips de cada VM es cuatro. Por lo tanto, si tu TPU_TOPOLOGY es 2x4x4 (TPU v4 con cuatro chips en cada VM), entonces el NUM_NODES es 32/4, lo que equivale a 8. Si omite esta marca, el número de nodos se calcula y se establece de forma predeterminada en función de la topología y el tipo de máquina.
  • RESERVATION_NAME: nombre de la reserva que usa GKE al crear el grupo de nodos. Si omite esta marca, GKE usará los grupos de nodos de segmentos de TPU disponibles. Para obtener más información sobre las reservas de TPU, consulta el artículo Reservas de TPU.
  • --spot: define el grupo de nodos que usará las VMs de Spot para los nodos de la porción de TPU. No se puede cambiar después de crear el grupo de nodos. Para obtener más información, consulta VMs de acceso puntual.
  • --flex-start: define que el grupo de nodos use el modo de aprovisionamiento flex-start. El inicio flexible se admite en GKE 1.33.0-gke.1712000 o versiones posteriores.
  • --enable-autoscaling: crea un grupo de nodos con el autoescalado habilitado. Cuando GKE escala un grupo de nodos de un segmento de TPU multianfitrión, escala el grupo de nodos de forma atómica de cero al tamaño máximo.

    • MAX_NODES: 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}) dividido entre el número de chips de cada máquina virtual.
  • --node-label=cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME, cloud.google.com/gke-workload-type=HIGH_AVAILABILITY: indica a GKE que el grupo de nodos de la porción de TPU multihost es una colección. Usa esta marca si se cumplen las siguientes condiciones:

    • El grupo de nodos ejecuta cargas de trabajo de inferencia en el nuevo grupo de nodos.
    • El grupo de nodos usa TPU Trillium.
    • Las máquinas virtuales de acceso puntual no admiten la programación de la recopilación.

    Para obtener más información sobre la gestión de la programación de la recogida, consulta Gestionar la programación de la recogida en slices de TPU de varios hosts.

  • --placement-type=COMPACT: crea un grupo de nodos con la colocación compacta habilitada. Esta opción debe usarse con la marca --tpu-topology. Para obtener más información, consulta Crear una política de colocación compacta y Topología de TPU.

Terraform

  1. Asegúrate de usar la versión 4.84.0 o una posterior del proveedor google.
  2. Añade el siguiente bloque a tu 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
        flex_start = false
      }
    
      placement_policy {
        type = "COMPACT"
        tpu_topology = TPU_TOPOLOGY
      }
    }
    

    Haz los cambios siguientes:

    • NODE_POOL_RESOURCE_NAME: el nombre del recurso de perfil de nodo en la plantilla de Terraform.
    • PROJECT_ID: tu ID de proyecto.
    • CLUSTER_NAME: el nombre del clúster al que quieres añadir el grupo de nodos.
    • POOL_NAME: nombre del grupo de nodos que se va a crear.
    • CLUSTER_LOCATION: ubicación de cálculo del clúster. Te recomendamos que tengas un clúster regional para aumentar la fiabilidad del plano de control de Kubernetes. También puedes usar un clúster zonal. Para obtener más información, consulta Seleccionar una versión y una topología de TPU.
    • NODE_ZONES: lista separada por comas de una o varias zonas en las que GKE crea el grupo de nodos.
    • NUM_NODES: número de nodos del grupo de nodos. Debe ser cero o el resultado de dividir el número de chips de TPU entre cuatro, ya que en las porciones de TPU multihost, cada nodo de la porción de TPU tiene cuatro chips. Por ejemplo, si TPU_TOPOLOGY es 4x8, hay 32 chips, lo que significa que NUM_NODES debe ser 8. Para obtener más información sobre las topologías de TPU, consulta la tabla de la sección Elegir la versión de TPU.
    • TPU_TOPOLOGY: indica la topología física deseada de la porción de TPU. El formato de la topología depende de la versión de TPU que estés usando. Para obtener más información sobre las topologías de TPU, consulta la tabla de la sección Elegir una topología.

    También puedes usar las siguientes variables:

    • RESERVATION_NAME: si usas una reserva de TPU, esta es la lista de etiquetas de los recursos de reserva que se usarán al crear el grupo de nodos. Para obtener más información sobre cómo rellenar el campo RESERVATION_LABEL_VALUES en reservation_affinity, consulta Proveedor de Terraform.
    • autoscaling: crea un grupo de nodos con el autoescalado habilitado. Cuando GKE escala un grupo de nodos de un segmento de TPU multianfitrión, escala el grupo de nodos de forma atómica de cero al 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}) dividido entre el número de chips de cada máquina virtual.
    • spot: permite que el grupo de nodos use máquinas virtuales de Spot para los nodos de la porción de TPU. No se puede cambiar después de crear el grupo de nodos. Para obtener más información, consulta Máquinas virtuales de acceso puntual.
    • flex_start: define que el grupo de nodos use el modo de aprovisionamiento flex-start. No se puede definir como true si spot está habilitado.

Consola

Para crear un grupo de nodos con TPUs, sigue estos pasos:

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

    Ir a Google Kubernetes Engine

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

  3. Haz clic en Añadir grupo de nodos.

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

  5. Selecciona el nombre de la zona en función de la versión de TPU que quieras usar. Para identificar una ubicación disponible, consulta Disponibilidad de las TPUs en GKE.

  6. En el panel de navegación, haga 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:

    • CT3P: para TPU v3.
    • CT4P: para TPU v4.
    • CT5LP: para la TPU v5e.
  9. En el menú desplegable Tipo de máquina, selecciona el nombre de la máquina que quieras usar para los nodos. Consulta la tabla Elegir la versión de TPU para saber cómo definir el tipo de máquina y la topología de TPU que crean un bloque de nodos de TPU multihost.

  10. En el menú desplegable Topología de TPU, selecciona la topología física del slice de TPU.

  11. En el cuadro de diálogo Cambios necesarios, haga clic en Hacer cambios.

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

  13. Si quieres, marca la casilla Habilitar nodos en VMs Spot para usar VMs Spot en los nodos del grupo de nodos.

  14. Haz clic en Crear.

Verificar el estado del grupo de nodos

  1. Obtén las credenciales para poder usar kubectl y acceder al clúster:

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

    Haz los cambios siguientes:

    • CLUSTER_NAME: el nombre del clúster.
    • PROJECT_ID: tu ID de proyecto.
    • CONTROL_PLANE_LOCATION: la ubicación de Compute Engine del plano de control de tu clúster. Proporciona una región para los clústeres regionales o una zona para los clústeres zonales.
  2. Usa kubectl en Cloud Shell para ver los nodos de tu segmento de TPU:

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

    Haz los cambios siguientes:

    • TPU_ACCELERATOR: el tipo de acelerador de TPU que has usado al crear los grupos de nodos. Por ejemplo, tpu-v4-podslice o tpu-v5-lite-podslice.
    • TPU_TOPOLOGY: la topología física del segmento de TPU.

    El resultado debería ser similar al siguiente:

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

Ejecutar una carga de trabajo de Multislice

En esta sección, se ejecuta una carga de trabajo de JAX que muestra el número global de chips de TPU en el segmento de TPU y, a continuación, se cierra.

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

  1. Crea el siguiente archivo de 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: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    - containerPort: 8431
                    command:
                    - bash
                    - -c
                    - |
                      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: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    - containerPort: 8431
                    securityContext:
                      privileged: true
                    command:
                    - bash
                    - -c
                    - |
                      python -c 'import jax; print("Global device count:", jax.device_count())'
                      sleep 60
                    resources:
                      limits:
                       google.com/tpu: NUM_CHIPS
    

    Haz los cambios siguientes:

    • NUM_SLICES: número de grupos de nodos de segmento de TPU. En este caso, NUM_SLICES es igual a 3.
    • ACCELERATOR_TYPE: el tipo de acelerador de TPU que has usado al crear los grupos de nodos. Por ejemplo, tpu-v4-podslice o tpu-v5-lite-podslice.
    • TPU_TOPOLOGY: la topología física del segmento de TPU. Por ejemplo, 4x4x4 o 2x2, según la versión de la TPU.
    • NUM_NODES: número de nodos del grupo de nodos. Debe ser cero o el producto de los valores definidos en TPU_TOPOLOGY ({A}x{B}x{C}) dividido entre el número de chips de TPU de cada máquina virtual. En el caso de las TPU v4 multihost, el número de chips de TPU de cada VM es cuatro. En el caso de las TPU v5e multihost, el número de chips de TPU de cada VM es uno, cuatro u ocho. Por lo tanto, si tu TPU_TOPOLOGY es 2x4x4 (TPU v4 con cuatro chips de TPU en cada VM), el NUM_NODES es 32/4, lo que equivale a 8.
    • NUM_CHIPS: En el caso de las TPUs v4 multihost, el número de chips de TPU en cada VM es cuatro. En el caso de las TPU v5e multihost, el número de chips de TPU de cada VM es uno, cuatro u ocho. Para obtener más información, consulta Chips de TPU en la VM de un segmento de TPU.

    En este manifiesto:

    • JobSet es un servicio sin encabezado con el mismo nombre que el nombre de JobSet. En este caso, es multislice-job.
    • La anotación alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool configura la afinidad de los pods para asegurarse de que todos los pods se programen en la misma porción.
    • El campo image se define como la imagen de IA de JAX más reciente. Para definir otra versión, consulta la lista de imágenes de IA de JAX actuales.
    • El valor maxRestarts: 4 indica el número máximo de veces que GKE reinicia el JobSet cuando falla un Job secundario. Si el número de reinicios de JobSet alcanza el máximo definido, JobSet se marca como fallido.
    • Los campos parallelism y completions son iguales al número de nodos de cada grupo de nodos.
    • El valor de backoff es 0 porque Multislice solo admite el entrenamiento síncrono de varios controladores. Debe tener el valor 0. Falla la tarea cuando falle algún pod.
    • Los valores de la sección de afinidad aseguran que solo se ejecute una carga de trabajo de TPU Multislice en un grupo de Multislices.
    • containerPort: 8080 es el puerto del coordinador de MXLA
    • containerPort: 8431 es el puerto para exportar las métricas de uso de la TPU.
    • El securityContext: privileged: true indica que los nodos tienen el modo privilegiado habilitado para acceder a las TPUs. Los nodos de GKE 1.28 o versiones posteriores no necesitan tener habilitado el modo con privilegios para acceder a las TPUs. Para obtener más información, consulta Ejecutar contenedores sin el modo con privilegios.
  2. Aplica el archivo de manifiesto:

    kubectl apply -f tpu-multislice.yaml
    
  3. Confirma que se ha admitido la carga de trabajo:

    kubectl get jobsets
    

    El resultado debería ser similar al siguiente:

    NAME            RESTARTS   COMPLETED   AGE
    multislice-job                         3s
    
  4. Monitoriza el estado de los pods aprovisionados:

    kubectl get pods
    

    El resultado debería ser 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 multislice-job JobSet programa, crea y, a continuación, ejecuta los pods hasta que se completan. Los nombres de los pods tienen el formato <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>. El prefijo jobsetName determina el JobSet al que pertenece el Pod.

  5. Opcional: Elimina la carga de trabajo de JAX:

    kubectl delete -f tpu-multislice.yaml
    

Configurar ajustes adicionales

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

Mejorar el rendimiento de la red con hostNetwork

Para mejorar el rendimiento de la red entre los slices de TPU, te recomendamos que actives hostNetworking. Usa hostNetwork: true en la especificación de tu pod para omitir toda la pila de redes de Kubernetes y permitir que tus pods de Kubernetes usen la red del host directamente para la comunicación entre máquinas virtuales.

Para activar hostNetworking, añade las dos líneas siguientes a tu especificación de pod:

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

Para seguir usando podHostnames para la detección de nodos de trabajador con hostNetwork, define dnsPolicy: ClusterFirstWithHostNet. Esto es importante cuando ejecutas trabajos de entrenamiento con reanudación automática y necesitas tener los mismos nombres para volver a cargar los mismos puntos de control.

Si usas TPU Trillium (v6e) y tus pods usan hostNetworking, instala el siguiente DaemonSet para ajustar /proc/sys/net/ipv4/tcp_rmem en el nodo.

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/ai-on-gke/51bf3dcab6ff658cf62cc32867f96860bf58dfdc/scripts/network-setup/v6e-increase-rmem.yaml

Mejorar el rendimiento de la red sin hostNetwork en TPU Trillium

Si usas TPU Trillium y tus pods no pueden usar hostNetworking, habilita la multiconexión en red con el modo netdevice para obtener el mejor rendimiento de la red. La compatibilidad con NIC en modo netdevice con varias redes transfiere la NIC de la VM directamente al pod, sin pasar por Kubernetes ni GKE Dataplane V2.

El tipo de máquina ct6e-standard-4t se basa en dos NICs físicas. Kubernetes requiere una vNIC que no se puede transferir a los pods. Por lo tanto, cada nodo debe tener tres vNICs para que los pods tengan acceso directo a dos vNICs y así conseguir el mejor rendimiento de ambas NICs físicas.

Para habilitar el modo netdevice en ct6e-standard-4t, sigue estos pasos:

  1. Crea dos VPCs adicionales que admitan el modo netdevice.
  2. Crear un clúster de GKE con funciones de varias redes
  3. Configura dos netdevice redes. Por ejemplo, puedes usar los siguientes objetos GKENetworkParamSet y Network (SECOND_VPC y THIRD_VPC son las VPCs creadas en el paso anterior):

    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: tpu-second
    spec:
      vpc: SECOND_VPC
      vpcSubnet: SECOND_VPC_SUBNET
      deviceMode: NetDevice
    ---
    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: tpu-third
    spec:
      vpc: THIRD_VPC
      vpcSubnet: SECOND_VPC_SUBNET
      deviceMode: NetDevice
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: tpu-second
    spec:
      provider: "GKE"
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: tpu-second
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: tpu-third
    spec:
      provider: "GKE"
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: tpu-third
    
  4. Conecta tus Pods con tres redes. Por ejemplo, puede usar las siguientes anotaciones en la especificación de su pod:

    metadata:
      annotations:
        networking.gke.io/default-interface: 'eth0'
        networking.gke.io/interfaces: |
          [
            {"interfaceName":"eth0","network":"default"},
            {"interfaceName":"eth1","network":"tpu-second"},
            {"interfaceName":"eth2","network":"tpu-third"},
          ]
    
  5. Aplica sysctls de red en un pod, en el contenedor init o en el contenedor de la aplicación. Por ejemplo, puedes añadir el siguiente contenedor init a la especificación del pod:

    initContainers:
    - name: "network-optimization-sysctls"
      image: "busybox"
      securityContext:
        privileged: true
      command:
      - bash
      - -c
      - |
        echo 5000 > /proc/sys/net/ipv4/tcp_rto_min_us
        echo 1 > /proc/sys/net/ipv4/tcp_no_metrics_save
        echo 0 > /proc/sys/net/ipv4/tcp_slow_start_after_idle
        echo 131072 > /proc/sys/net/core/optmem_max
        echo "4096 41943040 314572800" > /proc/sys/net/ipv4/tcp_rmem
    
Práctica recomendada:

Usa las interfaces eth1 y eth2 para mejorar el rendimiento de la red en lugar de la interfaz eth0. Para ello, añade export LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS --megascale_grpc_interface_prefixes=eth1,eth2,lo" a la especificación de la carga de trabajo.

Activar el almacenamiento de registros

Los registros emitidos por los contenedores que se ejecutan en los nodos de GKE, incluidos los nodos de slice de TPU, se pueden ver en el explorador de registros si tienes habilitado el registro del sistema de GKE en tu clúster.

Puede ver sus registros de GKE con el Explorador de registros. Para ver los registros de los contenedores de su carga de trabajo, utilice el siguiente filtro:

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 las porciones 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 Ver registros de TPU de GKE.

Habilitar métricas adicionales

Además de las métricas de TPU generales, hay cuatro métricas de tiempo de ejecución de TPU específicas de Multislice. Estas métricas están disponibles en GKE 1.29.1-gke.1016000 o versiones posteriores. La carga de trabajo de TPU debe usar la versión 0.4.24 de JAX

Estas son las métricas de multirebanada disponibles:

  • Latencias de transferencia de la red del centro de datos (DCN): distribución de las latencias de transferencia de la red para el tráfico de varios segmentos.
  • Latencias colectivas: distribución de la latencia colectiva de extremo a extremo del tráfico de varios segmentos.
  • Latencias de transferencia de host a dispositivo: distribución de la latencia de transferencia de host a dispositivo de cada fragmento de datos del tráfico de varios segmentos.
  • Latencias de transferencia de dispositivo a host: distribución de la latencia de transferencia de dispositivo a host de cada fragmento de datos del tráfico de varios segmentos.

Estas métricas se encuentran en el esquema de 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

Slice de TPU frente a Multislice

En la siguiente tabla se diferencian la organización arquitectónica de una porción de TPU y de una Multislice:

Slice de TPU Multicorte
Interconectividad La carga de trabajo se ejecuta en un solo segmento de TPU. Todos los chips de TPU de un slice están conectados con ICI. La carga de trabajo se ejecuta en varios segmentos de TPU. La comunicación dentro de un slice se produce a través de ICI. La comunicación entre las porciones se produce a través de la DCN.
Grupos de nodos admitidos Slice de TPU de un solo host y slice de TPU de varios hosts Grupos de slices de TPU de varios hosts
Tipo de carga de trabajo recomendado IndexedJob o JobSet JobSet

Eliminar los recursos

La forma más fácil de evitar que te cobren es eliminar el Google Cloud proyecto que has creado para el tutorial. También puedes eliminar los recursos concretos.

Eliminar el proyecto

    Delete a Google Cloud project:

    gcloud projects delete PROJECT_ID

Eliminar recursos concretos

Elimina el clúster de GKE:

```sh
gcloud container clusters delete  CLUSTER_NAME \
   --project=PROJECT_ID  \
   --location=CONTROL_PLANE_LOCATION
```

Siguientes pasos