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:
¿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ía2x2x1
o unct5lp-hightpu-4t
con una topología2x2
, 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.
- Crea un clúster Standard o un clúster de Autopilot que ejecute una versión que admita Multislice. Para conocer las versiones compatibles, consulta Disponibilidad de Multislice en GKE.
- Asegúrate de que tu proyecto tenga la cuota suficiente para Cloud TPU en GKE.
- Instala JobSet v0.2.3 o una versión posterior.
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:
- Crea tres grupos de nodos de porciones de TPU de varios hosts
- 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-
, usaus-west1-c
,us-west4-a
,us-west4-b
,us-central1-a
,us-east1-c
,us-east5-b
oeurope-west4-a
. - Para los tipos de máquinas TPU v5p que comienzan con
ct5p-
, usaus-east1-d
,us-east5-a
ous-east5-c
.
Para obtener más información, consulta la disponibilidad de TPU en GKE.
- Para TPU v4, usa
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.
- TPU v4 o v5p: Define la topología en 3 tuplas (
NUM_NODES
: la cantidad de nodos en el grupo de nodos. Debe ser cero o el producto de los valores definidos enTPU_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 tuTPU_TOPOLOGY
es2x4x4
(TPU v4 con cuatro chips en cada VM), laNUM_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 enTPU_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 proyectoCLUSTER_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, siTPU_TOPOLOGY
es4x8
, entonces hay 32 chips, lo que significa queNUM_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 ejemplo4x4x4
. - Para TPU v5e: define la topología en 2 tuplas (
{A}x{B}
), por ejemplo,2x2
.
- Para TPU v4: define la topología en 3 tuplas (
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 completarRESERVATION_LABEL_VALUES
en el camporeservation_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 enTPU_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:
Ve a la página de Google Kubernetes Engine en la consola de Google Cloud.
En la lista de clústeres, haz clic en el nombre del clúster que deseas modificar.
Haz clic en add_box 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-
, usaus-west1-c
,us-west4-a
,us-west4-b
,us-central1-a
,us-east1-c
,us-east5-b
oeurope-west4-a
. - Para los tipos de máquinas TPU v5p que comienzan con
ct5p-
, usaus-east1-d
,us-east5-a
ous-east5-c
.
- Para TPU v4, usa
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 de porción de 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.
Verifica el estado del grupo de nodos
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
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
otpu-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:
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 a3
.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
otpu-v5-lite-podslice
.TPU_TOPOLOGY
: la topología física para la porción de TPU. Por ejemplo,4x4x4
o2x2
, 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 enTPU_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 tuTPU_TOPOLOGY
es2x4x4
(TPU v4 con cuatro chips TPU en cada VM), laNUM_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
ycompletions
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 MXLAcontainerPort: 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.
Aplica el manifiesto
kubectl apply -f tpu-multislice.yaml
Confirma que se admita la carga de trabajo:
kubectl get jobsets
El resultado es similar al siguiente:
NAME RESTARTS COMPLETED AGE multislice-job 3s
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 |