Agendar cargas de trabalho do GKE com o TAS

As cargas de trabalho de IA e ML exigem uma comunicação significativa entre pods. Devido a esse requisito, a largura de banda de rede entre pods afeta diretamente o tempo e o custo de execução da carga de trabalho. Essa largura de banda depende do posicionamento das instâncias de máquina virtual (VM) no cluster.

Neste documento, explicamos como otimizar a programação de cargas de trabalho de IA ou ML em grande escala em um cluster do Google Kubernetes Engine (GKE) para melhorar o desempenho e a confiabilidade. Especificamente, você configura o cluster para usar o Topology Aware Scheduling (TAS) para comunicação de baixa latência. Essa abordagem minimiza a sobrecarga de comunicação e ajuda a maximizar o desempenho das cargas de trabalho.

O que é o Topology Aware Scheduling (TAS)?

A TAS pode melhorar significativamente a eficiência do treinamento de modelos de linguagem grandes (LLMs). A TAS posiciona estrategicamente os workers na topologia de rede para minimizar o overhead de comunicação durante a agregação de gradientes, o que exige que os workers se comuniquem em uma ordem de classificação específica. Ao minimizar os saltos de rede entre workers que se comunicam sequencialmente, o TAS reduz a disputa por recursos de rede e otimiza o uso da largura de banda, o que resulta em convergência mais rápida e tempos de treinamento mais curtos. Com modelos de LLM cada vez maiores, o TAS é essencial para maximizar o desempenho e a escalonabilidade do treinamento distribuído.

O TAS funciona melhor com capacidade densamente posicionada, que pode ser obtida por meio de capacidade vinculada à reserva. Com as VMs spot e outras opções de consumo, é menos provável que sua capacidade seja alocada de forma próxima. Portanto, o TAS pode não funcionar bem nesse cenário.

Antes de começar

Antes de começar, verifique se você realizou as tarefas a seguir:

  • Ative a API Google Kubernetes Engine.
    • Ativar a API Google Kubernetes Engine
  • Para usar a Google Cloud CLI nesta tarefa, instale e inicialize a CLI gcloud. Se você instalou a CLI gcloud anteriormente, instale a versão mais recente executando o comando gcloud components update. Talvez as versões anteriores da CLI gcloud não sejam compatíveis com a execução dos comandos neste documento.

  • Para se conectar ao cluster, execute o seguinte comando:

    gcloud container clusters get-credentials CLUSTER_NAME

    Substitua CLUSTER_NAME pelo nome do cluster.

Preparar o cluster do GKE

Para preparar o cluster do GKE para executar cargas de trabalho com o TAS, siga estas etapas:

  1. Instalar o Kueue com o TAS ativado

  2. Ver a topologia do cluster do GKE

  3. Configurar o Kueue

Instalar o Kueue com o TAS ativado

Recomendamos que você use a TAS com o Kueue, um sistema nativo do Kubernetes que gerencia cotas e como os jobs devem consumi-las. A TAS exige a versão 0.10.0 ou mais recente do Kueue, e você precisa ativá-la explicitamente.

Para instalar o Kueue e ativar o TAS, selecione uma das seguintes opções:

Manifesto do Kueue

  1. Instale o Kueue:

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

  2. Ative a TAS no 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 do Helm

Instale o Kueue com o TAS ativado usando um gráfico do 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"

Depois de instalar o Kueue, configure-o para entender a infraestrutura que ele está gerenciando, conforme explicado na próxima seção.

Ver a topologia do cluster do GKE

Antes de conferir a topologia dos nós A4X, A4, A3 Ultra, A3 Mega e A3 High (8 GPUs) provisionados como VMs spot, é necessário definir o posicionamento compacto nos nós do GKE para expor a topologia física deles para o TAS. Caso contrário, você vai encontrar erros.

Para ver a topologia dos nós do cluster do GKE em um pool de nós específico, execute o seguinte 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

Substitua NODE_POOL_NAME pelo nome do pool de nós.

Para entender a topologia física dos nós do GKE nas VMs na saída, consulte os seguintes rótulos de nó:

  • cloud.google.com/gce-topology-block: o ID específico da organização do bloco reservado em que a VM está localizada.

  • cloud.google.com/gce-topology-subblock: o ID específico da organização do sub-bloco em que a VM está localizada.

  • cloud.google.com/gce-topology-host: o ID do host em que a VM está localizada.

  • kubernetes.io/hostname: o nome do host do nó do Kubernetes. Esse nome do host normalmente também é o nome do nó do GKE.

