Programa cargas de trabajo de GKE con la Programación consciente de la topología (TAS)

Las cargas de trabajo de IA y AA requieren una comunicación significativa de pod a pod. Debido a este requisito, el ancho de banda de la red entre los pods afecta directamente el tiempo de ejecución y el costo de la carga de trabajo. Este ancho de banda depende de la ubicación de las instancias de máquina virtual (VM) en el clúster.

En este documento, se explica cómo optimizar la programación de tus cargas de trabajo de IA o AA a gran escala en un clúster de Google Kubernetes Engine (GKE) para el rendimiento y la confiabilidad. En específico, debes configurar tu clúster para que use la programación consciente de la topología (TAS) para la comunicación de baja latencia. Este enfoque minimiza la sobrecarga de comunicación y ayuda a maximizar el rendimiento de tus cargas de trabajo.

¿Qué es la programación consciente de la topología (TAS)?

TAS puede mejorar significativamente la eficiencia del entrenamiento de modelos de lenguaje grandes (LLM). La TAS coloca estratégicamente a los trabajadores en la topología de red para minimizar la sobrecarga de comunicación durante la agregación de gradientes, lo que requiere que los trabajadores se comuniquen en un orden de clasificación específico. Al minimizar los saltos de red entre los trabajadores que se comunican de forma secuencial, la TAS reduce la contención de la red y optimiza el uso del ancho de banda, lo que lleva a una convergencia más rápida y a tiempos de entrenamiento más cortos. Con modelos de LLM cada vez más grandes, la TAS es esencial para maximizar el rendimiento y la escalabilidad del entrenamiento distribuido.

La TAS funciona mejor con la capacidad colocada de forma densa, que se puede obtener a través de las reservas. Con las VMs de inicio flexible o las VMs Spot, es menos probable que tu capacidad se asigne de forma cercana, por lo que es posible que la TAS no funcione bien en esta situación.

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 update para obtener la versión más reciente. Es posible que las versiones anteriores de gcloud CLI no admitan la ejecución de los comandos de este documento.

  • Para conectarte a tu clúster, ejecuta el siguiente comando:

    gcloud container clusters get-credentials CLUSTER_NAME

    Reemplaza CLUSTER_NAME por el nombre del clúster.

Prepara tu clúster de GKE

Para preparar tu clúster de GKE para ejecutar cargas de trabajo con TAS, completa los siguientes pasos:

  1. Instala Kueue con la TAS habilitada

  2. Visualiza la topología de tu clúster de GKE

  3. Configura Kueue

Instala Kueue con la TAS habilitada

Te recomendamos que uses la TAS con Kueue, un sistema nativo de Kubernetes que administra las cuotas y cómo los trabajos deberían consumirlas. La TAS requiere la versión 0.10.0 de Kueue o una posterior, y debes habilitarla de forma explícita.

Para instalar Kueue y habilitar la TAS, selecciona una de las siguientes opciones:

Manifiesto de Kueue

  1. Instala Kueue:

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/v0.10.0/manifests.yaml

  2. Habilita la TAS en Kueue:

    kubectl -n kueue-system patch deployment kueue-controller-manager \
    --type json -p='[{"op": "add", "path": "/spec/template/spec/containers/0/args/-", "value": "--feature-gates=TopologyAwareScheduling=true"}]'

Gráfico de Helm

Instala Kueue con la TAS habilitada mediante un gráfico de Helm:

helm install kueue oci://us-central1-docker.pkg.dev/k8s-staging-images/charts/kueue \
    --version="v0.10.0" \
    --create-namespace \
    --namespace=kueue-system \
    --set="controllerManager.featureGates[0].name=TopologyAwareScheduling,controllerManager.featureGates[0].enabled=true"

Después de instalar Kueue, debes configurarlo para que comprenda la infraestructura que administra, como se explica en la siguiente sección.

Visualiza la topología de tu clúster de GKE

Antes de visualizar la topología de los nodos A4X, A4, A3 Ultra, A3 Mega y A3 High (8 GPUs) aprovisionados como VMs Spot, debes definir la posición compacta en los nodos de GKE para exponer su topología física para la TAS. De lo contrario, tendrás errores.

Para visualizar la topología de los nodos de tu clúster de GKE en un grupo de nodos específico, ejecuta el siguiente comando:

kubectl get nodes -l cloud.google.com/gke-nodepool=NODE_POOL_NAME \
    -ocustom-columns='NAME:.metadata.name,BLOCK:.metadata.labels.cloud\.google\.com/gce-topology-block,SUBBLOCK:.metadata.labels.cloud\.google\.com/gce-topology-subblock,HOST:.metadata.labels.cloud\.google\.com/gce-topology-host' | sort -k2,4

