Visão geral de multislices do Cloud TPU

O Cloud TPU Multislice é uma tecnologia de escalonamento de desempenho de pilha completa que permite que um job de treinamento use várias frações de TPU em um único pod ou em frações em vários pods com paralelismo de dados simples. Com chips TPU v4, isso significa que os jobs de treinamento podem usar mais de 4.096 chips em uma única execução. Para jobs de treinamento que exigem menos de 4.096 chips, uma única fração pode oferecer o melhor desempenho. No entanto, várias frações menores ficam mais prontamente disponíveis, permitindo um tempo de inicialização mais rápido quando o Multislice é usado com frações menores.

Quando implantados em configurações de multislice, os chips de TPU em cada fração se comunicam por meio de interconexão entre chips (ICI, na sigla em inglês). Os chips de TPU em diferentes frações se comunicam transferindo dados para CPUs (hosts), que, por sua vez, transmitem os dados pela rede de data center (DCN).

Os desenvolvedores não precisam escrever código para implementar a comunicação DCN entre fatias. O compilador XLA gera esse código e sobrepõe a comunicação com a computação para garantir o desempenho máximo.

conceitos

Tipo de acelerador

O formato de cada fração de TPU que compreende um Multislice. Cada fração em uma solicitação de várias fatias é do mesmo tipo de acelerador. Um tipo de acelerador consiste em um tipo de TPU (v4 ou v5e) seguido pelo número de TensorCores. Por exemplo, v4-128 especifica uma TPU v4 com 128 TensorCores.

Reparo automático

Quando uma fração encontra um evento de manutenção, preempção ou falha de hardware, o Cloud TPU cria uma nova fração. No caso raro em que há recursos insuficientes para criar uma nova fração, a criação não será concluída até que o hardware esteja disponível. Depois que a nova fração for criada, todas as outras frações no ambiente Multislice serão reiniciadas para que o treinamento possa continuar.Com um script de inicialização configurado corretamente, o script pode ser reiniciado automaticamente sem intervenção do usuário, carregar e retomar a partir do checkpoint mais recente.

Conjunto de dados

Os dados usados por um modelo para treinamento ou inferência.

Rede de data center (DCN)

Uma rede de menor latência e capacidade de processamento (em comparação com a ICI) que conecta frações de TPU em uma configuração de multislices.

Programação de gangues

Quando todas as frações de TPU são provisionadas juntas, ao mesmo tempo, garantindo que todas ou nenhuma das frações sejam provisionadas com sucesso.

Host

Um host é um computador físico que executa VMs. Um host pode executar no máximo quatro VMs ao mesmo tempo. Cada VM tem uma TPU dedicada.

Inferência

Carregar um modelo de machine learning pré-treinado em um host e fazer previsões nos dados.

Interconexão de interconexão (ICI)

Links internos de alta velocidade e baixa latência que conectam TPUs dentro de um pod de TPU.

Multislice (em inglês)

Duas ou mais frações de chip de TPU que podem se comunicar por DCN.

Nó

No contexto de Multislice, o nó se refere a uma única fração de TPU. Cada fração de TPU em um Multislice recebe um ID de nó.

Pod

Uma coleção de chips de TPU conectados por interfaces de rede ICI dedicadas. Um pod permite distribuir a carga de processamento entre várias TPUs.

Recurso em fila (QR)

Uma representação de recursos da TPU, usada para enfileirar e gerenciar uma solicitação para um ambiente de TPU de fração única ou multislice.

Script de inicialização

Um script de inicialização padrão do Compute Engine executado sempre que uma VM é inicializada ou reinicializada. Para o Multislice, ele é especificado na solicitação de criação de QR code. Para mais informações sobre os scripts de inicialização do Cloud TPU, consulte Gerenciar recursos da TPU.

Fração da TPU

Uma subseção lógica de um pod de TPU, que consiste em chips de TPU. Todos os ícones em uma fatia se comunicam usando a rede ICI.

VM da TPU

Uma máquina virtual que executa Linux e tem acesso às TPUs subjacentes. Para TPUs v4, cada VM da TPU tem acesso direto a quatro chips. Às vezes, chamamos uma VM da TPU de worker.

Tensor (link em inglês)

Uma estrutura de dados usada para representar dados multidimensionais em um modelo de machine learning.

Unidade de Processamento de Tensor (TPU)

Chip de aceleração de ML desenvolvido internamente pelo Google. Elas foram projetadas para oferecer computação rápida e eficiente em termos de energia para as principais tarefas de machine learning, como multiplicação de matrizes.

Tipos de capacidade do Cloud TPU

TPUs podem ser criadas a partir de diferentes tipos de capacidade. Consulte "Opções de uso" em Como funcionam os preços de TPU:

• Reserva: segmenta a cota reservada. Para usar a cota reservada, você precisa ter um contrato de reserva com o Google. Use a sinalização --reserved ao criar os recursos.
• Spot: segmenta a cota preemptiva usando VMs spot. Seus recursos podem ser interrompidos para dar espaço a solicitações de jobs de maior prioridade. Use a sinalização --spot ao criar os recursos.
• Sob demanda: segmenta a cota sob demanda, que não precisa de reserva e não será interrompida. A solicitação de TPU será enfileirada para uma fila de cota sob demanda oferecida pelo Cloud TPU. A disponibilidade dos recursos não é garantida. Selecionado por padrão, nenhuma sinalização necessária.