Quanto mais valores de rótulo duas VMs compartilharem, mais próximas elas estarão fisicamente uma da outra. Para mais informações sobre esses termos, consulte Terminologia.

Configurar o Kueue

Depois de instalar o Kueue, é preciso configurá-lo para especificar a infraestrutura que ele está gerenciando. Normalmente, o Kueue requer uma definição de cota de recursos de ClusterQueue de infraestrutura estática ou dinâmica com o escalonamento automático de cluster ativado. O ClusterQueue só aceita uma carga de trabalho se os recursos solicitados por ela forem menores ou iguais ao pool de recursos definido no ClusterQueue. Depois de configurar o Kueue conforme descrito nesta seção, ele vai admitir cargas de trabalho usando o TAS da seguinte maneira:

  • Cargas de trabalho de TAS: o Kueue verifica a topologia da infraestrutura física e o uso atual dela.

  • Cargas de trabalho não TAS: o Kueue não verifica a topologia da infraestrutura física. O Kueue gerencia toda a cota definida na configuração e deixa a atribuição de nós para o kube-scheduler.

Para entender como fornecer uma definição de cota de recursos do ClusterQueue para o Kueue, confira os exemplos a seguir:

  • Cota muito alta: o Kueue praticamente nunca interrompe a admissão de uma carga de trabalho com base nos recursos solicitados. Com base nas definições de TAS, o Kueue pode ou não admitir cargas de trabalho com base na topologia da infraestrutura. Para mais informações, consulte Cota de recursos muito alta.

  • Cota realista: o Kueue só aceita a carga de trabalho se os recursos que ela solicita estiverem dentro desses limites de cota. Com base nas definições de TAS, o Kueue verifica a topologia da infraestrutura antes de admitir a carga de trabalho. Para mais informações, consulte Cota de recursos realista.

Todas as referências à cota de recursos nas seções a seguir se referem à cota de recursos do ClusterQueue.

Cota de recursos muito alta

O exemplo a seguir usa uma cota de recursos muito alta, de modo que o Kueue nunca interrompe uma carga de trabalho com base na cota de recursos disponível. Em vez disso, o Kueue usa as informações de topologia dos nós disponíveis para tentar corresponder a topologia aos requisitos da carga de trabalho.

Para usar a seguinte definição de cota de recursos, siga estas etapas:

  1. Abra um editor de arquivos de sua preferência. Em seguida, inclua a seguinte definição de cota em um arquivo YAML chamado 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"

    Substitua NODE_POOL_NAME pelo nome do pool de nós.

  2. Crie e aplique a configuração de cota de recursos para o sistema de enfileiramento de jobs do Kueue:

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

Cota de recursos realista

O exemplo anterior só configurou recursos de GPU. No entanto, o Kueue pode gerenciar todos os recursos compatíveis com o Kubernetes.

O exemplo a seguir define uma cota de recursos mais realista, incluindo CPU, memória e GPU. Isso é para 100 máquinas a3-ultragpu-8g. Uma única máquina tem 224 vCPUs, 2.944 GB de memória e 8 GPUs. O mesmo princípio se aplica aos tipos de máquina de GPU geral (A2, A3 Edge, G2, G4 e N1), usando as respectivas especificações de CPU e memória.

Para usar a seguinte definição de cota de recursos, siga estas etapas:

  1. Abra um editor de arquivos de sua preferência. Em seguida, inclua a seguinte definição de cota em um arquivo YAML chamado 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"

    Substitua NODE_POOL_NAME pelo nome do pool de nós.

  2. Crie e aplique uma configuração de cota de recursos para o sistema de enfileiramento de jobs do Kueue:

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

    O resultado será o seguinte:

    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

Programar cargas de trabalho com TAS usando o Kueue

Os cenários a seguir mostram como instruir o Kueue e o TAS a gerenciar combinações comuns de carga de trabalho e infraestrutura usando tipos e níveis de solicitação de topologia:

  • Confira a seguir os tipos de solicitação de topologia disponíveis (preferenciais ou obrigatórios):

    • kueue.x-k8s.io/podset-preferred-topology: o Kueue prioriza o agendamento de toda a carga de trabalho em um determinado nível de topologia, mas ainda admite uma carga de trabalho que não se encaixa nesse nível de topologia. Para uma carga de trabalho que poderia caber em um único nível de topologia, o Kueue pode programar essa carga de trabalho em várias instâncias desse nível de topologia.

    • kueue.x-k8s.io/podset-required-topology: o Kueue continua tentando admitir essa carga de trabalho até que ela caiba inteiramente no nível de topologia escolhido.

  • Confira a seguir os níveis de solicitação de topologia disponíveis, que permitem que você seja mais ou menos específico sobre a infraestrutura física em que prefere ou exige que o job seja executado:

    • 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 trabalho usando esses valores, use o seguinte arquivo YAML de job:

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

