Implemente multislices de TPUs no GKE


Esta página mostra como implementar cargas de trabalho no Google Kubernetes Engine (GKE) usando a configuração de vários fragmentos da Cloud TPU para uma preparação em grande escala e rentável.

Este tutorial destina-se a engenheiros de aprendizagem automática (ML) e administradores e operadores de plataformas que querem usar a orquestração de contentores do Kubernetes para gerir cargas de trabalho de preparação, otimização e inferência de modelos em grande escala usando TPUs. Para saber mais sobre as funções comuns e as tarefas de exemplo referenciadas no conteúdo, consulte o artigo Funções e tarefas comuns do utilizador do GKE. Google Cloud

Antes de configurar o Multislice no GKE, certifique-se de que conhece os seguintes conceitos:

  1. Introdução à Cloud TPU
  2. Arquitetura do sistema Cloud TPU
  3. Acerca das TPUs no GKE

O que é o TPU Multislice

O Multislice da TPU é a organização arquitetónica das VMs numa fatia da TPU, em que duas ou mais fatias da TPU do Google Cloud comunicam através da rede do centro de dados (DCN). A segmentação múltipla permite uma preparação de grande escala, rentável e de pilha completa com um aumento quase linear até dezenas de milhares de chips de TPU. Numa configuração de várias fatias, o GKE implementa uma carga de trabalho de várias fatias em várias fatias de TPU. A comunicação entre os chips da TPU numa fatia ocorre através de interconexões entre chips (ICI). A comunicação entre as divisões ocorre através da DCN.

Recomendamos que use o Multislice se o seu trabalho for demasiado grande para caber numa única fatia de TPU.

Disponibilidade de vários segmentos no GKE

  • O Standard suporta o Multislice na versão 1.27.4-gke.900 e posteriores.
  • O Autopilot suporta o Multislice na versão 1.29.2-gke.1521000 e posteriores.
  • O Multislice suporta as frameworks JAX e PyTorch. A versão mínima suportada do JAX é a 2.1.
  • O multislice só suporta pools de nós de fatias de TPU com vários anfitriões. Por exemplo, não pode usar o Multislice com uma topologia ct4p-hightpu-4t com uma topologia 2x2x1 ou uma topologia ct5lp-hightpu-4t com uma topologia 2x2, porque estes são pools de nós de fatias de TPU de host único.
  • O Multislice só suporta a preparação de vários controladores síncronos.
  • As cargas de trabalho de divisão múltipla só podem ser executadas em divisões de TPU que partilham o mesmo tipo, tamanho e topologia de TPU.
  • O Multislice não suporta a TPU v3.

Antes de começar

Antes de começar, certifique-se de que realizou as seguintes tarefas:

  • Ative a API Google Kubernetes Engine.
  • Ative a API Google Kubernetes Engine
  • Se quiser usar a CLI gcloud para esta tarefa, instale-a e, em seguida, inicialize-a. Se instalou anteriormente a CLI gcloud, execute gcloud components update para obter a versão mais recente.

Execute uma carga de trabalho numa Multislice

Esta secção mostra como executar uma carga de trabalho num Multislice. Se usar o modo GKE Autopilot, avance para a secção Execute uma carga de trabalho com vários fragmentos. Os clusters do Autopilot que executam a versão 1.29.2-gke.1521000 ou posterior ativam as TPUs por predefinição.

Prepare um node pool do modo padrão

Esta secção aborda os seguintes passos:

  1. Crie três node pools de fatias de TPU com vários anfitriões
  2. Valide o estado do conjunto de nós

Crie o node pool da fatia de TPU

Pode criar mais do que um conjunto de nós de segmentação de TPU multi-anfitrião. Para efeitos deste guia, crie três pools de nós de fatia de TPU com vários anfitriões para executar uma carga de trabalho de várias fatias. Pode criar um nó de segmento de TPU com vários anfitriões usando a CLI Google Cloud, o Terraform ou a Google Cloud consola.

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 o seguinte:

  • POOL_NAME: o nome do novo node pool.
  • LOCATION: o nome da zona com base na versão da TPU que quer usar. Para identificar uma localização disponível, consulte o artigo Disponibilidade de TPUs no GKE.
  • CLUSTER_NAME: o nome do cluster.
  • NODE_ZONES: a lista separada por vírgulas de uma ou mais zonas onde o GKE cria o conjunto de nós.
  • MACHINE_TYPE: o tipo de máquina a usar para os nós. Para saber mais sobre os tipos de máquinas disponíveis, consulte o artigo Escolha a versão da TPU.
  • TPU_TOPOLOGY: a topologia física da fatia da TPU. O formato da topologia depende da versão da TPU. Para saber mais sobre as topologias de TPUs, use a tabela em Escolha uma topologia.

    Para saber mais, consulte o artigo Topologia.