Observação: o Cloud TPU limita o tamanho de fração para todos os tipos de acelerador. Para saber os limites de tamanho de fração para solicitações sob demanda, consulte Tamanhos de fração padrão para QRs sob demanda.

Começar

Se você nunca usou TPUs antes, comece instalando a Google Cloud CLI e configure seu ambiente do Cloud TPU. Para usar o multislice, seus recursos de TPU precisam ser gerenciados como recursos na fila.

Se você for um usuário existente da TPU v4 e tiver uma reserva, talvez seja necessário migrar sua reserva para um novo sistema de reservas. Para mais informações, entre em contato com seu representante de conta do Google Cloud.

Exemplo introdutório

Neste tutorial, usamos o código do repositório MaxText no GitHub (link em inglês). O MaxText é um LLM básico bem testado, de alto desempenho, escalonável arbitrariamente e de código aberto, escrito em Python e Jax. O MaxText foi projetado para treinar de maneira eficiente no Cloud TPU.

O código em shardings.py foi projetado para ajudar você a começar a testar diferentes opções de paralelismo. Por exemplo, paralelismo de dados, paralelismo de dados totalmente fragmentado (FSDP, na sigla em inglês) e paralelismo de tensor. O código é escalonado de uma única fração para ambientes multislice.

paralelismo da ICI

ICI refere-se à interconexão de alta velocidade que conecta as TPUs em uma única fração. A fragmentação de ICI corresponde à fragmentação dentro de uma fração. shardings.py fornece três parâmetros de paralelismo do ICI:

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

Os valores especificados para esses parâmetros determinam o número de fragmentos para cada método de paralelismo.

Essas entradas precisam ser restritas para que ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism seja igual ao número de ícones na fatia.

A tabela a seguir mostra exemplos de entradas do usuário para paralelismo ICI para os quatro chips disponíveis na v4-8:

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
FSDP de quatro vias	1	4	1
Paralelismo do tensor de quatro direções	1	1	4
FSDP bidirecional + paralelismo do tensor bidirecional	1	2	2

Observe que ici_data_parallelism precisa ser deixado como 1 na maioria dos casos porque a rede ICI é rápida o suficiente para quase sempre preferir o FSDP ao paralelismo de dados.

Neste exemplo, presumimos que você esteja familiarizado com a execução de código em uma única fração de TPU, como em Executar um cálculo em uma VM do Cloud TPU usando JAX. Este exemplo mostra como executar shardings.py em uma única fração.

1. Configure o ambiente:

   $ gcloud auth login
   $ gcloud config set project your-project-id
   $ gcloud config set compute/zone your-zone
   

2. Crie chaves SSH para gcloud. Recomendamos deixar uma senha em branco (pressione Enter duas vezes depois de executar o comando a seguir). Se você for solicitado que o arquivo google_compute_engine já existe, substitua a versão atual.

   $ ssh-keygen -f ~/.ssh/google_compute_engine
   

