Descripción general de las Cloud TPU de Cloud en la GKE [versión preliminar pública]

Resumen

En este documento, se explica cómo administrar Cloud TPU en una configuración de Slices múltiples con GKE. Suponemos que estás familiarizado con el uso de únicas porciones de Cloud TPU en GKE y que tienes experiencia general con las porciones múltiples de Cloud TPU. Te recomendamos que modeles cada porción de Cloud TPU como un grupo de nodos y envíes trabajos que impliquen varios grupos de nodos mediante la API de JobSet.

Comentarios

Si tienes preguntas o comentarios, envíanos un correo electrónico a cloudtpu-multislice-preview@google.com.

Conceptos

Reparación automática
Cuando una sección encuentra un evento de mantenimiento, interrupción o falla de hardware, Google Kubernetes Engine creará una porción nueva. Si no hay capacidad suficiente para crear una porción nueva, esta no se completará hasta que el hardware esté disponible. Luego, se reiniciarán todas las porciones del entorno de varias porciones para que el entrenamiento pueda continuar. La secuencia de comandos de entrenamiento debe definir una secuencia de comandos de inicio que verifique los puntos de control y los cargue antes de reiniciar el entrenamiento.
Redes del centro de datos (DCN)
Una red de mayor capacidad de procesamiento y latencia más alta que conecte segmentos de Cloud TPU en una configuración de varias porciones.
Conjunto de datos
Los datos que usa un modelo para el entrenamiento o la inferencia.
Host
Un host es una computadora física que ejecuta VM. Un host puede ejecutarse como máximo en la VM a la vez. Cada VM tiene una Cloud TPU dedicada.
Inferencia
Cargar un modelo de aprendizaje automático previamente entrenado en un host y realizar predicciones sobre los datos
Interchip Interconnect (ICI)
Vínculos internos de alta velocidad y baja latencia que conectan las Cloud TPU dentro de un pod de Cloud TPU.
Pod de Kubernetes
Un grupo de uno o más contenedores, con almacenamiento compartido y recursos de red, y una especificación sobre cómo ejecutar los contenedores.
Varios segmentos
Dos o más porciones de Cloud TPU que se pueden comunicar a través de la red de centro de datos (DCN)
Grupos de nodos
En el contexto de GKE, un grupo de nodos es un grupo de nodos (VM) que el programador de Kubernetes puede programar con cargas de trabajo (por ejemplo, un trabajo de entrenamiento). Para obtener más información sobre la arquitectura y los conceptos del clúster de GKE, consulta Arquitectura del clúster.
Recurso en cola
Una representación de los recursos de Cloud TPU, que se usa para poner en cola y administrar una solicitud de un entorno de Cloud TPU de una sola porción o de varias porciones.
Tensor
Estructura de datos que se usa para representar datos multidimensionales en un modelo de aprendizaje automático.
Unidad de procesamiento tensorial (TPU)
El chip de aceleración de AA desarrollado por Google. Están diseñados para ofrecer el procesamiento más rápido y eficiente para las tareas clave de aprendizaje automático, como la multiplicación de matrices.
pod de Cloud TPU
Un conjunto de chips de Cloud TPU conectados por interfaces de red ICI dedicadas. Un pod de Cloud TPU te permite distribuir la carga de procesamiento en varias Cloud TPU.
Porción de Cloud TPU
Una subsección lógica de un pod de Cloud TPU que consta de chips de Cloud TPU. Todos los chips de una porción se comunican entre sí mediante la red ICI.
VM de Cloud TPU
Una máquina virtual que ejecuta Linux que tiene acceso a las Cloud TPU subyacentes. En el caso de las Cloud TPU v4, cada VM de Cloud TPU tiene acceso directo a 4 chips. A veces, las VM de Cloud TPU se denominan trabajadores.

Requisitos previos

En este documento, se supone que estás familiarizado con la descripción general de Cloud TPU multifragmento y la introducción a Cloud TPU en GKE.

Administra Cloud TPU de varias porciones en GKE

Aprovisiona la capacidad de Cloud TPU para usarla con múltiples secciones de GKE