Opcionalmente, também pode usar as seguintes flags:

  • NUM_NODES: o número de nós no node pool. Tem de ser zero ou o produto dos valores definidos em TPU_TOPOLOGY ({A}x{B}x{C}) dividido pelo número de chips em cada MV. Para a TPU v4 e a TPU v5e com vários anfitriões, o número de chips em cada VM é quatro. Por conseguinte, se o seu TPU_TOPOLOGY for 2x4x4 (TPU v4 com quatro chips em cada VM), o NUM_NODES é 32/4, o que equivale a 8. Se omitir esta flag, o número de nós é calculado e definido por predefinição com base na topologia e no tipo de máquina.
  • RESERVATION_NAME: o nome da reserva que o GKE usa quando cria o conjunto de nós. Se omitir este sinalizador, o GKE usa os conjuntos de nós de fatias de TPU disponíveis. Para saber mais acerca das reservas de TPUs, consulte o artigo Reserva de TPUs.
  • --spot: define o node pool para usar VMs do Spot para os nós da fatia de TPU. Não é possível alterar esta opção após a criação do conjunto de nós. Para mais informações, consulte VMs do Spot.
  • --flex-start: define o node pool para usar o modo de aprovisionamento flex-start. O início flexível é suportado na versão 1.33.0-gke.1712000 ou posterior do GKE.
  • --enable-autoscaling: crie um node pool com o dimensionamento automático ativado. Quando o GKE dimensiona um conjunto de nós de fatia de TPU com vários anfitriões, dimensiona atomicamente o conjunto de nós de zero para o tamanho máximo.

    • MAX_NODES: o tamanho máximo do node pool. A flag --max-nodes é obrigatória se --enable-autoscaling for fornecido e tem de ser igual ao produto dos valores definidos em TPU_TOPOLOGY ({A}x{B}x{C}) dividido pelo número de chips em cada MV.
  • --node-label=cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME, cloud.google.com/gke-workload-type=HIGH_AVAILABILITY: indica ao GKE que o conjunto de nós da fatia de TPU com vários anfitriões é uma coleção. Use esta flag se as seguintes condições se aplicarem:

    • O node pool executa cargas de trabalho de inferência no novo node pool.
    • O conjunto de nós usa a TPU Trillium.
    • As VMs do Spot não suportam o agendamento de recolhas.

    Para saber mais sobre a gestão do agendamento da recolha, consulte o artigo Faça a gestão do agendamento da recolha em fatias de TPU com vários anfitriões.

  • --placement-type=COMPACT: crie um node pool com o posicionamento compacto ativado. Esta opção tem de ser usada com a flag --tpu-topology. Para mais informações, consulte os artigos Crie uma política de posicionamento compacta e Topologia da TPU.