3. Provisione suas TPUs com o seguinte comando:

   $ gcloud compute tpus queued-resources \
   create your-qr-id \
   --accelerator-type your-accelerator-type \
   --runtime-version tpu-ubuntu2204-base \
   --node-id qr-id \
   [--reserved |--spot]
   

   Descrições de sinalizações de comando

   your-qr-id

   Uma string definida pelo usuário que identifica a solicitação de QR code.

   accelerator-type

   O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.

   runtime-version

   A [versão do software do Cloud TPU](/tpu/docs/supported-tpu-configurations#tpu_software_versions).

   node-id

   O ID dos recursos de TPU que serão criados em resposta à solicitação de QR.

   reserved

   Use a cota reservada ao criar as frações.

   spot

   Usar a cota de VMs spot ao criar as frações.

   A CLI do Google Cloud não é compatível com todas as opções de criação de QR code, como tags. Para mais informações, consulte Criar QRs.

4. Aguarde até que o QR esteja no estado ACTIVE, o que significa que os nós de trabalho estão no estado READY. Depois que o provisionamento de QR for iniciado, pode levar de um a cinco minutos para ser concluído, dependendo do tamanho do QR. É possível verificar o status de uma solicitação de QR code usando o seguinte comando:

   $ gcloud compute tpus queued-resources \
     list --filter=your-qr-id
   

5. Uma fração v4-8 tem uma única VM de TPU. Conecte-se à VM da TPU usando SSH:

   $ gcloud compute tpus tpu-vm ssh your-qr-id
   

6. Clone o MaxText (que inclui shardings.py) na VM da TPU.

7. No diretório do repositório MaxText, execute o script de configuração para instalar o JAX e outras dependências na fração de TPU. O script de configuração leva alguns minutos para ser executado.

   $ bash setup.sh
   

8. Execute o comando a seguir para executar shardings.py na sua fração de TPU.

   $ python3 pedagogical_examples/shardings.py \
     --ici_fsdp_parallelism 4 \
     --batch_size 131072 \
     --embedding_dimension 2048
   

   Confira os resultados nos registros. Suas TPUs devem alcançar cerca de 260 TFLOP por segundo ou uma utilização impressionante de mais de 90%de FLOP. Nesse caso, selecionamos aproximadamente o lote máximo que se encaixa na memória de alta largura de banda (HBM, na sigla em inglês) da TPU.

9. Fique à vontade para explorar outras estratégias de fragmentação em relação à ICI. Por exemplo, você pode tentar a seguinte combinação:

   $ python3 pedagogical_examples/shardings.py \
     --ici_tensor_parallelism 4 \
     --batch_size 131072 \
     --embedding_dimension 2048
   

10. Exclua a fração de QR e TPU quando terminar. Execute essas etapas de limpeza no ambiente em que a fração foi configurada. Primeiro, execute exit para sair da sessão SSH. A exclusão levará de dois a cinco minutos para ser concluída e pode ser executada em segundo plano com a sinalização --async opcional.

    Importante: os QRs precisam ser excluídos após o uso. Se você não fizer isso, eles continuarão consumindo sua cota de QR code, que é rastreada separadamente da cota de TPU.

    $ gcloud compute tpus queued-resources
      delete your-qr-id --force (--async)
    

Fragmentação de várias fatias usando paralelismo DCN

O script shardings.py usa três parâmetros que especificam o paralelismo DCN, correspondente ao número de fragmentos de cada tipo de paralelismo de dados:

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

Os valores desses parâmetros precisam ser restritos para que dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism seja igual ao número de frações.

Como exemplo para duas frações, use --dcn_data_parallelism = 2.

	dcn_data_parallelism	dcn_fsdp_parallelism	dcn_tensor_parallelism	No de fatias
Paralelismo de dados bidirecional	2	1	1	2

dcn_tensor_parallelism precisa ser sempre definido como 1, porque a DCN não é adequada para essa fragmentação. Para cargas de trabalho LLM típicas em chips v4, dcn_fsdp_parallelism também precisa ser definido como 1 e, portanto, dcn_data_parallelism precisa ser definido como o número de frações, mas isso depende do aplicativo.

À medida que você aumenta o número de frações (supondo que o tamanho da fração e o lote por fração sejam constantes), aumenta a quantidade de paralelismo de dados.

Como executar shardings.py em um ambiente multislice

É possível executar shardings.py em um ambiente multislice usando multihost_runner.py ou executando shardings.py em cada VM da TPU. Aqui, usamos multihost_runner.py. As etapas a seguir são muito semelhantes às Introdução: experimentos rápidos em várias fatias do repositório MaxText, mas aqui executamos shardings.py em vez do LLM mais complexo em train.py.

A ferramenta multihost_runner.py é otimizada para experimentos rápidos, reutilizando repetidamente as mesmas TPUs. Como o script multihost_runner.py depende de conexões SSH de longa duração, não o recomendamos para jobs de longa duração. Se você quiser executar um job mais longo (por exemplo, horas ou dias), recomendamos usar multihost_job.py.

Neste tutorial, usamos o termo runner para indicar a máquina em que você executa o script multihost_runner.py. Usamos o termo workers para indicar as VMs de TPU que compõem as frações. É possível executar multihost_runner.py em uma máquina local ou em qualquer VM do Compute Engine no mesmo projeto das frações. Não é possível executar multihost_runner.py em um worker.

multihost_runner.py se conecta automaticamente a workers da TPU usando SSH.

Neste exemplo, executamos shardings.py em duas frações v4-16, um total de quatro VMs e 16 chips de TPU. É possível modificar o exemplo para execução em mais TPUs.

Configure seu ambiente

1. Clone MaxText na máquina do executor.

2. Acesse o diretório do repositório.

3. Crie chaves SSH para gcloud. Recomendamos deixar uma senha em branco (pressione Enter duas vezes depois de executar o comando a seguir). Se você receber a solicitação de que o arquivo google_compute_engine já existe, selecione para não manter a versão atual.

     $ ssh-keygen -f ~/.ssh/google_compute_engine
     

4. Adicione uma variável de ambiente para definir a contagem de frações de TPU como 2.

     $ export SLICE_COUNT=2
     

5. Crie um ambiente multislice usando queued-resources create.

   Observação: verifique se você tem a respectiva cota antes de selecionar --reserved, --spot ou a cota sob demanda padrão. Para informações sobre os tipos de cotas, consulte a Política de cotas.

   O comando a seguir mostra como criar uma TPU Multislice v4. Para usar a v5e, especifique um accelerator-type (por exemplo, v5litepod-16) e o runtime-version (v2-alpha-tpuv5-lite) da v5e.

     $ gcloud compute tpus queued-resources 
   
     create your-qr-id 
   
     --accelerator-type=your-accelerator-type 
   
     --runtime-version=tpu-vm-runtime-version 
   
     --node-count=node-count 
   
     --node-prefix=your-qr-id 
   
     [--reserved|--spot]
     

   Descrições de sinalizações de comando

   your-qr-id

   Uma string definida pelo usuário que identifica a solicitação de QR code.

   accelerator-type

   O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.

   runtime-version

   A versão do software do Cloud TPU.

   node-count

   O número de frações a serem criadas.

   node-prefix

   O prefixo usado para gerar nomes para cada fatia. Um número é anexado ao prefixo de cada fatia. Por exemplo, se você definir node-prefix como mySlice, as fatias serão nomeadas: mySlice-0, mySlice-1, continuando numericamente para cada uma delas.

   reserved

   Use a cota reservada ao criar as frações.

   spot

   Usar a cota de VMs spot ao criar as frações.

6. Quando o provisionamento de QR code for iniciado, poderá levar até cinco minutos para ser concluído, dependendo do tamanho do QR. Aguarde até que o recurso na fila (QR, na sigla em inglês) esteja no estado ACTIVE. Para verificar o status de uma solicitação de QR code, use o seguinte comando:

   $ gcloud compute tpus queued-resources list \
   --filter=your-qr-id
   

   Isso vai gerar uma saída semelhante a esta:

   NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
   ...
   que-res-id  us-central2-b  4           v4-16             ACTIVE
   ...
   

   Entre em contato com o representante da sua conta do Google Cloud se o status do QR code ficar no estado WAITING_FOR_RESOURCES ou PROVISIONING por mais de 15 minutos.

7. Instale as dependências.

   $ python3 multihost_runner.py \
     --TPU_PREFIX=your-qr-id \
     --COMMAND="bash setup.sh"
   

8. Execute shardings.py em cada worker usando multihost_runner.py.

   $ python3 multihost_runner.py \
     --TPU_PREFIX=your-qr-id \
     --COMMAND="python3 pedagogical_examples/shardings.py \
     --dcn_data_parallelism $SLICE_COUNT \
     --ici_fsdp_parallelism 8 \
     --batch_size 131072 \
     --embedding_dimension 2048"
   

   Você verá aproximadamente 230 TFLOPs por segundo de desempenho nos arquivos de registro.

9. Limpe as TPUs e o QR code quando terminar. A exclusão levará de dois a cinco minutos para ser concluída e pode ser executada em segundo plano com a sinalização --async opcional.

Como escalonar uma carga de trabalho para multislices

Antes de executar o modelo em um ambiente multislice, faça as mudanças de código a seguir:

• Use jax.experimental.mesh_utils.create_hybrid_device_mesh em vez de jax.experimental.mesh_utils.create_device_mesh ao criar a malha.

Essas são as únicas mudanças de código necessárias ao migrar para o multislice. Para alcançar alto desempenho, a DCN precisa ser mapeada para dados paralelos, dados totalmente fragmentados ou eixos paralelos de pipeline. Considerações de desempenho e estratégias de fragmentação são discutidas mais detalhadamente em Fragmentação com multislices para desempenho máximo.

Para validar que seu código pode acessar todos os dispositivos, declare que len(jax.devices()) é igual ao número de ícones no ambiente Multislice. Por exemplo, se você estiver usando quatro frações de v4-16, terá oito ícones por fatia * 4 fatias, portanto, len(jax.devices()) retornará 32.

Como escolher tamanhos de fatia para ambientes multislice

Para aumentar a velocidade linear, adicione novas fatias do mesmo tamanho da fatia atual. Por exemplo, se você usar uma fração v4-512, o Multislice terá aproximadamente o dobro do desempenho ao adicionar uma segunda fração de v4-512 e dobrar o tamanho do lote global. Para mais informações, consulte Fragmentação com multislices para desempenho máximo.

Como executar o job em várias frações

Há três abordagens diferentes para executar a carga de trabalho personalizada em um ambiente Multislice:

1. Usando o script do executor de experimentos, multihost_runner.py
2. Usando o script do executor de produção, multihost_job.py
3. Usar uma abordagem manual

Script do executor do experimento

O script multihost_runner.py distribui o código para um ambiente Multislice atual e executa o comando em cada host, copia os registros e rastreia o status de erro de cada comando. O script multihost_runner.py está documentado no README de MaxText.

Como multihost_runner.py mantém conexões SSH persistentes, ele é adequado apenas para experimentos de tamanho modesto e relativamente curta. É possível adaptar as etapas do tutorial multihost_runner.py à sua configuração de carga de trabalho e hardware.

Script do executor Production

Para jobs de produção que precisam de resiliência contra falhas de hardware e outras preempções, é melhor integrar-se diretamente à API Create Queued Resource. Como exemplo funcional, fornecemos multihost_job.py, que aciona a chamada da API Created Queued Resource com o script de inicialização apropriado para executar o treinamento e retomar a preempção. O script multihost_job.py está documentado no README de MaxText.

Como multihost_job.py precisa provisionar recursos para cada execução, ele não fornece um ciclo de iteração tão rápido quanto multihost_runner.py.

Abordagem manual

Recomendamos que você use ou adapte multihost_runner.py ou multihost_job.py para executar a carga de trabalho personalizada na configuração do Multislice. No entanto, se você preferir provisionar e gerenciar seu ambiente usando comandos QR diretamente, consulte Gerenciar um ambiente de multislices.

Gerenciar um ambiente multislice

Para provisionar e gerenciar QRs manualmente sem usar as ferramentas fornecidas no repositório MaxText, leia as seções a seguir.

Criar QR codes

Defina as variáveis de ambiente abaixo antes de provisionar a capacidade:

  $ export your-qr-id=your-queued-resource-id
  $ export PROJECT=your-project-name
  $ export ZONE=us-central2-b
  $ export NETWORK_NAME=your-network-name
  $ export SUBNETWORK_NAME=your-subnetwork-name
  $ export RUNTIME_VERSION=tpu-ubuntu2204-base
  $ export ACCELERATOR_TYPE=v4-16
  $ export SLICE_COUNT=4
  $ export STARTUP_SCRIPT="#!/bin/bash\n ..."
  $ gcloud config set project project-name
  $ gcloud config set compute/zone zone

Salvaguardas	Descrição
your-qr-id	O ID atribuído pelo usuário do QR.
PROJETO	Nome do projeto do Google Cloud
ZONA	us-central2-b
NETWORK_NAME	Nome das redes VPC.
SUBNETWORK_NAME	Nome da sub-rede em redes VPC
RUNTIME_VERSION	tpu-ubuntu2204-base
ACCELERATOR_TYPE	v4-16
EXAMPLE_TAG_1, EXAMPLE_TAG_2...	Tags usadas para identificar origens ou destinos válidos em firewalls de rede
SLICE_COUNT	Número de frações. Limitado a um máximo de 256 frações.
STARTUP_SCRIPT	Se adicionado à solicitação de criação, um script de inicialização poderá ser executado sempre que uma fração de TPU for provisionada ou reiniciada e se a fração de TPU for reparada ou redefinida.

Criar uma solicitação de QR code usando gcloud

$ gcloud compute tpus queued-resources \
  create ${your-qr-id} \
  --project your-project-id \
  --zone your-zone \
  --node-count ${SLICE_COUNT} \
  --accelerator-type ${ACCELERATOR_TYPE} \
  --runtime-version ${RUNTIME_VERSION} \
  --network ${NETWORK_NAME} \
  --subnetwork ${SUBNETWORK_NAME} \
  --tags ${EXAMPLE_TAG_1},${EXAMPLE_TAG_2} \ --metadata=startup-script='${STARTUP_SCRIPT}'
  [--reserved|--spot]
  

Descrições de sinalizações de comando

your-qr-id

Uma string definida pelo usuário que identifica a solicitação de QR code.

project

Uma string definida pelo usuário que identifica a solicitação de QR code.

zone

A zona do Google Cloud em que o QR code será criado.

node-count

O número de frações a serem criadas.

accelerator-type

O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.

runtime-version

A versão do software do Cloud TPU.

network

O nome de uma rede VPC a que o recurso de TPU será anexado.

subnetwork

O nome de uma sub-rede VPC a que o recurso de TPU será anexado.

reserved

Use a cota reservada ao criar as frações.

spot

Usar a cota de VMs spot ao criar as frações.

Verifique se você tem a respectiva cota antes de selecionar --reserved, --spot ou a cota padrão sob demanda. Para informações sobre os tipos de cotas, consulte a Política de cotas.

Criar uma solicitação de QR code usando curl

Crie um arquivo chamado queued-resource-req.json e copie o JSON a seguir nele.

{
  "guaranteed": { "reserved": true },
  "tpu": {
    "node_spec": [
    {
      "parent": "projects/your-project-number/locations/your-zone",
        "node": {
          "accelerator_type": "accelerator-type",
          "runtime_version": "tpu-vm-runtime-version",
          "network_config": {
            "network": "your-network-name",
            "subnetwork": "your-subnetwork-name",
            "enable_external_ips": true
          },
          "tags" : ["example-tag-1"]
          "metadata": {
            "startup-script": "your-startup-script"
          }
      },
      "multi_node_params": {
        "node_count": slice-count,
        "node_id_prefix": "your-queued-resource-id"
      }
    }
    ]
  }
}

Observação: substitua os seguintes valores:

• your-project-number: o número do projeto do Google Cloud
• your-zone: a zona em que você quer criar o QR code
• accelerator-type: a versão e o tamanho de uma única fração
• tpu-vm-runtime-version: as versões do ambiente de execução da VM da TPU
• your-network-name: (opcional) uma rede à qual o QR code será anexado
• your-subnetwork-name: (opcional) uma sub-rede à qual o QR code será anexado
• example-tag-1: opcional, uma string de tag arbitrária
• your-startup-script: um script de inicialização que será executado quando o QR code for alocado
• slice-count: o número de frações de TPU no ambiente Multislice
• your-qr-id: o ID fornecido pelo usuário para o QR code

Para mais informações, consulte a documentação da API REST Queued Resource para ver todas as opções disponíveis.

Observação: embora node_spec seja uma lista, apenas uma entrada node_spec é aceita.

Para usar a capacidade do Spot, substitua:

"guaranteed": { "reserved": true } com "spot": {}

Remova a linha para usar a capacidade padrão sob demanda.

Envie a solicitação de criação de QR code com o payload JSON:

  $ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d @queuedresourcereq.json https://tpu.googleapis.com/v2alpha1/projects/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=your-qr-id

Observação: substitua os seguintes valores:

• your-project-id: ID do projeto do Google Cloud
• your-zone: a zona em que você quer criar o QR code
• your-qr-id: o ID fornecido pelo usuário para o QR code

A resposta será semelhante a esta:

{
  "name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qr-guid>",
  "metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-11-01T00:17:05.742546311Z",
    "target": "projects/<your-project-id>/locations/<your-zone>/queuedResources/<your-qa-id>",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
  },
  "done": false
}

