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 4096 ícones em uma única execução. Para trabalhos de treinamento que exigem menos de 4.096 chips, uma única fatia pode oferecer o melhor desempenho. No entanto, várias frações menores ficam disponíveis mais facilmente, permitindo um tempo de inicialização mais rápido quando o Multislice é usado com fatias

Quando implantados em configurações de multislices, os chips de TPU em cada fração se comunicarem pela Inter-chip-interconnect (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 do data center (DCN, na sigla em inglês).

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 de computação para o melhor 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 multislices é do mesmo tipo de acelerador. Uma aceleradora 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 um slice encontra um evento de manutenção, preempção ou falha de hardware, o Cloud TPU cria um novo slice. 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 fatia é criada, todas as outras fatias no ambiente de fatias múltiplas são reiniciadas para que o treinamento continue. 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 ponto de verificação 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 mais alta e capacidade de processamento menor (quando comparada com a ICI) que conecta frações de TPU em uma configuração Multislice.

Programação de gangues

Quando todas as frações de TPU são provisionadas juntas, ao mesmo tempo, garantindo ou nenhuma das frações foi provisionada.

Host

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

Inferência

Carregar um modelo de machine learning pré-treinado em um host e fazer previsões nele 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 pela DCN.

Nó

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

Pod

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

Recurso em fila (QR)

Uma representação de recursos de TPU, usada para enfileirar e gerenciar uma solicitação para um ambiente de TPU de uma ou várias fatias.

Script de inicialização

Um script de inicialização padrão do Compute Engine executada sempre que uma VM é inicializada ou reinicializada. Para multislices, 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 de 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 um fração se comunicam usando a rede ICI.

VM de TPU

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

Tensor

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 aprendizado de máquina, 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 da 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 seus recursos.
• Spot: segmenta a cota preemptiva usando VMs spot. Seu os recursos podem ser interrompidos para dar espaço a solicitações de um job prioritário. Use a sinalização --spot ao criar os recursos.
• Sob demanda: segmenta uma cota sob demanda, que não precisa de reserva. e não será interrompido. A solicitação de TPU será enfileirada em uma fila de cota sob demanda oferecida pela Cloud TPU. A disponibilidade de recursos não é garantida. Selecionado por padrão, sem flags necessárias.

Primeiros passos

Se você ainda não usou TPUs, comece instalando a CLI do Google Cloud e configurando seu ambiente de TPU do Cloud. Para usar Multislice, os 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 migrar sua reserva para um novo sistema de reservas. Para mais informações, entre em contato com o representante da sua 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 de alto desempenho, escalonável de forma arbitrária, de código aberto e bem testado, escrito em Python e Jax. O MaxText foi projetado para treinar de forma eficiente Cloud TPU:

O código em shardings.py foi desenvolvido para ajudar você a começar a testar diferentes modelos de paralelismo . Por exemplo, paralelismo de dados, paralelismo de dados totalmente fragmentados (FSDP), e paralelismo de tensores. O código é dimensionado de uma única fatia para ambientes de várias fatias.

paralelismo da ICI

ICI refere-se à interconexão de alta velocidade que conecta as TPUs em um único fatia 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 limitadas 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 o paralelismo do ICI dos quatro chips disponíveis na v4-8:

Observe que ici_data_parallelism deve ser deixado como 1 na maioria dos casos, porque o A rede ICI é rápida o suficiente para preferir quase sempre 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 digite duas vezes depois de executar o comando a seguir). Se aparecer a mensagem 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 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.

   accelerator-type

   O tipo de acelerador especifica a versão e o tamanho da 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 da TPU que serão criados em resposta à solicitação de QR.

   reserved

   Use a cota reservada ao criar as fatias.

   spot

   Use a cota de VMs Spot ao criar as fatias.

   A CLI do Google Cloud não oferece suporte a todas as opções de criação de QR code, como tags. Para mais informações, consulte Criar códigos QR.

4. Aguarde até que o código QR esteja no estado ACTIVE, o que significa que os nós de worker estão no estado READY. Depois que o provisionamento de QR code for iniciado, poderá levar de um a cinco minutos, dependendo do tamanho do QR code. É 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 do TPU usando SSH:

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

6. Clone o MaxText (que inclui shardings.py) para a VM do TPU.

