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 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 mejorar el rendimiento y la confiabilidad. Específicamente, configurarás tu clúster para que use la programación adaptable a 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 que tiene en cuenta la topología (TAS)?

TAS puede mejorar significativamente la eficiencia del entrenamiento de modelos de lenguaje grandes (LLM). TAS coloca estratégicamente a los trabajadores en la topología de la 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, TAS reduce la contención de la red y optimiza la utilización del ancho de banda, lo que genera una convergencia más rápida y tiempos de entrenamiento más cortos. Con los modelos de LLM cada vez más grandes, TAS es esencial para maximizar el rendimiento y la escalabilidad del entrenamiento distribuido.

TAS funciona mejor con capacidad ubicada de forma densa, lo que se puede obtener a través de 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 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 que se describen en 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 TAS habilitado

  2. Cómo ver la topología de tu clúster de GKE

  3. Configura Kueue

Instala Kueue con TAS habilitado

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

Para instalar Kueue y habilitar 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 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 TAS habilitado usando 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 Kueke, 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 ver 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 TAS. De lo contrario, tendrás errores.

Para ver 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 nodos:

  • 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 definición de cuota de recursos de ClusterQueue de infraestructura estática o infraestructura dinámica con el escalamiento 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 el TAS de la siguiente manera:

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

  • Cargas de trabajo que no son de 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: Prácticamente nunca detiene la admisión de una carga de trabajo en función de los recursos solicitados. Según las definiciones de TAS, Kueue puede admitir o no cargas de trabajo según 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 se encuentran dentro de estos límites de cuota de recursos. Según las definiciones de 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 según 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 la 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 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 CPUs 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 lo siguiente:

    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 TAS usando Kueue

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

  • A continuación, se indican los tipos de solicitudes de topología disponibles (preferidos o requeridos):

    • 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 aun así admite una carga de trabajo que no se ajusta a 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 sigue intentando admitir esta carga de trabajo hasta que toda la carga de trabajo quepa en el nivel de topología elegido.

  • A continuación, se indican los niveles disponibles de la solicitud de topología, que te permiten 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 y TAS a través de 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 PodGroup puede requerir una GPU o no. Debido a una limitación de Kueue, estos casos deben controlarse 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, 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, debe tener el último índice en 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á las asignaciones de rango.

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?