Use o valor de GUID no final do valor da string do atributo name para acessar informações sobre a solicitação de QR code.

Recuperar o status de um QR code

Para ver o status da solicitação de QR code, use o seguinte comando:

  $ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/operations/operation-your-qr-guid

Observação: no comando anterior, substitua o seguinte:

• your-project-id: ID do projeto do Google Cloud
• your-zone: a zona em que o QR code será criado
• your-qr-guid: o GUID após name na saída da solicitação de criação de QR code.

A resposta desse comando contém o status da operação:

{
  "name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qa-guid>,
  "metadata": {...},
  "done": true,
  "response": {
    "@type": "type.googleapis.com/google.cloud.tpu.v2.QueuedResource",
    ...
    "state": {
      "state": "WAITING_FOR_RESOURCES"
    }
  }
}

Se o QR code for criado corretamente como ("done = true"), o estado no campo response será WAITING_FOR_RESOURCES ou FAILED. Se o QR estiver no estado WAITING_FOR_RESOURCES, significa que ele foi enfileirado e iniciará o provisionamento quando houver recursos suficientes. Se o QR estiver no estado FAILED, o motivo da falha estará na saída. Para mais informações sobre outros estados possíveis, consulte o Guia do usuário de recursos em fila.

Quando a operação for concluída, use describe QRs para monitorar os estágios do QR.