En Implementa cargas de trabajo de Cloud TPU en GKE, se detalla cómo puedes reservar Cloud TPU. Sin embargo, si usas la capacidad de Cloud TPU directamente a través de la API de Cloud TPU y deseas usar la misma capacidad para GKE Multislice, comunícate con gke-tpu-support@google.com a fin de transferir parte de tu capacidad a GKE.

Crea grupos de nodos en un entorno de varias porciones

Para crear un entorno de varias secciones, debes agrupar varios grupos de nodos de GKE en un grupo de varias porciones. Debes crear un grupo de nodos independiente para cada porción que sea parte de tu entorno de varias secciones. Debes agregar cada grupo de nodos a un grupo de varias secciones. Los grupos de varios fragmentos son una forma de agrupar y administrar varios grupos de nodos. Los grupos de Slices múltiples se definen mediante la aplicación de etiquetas (pares nombre/valor) a cada grupo de nodos que será parte de tu entorno Multislice. En este instructivo, usamos etiquetas llamadas MultisliceGroup y MultisliceGroupSize, pero puedes usar el nombre que desees. Todos los grupos de nodos de un grupo de varias porciones deben tener las mismas etiquetas (pares nombre/valor) y la misma topología de Cloud TPU.

Puedes agregar la etiqueta de varios segmentos en el momento de la creación. En este ejemplo, usamos MultisliceGroup y MultisliceGroupSize para los nombres de etiquetas:

$ gcloud beta container node-pools create pool-name \
  --region=cluster-region \
  --cluster=cluster-name \
  --node-locations=node-zone \
  --machine-type=machine-type \
  --tpu-topology=tpu-topology \
  --node-labels=MultisliceGroup=multislice-group-name,MultisliceGroupSize=num-slices \
  --num-nodes=num-nodes \
  [--reservation-affinity=specific \
     --reservation=reservation-name]

Notas:

  • Solo se puede realizar una operación de grupo de nodos a la vez, por lo que tendrás que esperar a que se cree o borre cada grupo de nodos para poder crear o borrar el siguiente.

  • No se admiten los entornos de varias secciones compuestos por porciones de host único. Por ejemplo, no puedes crear una multisección con varias v4-16, v5e con topología (1,4) ni v5e-8.

También puedes agregar las etiquetas MultisliceGroup y MultisliceGroupSize a un grupo de nodos existente:

$ gcloud beta container node-pools update pool-name \
  --region=cluster-region \
  --cluster=cluster-name \
  --node-labels=MultisliceGroup=multislice-group-name,MultisliceGroupSize=num-slices

Verifica el estado del grupo de nodos

Puedes verificar el estado de tus grupos de nodos con kubectl:

$ kubectl get nodes -l MultisliceGroup=multislice-group-1 | grep "Ready" | wc -l

Notas:

  1. Debes conectarte al clúster de GKE antes de ejecutar el comando. Para obtener más información, consulta Accede a clústeres.

  2. Cuando ejecutes el comando get nodes, es posible que veas el siguiente error: 

    E0608 13:25:01.586645 21293 memcache.go:255] couldn't get resource list for
http://metrics.k8s.io/v1beta1 : the server is currently unable to handle the
request

    Por lo general, este error significa que no tienes un grupo de nodos que no sea de Cloud TPU en el clúster de GKE. Agrega un grupo de nodos que no sea de Cloud TPU a tu clúster y vuelve a probar el comando.

Ejecuta cargas de trabajo de varios segmentos

API de JobSet

La API de JobSet es una API de código abierto para administrar un grupo de trabajos en conjunto. Es una gran opción para las cargas de trabajo de varias secciones; como se describe en la introducción a las Cloud TPU en GKE, usamos una sola IndexedJob de Kubernetes para ejecutar cargas de trabajo en una sola porción de Cloud TPU. De manera similar, para los entornos de varias secciones, necesitamos que varios IndexedJobs de Kubernetes trabajen juntos y JobSet ayuda a lograrlo.

imagen