Terraform

  1. Certifique-se de que usa a versão 4.84.0 ou posterior do fornecedor 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 o seguinte:

    • NODE_POOL_RESOURCE_NAME: o nome do recurso do conjunto de nós no modelo do Terraform.
    • PROJECT_ID: o ID do seu projeto.
    • CLUSTER_NAME: O nome do cluster existente ao qual adicionar o conjunto de nós.
    • POOL_NAME: o nome do node pool a criar.
    • CLUSTER_LOCATION: calcular a localização do cluster. Recomendamos que tenha um cluster regional para uma maior fiabilidade do plano de controlo do Kubernetes. Também pode usar um cluster zonal. Para saber mais, consulte o artigo Selecione uma versão e uma topologia da TPU.
    • NODE_ZONES: a lista separada por vírgulas de uma ou mais zonas onde o GKE cria o conjunto de nós.
    • NUM_NODES: o número de nós no node pool. Tem de ser zero ou o produto do número de chips da TPU dividido por quatro, porque nas fatias de TPU com vários anfitriões, cada nó de fatia de TPU tem 4 chips. Por exemplo, se TPU_TOPOLOGY for 4x8, existem 32 chips, o que significa que NUM_NODES tem de ser 8. Para saber mais sobre as topologias de TPU, use a tabela em Escolha a versão do TPU.
    • TPU_TOPOLOGY: isto indica a topologia física desejada para a fatia da TPU. O formato da topologia depende da versão da TPU que está a usar. Para saber mais sobre as topologias de TPUs, use a tabela em Escolha uma topologia.

    Opcionalmente, também pode usar as seguintes variáveis:

    • RESERVATION_NAME: se usar a reserva de TPU, esta é a lista de etiquetas dos recursos de reserva a usar quando criar o conjunto de nós. Para saber como preencher o campo RESERVATION_LABEL_VALUES no campo reservation_affinity, consulte o fornecedor do Terraform.
    • autoscaling: crie um node pool com o dimensionamento automático ativado. Quando o GKE dimensiona um conjunto de nós de fatia de TPU com vários anfitriões, dimensiona atomicamente o conjunto de nós de zero para o tamanho máximo.
      • MAX_NODES: é o tamanho máximo do node pool. Tem de ser igual ao produto dos valores definidos em TPU_TOPOLOGY ({A}x{B}x{C}) dividido pelo número de chips em cada MV.
    • spot: permite que o node pool use VMs do Spot para os nós de fatia da TPU. Não é possível alterar esta opção após a criação do conjunto de nós. Para mais informações, consulte VMs de spot.
    • flex_start: define o node pool para usar o modo de aprovisionamento flex-start. Não é possível definir como true se spot estiver ativado.

Consola

Para criar um node pool com TPUs:

  1. Aceda à página do Google Kubernetes Engine na Google Cloud consola.

    Aceda ao Google Kubernetes Engine

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

  3. Clique em Adicionar conjunto de nós.

  4. Na secção Detalhes do conjunto de nós, selecione a caixa Especificar localizações dos nós.

  5. Selecione o nome da zona com base na versão da TPU que quer usar. Para identificar uma localização disponível, consulte o artigo Disponibilidade de TPUs no GKE.

  6. No painel de navegação, clique em Nodes.

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

  8. No menu pendente 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 pendente Tipo de máquina, selecione o nome da máquina a usar para os nós. Use a tabela Escolha a versão do TPU para saber como definir o tipo de máquina e a topologia do TPU que criam um conjunto de nós de fatia do TPU com vários anfitriões.

  10. No menu pendente 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 arranque é um Disco persistente padrão ou um Disco persistente SSD.

  13. Opcionalmente, selecione a caixa de verificação Ativar nós em VMs spot para usar VMs spot para os nós no conjunto de nós.

  14. Clique em Criar.

Valide o estado do conjunto de nós

  1. Obtenha credenciais para poder usar o kubectl para aceder ao cluster:

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

    Substitua o seguinte:

    • CLUSTER_NAME: o nome do cluster.
    • PROJECT_ID: o ID do seu projeto.
    • CONTROL_PLANE_LOCATION: a localização do Compute Engine do plano de controlo do seu cluster. Indique uma região para clusters regionais ou uma zona para clusters zonais.
  2. Use kubectl no Cloud Shell para ver os nós da fatia de TPU:

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

    Substitua o seguinte:

    • TPU_ACCELERATOR: o tipo de acelerador de TPU que usou quando criou os conjuntos de nós. Por exemplo, tpu-v4-podslice ou tpu-v5-lite-podslice.
    • TPU_TOPOLOGY: A topologia física da fatia da TPU.

    O resultado é semelhante ao seguinte:

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

Execute uma carga de trabalho de vários fragmentos