Em um cenário raro, você pode encontrar o QR no estado FAILED enquanto algumas frações são ACTIVE. Se isso acontecer, exclua os recursos criados e tente novamente em alguns minutos ou entre em contato com a equipe do Cloud TPU para resolver o problema.

SSH e instalar dependências

Executar o código JAX em frações do Pod de TPU TPU descreve como se conectar às VMs da TPU usando SSH em uma única fração. Para se conectar a todas as VMs da TPU no ambiente Multislice por SSH e instalar dependências, use o seguinte comando gcloud:

  $ gcloud compute tpus queued-resources ssh ${your-qr-id} \
    --zone your-zone \
    --node=all \
    --worker=all \
    --command="command-to-run"
    --batch-size=4

O comando gcloud envia o comando especificado para todos os workers e nós no QR usando SSH. O comando é reunido em grupos de quatro e enviado simultaneamente. O próximo lote de comandos é enviado quando o lote atual conclui a execução. Em caso de falha em um dos comandos, o processamento é interrompido e nenhum outro lote é enviado. Para mais informações, consulte a referência da API de recursos em fila. Se o número de frações que você usa exceder o limite de linhas de execução do computador local (também chamado de limite de lotes), ocorrerá um impasse. Por exemplo, suponha que o limite de lotes na máquina local seja 64. Se você tentar executar um script de treinamento em mais de 64 frações, digamos 100, o comando SSH dividirá as fatias em lotes. Ele executará o script de treinamento no primeiro lote de 64 frações e aguardará a conclusão dos scripts antes de executá-lo no lote restante de 36 frações. No entanto, o primeiro lote de 64 frações não será concluído até que as 36 frações restantes comecem a executar o script, causando um impasse.

