Implante cargas de trabalho de TPU no Autopilot do GKE

Nesta página, descrevemos como acelerar as cargas de trabalho de machine learning (ML) usando aceleradores do Cloud TPU (TPUs) em clusters do Autopilot do Google Kubernetes Engine (GKE). Antes de ler esta página, você deve conhecer os seguintes conceitos:

  1. Introdução ao Cloud TPU
  2. Arquitetura do sistema do Cloud TPU
  3. Sobre TPUs no GKE

Como as TPUs funcionam no Autopilot

Para usar TPUs em cargas de trabalho do Autopilot, é preciso solicitar uma versão da TPU e uma topologia compatível para ela no manifesto da carga de trabalho. Em seguida, use os campos resources.requests e resources.limits do Kubernetes para especificar o número de chips de TPU para a carga de trabalho. Quando você implanta a carga de trabalho, o GKE provisiona nós que têm a configuração de TPU solicitada e programa seus pods nos nós. O GKE coloca cada carga de trabalho no próprio nó para que cada pod possa acessar todos os recursos do nó com risco minimizado de interrupção.

As TPUs no Autopilot são compatíveis com os seguintes recursos:

  1. Pods do Spot
  2. Reservas de capacidade específicas
  3. Pods de tempo de execução estendido

Planejar a configuração da TPU

Antes de solicitar TPUs, defina a configuração com base nos requisitos de CPU e memória da carga de trabalho. Você precisa decidir o seguinte:

  • Versão do TPU: a versão específica do Cloud TPU, como v5e.
  • Topologia da versão selecionada da TPU: a disposição e o número de TPUs.

A versão da TPU e a topologia selecionadas determinam se o GKE provisiona nós como frações de host único ou frações de vários hosts. Em frações de host único, cada nó é independente de outros nós da TPU. Em frações de vários hosts, o GKE cria um grupo de nós que têm VMs interconectadas em uma fração de TPU. As frações de vários hosts são atômicas, o que significa que o GKE aumenta ou diminui todo o grupo interconectado de nós como uma única unidade.

Para informações sobre as versões de TPU disponíveis, as topologias correspondentes, a CPU e a capacidade de memória e o tipo de fração resultante, consulte Escolher uma configuração de TPU do Autopilot.

Preços

Saiba mais sobre os preços em Preços do Autopilot.

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 gcloud components update.
  • Verifique se você tem um cluster do Autopilot em execução na versão 1.29.2-gke.1521000 ou mais recente do GKE.
  • Para usar TPUs reservadas, verifique se você tem uma reserva de capacidade específica. Para mais instruções, consulte Como consumir recursos por zona reservados.

Verifique se você tem uma cota de TPU

Para criar nós de fração de TPU, é necessário ter uma cota de TPU disponível, a menos que você esteja usando uma reserva de capacidade atual. Se você estiver usando TPUs reservadas, pule esta seção.

Para criar nós de fração de TPU no GKE, é preciso ter uma cota da API Compute Engine (compute.googleapis.com), e não uma cota da API Cloud TPU (tpu.googleapis.com). O nome da cota é diferente nos pods normais do Autopilot e nos pods do Spot.

Para verificar o limite e o uso atual da sua cota da API Compute Engine para TPUs, siga estas etapas:

  1. Acesse a página Cotas no console do Google Cloud.

    Acessar "Cotas"

  2. Na caixa Filtro , faça o seguinte:

    1. Selecione a propriedade Serviço, insira API Compute Engine e pressione Enter.

    2. Selecione a propriedade Tipo e escolha Cota.

    3. Selecione a propriedade Nome e digite um nome de cota com base no tipo de TPU que você quer, da seguinte maneira:

      • TPU v5p (tpu-v5p-slice): chips TPU v5p
      • TPU v5e (tpu-v5-lite-podslice): chips TPU v5 Lite PodSlice
      • TPU v5e (tpu-v5-lite-device): chips de dispositivo TPU v5 Lite
      • TPU v4 (tpu-v4-podslice): chips PodSlice de TPU v4

      Para pods do Spot, selecione a cota "Preemptiva" correspondente.

    4. Selecione a propriedade Dimensões (por exemplo, locais) e insira region: seguido do nome da região em que você planeja criar TPUs no GKE. Por exemplo, insira region:us-west4 se planeja criar nós de fração da TPU na zona us-west4-a. A cota de TPU é regional, portanto, todas as zonas na mesma região consomem a mesma cota de TPU.

