Implantar Multislices de TPU no GKE

Nesta página, apresentamos a configuração Multislice de TPU no Google Kubernetes Engine (GKE). Antes de configurar Multislice no GKE, é preciso conhecer os seguintes conceitos:

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

O que é Multislice de TPU

Multislice de TPU é a organização arquitetônica das VMs de TPU em que duas ou mais frações do Cloud TPU se comunicam pela rede de data center (DCN, na sigla em inglês). Multislice permite um treinamento de pilha completa, econômico e em grande escala com escalonamento vertical quase linear para até dezenas de milhares de chips de TPUs. Em uma configuração Multislice, o GKE implanta uma carga de trabalho Multislice em diversas frações de TPU. A comunicação entre chips em uma fração acontece por interconexões entre chips (ICI). A comunicação entre as frações acontece pela DCN.

Recomendamos o uso do recurso Multislice se o job for grande demais para caber em uma única fração de TPU.

Disponibilidade de Multislice no GKE

  • O GKE é compatível com Multislice na versão 1.27.4-gke.900 e posterior.
  • O Autopilot é compatível com o Multislice na versão 1.29.2-gke.1521000 e mais recente.
  • Multislice oferece suporte aos framework JAX e PyTorch. A versão mínima do JAX compatível é a 2.1.
  • Multislice oferece suporte apenas a pools de nós de frações de TPU de vários hosts. Por exemplo, não é possível usar Multislice com ct4p-hightpu-4t com uma topologia 2x2x1 ou um ct5lp-hightpu-4t com 2x2, porque são pools de nós de frações de TPU de host único.
  • Multislice oferece suporte apenas ao treinamento síncrono com vários controladores.
  • Cargas de trabalho Multislice só podem ser executadas em frações de TPU que compartilhem o mesmo tipo, tamanho e topologia de TPU.

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.

Executar uma carga de trabalho em um Multislice

Nesta seção, mostramos como executar uma carga de trabalho em um multislice. Se você usar o modo Autopilot do GKE, pule para a seção Executar uma carga de trabalho de multislice. Os clusters do Autopilot que executam a versão 1.29.2-gke.1521000 ou mais recente ativa as TPUs por padrão.

Preparar um pool de nós no modo Standard

Esta seção abrange as seguintes etapas:

  1. Criar três pools de nós de frações de TPU com vários hosts.
  2. Verificar o status do pool de nós

Criar o pool de nós da TPU

É possível criar mais de um pool de nós de TPU de vários hosts. Para os fins deste guia, crie três pools de nós de TPU de vários hosts para executar uma carga de trabalho Multislice. É possível criar um pool de nós de fração de TPU de vários hosts usando a CLI do Google Cloud, o Terraform ou o console do Google Cloud.

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONE \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    --num-nodes=NUM_NODES \
    [--spot \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME]

Substitua:

  • POOL_NAME: o nome do novo pool de nós.

  • LOCATION: o nome da zona com base na versão da TPU que você quer usar:

    • Para a TPU v4, use us-central2-b.
    • Os tipos de máquina TPU v5e que começam com ct5l- nunca são de vários hosts.
    • Para os tipos de máquina TPU v5e que começam com ct5lp-, use us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b, ou europe-west4-a.
    • Para os tipos de máquina TPU v5p que começam com ct5p-, use us-east1-d, us-east5-a ou us-east5-c.

    Para saber mais, consulte a disponibilidade da TPU no GKE.

  • CLUSTER_NAME: o nome do cluster.

  • NODE_ZONE: a lista separada por vírgulas de uma ou mais zonas em que o GKE cria o pool de nós.

  • MACHINE_TYPE: o tipo de máquina a ser usado para nós. Para saber mais sobre os tipos de máquina disponíveis, consulte Mapeamento da configuração de TPU.

  • TPU_TOPOLOGY: a topologia física da fatia de TPU. O formato da topologia depende da versão da TPU, da seguinte maneira:

    • TPU v4 ou v5p: defina a topologia em três tuplas ({A}x{B}x{C}), por exemplo, 4x4x4.
    • TPU v5e: defina a topologia em duas tuplas ({A}x{B}), por exemplo, 2x2.

    Para saber mais, consulte Topologia.

  • NUM_NODES: o número de nós no pool de nós. Precisa ser zero ou o produto dos valores definidos em TPU_TOPOLOGY ({A}x{B}x{C}) dividido pelo número de chips em cada VM. Para a TPU v4e de vários hosts e a TPU v5e, o número de chips em cada VM é quatro. Portanto, se a TPU_TOPOLOGY for 2x4x4 (TPU v4 com quatro chips em cada VM), a NUM_NODES será 32/4, que é igual a 8.