Para evitar esse cenário, execute o script de treinamento em segundo plano em cada VM anexando um E comercial (&) ao comando de script que você especificar com a sinalização --command. Ao fazer isso, depois de iniciar o script de treinamento no primeiro lote de frações, o controle retornará imediatamente ao comando SSH. O comando SSH pode começar a executar o script de treinamento no lote restante de 36 frações. Você precisa canalizar seus streams stdout e stderr adequadamente ao executar os comandos em segundo plano. Para aumentar o paralelismo no mesmo QR, é possível selecionar frações específicas usando o parâmetro --node.

Configuração da rede

Execute as etapas a seguir para verificar se as frações de TPU podem se comunicar entre si. Instale o JAX em cada uma das frações. Para mais informações, consulte Executar o código JAX em frações do Pod de TPU. Declare que len(jax.devices()) é igual ao número de ícones no ambiente Multislice. Para fazer isso, em cada fatia, execute:

  $ python3 -c 'import jax; print(jax.devices())'

Se você executar esse código em quatro fatias de v4-16, haverá oito ícones por fração e quatro frações, um total de 32 ícones (dispositivos) deverá ser retornado por jax.devices().

Listar QRs

É possível ver o estado dos seus QRs usando o comando queued-resources list:

$ gcloud compute tpus queued-resources list

NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central2-b  4           v4-16             ACTIVE
...

Descrever QRs

Para ver a configuração e o estado detalhados de um QR, use a API describe QR. É possível chamar essa API usando gcloud ou curl.

Utilizando gcloud:

$ gcloud compute tpus queued-resources describe ${your-qr-id}
...state:
 state: ACTIVE
...

Utilizando curl:

$ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/queuedResources/${your-qr-id}
{
  "name": your-queued-res,
  "tpu": {
    "nodeSpec": [
      {
        ... // node 1
      },
      {
        ... // node 2
      },
      ...
    ]
  },
  ...
  "state": "ACTIVE"
}

state representa o status de um QR code. Para mais informações sobre os possíveis estados de QRs, consulte Recursos na fila.

Iniciar o job em um ambiente provisionado

É possível executar cargas de trabalho manualmente conectando-se a todos os hosts em cada fração por SSH e executando o comando a seguir em todos os hosts.

$ gcloud compute tpus tpu-vm ssh your-qr-id \
  --zone=your-zone \
  --worker=all \
  --node=all \
  --command="command-to-run"

Redefinindo QRs

A API ResetQueuedResource pode ser usada para redefinir todas as VMs em um QR code ACTIVE. Isso força a limpeza da memória da máquina e a redefine o estado inicial da VM. Todos os dados armazenados localmente permanecerão intactos, e o script de inicialização será invocado após uma redefinição. A API ResetQueuedResource pode ser útil quando você quer reiniciar todas as TPUs. Por exemplo, quando o treinamento está travado e redefina todas as VMs do que a depuração.