Substitua as seguintes variáveis:

  • JOB_NAME: um nome para o job.

  • NUMBER_OF_REPLICAS: o número de pods em execução em paralelo.

  • ANNOTATIONS_STRING: consulte a tabela a seguir:

    Tipo e nível de topologia solicitados Descrição ANNOTATIONS_STRING
    Preferencial para execução em um nome do host (recomendado) Essa configuração vai aceitar sua carga de trabalho, desde que haja recursos suficientes disponíveis para atender aos requisitos dela, mesmo que a capacidade esteja fragmentada. O Kueue vai programar os pods da forma mais compacta possível. kueue.x-k8s.io/podset-preferred-topology: "kubernetes.io/hostname"
    Obrigatório para execução em um host

    Essa configuração vai admitir sua carga de trabalho se, e somente se, houver um host disponível com recursos suficientes para atender aos requisitos de recursos da carga de trabalho.

    Isso é útil quando há várias VMs por host (por exemplo, tipos de máquina menores) ou quando vários pods podem ser executados em um único nó. Nesses casos, se a carga de trabalho for admitida, ela será executada em um único host.

    		 kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-host"
    Preferível para execução em um host Essa configuração vai aceitar sua carga de trabalho, desde que haja recursos suficientes disponíveis para atender aos requisitos dela, mesmo que a capacidade esteja fragmentada. O Kueue vai tentar programar os pods em um host e vai usar outros hosts, se necessário. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-host"
    Obrigatório para execução em um sub-bloco Essa configuração vai admitir sua carga de trabalho se, e somente se, houver um sub-bloco disponível com recursos suficientes para atender aos requisitos de recursos da carga de trabalho. kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-subblock"
    Preferencial para execução em um sub-bloco Essa configuração vai aceitar sua carga de trabalho, desde que haja recursos suficientes disponíveis para atender aos requisitos dela, mesmo que a capacidade esteja fragmentada. O Kueue vai tentar programar os pods em um sub-bloco e vai usar outros sub-blocos, se necessário. Nesse caso, o Kueue vai classificar melhor um sub-bloco com mais capacidade disponível, mesmo que ele esteja fragmentado, em comparação com um sub-bloco com capacidade suficiente para atender aos requisitos. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-subblock"
    Obrigatório para execução em um bloco Essa configuração vai admitir sua carga de trabalho se, e somente se, os recursos disponíveis em um bloco atenderem aos requisitos de recursos da carga de trabalho. Se admitido, o Kueue vai minimizar o número de sub-blocos e hosts para programar a carga de trabalho. Isso pode resultar na fragmentação da capacidade disponível. kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
    Preferível para execução em um bloco Essa configuração vai aceitar sua carga de trabalho, desde que haja recursos suficientes disponíveis para atender aos requisitos dela, mesmo que a capacidade esteja fragmentada. O Kueue vai tentar programar os pods em um bloco e vai usar outros blocos, se necessário. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-block"

Programar cargas de trabalho usando o PodGroup com TAS usando o Kueue

Ao usar PodGroups, é necessário especificar três campos adicionais para cada pod em um PodGroup:

Dependendo do framework de ML usado, um líder de um PodGroup pode ou não exigir uma GPU. Devido a uma limitação do Kueue, esses casos precisam ser tratados de maneira diferente. Os exemplos a seguir demonstram como criar um PodGroup de três pods com um líder e dois workers.

Caso 1: o líder também é um worker e requer uma GPU

Se o líder for um dos workers e também precisar de uma GPU, ele poderá ter qualquer número no PodGroup. Para simplificar, no exemplo a seguir, o índice do líder é 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: o líder não é um worker e não requer uma GPU

Se o líder não for um dos workers devido à limitação do Kueue, ele precisará ter o último índice no PodGroup, devido à forma como o Kueue cria PodSets. Se o líder não tiver o último índice e o primeiro worker não usar o primeiro índice, o Kueue não vai aplicar as atribuições de classificação.

Veja o exemplo a seguir:

---
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"

A seguir