Agendar cargas de trabalho do GKE com o Agendamento com reconhecimento de topologia (TAS, na sigla em inglês)

As cargas de trabalho de IA e ML exigem uma comunicação significativa entre pods. Devido a esse requisito, a largura de banda da rede entre os 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 o agendamento das suas 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 agendamento com reconhecimento de topologia (TAS, na sigla em inglês) para comunicação de baixa latência. Essa abordagem minimiza a sobrecarga de comunicação e ajuda a maximizar a performance das suas cargas de trabalho.

O que é o agendamento com reconhecimento de topologia (TAS, na sigla em inglês)?

A TAS pode melhorar significativamente a eficiência do treinamento de modelos de linguagem grandes (LLMs). O TAS posiciona estrategicamente os workers na topologia de rede para minimizar a sobrecarga de comunicação durante a agregação de gradientes, 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 de rede e otimiza a utilização da largura de banda, resultando em uma convergência mais rápida e tempos de treinamento mais curtos. Com modelos de LLM cada vez maiores, o TAS é essencial para maximizar a performance e a escalonabilidade do treinamento distribuído.

A TAS funciona melhor com capacidade densamente posicionada, que pode ser obtida por reservas. Com VMs de início flexível ou do Spot, é menos provável que sua capacidade seja alocada muito próxima, então 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:

  • Ativar a API Google Kubernetes Engine.
    • Ativar a API Google Kubernetes Engine
  • Se você quiser usar a CLI do Google Cloud para essa tarefa, instale e inicialize a gcloud CLI. Se você instalou a gcloud CLI anteriormente, instale a versão mais recente executando o comando gcloud components update. Talvez as versões anteriores da gcloud CLI 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 seu 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 usar a TAS com o Kueue, um sistema nativo do Kubernetes que gerencia cotas e como os jobs devem consumi-las. O TAS requer a versão 0.10.0 ou mais recente do Kueue, e você precisa ativar 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 o 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 Kueke, 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, defina o posicionamento compacto nos nós do GKE para expor a topologia física deles ao TAS. Caso contrário, você vai encontrar erros.

Para conferir 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 suas VMs na saída, consulte os seguintes rótulos de nós:

  • 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 subbloco 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 geralmente também é o nome do nó do GKE.

Quanto mais valores de rótulo duas VMs compartilharem, mais próximas elas estarão fisicamente. 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 exige uma definição de cota de recursos ClusterQueue de infraestrutura estática ou dinâmica com o escalonamento automático de cluster ativado. O ClusterQueue aceita uma Workload somente se os recursos solicitados pela carga de trabalho forem menores ou iguais ao pool de recursos definidos no ClusterQueue. Depois de configurar o Kueue conforme descrito nesta seção, ele aceita cargas de trabalho usando o TAS da seguinte maneira:

  • Cargas de trabalho do 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 ClusterQueue ao Kueue, consulte 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 dos limites da cota de recursos. 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 a 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 definição de cota de recursos a seguir, conclua 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.

Para usar a definição de cota de recursos a seguir, conclua 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 o 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:

  • A seguir estão 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 ter se encaixado em um único nível de topologia, o Kueue pode programar essa carga em várias instâncias desse nível.

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

  • Estes são os níveis de solicitação de topologia disponíveis, que permitem ser mais ou menos específico sobre a infraestrutura física em que você prefere ou precisa executar seu job:

    • 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 de host (recomendado) Essa configuração vai admitir 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 seus 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 somente se houver um host disponível com recursos suficientes para atender aos requisitos dela.

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

    		 kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-host"
    Preferencial para execução em um host Essa configuração vai admitir sua carga de trabalho desde que haja recursos suficientes disponíveis para atender aos requisitos dela, mesmo que a capacidade esteja fragmentada. O Kueue tenta programar seus pods em um host e usa 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"
    Preferível para execução em um sub-bloco Essa configuração vai admitir 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 seus pods em um subbloco e usará outros subblocos, se necessário. Nesse caso, o Kueue vai classificar mais alto um subbloco com mais capacidade disponível, mesmo que ele esteja fragmentado, em comparação com um subbloco 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 aceitar 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 admitir 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 seus pods em um bloco e usará outros, se necessário. kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-block"

Programar cargas de trabalho usando PodGroup com TAS usando Kueue

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

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

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

Se o líder for um dos workers e também exigir uma GPU, ele poderá ter qualquer número dentro do 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 exige uma GPU

Se o líder não for um dos trabalhadores 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 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