Opcionalmente, é possível usar as seguintes sinalizações:

  • RESERVATION_NAME: o nome da reserva que o GKE usa ao criar o pool de nós. Se você omitir essa sinalização, o GKE usará pools de nós da TPU disponíveis. Para saber mais sobre reservas de TPU, consulte Reserva de TPU.
  • --spot: define o pool de nós para usar VMs do Spot nos nós da TPU. Isso não pode ser alterado após a criação do pool de nós. Para mais informações, consulte VMs spot.
  • --enable-autoscaling: adicionar um pool de nós com escalonamento automático ativado. Quando o GKE escalona um pool de nós de uma fração da TPU de vários hosts, ele escalona horizontalmente e de maneira atomizada esse pool de nós de zero até o tamanho máximo.
    • MAX_NODES: o tamanho máximo do pool de nós. A flag --max-nodes será obrigatória se --enable-autoscaling for fornecido e precisar ser igual ao produto dos valores definidos em TPU_TOPOLOGY ({A}x{B}x{C}) dividido pelo número de ícones para cada VM.

Terraform

  1. Use a versão 4.84.0 ou mais recente do provedor google.

  2. Adicione o seguinte bloco à configuração do Terraform:

    resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]
  initial_node_count = NUM_NODES

  autoscaling {
    max_node_count = MAX_NODES
    location_policy      = "ANY"
  }
  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    spot = true
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = TPU_TOPOLOGY
  }
}

    Substitua:

    • NODE_POOL_RESOURCE_NAME: o nome do recurso do pool de nós no modelo do Terraform.
    • PROJECT_ID: o ID do projeto.
    • CLUSTER_NAME: o nome do cluster atual ao qual o pool de nós será adicionado.
    • POOL_NAME: o nome do pool de nós a ser criado.
    • CLUSTER_LOCATION: local do Compute do cluster. Recomendamos ter um cluster regional para aumentar a confiabilidade do plano de controle do Kubernetes. Também é possível usar um cluster zonal. Para saber mais, consulte Selecionar uma versão de TPU e topologia.
    • NODE_ZONES: a lista separada por vírgulas de uma ou mais zonas em que o GKE cria o pool de nós.
    • NUM_NODES: o número de nós no pool de nós. Esse valor precisa ser zero ou o produto do número de chips de TPU dividido por quatro. Isso acontece porque, nas frações da TPU de vários hosts, cada nó de TPU tem quatro chips. Por exemplo, se TPU_TOPOLOGY for 4x8, haverá 32 chips, o que significa que NUM_NODES precisa ser 8. Para saber mais sobre topologias de TPU, use a tabela em Mapeamento da configuração de TPU.
    • TPU_TOPOLOGY: indica a topologia física desejada para a fração de TPU. O formato da topologia depende da versão da TPU usada:
      • Para a TPU v4: defina a topologia em três tuplas ({A}x{B}x{C}), por exemplo, 4x4x4.
      • Para a TPU v5e: defina a topologia em duas tuplas ({A}x{B}), por exemplo, 2x2.

    Também é possível usar as seguintes variáveis:

    • RESERVATION_NAME: se você usar a reserva de TPU, esta será a lista de rótulos dos recursos de reserva a serem usados ao criar o pool de nós. Para saber como preencher o RESERVATION_LABEL_VALUES no campo reservation_affinity, consulte Provedor do Terraform.
    • autoscaling: adicionar um pool de nós com escalonamento automático ativado. Quando o GKE escalona um pool de nós de uma fração da TPU de vários hosts, ele escalona horizontalmente e de maneira atomizada esse pool de nós de zero até o tamanho máximo.
      • MAX_NODES: o tamanho máximo do pool de nós. Ele precisa ser zero ou o produto dos valores definidos em TPU_TOPOLOGY ({A}x{B}x{C}) dividido pelo número de chips em cada VM.
    • spot: permite que o pool de nós use VMs spot para os nós de TPU. Isso não pode ser alterado após a criação do pool de nós. Para mais informações, consulte VMs spot.