JobSet proporciona lo siguiente:

  • Manejo de fallas: Recreación automática de todos los objetos Job secundarios en caso de que falle cualquiera de los objetos Job secundarios.
  • Indexación de Slice: JobSet indexa cada objeto Job con un ID y lo agrega como parte de anotaciones/etiquetas en el Job y el Pod.
  • Posición de la carga de trabajo: Se encarga automáticamente de asignar un solo trabajo a un dominio específico (en este caso, cada trabajo en un solo grupo de nodos).
  • Creación de Headless-Service: Ciclo de vida administrado de un headless-svc para todos los objetos Job del JobSet.

Instalación

  1. Instala la API de JobSet en tu clúster:

    $ kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.2.1/manifests.yaml

  2. Verifica que el controlador de JobSet esté en ejecución:

    $ kubectl get pods -n jobset-system

Ejecuta una carga de trabajo de varios segmentos

En el siguiente manifiesto YAML, se muestra cómo ejecutar una carga de trabajo de varias porciones en cuatro porciones de la versión 4-16. Copia el siguiente archivo YAML, guárdalo en un archivo y, luego, ejecuta kubectl apply -f <file-name.yaml>:

apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-job  # JobSet name (${JOBSET_NAME})
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool # 1:1 job replica to node pool assignment
spec:
  failurePolicy:
    maxRestarts: 4  # The set will be restarted on failures up to 4 times.
  replicatedJobs:
    - name: slice    # Part of the name of the child Jobs (<replicateJobName>)
      replicas: 4    # Number of slices
      template:
        spec:
          parallelism: 2   # Must be set to number of nodes in each node pool
          completions: 2   # Must be set to number of nodes in each node pool
          backoffLimit: 0   # Must be set to 0. Fail the job when any pod fails.
          template:
            spec:
              affinity: # The affinity section is to make sure there is only one Multislice job running in a single Multislice Group. More on this below.
                podAffinity:
                  requiredDuringSchedulingIgnoredDuringExecution:
                  - labelSelector:
                      matchExpressions:
                      - key: jobset.sigs.k8s.io/jobset-name
                        operator: In
                        values:
                        - multislice-job # JobSet name
                    topologyKey: MultisliceGroup
                podAntiAffinity:
                  requiredDuringSchedulingIgnoredDuringExecution:
                  - labelSelector:
                      matchExpressions:
                      - key: jobset.sigs.k8s.io/jobset-name
                        operator: NotIn
                        values:
                        - multislice-job # JobSet name
                    topologyKey: MultisliceGroup
                    namespaceSelector:
                      matchExpressions:
                      - key: jobset.sigs.k8s.io/jobset-name
                        operator: Exists
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
                cloud.google.com/gke-tpu-topology: 2x2x2
                MultisliceGroupSize: "4"
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080 # Port for MXLA coordinator
                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())'
                resources:
                  limits:
                    google.com/tpu: 4 # Number of Cloud TPU VMs per worker

  • JobSet creará 4 trabajos indexados de Kubernetes (cantidad de segmentos). Los trabajos indexados de Kubernetes seguirán el patrón de nomenclatura: <jobSetName>-<replicateJobName>-<job-index> (en este caso, es multislice-job-slice-{0,1,2,3}). Los trabajos creados aparecerán en Google Cloud CLI en Workloads.

  • JobSet es un servicio sin interfaz gráfica con el mismo nombre que JobSet (en este caso, es multislice-job).

  • Cada trabajo indexado crea 2 (# de trabajadores) Pods de Kubernetes. Los trabajos siguen el patrón de nombres: <jobSetName>-<replicateJobName>-<job-index>-<worker-index>-<5 letter suffix> (en este caso, es multislice-job-slice-{0,1,2,3}-{1,2}-{5 letter suffix}).

  • En la actualidad, las múltiples secciones en Google Kubernetes Engine solo admiten el entrenamiento síncrono de varios controladores. Para lograrlo, parallelism y completions deben configurarse en la cantidad de nodos en cada grupo de nodos, y backoff debe establecerse en 0. La cantidad de nodos es dividida por la cantidad de núcleos/nodo. En el caso de las TPU v4, esto es ocho para v4-16, configura parallelism y completions en 2.

  • La versión mínima de JAX es: v0.4.9.

Inserción de variables de entorno en GKE

El webhook de GKE inyecta de forma automática las variables de entorno a las especificaciones del Pod del trabajo indexado, que luego el Pod de Kubernetes que ejecuta la carga de trabajo de Multi Slice hereda. GKE inserta las siguientes variables de entorno:

Variable de entorno Valor
TPU_WORKER_ID metadata.annotations['batch.kubernetes.io/job-completion-index']
TPU_WORKER_HOSTNAMES Calcula la lista separada por comas de los índices de Pod del trabajo de la siguiente manera: <job-name>-0.<subdomain>,.......,<job-name>-<n-1>.<subdomain>
MEGASCALE_NUM_SLICES metadata.annotations['jobset.sigs.k8s.io/replicatedjob-replicas']
MEGASCALE_SLICE_ID metadata.annotations['jobset.sigs.k8s.io/job-index']
MEGASCALE_COORDINATOR_ADDRESS Generate the string using the jobSetName and the replicatedJobName:

<jobSetName>-<replicatedJobName>-0-0.<subdominio>

Opcional: Desactiva hostNetwork en tus Pods de GKE

Usa hostNetwork: true en la especificación de tu Pod para omitir toda la pila de red de Kubernetes y permite que tus Pods de Kubernetes usen la red del host directamente para la comunicación de VM a VM. Esto mejora el rendimiento de red entre porciones. Si deseas seguir usando podHostnames para el descubrimiento de trabajadores con hostNetwork, configura dnsPolicy: ClusterFirstWithHostNet. Esto es importante cuando se reanuda automáticamente el trabajo de entrenamiento, y debe tener los mismos nombres para volver a cargar los mismos puntos de control.

Para desactivar hostNetworking, quita las siguientes dos líneas de las especificaciones del Pod:

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

Organiza trabajos de varias secciones

Evita programar un interbloqueo

Si envías varios trabajos de varias porciones (JobSet A y B) del mismo MultisliceGroupSize, podrías tener una situación en la que ambos objetos JobSet estén programados de manera parcial (por ejemplo, 5 de 8 trabajos secundarios están programados desde JobSet A y 3 de 8 trabajos secundarios están programados desde JobSet B). Esto hará que los grupos de nodos permanezcan inactivos y que el JobSet A y B estén colgados.

Para evitarlo, puedes aprovechar podAffinity y podAntiAffinity. En las cargas de trabajo de varias secciones, cambia topologyKey para que sea la etiqueta (MultisliceGroup) que representa un grupo de varias secciones en lugar de la etiqueta cloud.google.com/gke-nodepool. Por ejemplo:

affinity:
  podAffinity:
    requiredDuringSchedulingIgnoredDuringExecution:
      - labelSelector:
          matchExpressions:
          - key: jobset.sigs.k8s.io/jobset-name
            operator: In
            values:
            - multislice-job # JobSet name
          topologyKey: MultisliceGroup
  podAntiAffinity:
    requiredDuringSchedulingIgnoredDuringExecution:
    - labelSelector:
        matchExpressions:
        - key: jobset.sigs.k8s.io/jobset-name
          operator: NotIn
          values:
          - multislice-job # JobSet name
      topologyKey: MultisliceGroup
      namespaceSelector:
        matchExpressions:
        - key: jobset.sigs.k8s.io/jobset-name
          operator: Exists

Observabilidad

Registros de Cloud Logging

Los Pods de GKE del JobSet se nombran con el siguiente patrón:

<jobSetName>-<replicateJobName>-<job-index>-<worker-index>-<5 letter suffix>

Puedes ver tus registros mediante el Explorador de registros de Cloud Logging con el siguiente filtro a fin de ver los registros de contenedores de tu carga de trabajo:

resource.type="k8s_container"
resource.labels.cluster_name=<cluster-name>
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=<jobSetName>

Para filtrar los registros de la porción x y el trabajador y, usa el siguiente filtro:

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