7. No diretório do repositório MaxText, execute o script de configuração para instalar 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

   Você pode conferir os resultados nos registros. As TPUs precisam alcançar cerca de 260 TFLOP por segundo ou uma impressionante utilização de FLOPs de mais de 90%. Neste caso, selecionamos aproximadamente o lote máximo que cabe 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 no 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 você configurou o slice (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 flag --async opcional.

    $ 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 de 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 dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism é igual a o número de fatias.

Como exemplo para duas fatias, use --dcn_data_parallelism = 2.

O dcn_tensor_parallelism precisa ser sempre definido como 1, porque o DCN não é adequado para esse tipo de fragmentação. Para cargas de trabalho LLMs típicas em chips v4, dcn_fsdp_parallelism também precisa ser definido como 1. dcn_data_parallelism deve ser definido como o número de fatias, dependente do aplicativo.

Conforme você aumenta o número de fatias (supondo que você mantenha o tamanho da fatia e o lote por fatia), você 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 Primeiros passos: experimentos rápidos com várias frações no repositório MaxText, mas aqui executamos shardings.py em vez do um LLM mais complexo em train.py.

A ferramenta multihost_runner.py é otimizada para experimentos rápidos, repetidamente e reutilizar as mesmas TPUs. Como o script multihost_runner.py depende 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 executor para indicar a máquina em que você execute o script multihost_runner.py. Usamos o termo trabalhadores para indicar VMs de TPU que compõem suas frações. É possível executar o multihost_runner.py em um máquina virtual ou qualquer VM do Compute Engine no mesmo projeto das frações. Não é possível executar multihost_runner.py em um worker.

O 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 quatro VMs e 16 chips de TPU. É possível modificar o exemplo para executar em mais TPUs.

Configurar o ambiente

1. Clone o 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 após executar o comando abaixo). Se uma mensagem informar que o arquivo google_compute_engine já existe, selecione a opção 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 fatias de TPU como 2.

     $ export SLICE_COUNT=2
     

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

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

     $ 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 fatias 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 são nomeadas: mySlice-0, mySlice-1, continuando numericamente para cada fatia.

   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 do QR Code começa, ele pode levar até cinco minutos para ser concluído, dependendo do tamanho do código. Aguardar até que o recurso em fila (QR, na sigla em inglês) entre o estado ACTIVE. Verifique o status de uma solicitação de QR usando 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 conta do Google Cloud se o status do QR 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"

   Você vai encontrar aproximadamente 230 TFLOPs por segundo de desempenho nos arquivos de registro.

9. Limpe os TPUs e o QR Code quando terminar. A exclusão vai levar de duas a cinco minutos para concluir e pode ser executado em segundo plano com o comando sinalização --async.

Como escalonar uma carga de trabalho para multislice

Antes de executar o modelo em um ambiente de várias fatias, faça as seguintes mudanças no código:

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

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

Como escolher tamanhos de fatia para ambientes multislice

Para acelerar de forma linear, adicione novas fatias do mesmo tamanho que a atual. Por exemplo, se você usar uma fração v4-512, o Multislice vai alcançar aproximadamente o dobro do desempenho adicionando uma segunda fração de v4-512 e dobra o tamanho do lote global. Para mais informações, consulte Como dividir com o Multislice para desempenho máximo.

Como executar seu job em várias fatias

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

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 de experimentação

O script multihost_runner.py distribui o código para um ambiente Multislice existente e 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 de MaxText.

Como multihost_runner.py mantém conexões SSH persistentes, ele só adequado para experimentos de tamanho modesto e de execução relativamente curta. Você pode Adapte as etapas do tutorial multihost_runner.py. à carga de trabalho e à configuração de hardware.

Script do executor Production

Para jobs de produção que precisam de resiliência contra falhas de hardware e outros preempções, o ideal é fazer a integração diretamente com a ferramenta Criar recurso em fila API. Como um exemplo funcional, fornecemos multihost_job.py, que aciona a chamada da API Created Queued Resource com a inicialização apropriada para executar o treinamento e retomar a preempção. O multihost_job.py está documentado na README de MaxText.

Como o multihost_job.py precisa provisionar recursos para cada execução, ele não fornecem 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 de Multislice. No entanto, se você preferir provisionar e gerenciar seu ambiente usando comandos QR diretamente, consulte Gerenciar um ambiente multislice.

Gerenciar um ambiente Multislice

Para provisionar e gerenciar QRs de forma manual sem usar as ferramentas fornecido no repositório MaxText, leia o nas 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

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]
  

Verifique se você tem a respectiva cota antes de selecionar --reserved. --spot ou a cota padrão sob demanda. Para informações sobre tipos de cota, 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"
      }
    }
    ]
  }
}

• 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 fração
• tpu-vm-runtime-version: as versões do ambiente de execução da VM da TPU
• your-network-name: opcional, uma rede a que o código QR será anexado
• your-subnetwork-name: opcional, uma sub-rede a que o código 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 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

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