Console

Para criar um pool de nós com TPUs:

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

    Acessar o Google Kubernetes Engine

  2. Na lista de clusters, clique no nome do cluster que você quer modificar.

  3. Clique em Adicionar pool de nós.

  4. Na seção Detalhes do pool de nós, marque a caixa Especificar locais do nó.

  5. Selecione a zona com base na versão de TPU que você quer usar:

    • Para a TPU v4, use us-central2-b.
    • Os tipos de máquina TPU v5e que começam com ct5l- nunca são de vários hosts.
    • Para os tipos de máquina TPU v5e que começam com ct5lp-, use us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b, ou europe-west4-a.
    • Para os tipos de máquina TPU v5p que começam com ct5p-, use us-east1-d, us-east5-a ou us-east5-c.

  6. No painel de navegação, clique em Nós.

  7. Na seção Configuração da máquina, selecione TPUs.

  8. No menu suspenso Série, selecione uma das seguintes opções:

    • CT4P: para a TPU v4.
    • CT5LP: para a TPU v5e.

  9. No menu suspenso Tipo de máquina, selecione o nome da máquina que será usada para os nós. Use a tabela Mapeamento da configuração de TPU para saber como definir o tipo de máquina e a topologia de TPU que criam um pool de nós de TPU de vários hosts.

  10. No menu suspenso Topologia da TPU, selecione a topologia física para a fração da TPU.

  11. Na caixa de diálogo Alterações necessárias, clique em Fazer alterações.

  12. Certifique-se de que o Tipo de disco de inicialização seja Disco permanente padrão ou Disco permanente SSD.

  13. Como opção, marque a caixa de seleção Ativar nós em VMs do Spot para usar VMs do Spot nos nós do pool de nós.

  14. Clique em Criar.

Verificar o status do pool de nós

  1. Receba as credenciais para poder usar kubectl para acessar o cluster:

    gcloud container clusters get-credentials CLUSTER_NAME \
    --project=PROJECT_ID

    Substitua:

    • CLUSTER_NAME: o nome do cluster.
    • PROJECT_ID: o ID do projeto.

  2. Use kubectl no Cloud Shell para ver os nós de TPU:

    kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=TPU_ACCELERATOR \
   -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGY

    Substitua:

    • TPU_ACCELERATOR: o tipo de acelerador de TPU que você usou quando criou os pools de nós. Por exemplo, tpu-v4-podslice, tpu-v5-lite-device ou tpu-v5-lite-podslice.
    • TPU_TOPOLOGY: a topologia física da fração de TPU.

    O resultado será assim:

     NAME                                    STATUS   ROLES    AGE    VERSION
 gke-tpu-20ee2cce-5tv6                   Ready    <none>   34h     v1.28.1-gke.1066000

Executar uma carga de trabalho Multislice

Nesta seção, você executará uma carga de trabalho JAX que mostra o número global de chips de TPU na fração de TPU e, em seguida, é encerrada.

