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.
Antes de configurar Multislice no GKE, é preciso conhecer os seguintes conceitos:
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
ct4p-hightpu-4t
com uma topologia2x2x1
ou umct5lp-hightpu-4t
com2x2
, 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
.
- Crie um cluster padrão ou de piloto automático que execute uma versão com suporte ao Multislice. Para versões com suporte, consulte Disponibilidade de multislices no GKE.
- Verifique se o projeto tem cota suficiente para o Cloud TPU no GKE.
- Instale o JobSet v0.2.3 ou mais recente.
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:
- Criar três pools de nós de frações de TPU com vários hosts.
- 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 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-
, useus-west1-c
,us-west4-a
,us-west4-b
,us-central1-a
,us-east1-c
,us-east5-b
, oueurope-west4-a
. - Para os tipos de máquina TPU v5p que começam com
ct5p-
, useus-east1-d
,us-east5-a
ouus-east5-c
.
Para saber mais, consulte a disponibilidade da TPU no GKE.
- Para a TPU v4, use
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.
- TPU v4 ou v5p: defina a topologia em três tuplas (
NUM_NODES
: o número de nós no pool de nós. Precisa ser zero ou o produto dos valores definidos emTPU_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 aTPU_TOPOLOGY
for2x4x4
(TPU v4 com quatro chips em cada VM), aNUM_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 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.--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 sinalização--max-nodes
será obrigatória se--enable-autoscaling
for fornecido e precisar ser igual ao produto dos valores definidos emTPU_TOPOLOGY
({A}x{B}x{C}
) dividido pelo número de ícones para cada VM.
Terraform
- Use a versão 4.84.0 ou mais recente do
provedor
google
. 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 de fração tem quatro chips. Por exemplo, seTPU_TOPOLOGY
for4x8
, haverá 32 chips, o que significa queNUM_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
.
- Para a TPU v4: defina a topologia em três tuplas (
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 oRESERVATION_LABEL_VALUES
no camporeservation_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 emTPU_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.
Console
Para criar um pool de nós com TPUs:
Acesse a página Google Kubernetes Engine no console do Google Cloud.
Na lista de clusters, clique no nome do cluster que você quer modificar.
Clique em add_box Adicionar pool de nós.
Na seção Detalhes do pool de nós, marque a caixa Especificar locais do nó.
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-
, useus-west1-c
,us-west4-a
,us-west4-b
,us-central1-a
,us-east1-c
,us-east5-b
, oueurope-west4-a
. - Para os tipos de máquina TPU v5p que começam com
ct5p-
, useus-east1-d
,us-east5-a
ouus-east5-c
.
- Para a TPU v4, use
No painel de navegação, clique em Nós.
Na seção Configuração da máquina, selecione TPUs.
No menu suspenso Série, selecione uma das seguintes opções:
- CT4P: para a TPU v4.
- CT5LP: para a TPU v5e.
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 fração de TPU de vários hosts.
No menu suspenso Topologia da TPU, selecione a topologia física para a fatia da TPU.
Na caixa de diálogo Alterações necessárias, clique em Fazer alterações.
Certifique-se de que o Tipo de disco de inicialização seja Disco permanente padrão ou Disco permanente SSD.
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.
Clique em Criar.
Verificar o status do pool de nós
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.
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=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
outpu-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:
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: 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 fração de TPU. Nesse caso, oNUM_SLICES
é igual a3
.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
outpu-v5-lite-podslice
.TPU_TOPOLOGY
: a topologia física da fração de TPU. Por exemplo,4x4x4
ou2x2
, 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 emTPU_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 aTPU_TOPOLOGY
for2x4x4
(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
. - 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
ecompletions
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 MXLAcontainerPort: 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.
Aplique o manifesto:
kubectl apply -f tpu-multislice.yaml
Confirme que a carga de trabalho foi permitida:
kubectl get jobsets
O resultado será assim:
NAME RESTARTS COMPLETED AGE multislice-job 3s
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 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.
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 |