Para usar a capacidade do Spot, substitua:

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

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

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

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

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 receber informações sobre a solicitação de QR code.

Recuperar o status de um código QR

Para conferir 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

• your-project-id: ID do projeto do Google Cloud
• your-zone: a zona em que o QR será criado
• your-qr-guid: o GUID que segue 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 código QR for criado com sucesso ("done = true"), o estado no campo response será WAITING_FOR_RESOURCES ou FAILED. Se o QR estiver no estado WAITING_FOR_RESOURCES, ele foi será enfileirado e iniciará o provisionamento quando houver recursos suficientes. Se o código QR estiver no estado FAILED, o motivo da falha vai aparecer na saída. Para mais informações sobre outros estados possíveis, consulte a Guia do usuário sobre recursos na fila.

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

Em um cenário raro, seu QR pode estar no estado FAILED, enquanto algumas fatias estão ACTIVE. Se isso acontecer, exclua os recursos criados e tente de novo em alguns minutos ou entre em contato para a equipe do Cloud TPU resolver o problema.

SSH e instalar dependências

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

  $ 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 em QR usando SSH. O comando é reunido em grupos de quatro e enviado ao mesmo tempo. O próximo lote de comandos é enviado quando o lote atual conclui a execução. Se houver falha em um dos comandos, o processamento paradas, 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ê está usando exceder a linha de execução do seu computador local limite de lote (também chamado de limite de lotes), você encontrará um impasse. Como exemplo, suponha que o limite de lote na sua máquina local seja 64. Se você tentar executar uma script de treinamento em mais de 64 frações, digamos 100, o comando SSH quebrará as fatias em lotes. Ele vai executar o script de treinamento no primeiro lote de 64 frações e aguardar a conclusão dos scripts antes de executar o script na lote restante de 36 frações. No entanto, o primeiro lote de 64 frações não pode concluir até que as 36 frações restantes comecem a executar o script, o que resulta um impasse.

Para evitar esse cenário, execute o script de treinamento em segundo plano 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 fatias, o controle volta imediatamente para o comando SSH. O comando SSH pode começar a executar o script de treinamento no lote restante de 36 fatias. Você vai precisar canalizar as transmissões stdout e stderr adequadamente ao executar os comandos em segundo plano. Para aumentar paralelismo no mesmo QR, é possível selecionar frações específicas usando a função --node .

Configuração da rede

Siga as etapas abaixo para garantir que as frações do TPU possam se comunicar entre si. Instale o JAX em cada uma das fatias. Para mais informações, consulte Executar o código JAX em frações do pod de TPU. Verifique se len(jax.devices()) é igual ao número de chips 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 chips por fatiamento e quatro fatias, um total de 32 chips (dispositivos) 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 conferir a configuração e o estado detalhados de um QR code, use o descrever a API 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 em fila.

Iniciar o trabalho em um ambiente provisionado

É possível executar cargas de trabalho manualmente conectando-se a todos os hosts em cada fração por SSH e executar 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"

Como redefinir QRs

A API ResetQueuedResource pode ser usada para redefinir todas as VMs em um QR ACTIVE. A redefinição das VMs exclui permanentemente a memória da máquina e a redefine para o estado inicial. Todos os dados armazenados localmente permaneçam intactos, e o script de inicialização será invocado após uma redefinição. 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 do que depurar.

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 comando comando:

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

Como excluir códigos QR

Para liberar os recursos ao final da sessão de treinamento, exclua a fila recurso 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 flag --async opcional.

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

Recuperação automática de falhas

Em caso de interrupção, o Multislice oferece serviços sem intervenção reparação da fração afetada e redefinição de todas as frações posteriormente. Os afetados fração é substituída por uma nova e as fatias restantes saudáveis são redefinidas. Se não houver capacidade disponível para alocar uma fatia 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 pontos de verificação salvos. O script de inicialização é executado automaticamente sempre que uma fatia é realocada ou uma VM é redefinida. Especifique um script de inicialização no payload JSON enviado à API de criação de solicitações 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 checkpoints armazenados em um bucket do Cloud Storage durante o treinamento 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 tentar para fora.

Criação de perfil e depuração

O perfil é o mesmo em ambientes de fatia única e multifatia. Para 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

Alcançar o desempenho máximo em ambientes Multislice exige considerando como fragmentar entre 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 dividir as ativações nas dimensões do modelo (às vezes chamado de paralelismo de tensor) porque isso exige muita largura de banda entre as fatias. Para todas essas estratégias, você pode 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 puro. O uso de paralelismo de dados totalmente fragmentados é útil para liberar o uso da memória. A desvantagem é que a comunicação entre as fatias usa a rede DCN e diminui a carga de trabalho. Use o paralelismo do 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 funcionando bem, mas gostaria de melhorar o desempenho escalonando várias fatias.