As redefinições de todas as VMs são realizadas em paralelo e uma operação ResetQueuedResource leva de um a dois minutos para ser concluída. Para invocar a API, use o seguinte comando:

$ gcloud compute tpus queued-resources reset your-qr-id

Excluindo QRs

Para liberar recursos no final da sessão de treinamento, exclua o recurso na fila com a sinalização --force. A exclusão levará de dois a cinco minutos para ser concluída e pode ser executada em segundo plano com a sinalização --async opcional.

$ gcloud compute tpus queued-resources \
delete your-qr-id --force (--async)

Recuperação automática de falhas

No caso de uma interrupção, o Multislice oferece reparo sem intervenção da fração afetada e redefinição de todas as fatias posteriormente. A fração afetada é substituída por uma nova e as restantes são redefinidas. Se nenhuma capacidade estiver disponível para alocar uma fração de substituição, o treinamento será interrompido.

Para retomar o treinamento automaticamente após uma interrupção, especifique um script de inicialização que verifique e carregue os últimos checkpoints salvos. Esse script de inicialização é executado automaticamente sempre que uma fração é realocada ou uma VM é redefinida. Especifique um script de inicialização no payload JSON que você envia para a API de solicitação de QR code.

O script de inicialização a seguir (usado em Criar QRs) permite que você se recupere automaticamente de falhas e retome o treinamento em pontos de verificação armazenados em um bucket do Cloud Storage durante o treinamento do MaxText:

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 MaxText/train.py MaxText/configs/base.yml run_name=run_test_failure_recovery dcn_data_parallelism=4 ici_fsdp_parallelism=8 steps=10000 save_period=10 base_output_directory='gs://user1-us-central2'' EOF"
         }
     ...
     }
   ]
 }
}

Clone o repositório MaxText antes de testar isto.

Criação de perfil e depuração

A criação de perfil é a mesma nos ambientes de fração única e multislice. Para mais informações, consulte Como criar perfis de programas JAX.

Como otimizar o treinamento

Como fragmentar com o multislice para o desempenho máximo

Para alcançar o desempenho máximo em ambientes de multislices, é necessário considerar como fragmentar as várias fatias. Normalmente, há três opções (paralelismo de dados, paralelismo de dados totalmente fragmentado e paralelismo de pipeline). Não recomendamos fragmentar ativações entre as dimensões do modelo (às vezes chamado de paralelismo de tensor), porque exige muita largura de banda entre as fatias. Para todas essas estratégias, é possível manter a mesma estratégia de fragmentação em uma fração que funcionou para você no passado.

Recomendamos começar com o paralelismo puro de dados. O uso de paralelismo de dados totalmente fragmentado é útil para liberar o uso de memória. A desvantagem é que a comunicação entre frações usa a rede DCN e torna a carga de trabalho mais lenta. Use o paralelismo de pipeline somente quando necessário com base no tamanho do lote (conforme analisado abaixo).

Quando usar o paralelismo de dados

O paralelismo puro de dados funcionará bem nos casos em que você tem uma carga de trabalho que está funcionando bem, mas gostaria de melhorar o desempenho escalonando em várias frações.

Para alcançar um escalonamento forte em várias frações, o tempo necessário para executar a redução total na DCN precisa ser menor que o tempo necessário para executar uma passagem para trás. A DCN é usada para comunicação entre frações e é um fator limitante na capacidade da carga de trabalho.

Cada chip de TPU v4 tem um desempenho máximo de 275 * 10¹² FLOPS por segundo.

Há quatro chips por host de TPU, e cada host tem uma largura de banda de rede máxima de 50 Gbps.

Isso significa que a intensidade aritmética é 4 * 275 * 10¹² FLOPS / 50 Gbps = 22.000 FLOPS / bit.

Seu modelo usará de 32 a 64 bits de largura de banda DCN para cada parâmetro por etapa. Se você usar dois frações, seu modelo usará 32 bits de largura de banda DCN. Se você usar mais de duas frações, o compilador executará uma operação de redução total de embaralhamento completo e você usará até 64 bits de largura de banda DCN para cada parâmetro por etapa. A quantidade de FLOPS necessárias para cada parâmetro varia de acordo com seu modelo. Especificamente, para modelos de linguagem baseados em transformador, o número de FLOPS necessários para um passo para frente e para trás é aproximadamente 6 * B * P, em que:

• B é o tamanho do lote em tokens
• P é o número de parâmetros

O número de FLOPS por parâmetro é 6 * B, e o número de FLOPS por parâmetro durante o passe reverso é 4 * B.

Para garantir um escalonamento forte em várias frações, certifique-se de que a intensidade operacional exceda a intensidade aritmética do hardware da TPU. Para calcular a intensidade operacional, divida o número de FLOPS por parâmetro durante a passagem para trás pela largura de banda da rede (em bits) por parâmetro por etapa: Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

Portanto, para um modelo de linguagem baseado em Transformer, se você estiver usando duas frações: Operational intensity = 4 * B / 32

Se você estiver usando mais de duas frações: Operational intensity = 4 * B/64

Isso sugere um tamanho de lote mínimo entre 176 mil e 352 mil para modelos de linguagem baseados em Transformer. Como a rede DCN pode descartar pacotes brevemente, é melhor manter uma margem de erro significativa, implantando o paralelismo de dados somente se o tamanho do lote por pod for de pelo menos 350 mil (dois pods) a 700 mil (muitos pods).

