Agendar cargas de trabalho do GKE com o TAS

As cargas de trabalho de IA e ML exigem uma comunicação significativa de pod para pod. 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.

Este documento explica como otimizar o agendamento das cargas de trabalho de IA ou ML em grande escala em um cluster do Google Kubernetes Engine (GKE) para desempenho e 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 o desempenho das cargas de trabalho.

O que é o agendamento com reconhecimento de topologia (TAS)?

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 os workers que se comunicam sequencialmente, o TAS reduz a contenção de rede e otimiza a utilização da largura de banda, levando a uma 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 reservas. Com VMs de início flexível ou VMs 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 Google Cloud CLI para essa tarefa, instale e, em seguida, 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. Conferir a topologia do cluster do GKE

  3. Configurar o Kueue

Instalar o Kueue com o TAS ativado

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

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

Manifesto do Kueue

  1. Instalar o Kueue:

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

  2. Ativar 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 Kueue, você precisa configurá-lo para entender a infraestrutura que ele está gerenciando, conforme explicado na próxima seção.

Conferir 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 do Spot, defina 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 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 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 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, você precisa configurá-lo para especificar a infraestrutura que ele está gerenciando. Normalmente, o Kueue exige uma ClusterQueue definição de cota de recursos de infraestrutura estática ou infraestrutura dinâmica com 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 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 do ClusterQueue para o Kueue, revise 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 do TAS, o Kueue pode ou não aceitar 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 solicitados pela carga de trabalho estiverem dentro desses limites de cota de recursos. Com base nas definições do TAS, o Kueue verifica a topologia da infraestrutura antes de aceitar 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.

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á assim:

    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 (preferencial ou obrigatório):

    • 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 aceita 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 de trabalho em várias instâncias desse nível de topologia.

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

  • A seguir estão os níveis de solicitação de topologia disponíveis, permitindo 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 do 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 de recursos da carga de trabalho, 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 aceitar sua carga de trabalho 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áquinas menores) ou 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 aceitar sua carga de trabalho, desde que haja recursos suficientes disponíveis para atender aos requisitos de recursos da carga de trabalho, mesmo que a capacidade esteja fragmentada. O Kueue vai tentar programar seus 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 aceitar sua carga de trabalho 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 de recursos da carga de trabalho, mesmo que a capacidade esteja fragmentada. O Kueue vai tentar programar seus pods em um sub-bloco e vai usar outros sub-blocos, se necessário. Nesse caso, o Kueue vai classificar um sub-bloco com mais capacidade disponível, mesmo que 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 aceitar sua carga de trabalho somente se os recursos disponíveis em um bloco atenderem aos requisitos de recursos da carga de trabalho. Se aceito, 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"
    Preferencial 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 de recursos da carga de trabalho, mesmo que a capacidade esteja fragmentada. O Kueue vai tentar programar seus 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 PodGroup com o TAS usando o Kueue

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

Dependendo do framework de ML que você usa, um líder de um 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 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 exige uma GPU

Se o líder for um dos workers e também exigir 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 exige 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 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