Implantar Multislices de TPU no GKE

Nesta página, mostramos como implantar cargas de trabalho no Google Kubernetes Engine (GKE) usando a configuração Multislice do Cloud TPU para otimizar treinamento em grande escala.

Este tutorial é destinado a engenheiros de machine learning (ML) e administradores e operadores de plataforma que querem usar a orquestração de contêineres do Kubernetes para gerenciar cargas de trabalho de treinamento, ajuste e inferência de modelos em grande escala usando TPUs. Para saber mais sobre papéis comuns e exemplos de tarefas referenciados no conteúdo do Google Cloud , consulte Funções e tarefas comuns do usuário do GKE.

Antes de configurar o Multislice no GKE, confira se você conhece 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 de VMs em uma fração de TPU em que duas ou mais frações do Cloud TPU se comunicam pela rede do 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 TPU. Em uma configuração Multislice, o GKE implanta uma carga de trabalho Multislice em diversas frações de TPU. A comunicação entre chips TPU 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 um ct4p-hightpu-4t com uma topologia 2x2x1 ou um ct5lp-hightpu-4t com uma topologia 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.
  • O Multislice não é compatível com a TPU v3.

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

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 de fração da TPU

É possível criar mais de um pool de nós de TPU de fração de vários hosts. Para os fins deste guia, crie três pools de nós de TPU de fração 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 Google Cloud CLI, o Terraform ou o console do Google Cloud .

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    [--num-nodes=NUM_NODES] \
    [--spot \]
    [--flex-start \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME] \
    [--node-labels cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME,cloud.google.com/gke-workload-type=HIGH_AVAILABILITY]
    [--placement-type=COMPACT]

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 identificar um local disponível, consulte Disponibilidade da TPU no GKE.
  • CLUSTER_NAME: o nome do cluster.
  • NODE_ZONES: 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 Escolher a versão da TPU.
  • TPU_TOPOLOGY: a topologia física da fatia de TPU. O formato da topologia depende da versão da TPU. Para saber mais sobre topologias de TPU, use a tabela em Escolher uma topologia.

    Para saber mais, consulte Topologia.

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

  • 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. Se você omitir essa flag, o número de nós será calculado e definido como padrão com base na topologia e no tipo de máquina.
  • 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 de fração 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 de nós de fração. Isso não pode ser alterado após a criação do pool de nós. Para mais informações, consulte VMs spot.
  • --flex-start: define o pool de nós para usar VMs de início flexível. As VMs de início flexível são criadas usando a opção de consumo flex-start, que é compatível com a versão 1.33.0-gke.1712000 ou mais recente do GKE.
  • --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.
  • --node-label=cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME, cloud.google.com/gke-workload-type=HIGH_AVAILABILITY: informa ao GKE que o pool de nós de fração de TPU de vários hosts é uma coleção. Use essa flag se as seguintes condições forem verdadeiras:

    • O pool de nós executa cargas de trabalho de inferência no novo pool de nós.
    • O pool de nós usa a TPU Trillium.
    • As VMs spot não são compatíveis com o agendamento de coleta.

    Para saber mais sobre o gerenciamento da programação de coleta, consulte Gerenciar a programação de coleta em frações de TPU de vários hosts.

  • --placement-type=COMPACT: crie um pool de nós com o posicionamento compacto ativado. Essa opção precisa ser usada com a flag --tpu-topology. Para mais informações, consulte Criar uma política de posicionamento compacto e Topologia de TPU.

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
        flex_start = false
      }
    
      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 de fração 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 Escolher a versão da 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 saber mais sobre topologias de TPU, use a tabela em Escolher uma topologia.

    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 fração 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.
    • flex_start: define o pool de nós para usar a opção de consumo flex-start. Não pode ser definido como true se spot estiver ativado.

Console

Para criar um pool de nós com TPUs:

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

    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 o nome da zona com base na versão da TPU que você quer usar. Para identificar um local disponível, consulte Disponibilidade da TPU no GKE.

  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:

    • CT3P: para a TPU v3.
    • 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 Escolher a versão da TPU para saber como definir o tipo de máquina e a topologia de TPU que criam um pool de nós de fração de TPU de vários hosts.

  10. No menu suspenso Topologia da TPU, selecione a topologia física para a fatia 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 \
        --location=CONTROL_PLANE_LOCATION
    

    Substitua:

    • CLUSTER_NAME: o nome do cluster.
    • PROJECT_ID: o ID do projeto.
    • CONTROL_PLANE_LOCATION: o local do Compute Engine do plano de controle do cluster. Forneça uma região para clusters regionais ou uma zona para clusters zonais.
  2. Use kubectl no Cloud Shell para ver os nós de fração de TPU:

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

    Substitua:

    • TPU_ACCELERATOR: o tipo de acelerador de TPU que você usou ao criar os pools de nós. Por exemplo, tpu-v4-podslice 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:

    Piloto automático

    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: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    - containerPort: 8431
                    command:
                    - bash
                    - -c
                    - |
                      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: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    - containerPort: 8431
                    securityContext:
                      privileged: true
                    command:
                    - bash
                    - -c
                    - |
                      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 fração de TPU. Nesse caso, o NUM_SLICES é igual a 3.
    • ACCELERATOR_TYPE: o tipo de acelerador de TPU que você usou ao criar os pools de nós. Por exemplo, tpu-v4-podslice 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 de TPU 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 em uma fração de TPU.

    Nesse manifesto:

    • O JobSet é um serviço headless com o mesmo nome do JobSet. Nesse caso, é multislice-job.
    • A anotação alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool configura a afinidade de pods para garantir que todos os pods sejam programados na mesma fatia.
    • O campo image é definido como a imagem mais recente da IA do JAX. Para definir uma versão diferente, consulte a lista de imagens atuais da IA JAX.
    • 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.

  5. Opcional: remova a carga de trabalho JAX:

    kubectl delete -f tpu-multislice.yaml
    

