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.
Este instructivo está dirigido a ingenieros de aprendizaje automático (AA) y administradores y operadores de plataformas que deseen usar la organización de contenedores de Kubernetes para administrar cargas de trabajo de inferencia, ajuste y entrenamiento de modelos a gran escala con TPU. Para obtener más información sobre los roles comunes y las tareas de ejemplo a los que se hace referencia en el contenido de Google Cloud , consulta Roles y tareas comunes de los usuarios de GKE.
Antes de configurar Multislice en GKE, asegúrate de estar familiarizado 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-4tcon una topología2x2x1o unct5lp-hightpu-4tcon 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.
- Multislice no admite TPU v3.
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 el comando
gcloud components updatepara obtener la versión más reciente. Es posible que las versiones anteriores de gcloud CLI no admitan la ejecución de los comandos que se describen en este documento.
- 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 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]
Reemplaza lo siguiente:
POOL_NAME: el nombre del grupo de nodos nuevo.LOCATION: Es el nombre de la zona en función de la versión de TPU que deseas usar. Para identificar una ubicación disponible, consulta Disponibilidad de TPU en GKE.CLUSTER_NAME: Es el nombre del clúster.NODE_ZONES: 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 Elige la versió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. Para obtener más información sobre las topologías de TPU, usa la tabla de Elige una topología.Para obtener más información, consulta Topología.
De manera opcional, también puedes usar las siguientes marcas:
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_TOPOLOGYes2x4x4(TPU v4 con cuatro chips en cada VM), laNUM_NODESes 32/4, que es igual a 8. Si omites esta marca, la cantidad de nodos se calcula y se establece de forma predeterminada según la topología y el tipo de máquina.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.--flex-start: Configura el grupo de nodos para que use VMs de inicio flexible. Las VMs de inicio flexible se crean con la opción de consumo flex-start, que es compatible con la versión 1.33.0-gke.1712000 o posterior de GKE.--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-nodeses obligatoria si se proporciona--enable-autoscalingy 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.
--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 porción de TPU de varios hosts 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 grupo de nodos nuevo.
- El grupo de nodos usa TPU Trillium.
- Las VMs Spot no admiten la programación de la recopilación.
Para obtener más información sobre la administración de la programación de la recopilación, consulta Administra la programación de la recopilación en porciones de TPU de varios hosts.
--placement-type=COMPACT: Crea un grupo de nodos con la posición de compactación habilitada. Esta opción se debe usar con la marca--tpu-topology. Para obtener más información, consulta Crea una política de posición compacta y Topología de TPU.
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 flex_start = false } 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_TOPOLOGYes4x8, entonces hay 32 chips, lo que significa queNUM_NODESdebe ser 8. Para obtener más información sobre las topologías de TPU, usa la tabla de Elige la versió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 obtener más información sobre las topologías de TPU, usa la tabla de Elige una topología.
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_VALUESen 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.flex_start: Configura el grupo de nodos para usar la opción de consumo flex-start. No se puede establecer entruesispotestá habilitado.
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.
Selecciona el nombre de la zona según la versión de TPU que deseas usar. Para identificar una ubicación disponible, consulta Disponibilidad de TPU en GKE.
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:
- CT3P: Para TPU v3.
- 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 Elige la versió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
kubectlpara acceder al clúster:gcloud container clusters get-credentials CLUSTER_NAME \ --project=PROJECT_ID \ --location=CONTROL_PLANE_LOCATIONReemplaza lo siguiente:
CLUSTER_NAME: Es el nombre del clúster.PROJECT_ID: ID del proyectoCONTROL_PLANE_LOCATION: Es 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
kubectlen Cloud Shell para ver tus nodos de porción de TPU:kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=ACCELERATOR_TYPE \ -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGYReemplaza lo siguiente:
TPU_ACCELERATOR: El tipo de acelerador de TPU que usaste cuando creaste los grupos de nodos. Por ejemplo,tpu-v4-podsliceotpu-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: 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_CHIPSEstá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_CHIPSReemplaza lo siguiente:
NUM_SLICES: La cantidad de grupos de nodos de porción de TPU. En este caso,NUM_SLICESes igual a3.ACCELERATOR_TYPE: El tipo de acelerador de TPU que usaste cuando creaste los grupos de nodos. Por ejemplo,tpu-v4-podsliceotpu-v5-lite-podslice.TPU_TOPOLOGY: la topología física para la porción de TPU. Por ejemplo,4x4x4o2x2, 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_TOPOLOGYes2x4x4(TPU v4 con cuatro chips TPU en cada VM), laNUM_NODESes 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. - La anotación
alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepoolconfigura la afinidad de Pods para garantizar que todos los Pods se programen en la misma división. - El campo
imagese establece en la imagen de IA de JAX más reciente. Para establecer una versión diferente, consulta la lista de imágenes de JAX AI actuales. - El
maxRestarts: 4indica 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
parallelismycompletionsson iguales a la cantidad de nodos en cada grupo de nodos. backoffes 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: 8080es el puerto para el coordinador de MXLAcontainerPort: 8431es el puerto para exportar las métricas de uso de TPU- El
securityContext: privileged: trueindica 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.yamlConfirma que se admita la carga de trabajo:
kubectl get jobsetsEl resultado es similar al siguiente:
NAME RESTARTS COMPLETED AGE multislice-job 3sSupervisa el estado de los Pods aprovisionados:
kubectl get podsEl resultado es similar a este:
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 2m30sEl JobSet
multislice-jobprograma, crea y, luego, ejecuta los Pods hasta su finalización. Los nombres de los Pods tienen el formato<jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>. El prefijojobsetNamedetermina el JobSet al que pertenece el Pod.Opcional: Quita la carga de trabajo de JAX:
kubectl delete -f tpu-multislice.yaml
Configura los parámetros adicionales
En las siguientes secciones, se describen las configuraciones adicionales que puedes aplicar a tu Multislice.
Mejora el rendimiento de la red con hostNetwork
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, agrega las siguientes dos líneas a 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.
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
Mejora el rendimiento de la red sin hostNetwork en la TPU Trillium
Si usas la TPU Trillium y tus Pods no pueden usar hostNetworking, habilita la función de varias redes con el modo netdevice para obtener el mejor rendimiento de la red. La compatibilidad de NIC en modo netdevice con varias redes pasa la NIC de la VM directamente al Pod, lo que omite Kubernetes y GKE Dataplane V2.
El tipo de máquina ct6e-standard-4t está respaldado por dos NIC físicas. Kubernetes requiere una vNIC que no se puede pasar a los Pods. Por lo tanto, cada nodo debe tener tres vNIC para permitir que los Pods tengan acceso directo a dos vNIC y lograr el mejor rendimiento de ambas NIC físicas.
Para habilitar el modo netdevice para ct6e-standard-4t, completa los siguientes pasos:
- Crea dos VPC adicionales que admitan el modo
netdevice - Crea un clúster de GKE con capacidades de varias redes
Configura dos redes de
netdevice. Por ejemplo, puedes usar los siguientes objetosGKENetworkParamSetyNetwork(SECOND_VPCyTHIRD_VPCson las VPC 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-thirdConecta los Pods con tres redes. Por ejemplo, puedes usar las siguientes anotaciones en la especificación de tu 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 dentro del Pod, en el contenedor init o en el contenedor de la aplicación. Por ejemplo, puedes agregar el siguiente contenedor de inicialización a las especificaciones 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 obtener un mejor rendimiento de las redes, en lugar de la interfaz eth0. Para ello, agrega export LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS --megascale_grpc_interface_prefixes=eth1,eth2,lo" a la especificación de la carga de trabajo.
Habilitar registro
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.
Habilita métricas adicionales
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_latencieskubernetes.io/container/multislice/network/collective_end_to_end_latencieskubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencieskubernetes.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 | Porciones múltiples | |
|---|---|---|
| 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 |
Limpia los recursos
La manera más fácil de eliminar la facturación es borrar el Google Cloud proyecto que creaste para el instructivo. Como alternativa, puedes borrar los recursos individuales.
Borra el proyecto
Delete a Google Cloud project:
gcloud projects delete PROJECT_ID
Borra los recursos individuales
Borra el clúster de GKE:
```sh
gcloud container clusters delete CLUSTER_NAME \
--project=PROJECT_ID \
--location=CONTROL_PLANE_LOCATION
```