Nesta secção, executa uma carga de trabalho JAX que mostra o número global de chips de TPU na fatia de TPU e, em seguida, termina.

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
    

    Standard

    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 o seguinte:

    • NUM_SLICES: o número de pools de nós de fatias de TPU. Neste caso, o valor de NUM_SLICES é igual a 3.
    • ACCELERATOR_TYPE: o tipo de acelerador de TPU que usou quando criou os conjuntos de nós. Por exemplo, tpu-v4-podslice ou tpu-v5-lite-podslice.
    • TPU_TOPOLOGY: A topologia física da fatia da TPU. Por exemplo, 4x4x4 ou 2x2, consoante a versão da TPU.
    • NUM_NODES: o número de nós no node pool. Tem de ser zero ou o produto dos valores definidos em TPU_TOPOLOGY ({A}x{B}x{C}) dividido pelo número de chips TPU em cada VM. Para a TPU v4 com vários anfitriões, o número de chips de TPU em cada VM é quatro. Para a TPU v5e com vários anfitriões, o número de chips de TPU em cada VM é um, quatro ou oito. Por conseguinte, se o seu TPU_TOPOLOGY for 2x4x4 (TPU v4 com quatro chips de TPU em cada VM), o NUM_NODES é 32/4, o que equivale a 8.
    • NUM_CHIPS: para a TPU v4 multi-anfitrião, o número de chips da TPU em cada VM é quatro. Para a TPU v5e com vários anfitriões, o número de chips de TPU em cada VM é um, quatro ou oito. Para saber mais, consulte o artigo Chips de TPU na VM numa fatia de TPU.

    Neste manifesto:

    • O JobSet é um serviço sem interface com o mesmo nome que o nome do JobSet. Neste caso, é multislice-job.
    • A anotação alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool configura a afinidade do agrupamento para garantir que todos os agrupamentos são agendados na mesma fatia.
    • O campo image está definido para a imagem de IA JAX mais recente. Para definir uma versão diferente, consulte a lista de imagens de IA do JAX atuais.
    • O maxRestarts: 4 indica o número máximo de vezes que o GKE reinicia o JobSet quando uma tarefa secundária falha. Se os reinícios do JobSet atingirem o máximo definido, o JobSet é marcado como falhado.
    • Os campos parallelism e completions são iguais ao número de nós em cada conjunto de nós.
    • O valor de backoff é 0 porque o Multislice só suporta o treino síncrono com vários controladores. Tem de estar definido como 0. Falhar a tarefa quando qualquer pod falhar.
    • Os valores na secção de afinidade garantem que existe apenas uma carga de trabalho de vários fragmentos de TPU em execução num grupo de vários fragmentos.
    • O containerPort: 8080 é a porta do coordenador MXLA
    • O containerPort: 8431 é a porta para exportar as métricas de utilização da TPU
    • O ícone securityContext: privileged: true indica que os nós têm o modo privilegiado ativado para aceder a TPUs. Os nós no GKE versão 1.28 ou posterior não precisam de ter o modo privilegiado ativado para aceder às TPUs. Para saber mais, consulte o artigo Execute contentores sem o modo privilegiado.
  2. Aplique o manifesto:

    kubectl apply -f tpu-multislice.yaml
    
  3. Confirme se a carga de trabalho é admitida:

    kubectl get jobsets
    

    O resultado é semelhante ao seguinte:

    NAME            RESTARTS   COMPLETED   AGE
    multislice-job                         3s
    
  4. Monitorize o estado dos pods aprovisionados:

    kubectl get pods
    

    O resultado é semelhante ao seguinte:

     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 multislice-job JobSet agenda, cria e, em seguida, executa os pods até à conclusão. Os nomes dos agrupamentos estão no formato <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>. O prefixo jobsetName determina o JobSet ao qual o Pod pertence.

  5. Opcional: remova a carga de trabalho JAX:

    kubectl delete -f tpu-multislice.yaml
    

Configure definições adicionais

As secções seguintes descrevem as configurações adicionais que pode aplicar ao Multislice.

Melhore o desempenho da rede com hostNetwork

Para melhorar o desempenho da rede entre fatias de TPU, recomendamos que ative a opção hostNetworking. Use hostNetwork: true na especificação do pod para ignorar toda a pilha de rede do Kubernetes e permitir que os pods do Kubernetes usem a rede do anfitrião diretamente para a comunicação de VM para VM.

Para ativar hostNetworking, adicione as duas linhas seguintes à especificação do pod:

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

Para continuar a usar o podHostnames para a descoberta de nós trabalhadores com o hostNetwork, defina dnsPolicy: ClusterFirstWithHostNet. Isto é importante quando está a executar tarefas de preparação com retoma automática e precisa de ter os mesmos nomes para recarregar os mesmos pontos de verificação.