Reemplaza NODE_POOL_NAME por el nombre del grupo de nodos.

Para comprender la topología física de los nodos de GKE en tus VMs en el resultado, consulta las siguientes etiquetas de nodo:

  • cloud.google.com/gce-topology-block: Es el ID específico de la organización del bloque reservado en el que se encuentra la VM.

  • cloud.google.com/gce-topology-subblock: Es el ID específico de la organización del subbloque en el que se encuentra la VM.

  • cloud.google.com/gce-topology-host: Es el ID del host en el que se encuentra la VM.

  • kubernetes.io/hostname: Es el nombre de host del nodo de Kubernetes. Por lo general, este nombre de host también es el nombre del nodo de GKE.

Cuantos más valores de etiqueta compartan dos VMs, más cerca estarán físicamente entre sí. Para obtener más información sobre estos términos, consulta Terminología.

Configura Kueue

Después de instalar Kueue, debes configurarlo para especificar la infraestructura que administra. Por lo general, Kueue requiere una ClusterQueue definición de cuota de recursos de infraestructura estática o infraestructura dinámica con el ajuste de escala automático del clúster habilitado. ClusterQueue admite una carga de trabajo solo si los recursos que solicita la carga de trabajo son menores o iguales al grupo de recursos definidos en ClusterQueue. Después de configurar Kueue como se describe en esta sección, Kueue admite cargas de trabajo con la TAS de la siguiente manera:

  • Cargas de trabajo de la TAS: Kueue verifica la topología de la infraestructura física y su uso actual.

  • Cargas de trabajo que no son de la TAS: Kueue no verifica la topología de la infraestructura física. Kueue administra toda la cuota definida en la configuración y deja la asignación de nodos a kube-scheduler.

Para comprender cómo proporcionar una definición de cuota de recursos de ClusterQueue a Kueue, revisa los siguientes ejemplos:

  • Cuota muy alta: Kueue prácticamente nunca detiene la admisión de una carga de trabajo en función de los recursos solicitados. Según las definiciones de la TAS, Kueue puede admitir o no cargas de trabajo en función de la topología de la infraestructura. Para obtener más información, consulta Cuota de recursos muy alta.

  • Cuota realista: Kueue admite la carga de trabajo solo si los recursos que solicita la carga de trabajo están dentro de estos límites de cuota de recursos. Según las definiciones de la TAS, Kueue verifica la topología de la infraestructura antes de admitir la carga de trabajo. Para obtener más información, consulta Cuota de recursos realista.

Todas las referencias a la cuota de recursos en las siguientes secciones se refieren a la cuota de recursos de ClusterQueue.

Cuota de recursos muy alta

En el siguiente ejemplo, se usa una cuota de recursos muy alta, de modo que Kueue nunca detiene una carga de trabajo en función de la cuota de recursos disponible. En cambio, Kueue usa la información de topología de los nodos disponibles para intentar hacer coincidir la topología con los requisitos de la carga de trabajo.

Para usar la siguiente definición de cuota de recursos, completa los siguientes pasos:

  1. Abre el editor de archivos que prefieras. Luego, incluye la siguiente definición de cuota en un archivo YAML llamado kueue-tas-config-very-high-quota.yaml:

      apiVersion: kueue.x-k8s.io/v1alpha1
  kind: Topology
  metadata:
    name: "gke-default"
  spec:
    levels:
    - nodeLabel: "cloud.google.com/gce-topology-block"
    - nodeLabel: "cloud.google.com/gce-topology-subblock"
    - nodeLabel: "cloud.google.com/gce-topology-host"
    - nodeLabel: "kubernetes.io/hostname"
---
  kind: ResourceFlavor
  apiVersion: kueue.x-k8s.io/v1beta1
  metadata:
    name: "tas-flavor"
  spec:
    nodeLabels:
      cloud.google.com/gke-nodepool: "NODE_POOL_NAME"
    topologyName: "gke-default"
    tolerations:
    - key: "nvidia.com/gpu"
      operator: "Exists"
      effect: NoSchedule
---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: "tas-cluster-queue"
  spec:
    namespaceSelector: {}
    resourceGroups:
    - coveredResources: ["nvidia.com/gpu"]
      flavors:
      - name: "tas-flavor"
        resources:
        - name: "nvidia.com/gpu"
          nominalQuota: 10000000
---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    namespace: "default"
    name: "tas-user-queue"
  spec:
    clusterQueue: "tas-cluster-queue"

    Reemplaza NODE_POOL_NAME por el nombre del grupo de nodos.

  2. Crea y aplica la configuración de cuota de recursos para el sistema de puesta en cola de trabajos de Kueue:

    kubectl create -f kueue-tas-config-very-high-quota.yaml

