Visão geral de vários pedaços 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 fatias em vários pods com paralelismo de dados simples. Com chips de 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 estão mais prontamente disponíveis, permitindo um tempo de inicialização mais rápido quando o Multislice é usado com partes menores.

Quando implantados em configurações de multislice, os chips de TPU em cada fração se comunicam por meio da interconexão entre chips (ICI). 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 de DCN entre fatias. O compilador XLA gera esse código para você e sobrepõe a comunicação com a computação para otimizar o desempenho.

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 criará 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 de treinamento pode ser reiniciado automaticamente sem intervenção do usuário, carregando e retomando 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 latência maior e menor capacidade (quando comparada com ICI) que conecta frações de TPU em uma configuração Multislice.

Agendamento de grupo

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

Host

Um host é um computador físico que executa VMs. Um host pode executar no máximo 4 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 sobre os dados.

Interchip Interconnect (ICI)

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

Multislice

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

Nó

No contexto 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 em várias TPUs.

Recurso na fila (QR)

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

Script de inicialização

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

Fração de TPU

Uma subseção lógica de um Pod de TPU que consiste em chips de TPU. Todos os chips em uma fração se comunicam usando a rede ICI.

VM da TPU

Uma máquina virtual executando Linux que tem acesso às TPUs subjacentes. Para TPUs v4, cada VM de TPU tem acesso direto a quatro chips. Às vezes, chamamos uma VM de 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. Eles foram projetados para oferecer computação rápida e eficiente 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 três tipos de capacidade. Consulte "Opções de uso" em Como funcionam os preços de TPUs:

• Reserva: segmenta a cota reservada. Para usar a cota reservada, é preciso ter um contrato de reserva com o Google. Use a sinalização --reserved ao criar os recursos.
• Preemptivo: segmenta a cota Preemptiva. Os recursos podem ser interrompidos por preempção para abrir espaço para solicitações de um job de prioridade mais alta. Use a sinalização --best-effort 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 da TPU será colocada em uma fila de cotas sob demanda oferecida pelo Cloud TPU. A disponibilidade dos recursos não é garantida. Selecionado por padrão, nenhuma flag é necessária.

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

Começar

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

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

Exemplo introdutório

Neste tutorial, usamos código do repositório da MaxText no GitHub (link em inglês). O MaxText é um LLM básico de alto desempenho, arbitrariamente escalonável, de código aberto e bem testado, escrito em Python e Jax. O MaxText foi projetado para um treinamento eficiente na 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 fragmentados (FSDP, na sigla em inglês) e paralelismo de tensor. O código é escalonado de uma fração única para ambientes Multislice.

Paralelismo de ICI

ICI se refere à interconexão de alta velocidade que conecta as TPUs em uma única fatia. A fragmentação de ICI corresponde à fragmentação dentro de uma fração. shardings.py fornece três parâmetros de paralelismo de 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 fração.

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

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

Na maioria dos casos, ici_data_parallelism precisa ser deixado como 1, 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 do código em uma única fração de TPU, como em Executar um cálculo em uma VM do Cloud TPU usando o 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ê receber uma solicitação de que o arquivo google_compute_engine já existe, substitua a versão atual.

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

3. Provisione as TPUs com o seguinte comando:

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

   Descrições de sinalizações de comando

   your-qr-id

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

   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, acesse 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 code.

   reserved

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

   best-effort

   Use a cota de melhor esforço ao criar as frações [padrão].

   O Google Cloud CLI não oferece suporte a todas as opções de criação de QR, 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 do QR code começa, 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 usando o seguinte comando:

   $ gcloud alpha 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 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 fração de TPU.

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

   É possível ver os resultados nos registros. Suas TPUs precisam alcançar cerca de 260 TFLOP por segundo ou uma utilização impressionante de 90%ou mais. 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 conhecer outras estratégias de fragmentação (link em inglês) com o ICI. Por exemplo, tente 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 TPU e QR quando terminar. Execute essas etapas de limpeza no ambiente em que você configurou a fração. Primeiro, execute exit para sair da sessão SSH. A exclusão vai 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. Caso contrário, ele continuará consumindo sua cota de QR code, que é rastreada separadamente da cota de TPU.

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

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