Se nenhuma cota corresponder ao filtro inserido, isso significa que o projeto não recebeu nenhuma das cotas especificadas para a região desejada. Você precisará solicitar um aumento de cota de TPU.

Preparar seu aplicativo de TPU

As cargas de trabalho de TPU têm os requisitos de preparação a seguir.

  1. Frameworks como JAX, PyTorch e TensorFlow acessam VMs de TPU usando a biblioteca compartilhada libtpu. libtpu inclui o compilador XLA, o software do ambiente de execução da TPU e o driver da TPU. Cada versão do PyTorch e do JAX requer uma determinada versão de libtpu.so. Para usar TPUs no GKE, use as seguintes versões:
    Tipo de TPU Versão do libtpu.so
    TPU v5e
    tpu-v5-lite-podslice
    tpu-v5-lite-device
    TPU v5p
    tpu-v5p-slice
    • Versão recomendada do jax[tpu]: 0.4.19 ou mais recente.
    • Versão recomendada do torchxla[tpuvm]: sugestão de uso de um build de versão noturno em 23 de outubro de 2023.
    TPU v4
    tpu-v4-podslice
  2. Defina as variáveis de ambiente a seguir para o contêiner que solicita os recursos da TPU:
    • TPU_WORKER_ID: um número inteiro exclusivo para cada pod. Esse ID denota um código de worker exclusivo na fatia de TPU. Os valores aceitos para esse campo variam de zero ao número de pods menos um.
    • TPU_WORKER_HOSTNAMES: uma lista separada por vírgulas de nomes de host ou endereços IP de VM da TPU que precisam se comunicar entre si na fatia. É necessário que haja um nome do host ou endereço IP para cada VM da TPU na fatia. A lista de endereços IP ou nomes do host é ordenada e zero indexada pelo TPU_WORKER_ID.

      • O GKE injeta essas variáveis de ambiente automaticamente usando um webhook mutável quando um job é criado com completionMode: Indexed, subdomain e parallelism > 1 e solicitando as propriedades google.com/tpu. O GKE adiciona um serviço sem comando para que os registros DNS sejam adicionados aos pods de apoio do serviço.

Depois de concluir a preparação da carga de trabalho, é possível executar um job que usa TPUs.

Solicitar TPUs em uma carga de trabalho

Nesta seção, mostramos como criar um job que solicita TPUs no Autopilot. Em qualquer carga de trabalho que precise de TPUs, especifique o seguinte:

  • Seletores de nós para a versão e topologia da TPU
  • O número de chips de TPU para um contêiner na sua carga de trabalho

Para ver uma lista de versões de TPU, topologias e o número correspondente de chips e nós de TPU em uma fração, consulte Escolher uma configuração de TPU do Autopilot.

Considerações sobre solicitações de TPU em cargas de trabalho

Somente um contêiner em um pod pode usar TPUs. O número de chips de TPU que um contêiner solicita precisa ser igual ao número de chips anexados a um nó na fração. Por exemplo, se você solicitar a TPU v5e (tpu-v5-lite-podslice) com uma topologia 2x4, poderá solicitar qualquer um dos seguintes itens:

  • Chips 4, que cria dois nós de vários hosts com quatro chips de TPU cada.
  • Chips 8, que cria um nó de host único com oito chips de TPU