Definir outras configurações

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

Melhorar o desempenho da rede com hostNetwork

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, adicione as duas linhas a seguir à 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.

Se você usa a TPU Trillium (v6e) e seus pods usam hostNetworking, instale o DaemonSet a seguir para ajustar /proc/sys/net/ipv4/tcp_rmem no nó.

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/ai-on-gke/51bf3dcab6ff658cf62cc32867f96860bf58dfdc/scripts/network-setup/v6e-increase-rmem.yaml

Melhorar o desempenho da rede sem hostNetwork na TPU Trillium

Se você usa a TPU Trillium e seus pods não podem usar hostNetworking, ative várias redes com o modo netdevice para ter o melhor desempenho de rede. O suporte de placa de rede (NIC) do modo netdevice com várias redes transmite a placa de rede (NIC) da VM diretamente para o pod, ignorando o Kubernetes e o GKE Dataplane V2.

O tipo de máquina ct6e-standard-4t é compatível com duas NICs físicas. O Kubernetes exige uma vNIC que não pode ser transmitida para pods. Portanto, cada nó precisa ter três vNICs para permitir que os pods tenham acesso direto a duas vNICs e alcancem o melhor desempenho das duas NICs físicas.

Para ativar o modo netdevice para ct6e-standard-4t, siga estas etapas:

  1. Crie mais duas VPCs que ofereçam suporte ao modo netdevice
  2. Criar um cluster do GKE com recursos de várias redes
  3. Configure duas redes netdevice. Por exemplo, é possível usar os seguintes objetos GKENetworkParamSet e Network (SECOND_VPC e THIRD_VPC são as VPCs criadas na etapa anterior):

    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: tpu-second
    spec:
      vpc: SECOND_VPC
      vpcSubnet: SECOND_VPC_SUBNET
      deviceMode: NetDevice
    ---
    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: tpu-third
    spec:
      vpc: THIRD_VPC
      vpcSubnet: SECOND_VPC_SUBNET
      deviceMode: NetDevice
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: tpu-second
    spec:
      provider: "GKE"
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: tpu-second
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: tpu-third
    spec:
      provider: "GKE"
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: tpu-third
    
  4. Conecte seus pods com três redes. Por exemplo, é possível usar as seguintes anotações na especificação do pod:

    metadata:
      annotations:
        networking.gke.io/default-interface: 'eth0'
        networking.gke.io/interfaces: |
          [
            {"interfaceName":"eth0","network":"default"},
            {"interfaceName":"eth1","network":"tpu-second"},
            {"interfaceName":"eth2","network":"tpu-third"},
          ]
    
  5. Aplique sysctls de rede no pod, no contêiner de inicialização ou no contêiner de aplicativo. Por exemplo, você pode adicionar o seguinte contêiner de inicialização à especificação do pod:

    initContainers:
    - name: "network-optimization-sysctls"
      image: "busybox"
      securityContext:
        privileged: true
      command:
      - bash
      - -c
      - |
        echo 5000 > /proc/sys/net/ipv4/tcp_rto_min_us
        echo 1 > /proc/sys/net/ipv4/tcp_no_metrics_save
        echo 0 > /proc/sys/net/ipv4/tcp_slow_start_after_idle
        echo 131072 > /proc/sys/net/core/optmem_max
        echo "4096 41943040 314572800" > /proc/sys/net/ipv4/tcp_rmem
    
Prática recomendada:

Use as interfaces eth1 e eth2 para melhorar o desempenho da rede, em vez da interface eth0. Para isso, adicione export LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS --megascale_grpc_interface_prefixes=eth1,eth2,lo" à especificação da carga de trabalho.

Ativar a geração de registros

Os registros emitidos por contêineres em execução nos nós do GKE, incluindo nós de fração 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>

Para saber mais, consulte Acessar os registros de TPU do GKE.

Ativar métricas adicionais

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

Limpar recursos

O jeito mais fácil de evitar cobranças é excluir o Google Cloud projeto que você criou para o tutorial. A outra opção é excluir os recursos individuais.

Excluir o projeto

    Delete a Google Cloud project:

    gcloud projects delete PROJECT_ID

Excluir recursos individuais

Exclua o cluster do GKE:

```sh
gcloud container clusters delete  CLUSTER_NAME \
   --project=PROJECT_ID  \
   --location=CONTROL_PLANE_LOCATION
```

A seguir