O script shardings.py usa três parâmetros que especificam o paralelismo da DCN, correspondentes 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 frações
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 mantenha o tamanho da fração e o lote por fração 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 de TPU. Aqui, usamos multihost_runner.py. As etapas a seguir são muito semelhantes aos Primeiros passos: experimentos rápidos em várias frações do repositório MaxText, exceto aqui que 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 o script multihost_runner.py é executado. Usamos o termo trabalhadores para indicar as VMs da 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 que as frações. Não é possível executar multihost_runner.py em um worker.

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

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

configurar o ambiente

1. Clone MaxText na máquina de execução.

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 uma mensagem informando que o arquivo google_compute_engine já existe, escolha 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. Criar um ambiente Multislice usando queued-resources create.

   Observação :verifique se você tem a cota respectiva antes de selecionar --reserved ou --best_effort, ou a cota padrão sob demanda. Para informações sobre 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 da v5e (por exemplo, v5litepod-16) e a v5e runtime-version (v2-alpha-tpuv5-lite).

     $ gcloud alpha 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|--best-effort]
     

   Descrições de sinalizações de comando

   your-qr-id

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

   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, acesse 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 fração. Um número é anexado ao prefixo de cada fração. Por exemplo, se você definir node-prefix como mySlice, as frações serão nomeadas como mySlice-0, mySlice-1 e assim por diante.

   reserved

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

   best-effort

   Use a cota de melhor esforço ao criar as frações [padrão].

6. O provisionamento do QR code pode levar até cinco minutos para ser concluído, dependendo do tamanho dele. Aguarde até que o recurso em fila (QR, na sigla em inglês) esteja no estado ACTIVE. É possível verificar o status de uma solicitação de QR usando o seguinte comando:

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

   A saída será 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 estiver 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"
   

   Serão exibidos aproximadamente 230 TFLOPs por segundo de desempenho nos arquivos de registro.

9. Limpe as TPUs e o QR code quando terminar. A exclusão leva 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 multislice

Antes de executar o modelo em um ambiente Multislice, faça as seguintes alterações no código:

• 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 ter alto desempenho, a DCN precisa ser mapeada em eixos paralelos de dados, dados totalmente fragmentados ou paralelos de pipeline. As considerações de desempenho e as estratégias de fragmentação são discutidas em mais detalhes em Como fragmentar com multislices para obter o desempenho máximo.

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

Como escolher tamanhos de fração para ambientes Multislice

Para ter uma aceleração linear, adicione novas frações do mesmo tamanho que a atual. Por exemplo, se você usar uma fração v4-512, o Multislice atingirá aproximadamente o dobro do desempenho, adicionando uma segunda fatia v4-512 e dobrando o tamanho do lote global. Para mais informações, consulte Como fragmentar com vários pedaços para alcançar o 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 da experimentação, multihost_runner.py
2. Usando o script do executor de produção, multihost_job.py
3. Como usar uma abordagem manual

Script do executor de experimentos

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

Como multihost_runner.py mantém conexões SSH persistentes, ele é adequado apenas para experimentos de tamanho modesto e relativamente de curta duração. É possível adaptar as etapas do tutorial multihost_runner.py para 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 fazer a integração diretamente com a 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 na preempção. O script multihost_job.py está documentado no README do MaxText.

Como o 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 Multislice. No entanto, se você preferir provisionar e gerenciar o ambiente usando comandos QR diretamente, consulte Gerenciar um ambiente multislice.

Gerenciar um ambiente multislice

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

Criar respostas rápidas

Defina as variáveis de ambiente a seguir 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

Entrada	Descrição
your-qr-id	O ID do QR atribuído pelo usuário.
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 partes. Limitado a um máximo de 256 fatias.
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 alpha 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|--best-effort]
  

Descrições de sinalizações de comando

your-qr-id

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

project

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

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, acesse Versões de TPU.

runtime-version

A versão do software do Cloud TPU.

network

O nome de uma rede VPC à qual 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.

best-effort

Use a cota de melhor esforço ao criar as frações [padrão].

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