Para executar uma carga de trabalho JAX, faça o seguinte:

  1. Crie o seguinte manifesto tpu-multislice.yaml:

    Autopilot

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-job
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: NUM_SLICES
      template:
        spec:
          parallelism: NUM_NODES
          completions: NUM_NODES
          backoffLimit: 0
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                command:
                - bash
                - -c
                - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  python -c 'import jax; print("Global device count:", jax.device_count())'
                  sleep 60
                resources:
                 limits:
                    google.com/tpu: NUM_CHIPS

    Padrão

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: multislice-job
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
spec:
  failurePolicy:
    maxRestarts: 4
  replicatedJobs:
    - name: slice
      replicas: NUM_SLICES
      template:
        spec:
          parallelism: NUM_NODES
          completions: NUM_NODES
          backoffLimit: 0
          template:
            spec:
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
              containers:
              - name: jax-tpu
                image: python:3.8
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                securityContext:
                  privileged: true
                command:
                - bash
                - -c
                - |
                  pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                  python -c 'import jax; print("Global device count:", jax.device_count())'
                  sleep 60
                resources:
                  limits:
                   google.com/tpu: NUM_CHIPS

    Substitua:

    • NUM_SLICES: o número de pools de nós de TPU. Nesse caso, o NUM_SLICES é igual a 3.
    • ACCELERATOR_TYPE: o tipo de acelerador de TPU que você usou quando criou os pools de nós. Por exemplo, tpu-v4-podslice, tpu-v5-lite-device ou tpu-v5-lite-podslice.
    • TPU_TOPOLOGY: a topologia física da fração de TPU. Por exemplo, 4x4x4 ou 2x2, dependendo da versão da TPU.
    • NUM_NODES: o número de nós no pool de nós. Precisa ser zero ou o produto dos valores definidos em TPU_TOPOLOGY ({A}x{B}x{C}) dividido pelo número de chips em cada VM. Para a TPU v4 de vários hosts, o número de chips em cada VM é quatro. Para a TPU v5e de vários hosts, o número de chips em cada VM é um, quatro ou oito. Portanto, se a TPU_TOPOLOGY for 2x4x4 (TPU v4 com quatro chips em cada VM), NUM_NODES será 32/4, que é igual a 8.
    • NUM_CHIPS: para a TPU v4 de vários hosts, o número de chips em cada VM é quatro. Para a TPU v5e de vários hosts, o número de chips em cada VM é um, quatro ou oito. Para saber mais, consulte Chips de TPU na VM de TPU.

    Nesse manifesto:

    • O JobSet é um serviço headless com o mesmo nome do JobSet. Nesse caso, é multislice-job.
    • O maxRestarts: 4 indica o número máximo de vezes que o GKE reinicia o JobSet quando um job filho falha. Se as reinicializações do JobSet atingirem o limite máximo definido, uma falha será registrada.
    • Os campos parallelism e completions são iguais ao número de nós em cada pool.
    • O backoff é 0 porque o Multislice oferece suporte apenas ao treinamento síncrono de vários controladores. Precisa ser definido como 0. Falha em qualquer pod gera uma falha no job.
    • Os valores na seção de afinidade garantem que haja apenas uma carga de trabalho Multislice de TPU em execução em um grupo de Multislices.
    • containerPort: 8080 é a porta do coordenador do MXLA
    • containerPort: 8431 é a porta para exportar as métricas de uso da TPU.
    • O securityContext: privileged: true indica que os nós têm o modo privilegiado ativado para acessar TPUs. Os nós na versão 1.28 ou posterior do GKE não precisam que o modo privilegiado esteja ativado para acessar as TPUs. Para saber mais, consulte Executar contêineres sem modo privilegiado.

  2. Aplique o manifesto:

    kubectl apply -f tpu-multislice.yaml

  3. Confirme que a carga de trabalho foi permitida:

    kubectl get jobsets

    O resultado será assim:

    NAME            RESTARTS   COMPLETED   AGE