Como prática recomendada para maximizar o custo-benefício, consuma sempre toda a TPU na fração que você solicitou. Se você solicitar uma fração de vários hosts de dois nós com quatro chips de TPU cada, implante uma carga de trabalho que seja executada em ambos os nós e consuma todos os oito chips de TPU na fração.

Criar uma carga de trabalho que solicite TPUs

As etapas a seguir criam um job que solicita TPUs. Se você tiver cargas de trabalho executadas em frações de TPU de vários hosts, também será necessário criar um serviço sem comando que selecione a carga de trabalho por nome. Esse serviço sem comando permite que pods em diferentes nós na fração de vários hosts se comuniquem entre si, atualizando a configuração DNS do Kubernetes para apontar para os pods na carga de trabalho.

  1. Salve o seguinte manifesto como tpu-autopilot.yaml:

    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_TYPE
        cloud.google.com/gke-tpu-topology: TOPOLOGY
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS

    Substitua:

    • TPU_TYPE: o tipo de TPU a ser usado, como tpu-v4-podslice. Precisa ser um valor compatível com o GKE.
    • TOPOLOGY: a disposição dos chips de TPU na fração, como 2x2x4. Precisa ser uma topologia compatível com o tipo de TPU selecionado.
    • NUMBER_OF_CHIPS: o número de chips de TPU para o contêiner usar. Precisa ser o mesmo valor para limits e requests.

  2. Implante o job:

    kubectl create -f tpu-autopilot.yaml

Quando você cria esse job, o GKE faz automaticamente o seguinte:

  1. Provisiona nós para executar os pods. Dependendo do tipo de TPU, topologia e solicitações de recursos especificados, esses nós são frações de host único ou de vários hosts.
  2. Adiciona taints aos pods e tolerâncias aos nós para evitar que outras cargas de trabalho sejam executadas nos mesmos nós que as cargas de trabalho da TPU.

Exemplo: exibir o total de ícones de TPU em uma fração de vários hosts

A carga de trabalho a seguir retorna o número de chips de TPU em todos os nós em uma fração de TPU de vários hosts. Para criar uma fração de vários hosts, a carga de trabalho tem os seguintes parâmetros:

  • Versão do TPU: TPU v4
  • Topologia: 2x2x4

Essa seleção de versão e topologia resultam em uma fração de vários hosts.

  1. Salve o seguinte manifesto como available-chips-multihost.yaml: 
    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
  2. Implante o manifesto: 
    kubectl create -f available-chips-multihost.yaml

    O GKE executa uma fatia da TPU v4 com quatro VMs (fatia da TPU de vários hosts). A fração tem 16 chips de TPU interconectados.

  3. Verifique se o job criou quatro pods: 
    kubectl get pods

    O resultado será assim:

    NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
  4. Consiga os registros de um dos pods: 
    kubectl logs POD_NAME

    Substitua POD_NAME pelo nome de um dos pods criados. Por exemplo, tpu-job-podslice-0-5cd8r.

    O resultado será assim: 

    TPU cores: 16

Exemplo: exibir os chips da TPU em um único nó

A carga de trabalho a seguir é um pod estático que exibe o número de chips de TPU anexados a um nó específico. Para criar um nó de host único, a carga de trabalho tem os seguintes parâmetros:

  • Versão do TPU: TPU v5e
  • Topologia: 2x4

Essa seleção de versão e topologia resultam em uma fração de host único.

  1. Salve o seguinte manifesto como available-chips-singlehost.yaml: 
    apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8
  2. Implante o manifesto: 
    kubectl create -f available-chips-singlehost.yaml

    O GKE provisiona nós com oito frações de TPU de host único que usam a TPU v5e. Cada nó de TPU tem oito chips de TPU (fração de TPU de host único).

  3. Consiga os registros do pod: 
    kubectl logs tpu-job-jax-v5

    O resultado será assim: 

    Total TPU chips: 8

Observabilidade e métricas

Painel

Na página Clusters do Kubernetes no Console do Google Cloud, a guia Observabilidade exibe as métricas de observabilidade da TPU. Para mais informações, consulte Métricas de observabilidade do GKE.