Criar uma solicitação de QR code usando curl

Crie um arquivo chamado queued-resource-req.json e copie o seguinte JSON 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 seu 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 fatia
• 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 será anexado
• your-subnetwork-name: opcional, uma sub-rede em que o QR 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 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 todas as opções disponíveis.

Observação: embora node_spec seja uma lista, só é possível usar uma entrada node_spec.

Para usar a capacidade preemptiva, substitua:

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

Ou remova a linha para usar a capacidade sob demanda padrão.

Envie a solicitação de criação de QR 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 do GUID no final do valor da string para o atributo name para receber informações sobre a solicitação de QR.

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/v2alpha1/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 do QR.

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.v2alpha1.QueuedResource",
    ...
    "state": {
      "state": "WAITING_FOR_RESOURCES"
    }
  }
}

Se o QR code for criado ("done = true"), o estado no campo response será WAITING_FOR_RESOURCES ou FAILED. Se o QR estiver no estado WAITING_FOR_RESOURCES, ele foi colocado na fila e começará a provisionar 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 na fila.

Quando a operação for concluída, use a opção descrever QRs para monitorar os estágios deles.

Em um cenário raro, você pode encontrar seu QR no estado FAILED enquanto algumas partes estã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 instalação de dependências

Executar o código JAX em frações do pod de TPU descreve como se conectar às VMs de TPU usando SSH em uma única fração. Para se conectar a todas as VMs de 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

Este comando gcloud envia o comando especificado para todos os workers e nós no QR usando SSH. O comando é dividido em grupos de quatro e enviado simultaneamente. O próximo lote de comandos é enviado quando o lote atual conclui a execução. Se houver uma falha em um dos comandos, o processamento será interrompido e nenhum outro lote será enviado. Para mais informações, consulte a referência da API de recursos na fila. Se o número de frações que você está usando exceder o limite de linhas de execução do computador local (também chamado de limite de lotes), haverá um impasse. Por exemplo, suponha que o limite de lotes na máquina local seja de 64. Se você tentar executar um script de treinamento em mais de 64 frações, digamos 100, o comando SSH dividirá as frações 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 pode 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 do script especificado com a sinalização --command. Quando você fizer 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ê vai precisar canalizar os streams stdout e stderr corretamente ao executar os comandos em segundo plano. Para aumentar o paralelismo no mesmo QR code, é possível selecionar frações específicas usando o parâmetro --node.

Configuração da rede

Execute as etapas a seguir para garantir que as frações de TPU se comuniquem umas com as outras. 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 chips no ambiente Multislice. Para fazer isso, execute este comando em cada fração:

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

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

Listar QRs

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

$ gcloud alpha 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 conferir a configuração detalhada e o estado de um QR code, use a API de descrição do QR. É possível chamar essa API usando gcloud ou curl.

Utilizando gcloud:

$ gcloud alpha 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/v2alpha1/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 seguinte comando 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. A redefinição das VMs apaga à força a memória da máquina e redefine a VM para o estado inicial. 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 redefinir todas as VMs é mais fácil 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 chamar a API, use o seguinte comando:

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

Excluindo QR codes

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

$ gcloud alpha 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 a redefinição de todas as frações posteriormente. A fração afetada é substituída por uma nova, e as frações íntegras 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 verifica e carrega 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 enviado para a API de criação de solicitação de QR.

O script de inicialização a seguir (usado em Criar QRs) permite que você se recupere automaticamente de falhas e retome o treinamento dos checkpoints armazenados em um bucket do Cloud Storage durante o treinamento da 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.

Criação de perfil e depuração

A criação de perfil é a mesma em 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 multislice para conseguir o desempenho máximo