Para outras arquiteturas de modelo, você precisará estimar o tempo de execução de sua passagem para trás por fatia, seja cronometrando-a com um criador de perfil ou contando FLOPS. Em seguida, você pode comparar isso com o tempo de execução esperado para todos reduzirem na DCN e ter uma boa estimativa de se o paralelismo de dados fará sentido para você.

Quando usar o paralelismo de dados totalmente fragmentados (FSDP, na sigla em inglês)

O paralelismo de dados totalmente fragmentado (FSDP, na sigla em inglês) combina o paralelismo de dados (fragmentação de dados entre os nós) com a fragmentação dos pesos entre os nós. Para cada operação nas passagens de ida e volta, os pesos são todos coletados para que cada fração tenha os pesos necessários. Em vez de sincronizar os gradientes usando a redução total, eles são reduzidos e dispersos à medida que são produzidos. Dessa forma, cada fatia recebe apenas os gradientes dos pesos pelos quais é responsável.

Semelhante ao paralelismo de dados, o FSDP exigirá o escalonamento linear do tamanho do lote global com o número de frações. O FSDP diminui a pressão sobre a memória à medida que você aumenta o número de frações. Isso ocorre porque o número de pesos e o estado do otimizador por fração diminui, mas faz isso ao preço do aumento do tráfego de rede e da maior possibilidade de bloqueio devido a um atraso coletivo.

Na prática, o FSDP em frações é melhor se você estiver aumentando o lote por fatia, armazenando mais ativações para minimizar a rematerialização durante a passagem para trás ou aumentando o número de parâmetros na rede neural.

As operações de coleta total e redução no FSDP funcionam de maneira semelhante àquelas no DP. Assim, é possível determinar se a carga de trabalho do FSDP está limitada pelo desempenho do DCN da mesma maneira descrita na seção anterior.

Quando usar o paralelismo de pipeline

O paralelismo de pipelines se torna relevante ao alcançar alto desempenho com outras estratégias de paralelismo que exigem um tamanho de lote global maior que o tamanho máximo que você quer. O paralelismo de pipelines permite que frações de um pipeline "compartilhem" um lote. No entanto, o paralelismo de pipeline tem duas desvantagens significativas:

1. Isso causa o "bolha do pipeline", em que os ícones estão inativos porque estão aguardando dados.
2. Isso requer microlotes, que diminui o tamanho efetivo do lote, a intensidade aritmética e, por fim, modelam a utilização do FLOP.

O paralelismo de pipeline só deverá ser usado se as outras estratégias de paralelismo exigirem um tamanho de lote global muito grande. Antes de tentar o paralelismo de pipeline, convém testar empiricamente se a convergência por amostra diminui no tamanho do lote necessário para alcançar um FSDP de alto desempenho. O FSDP tende a alcançar maior utilização de FLOP de modelo, mas se a convergência por amostra diminuir à medida que o tamanho do lote aumenta, o paralelismo de pipeline ainda poderá ser a melhor escolha. A maioria das cargas de trabalho pode tolerar tamanhos de lote grandes o suficiente para não se beneficiar do paralelismo de pipeline, mas elas podem ser diferentes.

Se o paralelismo de pipeline for necessário, recomendamos combiná-lo com o paralelismo de dados ou FSDP. Isso permitirá que você minimize a profundidade do pipeline e aumente o tamanho do lote por pipeline até que a latência da DCN se torne menos importante na capacidade de processamento. Concretamente, se você tiver N frações, considere os pipelines de Isso minimiza a lentidão introduzida pelo paralelismo do pipeline, além de permitir o escalonamento além do limite global de tamanho do lote.

Práticas recomendadas de multislices

Carregamento de dados

Durante o treinamento, carregamos repetidamente lotes de um conjunto de dados para alimentar o modelo. Ter um carregador de dados assíncrono eficiente que fragmenta o lote entre os hosts é importante para evitar a privação das TPUs de trabalho. O carregador de dados atual em MaxText faz com que cada carregamento de host tenha um subconjunto igual dos exemplos. Essa solução é adequada para texto, mas requer uma refragmentação dentro do modelo. Além disso, o MaxText ainda não oferece snapshots determinísticos, o que permitiria ao iterador carregar os mesmos dados antes e depois da preempção.

Como estabelecer pontos de verificação

A biblioteca de checkpoint da Orbax fornece primitivos para o checkpoint do JAX PyTrees para o armazenamento local ou do Google Cloud. Fornecemos uma integração de referência com checkpoint síncrono no MaxText em checkpointing.py.

Configurações aceitas

Formas

Todos os cortes precisam ter a mesma forma (por exemplo, o mesmo AcceleratorType). Não é possível usar formatos heterogêneos de fatia.

Orquestração

O GKE oferece suporte à orquestração. Para mais informações, consulte TPUs no GKE.

Frameworks

O multislice só oferece suporte a cargas de trabalho JAX e PyTorch.

Paralelismo

Recomendamos que os usuários testem o Multislice com paralelismo de dados. Para saber mais sobre a implementação do paralelismo de pipeline com o Multislice, entre em contato com o representante da conta do Google Cloud.

Suporte e feedback

Seu feedback é muito bem-vindo! Para compartilhar feedback ou solicitar suporte, entre em contato usando o formulário de feedback ou suporte do Cloud TPU.