Se usar a TPU Trillium (v6e) e os seus pods usarem hostNetworking, instale o seguinte DaemonSet para otimizar /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

Melhore o desempenho da rede sem hostNetwork na TPU Trillium

Se usar a TPU Trillium e os seus pods não puderem usar hostNetworking, ative a multirrede com o modo netdevice para o melhor desempenho da rede. O suporte de NICs no modo netdevice com várias redes transfere a NIC da VM diretamente para o pod, ignorando o Kubernetes e o GKE Dataplane V2.

O tipo de máquina ct6e-standard-4t é suportado por duas NICs físicas. O Kubernetes requer uma vNIC que não pode ser transmitida para os pods. Por isso, cada nó tem de ter três vNICs para permitir que os pods tenham acesso direto a duas vNICs para alcançar o melhor desempenho de ambas as NICs físicas.

Para ativar o modo netdevice para o ct6e-standard-4t, conclua os seguintes passos:

  1. Crie duas VPCs adicionais que suportem o modo netdevice
  2. Crie um cluster do GKE com capacidades de várias redes
  3. Configure duas netdeviceredes. Por exemplo, pode usar os seguintes objetos GKENetworkParamSet e Network (SECOND_VPC e THIRD_VPC são as VPCs criadas no passo 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. Ligue os Pods a três redes. Por exemplo, pode 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 interior do pod, no contentor de inicialização ou no contentor de aplicações. Por exemplo, pode adicionar o seguinte contentor 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 um melhor desempenho de rede, em vez da interface eth0. Pode fazê-lo adicionando export LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS --megascale_grpc_interface_prefixes=eth1,eth2,lo" à especificação da carga de trabalho.

Ative o registo

Os registos emitidos por contentores em execução em nós do GKE, incluindo nós de fatias de TPU, são visíveis no Explorador de registos, se tiver o registo do sistema do GKE ativado no seu cluster.

Pode ver os registos do GKE através do Explorador de registos com o seguinte filtro para ver os registos de contentores da sua 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 fatias e trabalhadores de TPU:

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 o artigo Veja os registos do GKE TPU.

Ative métricas adicionais

Além das métricas de TPU gerais, existem 4 métricas de tempo de execução de TPU específicas de várias fatias adicionais. Estas métricas estão disponíveis na versão 1.29.1-gke.1016000 ou posterior do GKE. A carga de trabalho da TPU tem de usar a versão 0.4.24 do JAX

Seguem-se as métricas de vários segmentos disponíveis:

  • Latências de transferência da DCN (rede do centro de dados): distribuição das latências de transferência de rede para tráfego de várias divisões.
  • Latências coletivas: distribuição da latência coletiva ponto a ponto para tráfego de vários segmentos.
  • Latências de transferência do anfitrião para o dispositivo: distribuição da latência de transferência do anfitrião para o dispositivo para cada parte de dados para tráfego de vários segmentos.
  • Latências de transferência do dispositivo para o anfitrião: distribuição da latência de transferência do dispositivo para o anfitrião para cada bloco de dados para tráfego de vários segmentos.

Estas métricas encontram-se no esquema do contentor 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

Segmento de TPU versus Multislice

A tabela seguinte diferencia a organização arquitetónica de uma fatia de TPU e de uma multifatia:

Fatia de TPU Multislice
Interligação A carga de trabalho é executada numa única fatia de TPU. Todos os chips da TPU numa fatia estão ligados com a ICI. A carga de trabalho é executada em várias fatias de TPUs. A comunicação numa fatia ocorre através da ICI. A comunicação entre fatias ocorre através da DCN.
Node pools suportados Segmento de TPU de anfitrião único e segmento de TPU de vários anfitriões Grupos de fatias de TPU com vários anfitriões
Tipo de carga de trabalho recomendado IndexedJob ou JobSet JobSet

Limpe os recursos

A forma mais fácil de eliminar a faturação é eliminar o Google Cloud projeto que criou para o tutorial. Em alternativa, pode eliminar os recursos individuais.

Elimine o projeto

    Delete a Google Cloud project:

    gcloud projects delete PROJECT_ID

Elimine recursos individuais

Elimine o cluster do GKE:

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

O que se segue?