Alcançar o desempenho máximo em ambientes Multislice requer considerar como fragmentar as várias frações. Normalmente, há três opções: paralelismo de dados, paralelismo de dados totalmente fragmentados e paralelismo de pipeline. Não recomendamos a fragmentação de ativações nas dimensões do modelo (às vezes chamada de paralelismo de tensor) porque isso exige muita largura de banda entre 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 de dados puros. Usar o paralelismo de dados totalmente fragmentados é útil para liberar o uso da memória. A desvantagem é que a comunicação entre frações usa a rede DCN e atrasará a carga de trabalho. 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 de dados puro funcionará bem nos casos em que você tem uma carga de trabalho que está sendo executada corretamente, mas quer 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 pelo DCN precisa ser menor que o tempo necessário para executar uma passagem reversa. 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 pico 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 de DCN para cada parâmetro por etapa. Se você usar duas frações, o modelo vai usar 32 bits de largura de banda DCN. Se você usar mais de duas frações, o compilador vai executar uma operação de redução total aleatória e usar até 64 bits de largura de banda DCN para cada parâmetro por etapa. A quantidade de FLOPS necessários para cada parâmetro varia de acordo com o modelo. Especificamente para modelos de linguagem baseados em transformadores, o número de FLOPS necessários para um passe de frente e um retorno é de 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 para trás é 4 * B.

Para garantir um forte escalonamento em várias frações, verifique se a intensidade operacional excede 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 reversa 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 transformador, 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 transformadores. Como a rede DCN pode descartar pacotes brevemente, é melhor manter uma margem significativa para erro, implantando paralelismo de dados apenas 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 ambiente de execução da sua passagem reversa por fração, 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 reduzir toda a redução no DCN e ter uma boa estimativa de se o paralelismo de dados faz sentido para você.

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

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

Semelhante ao paralelismo de dados, o FSDP exigirá o escalonamento do tamanho do lote global linear com o número de frações. O FSDP diminui a pressão da 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 isso acontece devido ao aumento do tráfego de rede e à maior possibilidade de bloqueio devido a um coletivo atrasado.

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

As operações de coleta e redução total no FSDP funcionam de maneira semelhante às do DP, para que você possa determinar se a carga de trabalho do FSDP é limitada pelo desempenho do DCN da mesma maneira descrita na seção anterior.

Quando usar o paralelismo de pipelines

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

1. Ele gera a "bolha do pipeline", em que os chips estão inativos porque estão aguardando dados.
2. Ela requer microlotes, que diminui o tamanho efetivo do lote, a intensidade aritmética e, por fim, a utilização do FLOP do modelo.

O paralelismo de pipeline só será usado se as outras estratégias de paralelismo exigirem um tamanho de lote global muito grande. Antes de testar o paralelismo de pipeline, vale a pena fazer testes para verificar empiricamente se a convergência por amostra diminui no tamanho do lote necessário para atingir um FSDP de alto desempenho. O FSDP tende a alcançar uma maior utilização do FLOP do modelo, mas se a convergência por amostra diminuir à medida que o tamanho do lote aumenta, o paralelismo do pipeline ainda pode ser a melhor escolha. A maioria das cargas de trabalho pode tolerar tamanhos de lote suficientes para não se beneficiar do paralelismo de pipeline, mas sua carga de trabalho pode ser diferente.

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 um fator na capacidade. Concretamente, se você tiver N frações, considere pipelines de profundidade 2 e N/2 réplicas de paralelismo de dados, em seguida, pipelines com profundidade de 4 e N/4 réplicas de paralelismo de dados e assim por diante, até que o lote por pipeline fique grande o suficiente para que os coletivos de DCN possam ser ocultados por trás da aritmética na passagem de retorno. Isso minimizará a lentidão introduzida pelo paralelismo do pipeline, permitindo que você escale além do limite global de tamanho do lote.

Práticas recomendadas para multislice

Carregamento de dados

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

Como estabelecer pontos de verificação

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

Configurações compatíveis

Formas

Todas as fatias precisam ter a mesma forma (por exemplo, o mesmo AcceleratorType). Não há suporte para formas de fatia heterogêneas.

Orquestração

A orquestração é compatível com o GKE. Para mais informações, consulte TPUs no GKE.

Frameworks

O multislice oferece suporte apenas às cargas de trabalho JAX e PyTorch.

Paralelismo

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

Suporte e feedback

Gostaríamos de receber seu feedback. Para compartilhar feedback ou solicitar suporte, use o formulário de suporte ou feedback do Cloud TPU.