O painel da TPU será preenchido apenas se você tiver métricas do sistema ativadas no cluster do GKE.

Métricas do ambiente de execução

No GKE 1.27.4-gke.900 ou mais recente, as cargas de trabalho da TPU que usam o JAX 0.4.14 ou mais recente e especificam containerPort: 8431 exportam as métricas de utilização da TPU como métricas do sistema do GKE. As métricas a seguir estão disponíveis no Cloud Monitoring para monitorar o desempenho do ambiente de execução da carga de trabalho da TPU

  • Ciclo de trabalho: porcentagem de tempo durante o período de amostragem anterior (60 segundos) em que os TensorCores estavam ativamente em processamento em um chip de TPU. Uma porcentagem maior significa melhor uso da TPU.
  • Uso da memória: quantidade de memória do acelerador alocada em bytes. Amostras coletadas a cada 60 segundos.
  • Capacidade de memória: memória total do acelerador em bytes. Amostras coletadas a cada 60 segundos.

Essas métricas estão localizadas no esquema de nó (k8s_node) e contêiner de contêiner (k8s_container) do Kubernetes.

Contêiner do Kubernetes:

  • kubernetes.io/container/accelerator/duty_cycle
  • kubernetes.io/container/accelerator/memory_used
  • kubernetes.io/container/accelerator/memory_total

Nó do Kubernetes:

  • kubernetes.io/node/accelerator/duty_cycle
  • kubernetes.io/node/accelerator/memory_used
  • kubernetes.io/node/accelerator/memory_total

Métricas de hospedagem

No GKE 1.28.1-gke.1066000 ou mais recente, as VMs em uma fração da TPU exportam as métricas de utilização da TPU como métricas do sistema do GKE. As seguintes métricas estão disponíveis no Cloud Monitoring para monitorar o desempenho do host da TPU:

  • Uso do TensorCore: porcentagem atual do TensorCore que está sendo utilizada. O valor do TensorCore é igual à soma das unidades de multiplicação de matriz (MXUs) mais a unidade vetorial. O valor de utilização do TensorCore é a divisão das operações do TensorCore que foram realizadas no período de amostragem anterior (60 segundos) pelo número de operações do TensorCore com suporte no mesmo período. Um valor maior significa melhor utilização.
  • Uso de largura de banda da memória: porcentagem atual da largura de banda da memória do acelerador que está sendo usada. Calculada pela divisão da largura de banda da memória usada durante um período de amostragem (60 segundos) pela largura de banda máxima aceita no mesmo período.

Essas métricas estão localizadas no esquema de nó (k8s_node) e contêiner de contêiner (k8s_container) do Kubernetes.

Contêiner do Kubernetes:

  • kubernetes.io/container/accelerator/tensorcore_utilization
  • kubernetes.io/container/accelerator/memory_bandwidth_utilization

Nó do Kubernetes:

  • kubernetes.io/container/node/tensorcore_utilization
  • kubernetes.io/container/node/memory_bandwidth_utilization

Para mais informações, consulte Métricas do Kubernetes e Métricas do sistema do GKE.

Geração de registros

Os registros emitidos por contêineres em execução em nós do GKE, incluindo VMs de TPU, sãocoletou pelo agente do Logging do GKE, enviado para o Logging e sãovisível no Logging do Google Analytics.

Recomendações para cargas de trabalho de TPU no Autopilot

As recomendações a seguir podem melhorar a eficiência das cargas de trabalho da TPU:

  • Use pods de tempo de execução estendido por um período de carência de até sete dias antes que o GKE encerre seus pods para redução de escalonamento ou upgrade de nós. É possível usar janelas de manutenção e exclusões com pods de tempo de execução estendidos para atrasar ainda mais os upgrades automáticos de nós.
  • Use reservas de capacidade para garantir que as cargas de trabalho recebam as TPUs solicitadas sem serem colocadas em uma fila para disponibilidade.