Cuota de recursos realista

En el ejemplo anterior, solo se configuraron los recursos de GPU. Sin embargo, Kueue puede administrar todos los recursos compatibles con Kubernetes.

En el siguiente ejemplo, se define una cuota de recursos más realista, que incluye CPU, memoria y GPU. Esto es para 100 máquinas a3-ultragpu-8g. Una sola máquina tiene 224 CPU virtuales, 2, 944 GB de memoria y 8 GPUs.

Para usar la siguiente definición de cuota de recursos, completa los siguientes pasos:

  1. Abre el editor de archivos que prefieras. Luego, incluye la siguiente definición de cuota en un archivo YAML llamado kueue-tas-config-real-quota.yaml:

      apiVersion: kueue.x-k8s.io/v1alpha1
  kind: Topology
  metadata:
    name: "gke-default"
  spec:
    levels:
    - nodeLabel: "cloud.google.com/gce-topology-block"
    - nodeLabel: "cloud.google.com/gce-topology-subblock"
    - nodeLabel: "cloud.google.com/gce-topology-host"
    - nodeLabel: "kubernetes.io/hostname"
---
  kind: ResourceFlavor
  apiVersion: kueue.x-k8s.io/v1beta1
  metadata:
    name: "tas-flavor"
  spec:
    nodeLabels:
      cloud.google.com/gke-nodepool: "NODE_POOL_NAME"
    topologyName: "gke-default"
    tolerations:
    - key: "nvidia.com/gpu"
      operator: "Exists"
      effect: NoSchedule
---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: "tas-cluster-queue"
  spec:
    namespaceSelector: {} # match all
    resourceGroups:
    - coveredResources: ["cpu", "memory", "nvidia.com/gpu"]
      flavors:
      - name: "tas-flavor"
        resources:
        # numbers below represent quota of 100 a3-ultragpu-8g machines
        - name: "cpu"
          nominalQuota: 22400
        - name: "memory"
          nominalQuota: 294400Gi
        - name: "nvidia.com/gpu"
          nominalQuota: 800
---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    namespace: "default"
    name: "tas-user-queue"
  spec:
    clusterQueue: "tas-cluster-queue"

    Reemplaza NODE_POOL_NAME por el nombre del grupo de nodos.

  2. Crea y aplica una configuración de cuota de recursos para el sistema de puesta en cola de trabajos de Kueue:

    kubectl create -f kueue-tas-config-real-quota.yaml

    El resultado es similar a este:

    topology.kueue.x-k8s.io/gke-default created
resourceflavor.kueue.x-k8s.io/tas-flavor created
clusterqueue.kueue.x-k8s.io/tas-cluster-queue created
localqueue.kueue.x-k8s.io/tas-user-queue created

Programa cargas de trabajo con la TAS mediante Kueue

En las siguientes situaciones, se muestra cómo puedes indicarle a Kueue y a la TAS que administren combinaciones comunes de cargas de trabajo y de infraestructura con tipos de solicitudes de topología y niveles de solicitudes de topología:

  • Los siguientes son los tipos de solicitudes de topología disponibles (preferidos o obligatorios):

    • kueue.x-k8s.io/podset-preferred-topology: Kueue prioriza la programación de toda la carga de trabajo dentro de un nivel de topología determinado, pero aún admite una carga de trabajo que no cabe dentro de este nivel de topología. Para una carga de trabajo que podría haber cabido en un solo nivel de topología, Kueue podría programar esa carga de trabajo en varias instancias de ese nivel de topología.

    • kueue.x-k8s.io/podset-required-topology: Kueue continúa intentando admitir esta carga de trabajo hasta que toda la carga de trabajo pueda caber dentro del nivel de topología elegido.

  • Los siguientes son los niveles de solicitudes de topología disponibles, lo que te permite ser más o menos específico sobre la infraestructura física en la que prefieres o necesitas que se ejecute tu trabajo:

    • cloud.google.com/gce-topology-block

    • cloud.google.com/gce-topology-subblock

    • cloud.google.com/gce-topology-host

    • kubernetes.io/hostname

Para programar cargas de trabajo con estos valores, usa el siguiente archivo YAML de trabajo:

apiVersion: batch/v1
kind: Job
metadata:
  generateName: JOB_NAME
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
spec:
  parallelism: NUMBER_OF_REPLICAS
  completions: NUMBER_OF_REPLICAS
  completionMode: Indexed
  template:
    metadata:
      annotations:
        ANNOTATIONS_STRING
    spec:
      containers:
      - name: dummy-job
        image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
        args: ["60s"]
        resources:
          requests:
            nvidia.com/gpu: "1"
          limits:
            nvidia.com/gpu: "1"
      restartPolicy: Never