multislice-job                         3s

  4. Monitore o status dos pods provisionados:

    kubectl get pods

    O resultado será assim:

     NAME                                READY   STATUS      RESTARTS   AGE
 multislice-job-slice-0-0-wzq9t      0/1     Completed   0          2m31s
 multislice-job-slice-0-1-zf4dp      0/1     Completed   0          2m30s
 multislice-job-slice-1-0-hbfn5      0/1     Completed   0          2m31s
 multislice-job-slice-1-1-45fgl      0/1     Completed   0          2m30s
 multislice-job-slice-2-0-wjbp4      0/1     Completed   0          2m30s
 multislice-job-slice-2-1-lwnvs      0/1     Completed   0          2m30s

O JobSet multislice-job programa, cria e executa os pods até a conclusão. Os nomes dos pods estão no formato <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>. O prefixo jobsetName determina o JobSet a que o pod pertence.

Outras configurações

As seções a seguir descrevem as configurações adicionais que você pode aplicar ao Multislice.

Ativar a hostNetwork nos pods padrão do GKE

Para melhorar o desempenho da rede entre frações de TPU, recomendamos ativar hostNetworking. Use hostNetwork: true na especificação do pod para pular toda a pilha de rede do Kubernetes e permitir que os pods do Kubernetes usem a rede do host diretamente para a comunicação entre VMs.

Para ativar hostNetworking, remova as duas linhas a seguir da especificação do pod:

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

Para continuar usando o podHostnames para a descoberta de nós de trabalho com hostNetwork, defina dnsPolicy: ClusterFirstWithHostNet. Isso é importante quando você está executando jobs de retomada automática de treinamento e precisa ter os mesmos nomes para recarregar os mesmos checkpoints.

Geração de registros

Os registros emitidos por contêineres em execução nos nós do GKE, incluindo VMs de TPU, ficam visíveis na Análise de registros, se a geração de registros do sistema GKE estiver ativada no seu cluster.

É possível visualizar os registros do GKE usando a Análise de registros com o seguinte filtro para visualizar os registros de contêiner da carga de trabalho:

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME

Use o seguinte filtro para a fração de TPU e os workers:

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME
resource.labels.pod_name:<jobSetName>-<replicateJobName>-<job-index>-<worker-index>

Observabilidade e métricas

Além das métricas gerais de TPU, há quatro outras métricas específicas de ambiente de execução de TPU em várias partes. Essas métricas estão disponíveis na versão 1.29.1-gke.1016000 ou posterior do GKE. A carga de trabalho da TPU precisa usar a versão 0.4.24 do JAX

Estas são as métricas de multislice disponíveis:

  • Latências de transferência da rede de data center (DCN, na sigla em inglês): distribuição das latências de transferência de rede para o tráfego multislice.
  • Latências coletivas: distribuição da latência coletiva de ponta a ponta para o tráfego multislice.
  • Latências de transferência de host para dispositivo: distribuição da latência de transferência de host para dispositivo em cada bloco de dados para tráfego multislice.
  • Latências de transferência de dispositivo para host: distribuição da latência de transferência do dispositivo para o host em cada bloco de dados para tráfego multislice.

Essas métricas estão localizadas no esquema do contêiner do Kubernetes (k8s_container):

  • kubernetes.io/container/multislice/network/dcn_transfer_latencies
  • kubernetes.io/container/multislice/network/collective_end_to_end_latencies
  • kubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencies
  • kubernetes.io/container/multislice/accelerator/device_to_host_transfer_latencies

Fração de TPU x Multislice

A tabela a seguir diferencia a organização arquitetônica de uma fração de TPU e um Multislice:

Fração de TPU Multislice
Interconectividade A carga de trabalho é executada em uma única fração da TPU. Todos os chips de TPU em uma fração são conectados com o ICI. A carga de trabalho é executada em várias frações de TPU. A comunicação de uma fração acontece por ICI. A comunicação entre frações ocorre por DCN.
Pools de nós compatíveis Fração de TPU de host único e fração de TPU de vários hosts Grupos de frações de TPU de vários hosts
Tipo de carga de trabalho recomendado IndexedJob ou JobSet JobSet

A seguir