Para alcançar um escalonamento forte em vários frações, o tempo necessário para realizar a redução total em DCN precisa ser menor que o tempo necessário para dar um passe para trás. A DCN é usada para a comunicação entre as frações e é um fator limitante na taxa de transferência de carga de trabalho.

Cada chip 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 máxima de rede 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 duas fatias, seu modelo usará 32 bits de largura de banda DCN. Se você Se você usar mais de duas frações, o compilador realizará uma operação de embaralhamento total na redução. operação e você 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 sua um modelo de machine learning. Especificamente, para modelos de linguagem com base em Transformer, o número de FLOPS necessárias para um passe para frente e para trás são 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 fatias, verifique se a intensidade operacional excede a intensidade aritmética do hardware de TPU. Para calcular a intensidade operacional, divida o número de FLOPS por parâmetro durante o passa 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 o Transformer de linguagem grandes e baseados em machine learning. Como a rede DCN pode descartar pacotes brevemente, é melhor manter uma margem de erro significativa, implantando apenas o paralelismo de dados se o tamanho do lote por pod for de pelo menos 350 mil (dois pods) a 700 mil (vários pods).

Para outras arquiteturas de modelo, você precisará estimar o tempo de execução da transmissão reversa por fatia (cronometrando-a usando um perfilador ou contando FLOPS). Em seguida, você pode comparar isso com o tempo de execução esperado para todos DCN e ter uma boa estimativa de se o paralelismo de dados vai fazer 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 o paralelismo de dados (fragmentação do dados entre nós) com a fragmentação dos pesos entre os nós. Para cada operação no os passos para frente e para trás, os pesos são todos agrupados de modo que cada fatia tem os pesos necessários. Em vez de sincronizar os gradientes usando a redução total, os gradientes serão reduzidos e dispersos à medida que forem produzidos. Dessa forma, cada fatia recebe apenas os gradientes dos pesos pelos quais é responsável.

Assim como o paralelismo de dados, o FSDP vai exigir o dimensionamento do tamanho do lote global de forma linear com o número de fatias. O FSDP diminui a pressão sobre a memória você aumenta o número de fatias. Isso ocorre porque o número de pesos e o estado do otimizador por fatia diminuem, mas isso acontece 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 fatias é melhor se você aumentar o lote por fatia, armazenando mais ativações para minimizar a rematerialização durante a transmissão reversa ou aumentando o número de parâmetros na rede neural.

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

Quando usar o paralelismo de pipeline

O paralelismo de pipelines se torna relevante ao alcançar alto desempenho com outros estratégias de paralelismo que exigem um tamanho de lote global maior que o tamanho máximo do lote preferencial. O paralelismo do pipeline permite que as fatias que compõem um pipeline "compartilhem" um lote. No entanto, o paralelismo do pipeline tem duas desvantagens significativas:

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

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

Se o paralelismo de pipeline for necessário, recomendamos combiná-lo com paralelismo de dados ou FSDP. Isso permite minimizar a profundidade do pipeline e, aumentar o tamanho de lote por pipeline até que a latência de DCN se torne menor na capacidade de processamento. Se você tiver N frações, considere pipelines réplicas de profundidade 2 e N/2 de paralelismo de dados, e os pipelines de profundidade 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 DCN possam ficar escondidos por trás da aritmética na passe para trás. Isso vai minimizar a lentidão introduzida pelo paralelismo do pipeline, permitindo que você dimensione além do limite de tamanho de lote global.

Práticas recomendadas para fatias múltiplas

Carregamento de dados

Durante o treinamento, carregamos repetidamente lotes de um conjunto de dados para alimentar um modelo de machine learning. Ter um carregador de dados assíncrono e eficiente que fragmenta o lote é importante para evitar a fome 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 reshard no modelo. Além disso, o MaxText ainda não oferece snapshots determinísticos, o que permitiria ao iterador de dados para carregar os mesmos dados antes e depois da preempção.

Como estabelecer pontos de verificação

A biblioteca de controle de ponto de verificação Orbax fornece primitivas para o controle de ponto de verificação de PyTrees JAX para 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

Todas as fatias precisam ter a mesma forma (por exemplo, o mesmo AcceleratorType). Não é possível usar formas de fatias 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 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 seu representante da conta do Google Cloud.

Suporte e feedback

Queremos saber sua opinião. Para compartilhar feedback ou solicitar suporte, entre em contato com nossa equipe usando o formulário de feedback ou suporte do Cloud TPU.