Reemplaza las siguientes variables:

  • JOB_NAME: Es un nombre para el trabajo.

  • NUMBER_OF_REPLICAS: Es la cantidad de pods que se ejecutan en paralelo.

  • ANNOTATIONS_STRING: Consulta la siguiente tabla:

    Tipo y nivel de topología solicitados Descripción ANNOTATIONS_STRING
    Preferido para ejecutarse dentro de un nombre de host (recomendado) Esta configuración admitirá tu carga de trabajo siempre que haya suficientes recursos disponibles para satisfacer los requisitos de recursos de tu carga de trabajo, incluso si la capacidad está fragmentada. Kueue programará tus pods de la manera más compacta posible. kueue.x-k8s.io/podset-preferred-topology: "kubernetes.io/hostname"
    Obligatorio para ejecutarse dentro de un host

    Esta configuración admitirá tu carga de trabajo solo si hay un host disponible con suficientes recursos para satisfacer los requisitos de recursos de tu carga de trabajo.

    Esto es útil cuando hay varias VMs por host (por ejemplo, tipos de máquinas más pequeños) o cuando se pueden ejecutar varios pods en un solo nodo. En esos casos, si se admite la carga de trabajo, se ejecutará en un solo host.

    		 kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-host"
    Preferido para ejecutarse dentro de un host Esta configuración admitirá tu carga de trabajo siempre que haya suficientes recursos disponibles para satisfacer los requisitos de recursos de tu carga de trabajo, incluso si la capacidad está fragmentada. Kueue intentará programar tus pods dentro de un host y usará hosts adicionales si es necesario. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-host"
    Obligatorio para ejecutarse dentro de un subbloque Esta configuración admitirá tu carga de trabajo solo si hay un subbloque disponible con suficientes recursos para satisfacer los requisitos de recursos de tu carga de trabajo. kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-subblock"
    Preferido para ejecutarse dentro de un subbloque Esta configuración admitirá tu carga de trabajo siempre que haya suficientes recursos disponibles para satisfacer los requisitos de recursos de tu carga de trabajo, incluso si la capacidad está fragmentada. Kueue intentará programar tus pods dentro de un subbloque y usará subbloques adicionales si es necesario. En este caso, Kueue clasificará más alto un subbloque con más capacidad disponible, incluso si está fragmentado en comparación con un subbloque con la capacidad suficiente para satisfacer los requisitos. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-subblock"
    Obligatorio para ejecutarse dentro de un bloque Esta configuración admitirá tu carga de trabajo solo si los recursos disponibles dentro de un bloque satisfacen los requisitos de recursos de tu carga de trabajo. Si se admite, Kueue minimizará la cantidad de subbloques y hosts para programar la carga de trabajo. Esto puede provocar la fragmentación de tu capacidad disponible. kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
    Preferido para ejecutarse dentro de un bloque Esta configuración admitirá tu carga de trabajo siempre que haya suficientes recursos disponibles para satisfacer los requisitos de recursos de tu carga de trabajo, incluso si la capacidad está fragmentada. Kueue intentará programar tus pods dentro de un bloque y usará bloques adicionales si es necesario. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-block"

Programa cargas de trabajo con PodGroup con la TAS mediante Kueue

Cuando usas PodGroups, debes especificar tres campos adicionales para cada pod en un PodGroup:

Según el framework de AA que uses, un líder de un PodGroup puede requerir una GPU o no. Debido a una limitación de Kueue, estos casos deben manejarse de manera diferente. En los siguientes ejemplos, se muestra cómo crear un PodGroup de tres pods con un líder y dos trabajadores.

Caso 1: El líder también es un trabajador y requiere una GPU

Si el líder es uno de los trabajadores y también requiere una GPU, el líder puede tener cualquier número dentro del PodGroup. Para simplificar, en el siguiente ejemplo, el índice del líder es 0:

apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-leader-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "0"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  containers:
  - name: leader
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
  restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-1-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "1"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-2-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "2"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

Caso 2: El líder no es un trabajador y no requiere una GPU

Si el líder no es uno de los trabajadores debido a la limitación de Kueue, el líder debe tener el último índice en el PodGroup, debido a la forma en que Kueue crea PodSets. Si el líder no tiene el último índice y el primer trabajador no usa el primer índice, Kueue no aplicará asignaciones de clasificación.

Consulta el siguiente ejemplo:

---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-leader-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "2"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  containers:
  - name: leader
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        cpu: "1"
      limits:
        cpu: "1"
  restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-0-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "0"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-1-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "1"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

¿Qué sigue?