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:
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ía2x2x1
o unact5lp-hightpu-4t
con una topología2x2
, 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
.
- Crea un clúster Standard o un clúster Autopilot que ejecute una versión compatible con Multislice. Para ver las versiones compatibles, consulta Disponibilidad de Multislice en GKE.
- Asegúrate de que tu proyecto tiene suficiente cuota para TPU de Cloud en GKE.
- Instala JobSet v0.2.3 o una versión posterior.
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:
- Crea tres grupos de nodos de segmento de TPU multihost
- 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 enTPU_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 tuTPU_TOPOLOGY
es2x4x4
(TPU v4 con cuatro chips en cada VM), entonces elNUM_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 enTPU_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
- Asegúrate de usar la versión 4.84.0 o una posterior del proveedor
google
. 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, siTPU_TOPOLOGY
es4x8
, hay 32 chips, lo que significa queNUM_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 campoRESERVATION_LABEL_VALUES
enreservation_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 enTPU_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 comotrue
sispot
está habilitado.
Consola
Para crear un grupo de nodos con TPUs, sigue estos pasos:
Ve a la página Google Kubernetes Engine en la Google Cloud consola.
En la lista de clústeres, haga clic en el nombre del clúster que quiera modificar.
Haz clic en add_box Añadir grupo de nodos.
En la sección Detalles del grupo de nodos, marca la casilla Especificar ubicaciones de nodos.
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.
En el panel de navegación, haga 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:
- CT3P: para TPU v3.
- CT4P: para TPU v4.
- CT5LP: para la TPU v5e.
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.
En el menú desplegable Topología de TPU, selecciona la topología física del slice de TPU.
En el cuadro de diálogo Cambios necesarios, haga clic en Hacer cambios.
Asegúrate de que Tipo de disco de arranque sea Disco persistente estándar o Disco persistente SSD.
Si quieres, marca la casilla Habilitar nodos en VMs Spot para usar VMs Spot en los nodos del grupo de nodos.
Haz clic en Crear.
Verificar el estado del grupo de nodos
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.
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
otpu-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:
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 a3
.ACCELERATOR_TYPE
: el tipo de acelerador de TPU que has usado al crear los grupos de nodos. Por ejemplo,tpu-v4-podslice
otpu-v5-lite-podslice
.TPU_TOPOLOGY
: la topología física del segmento de TPU. Por ejemplo,4x4x4
o2x2
, 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 enTPU_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 tuTPU_TOPOLOGY
es2x4x4
(TPU v4 con cuatro chips de TPU en cada VM), elNUM_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
ycompletions
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 MXLAcontainerPort: 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.
Aplica el archivo de manifiesto:
kubectl apply -f tpu-multislice.yaml
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
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 prefijojobsetName
determina el JobSet al que pertenece el Pod.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:
- Crea dos VPCs adicionales que admitan el modo
netdevice
. - Crear un clúster de GKE con funciones de varias redes
Configura dos
netdevice
redes. Por ejemplo, puedes usar los siguientes objetosGKENetworkParamSet
yNetwork
(SECOND_VPC
yTHIRD_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
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"}, ]
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
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
- Consulta cómo orquestar cargas de trabajo de Multislice con slices de TPU.