Visão geral do Multislice da 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 uma única fração ou em frações em vários pods com paralelismo de dados padrão. Com os chips TPU v4, isso significa que os trabalhos 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 fatias menores ficam mais disponíveis, permitindo um tempo de inicialização mais rápido quando o Multislice é usado com fatias menores.

Quando implantados em configurações Multislice, os chips de TPU em cada fração se comunicam por interconexão entre chips (ICI). Os chips de TPU em diferentes slices se comunicam transferindo dados para CPUs (hosts), que transmitem os dados pela rede de data center (DCN). Para mais informações sobre escalonamento com Multislice, consulte Como escalonar o treinamento de IA para até dezenas de milhares de chips do Cloud TPU com Multislice.

Os desenvolvedores não precisam escrever código para implementar a comunicação DCN entre segmentos. O compilador XLA gera esse código para você e sobrepõe a comunicação com a computação para maximizar 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 mais recente) seguido pelo número de TensorCores. Por exemplo, v5litepod-128 especifica uma TPU v5e 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. Se não houver recursos suficientes para criar uma nova fração, a criação não será concluída até que o hardware fique disponível. Depois que a nova fração é criada, todas as outras no ambiente Multifragmento sã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 do último ponto de verificação.

Rede de data center (DCN)

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

Agendamento em grupo

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.

Interconexão entre chips (ICI)

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

Multislice

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

Nó

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

Script de inicialização

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

Tensor

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

Tipos de capacidade da Cloud TPU

As TPUs podem ser criadas com diferentes tipos de capacidade. Consulte "Opções de uso" em Como funciona o preço das TPUs:

• Reserva: para consumir uma reserva, você precisa ter um contrato de reserva com o Google. Use a flag --reserved ao criar seus recursos.

• Spot: segmenta a cota preemptiva usando VMs spot. Seus recursos podem ser interrompidos para abrir espaço para solicitações de um job de maior prioridade. Use a sinalização --spot ao criar seus recursos.

• On demand: segmenta a cota on demand, que não precisa de uma reserva e não será interrompida. A solicitação de TPU será enfileirada em uma fila de cota sob demanda oferecida pelo Cloud TPU, mas a disponibilidade de recursos não é garantida. Selecionado por padrão, sem necessidade de flags.

Primeiros passos

1. Configure seu ambiente do Cloud TPU.

2. In the Google Cloud console, activate Cloud Shell.

   Activate Cloud Shell

   At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

Para usar o multislice, os recursos de TPU precisam ser gerenciados como recursos enfileirados.

Exemplo introdutório

Este tutorial usa código do repositório do GitHub MaxText. 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 com eficiência na Cloud TPU.

O código em shardings.py foi criado para ajudar você a começar a testar diferentes opções de paralelismo. Por exemplo, paralelismo de dados, paralelismo de dados totalmente fragmentados (FSDP) e paralelismo de tensores. O código é dimensionado de um único ambiente de corte para ambientes Multislice.

Paralelismo de ICI

ICI se refere à interconexão de alta velocidade que conecta as TPUs em uma única fração. O sharding de ICI corresponde ao sharding em uma fração. 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 chips na fração.

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

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

Observe que ici_data_parallelism deve 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.

Este exemplo pressupõe que você já sabe como executar código em uma única fração de TPU, como em Executar um cálculo em uma VM da Cloud TPU usando o JAX. Este exemplo mostra como executar shardings.py em uma única fração.

1. Configure o ambiente:

   $ gcloud auth login
   $ export QR_ID=your-queued-resource-id
   $ export TPU_NAME=your-tpu-name
   $ export PROJECT=your-project-name
   $ export ZONE=us-central1-a
   $ export NETWORK_NAME=your-network-name
   $ export SUBNETWORK_NAME=your-subnetwork-name
   $ export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   $ export ACCELERATOR_TYPE=v5litepod-16
   $ export EXAMPLE_TAG_1=your-tag-1
   $ export EXAMPLE_TAG_2=your-tag-2
   $ export SLICE_COUNT=4
   $ export STARTUP_SCRIPT='#!/bin/bash\n'

   Descrições de variáveis

   Entrada	Descrição
   QR_ID	O ID atribuído pelo usuário do recurso enfileirado.
   TPU_NAME	O nome atribuído pelo usuário da TPU.
   PROJETO	Nome do projetoGoogle Cloud
   ZONA	Especifica a zona em que os recursos serão criados.
   NETWORK_NAME	Nome das redes VPC.
   SUBNETWORK_NAME	Nome da sub-rede em redes VPC
   RUNTIME_VERSION	A versão do software da Cloud TPU.
   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 fatias. Limitado a um máximo de 256 intervalos.
   STARTUP_SCRIPT	Se você especificar um script de inicialização, ele será executado quando a fração de TPU for provisionada ou reiniciada.

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 mensagem informando que o arquivo google_compute_engine já existe, substitua a versão atual.

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

3. Provisione as TPUs:

   gcloud

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --node-id=${TPU_NAME} \
       --zone=${ZONE} \
       [--reserved |--spot]

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

   Console

   1. No console Google Cloud , acesse a página TPUs:

      Acessar TPUs

   2. Clique em Criar TPU.

   3. No campo Nome, insira um nome para a TPU.

   4. Na caixa Zona, selecione a zona em que você quer criar a TPU.

   5. Na caixa Tipo de TPU, selecione um tipo de acelerador. 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 acelerador compatíveis com cada versão de TPU, consulte Versões de TPU.

   6. Na caixa Versão do software de TPU, selecione uma versão. Ao criar uma VM do Cloud TPU, a versão do software da TPU especifica a versão do ambiente de execução da TPU a ser instalada. Para mais informações, consulte Versões de software da TPU.

   7. Clique no botão Ativar enfileiramento.

   8. No campo Nome do recurso na fila, digite um nome para sua solicitação de recurso na fila.

   9. Clique em Criar para criar a solicitação de recurso em fila.

4. Aguarde até que o recurso na fila esteja no estado ACTIVE, o que significa que os nós de worker estão no estado READY. Depois que o provisionamento de recursos enfileirados começa, pode levar de um a cinco minutos para ser concluído, dependendo do tamanho do recurso enfileirado. É possível verificar o status de uma solicitação de recurso em fila usando a CLI gcloud ou o console Google Cloud :

   gcloud

   $ gcloud compute tpus queued-resources \
       list --filter=${QR_ID} --zone=${ZONE}

   Console

   1. No console Google Cloud , acesse a página TPUs:

      Acessar TPUs

   2. Clique na guia Recursos na fila.

   3. Clique no nome da solicitação de recurso na fila.

5. Conecte-se à VM da TPU usando SSH:

   $ gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE}

6. Clone o MaxText (que inclui o shardings.py) na sua VM de TPU:

   $ git clone https://github.com/AI-Hypercomputer/maxtext && cd maxtext

7. Instale o Python 3.10:

   $ sudo apt-get update
   $ sudo apt install python3.10
   $ sudo apt install python3.10-venv

8. Crie e ative um ambiente virtual:

   $ python3 -m venv your-venv-name
   $ source your-venv-name/bin/activate

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

   $ bash setup.sh

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

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

    Você pode conferir os resultados nos registros. As TPUs devem atingir cerca de 260 TFLOP por segundo ou uma impressionante utilização de FLOP de mais de 90%. Neste caso, selecionamos aproximadamente o lote máximo que cabe na memória de alta largura de banda (HBM) da TPU.

11. Você pode explorar outras estratégias de fragmentação no ICI. Por exemplo, tente a seguinte combinação:

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

12. Exclua o recurso na fila e a fração de TPU 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 leva de dois a cinco minutos. Se você estiver usando a CLI gcloud, execute esse comando em segundo plano com a flag opcional --async.

    gcloud

    $ gcloud compute tpus queued-resources \
        delete ${QR_ID} --force (--async)

    Console

    1. No console Google Cloud , acesse a página TPUs:

       Acessar TPUs

    2. Clique na guia Recursos na fila.

    3. Marque a caixa de seleção ao lado da solicitação de recurso na fila.

    4. Clique em deleteExcluir.

Fragmentação multislice 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 para que dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism seja igual ao número de fatias.

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

	dcn_data_parallelism	dcn_fsdp_parallelism	dcn_tensor_parallelism	# de intervalos
Paralelismo de dados bidirecional	2	1	1	2

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

Ao aumentar o número de intervalos (supondo que você mantenha o tamanho do intervalo e o lote por intervalo constantes), você aumenta a quantidade de paralelismo de dados.

Executar shardings.py em um ambiente Multislice

É possível executar shardings.py em um ambiente de várias fatias 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 às de Como começar: experimentos rápidos em várias partes do repositório MaxText. A diferença é que 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 executor 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 suas frações. É possível executar multihost_runner.py em uma máquina local ou em qualquer VM do Compute Engine no mesmo projeto que suas partiçõ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, você executa shardings.py em duas frações v5e-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:

   $ git clone https://github.com/AI-Hypercomputer/maxtext

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

   $ cd maxtext

3. Crie chaves SSH para gcloud. Recomendamos deixar uma senha em branco. Para isso, pressione Enter duas vezes depois de executar o comando a seguir. Se aparecer uma mensagem informando 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 do Multislice usando o comando queued-resources create ou o console Google Cloud .

   gcloud

   O comando a seguir mostra como criar uma TPU Multislice v5e. Para usar uma versão diferente da TPU, especifique um accelerator-type e um runtime-version diferentes.

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --node-count=${SLICE_COUNT} \
       --node-prefix=${TPU_NAME} \
       --zone=${ZONE} \
       [--reserved|--spot]

   Console

   1. No console Google Cloud , acesse a página TPUs:

      Acessar TPUs

   2. Clique em Criar TPU.

   3. No campo Nome, insira um nome para a TPU.

   4. Na caixa Zona, selecione a zona em que você quer criar a TPU.

   5. Na caixa Tipo de TPU, selecione um tipo de acelerador. O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. O Multislice só é compatível com o Cloud TPU v4 e versões posteriores. Para mais informações sobre as versões da TPU, consulte Versões da TPU.

   6. Na caixa Versão do software de TPU, selecione uma versão. Ao criar uma VM do Cloud TPU, a versão do software da TPU especifica a versão do ambiente de execução da TPU a ser instalada nas VMs da TPU. Para mais informações, consulte Versões de software da TPU.

   7. Clique no botão Ativar enfileiramento.

   8. No campo Nome do recurso na fila, digite um nome para sua solicitação de recurso na fila.

   9. Clique na caixa de seleção Usar TPU com várias frações.

   10. No campo Contagem de intervalos, insira o número de intervalos que você quer criar.

   11. Clique em Criar para criar a solicitação de recurso em fila.

6. Quando o provisionamento de recursos enfileirados começa, ele pode levar até cinco minutos para ser concluído, dependendo do tamanho do recurso. Aguarde até que o recurso enfileirado esteja no estado ACTIVE. É possível verificar o status de uma solicitação de recurso em fila usando a CLI gcloud ou o console Google Cloud :

   gcloud

   $ gcloud compute tpus queued-resources list \
       --filter=${QR_ID} --zone=${ZONE} --project=${PROJECT}

   Isso vai gerar uma saída semelhante a esta:

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

   Console

   1. No console Google Cloud , acesse a página TPUs:

      Acessar TPUs

   2. Clique na guia Recursos na fila.

   3. Clique no nome da solicitação de recurso na fila.

   Entre em contato com seu representante de conta do Google Cloud se o status do QR ficar no estadoWAITING_FOR_RESOURCESouPROVISIONINGpor mais de 15 minutos.

7. Instale as dependências.

   $ python3 multihost_runner.py \
       --TPU_PREFIX=${TPU_NAME} \
       --ZONE=${ZONE} \
       --COMMAND="bash setup.sh"

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

   $ python3 multihost_runner.py \
       --TPU_PREFIX=${TPU_NAME} \
       --ZONE=${ZONE} \
       --COMMAND="python3 -m pedagogical_examples.shardings \
       --dcn_data_parallelism ${SLICE_COUNT} \
       --ici_fsdp_parallelism 16 \
       --batch_size 131072 \
       --embedding_dimension 2048"

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

   Para mais informações sobre como configurar o paralelismo, consulte Fragmentação multislice usando paralelismo DCN e shardings.py.

9. Limpe as TPUs e o recurso na fila quando terminar. A exclusão leva de dois a cinco minutos para ser concluída. Se você estiver usando a CLI gcloud, execute esse comando em segundo plano com a flag opcional --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:

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

Essas devem ser as únicas mudanças de código necessárias ao migrar para o Multislice. Para alcançar alto desempenho, a DCN precisa ser mapeada em eixos de paralelismo de dados, paralelismo de dados totalmente fragmentados ou paralelismo de pipeline. As considerações sobre performance e as estratégias de fragmentação são discutidas em mais detalhes em Fragmentação com multislice para performance máxima.

Para validar se o código pode acessar todos os dispositivos, confirme se 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 chips por fatia * 4 fatias, então len(jax.devices()) vai retornar 32.

Como escolher tamanhos de fração para ambientes de várias fatias

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

Executar seu job em várias frações

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

1. Usando o script de execução 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 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 no README do MaxText.

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

Script do executor Production

Para jobs de produção que precisam de capacidade de recuperação contra falhas de hardware e outras interrupções, é melhor fazer a integração diretamente com a API Create Queued Resource. Use multihost_job.py como um exemplo funcional que aciona a chamada da API de recurso enfileirado criado com o script de inicialização adequado para executar o treinamento e retomar a execução após uma interrupçã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 oferece um ciclo de iteração tão rápido quanto o multihost_runner.py.

Abordagem manual

Recomendamos usar ou adaptar multihost_runner.py ou multihost_job.py para executar sua 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 multislice.

Gerenciar um ambiente Multislice

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

Criar recursos na fila

gcloud

1. Crie uma solicitação de recurso em fila usando o comando a seguir:

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --project=${PROJECT} \
       --zone=${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 cota respectiva antes de selecionar --reserved, --spot ou a cota padrão sob demanda. Para informações sobre tipos de cota, consulte Política de cotas.

curl

1. 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"
       }
       }
       ]
   }
   }

   Substitua os seguintes valores:

   • your-project-number: o Google Cloud número do projeto
   • your-zone: a zona em que você quer criar o recurso enfileirado
   • accelerator-type: a versão e o tamanho de uma única fatia. O Multislice só é compatível com o Cloud TPU v4 e versões posteriores.
   • tpu-vm-runtime-version: a versão do ambiente de execução da VM da TPU que você quer usar.
   • your-network-name: opcional, uma rede a que o recurso enfileirado será anexado
   • your-subnetwork-name: opcional, uma sub-rede a que o recurso enfileirado 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 recurso enfileirado for alocado.
   • slice-count: o número de frações de TPU no seu ambiente Multislice
   • your-queued-resource-id: o ID fornecido pelo usuário para o recurso enfileirado.

   Para mais informações, consulte a documentação da API REST de recursos enfileirados com 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 on demand padrão.

2. Envie a solicitação de criação de recurso na fila 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-queued-resource-id

   Substitua os seguintes valores:

   • your-project-id: o ID do projeto Google Cloud
   • your-zone: a zona em que você quer criar o recurso enfileirado
   • your-queued-resource-id: o ID fornecido pelo usuário para o recurso enfileirado.

Ela será parecida com o exemplo a seguir:

{
"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 GUID no final do valor da string para o atributo name e receba informações sobre a solicitação de recurso na fila.

Console

1. No console Google Cloud , acesse a página TPUs:

   Acessar TPUs

2. Clique em Criar TPU.

3. No campo Nome, insira um nome para a TPU.

4. Na caixa Zona, selecione a zona em que você quer criar a TPU.

5. Na caixa Tipo de TPU, selecione um tipo de acelerador. O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. O Multislice só é compatível com a Cloud TPU v4 e versões posteriores. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.

6. Na caixa Versão do software de TPU, selecione uma versão. Ao criar uma VM do Cloud TPU, a versão do software da TPU especifica a versão do ambiente de execução da TPU a ser instalada. Para mais informações, consulte Versões de software da TPU.

7. Clique no botão Ativar enfileiramento.

8. No campo Nome do recurso na fila, digite um nome para sua solicitação de recurso na fila.

9. Clique na caixa de seleção Usar TPU com várias frações.

10. No campo Contagem de intervalos, insira o número de intervalos que você quer criar.

11. Clique em Criar para criar a solicitação de recurso em fila.

Recuperar o status de um recurso na fila

gcloud

$ gcloud compute tpus queued-resources describe ${QR_ID} --zone=${ZONE}

Para um recurso enfileirado no estado ACTIVE, a saída é semelhante a esta:

...
state:
    state: ACTIVE
...

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}

Para um recurso enfileirado no estado ACTIVE, a saída é semelhante a esta:

{
"name": your-queued-res,
"tpu": {
    "nodeSpec": [
    {
        ... // node 1
    },
    {
        ... // node 2
    },
    ...
    ]
},
...
"state": "ACTIVE"
}

Console

1. No console Google Cloud , acesse a página TPUs:

   Acessar TPUs

2. Clique na guia Recursos na fila.

3. Clique no nome da solicitação de recurso na fila.

Depois que a TPU for provisionada, você também poderá conferir detalhes sobre a solicitação de recurso enfileirado acessando a página "TPUs", encontrando sua TPU e clicando no nome da solicitação de recurso enfileirado correspondente.

Em um cenário raro, você pode encontrar o recurso na fila no estado FAILED enquanto algumas frações estão ACTIVE. Se isso acontecer, exclua os recursos criados e tente novamente em alguns minutos ou entre em contato com o suporte doGoogle Cloud .

SSH e instalação de dependências

Executar código JAX em frações 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 seu ambiente de várias fatias por SSH e instalar dependências, use o seguinte comando gcloud:

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

Esse comando gcloud envia o comando especificado para todos os trabalhadores e nós no QR usando SSH. O comando é agrupado em lotes de quatro e enviado simultaneamente. O próximo lote de comandos é enviado quando o lote atual conclui a execução. Se um dos comandos falhar, o processamento será interrompido, e nenhum outro lote será enviado. Para mais informações, consulte a referência da API de recursos enfileirados. Se o número de slices que você está usando exceder o limite de encadeamento do computador local (também chamado de limite de lote), ocorrerá um deadlock. Por exemplo, suponha que o limite de agrupamento em lote na sua máquina local seja 64. Se você tentar executar um script de treinamento em mais de 64 intervalos, digamos 100, o comando SSH vai dividir os intervalos em lotes. Ele vai executar o script de treinamento no primeiro lote de 64 fatias e aguardar a conclusão dos scripts antes de executar o script no lote restante de 36 fatias. No entanto, o primeiro lote de 64 partes não pode ser concluído até que as 36 partes restantes comecem a executar o script, causando um deadlock.

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 flag --command. Quando você faz isso, depois de iniciar o script de treinamento no primeiro lote de intervalos, o controle retorna imediatamente para o comando SSH. O comando SSH pode começar a executar o script de treinamento no restante do lote de 36 slices. Você precisará redirecionar os fluxos stdout e stderr de maneira adequada ao executar os comandos em segundo plano. Para aumentar o paralelismo no mesmo QR, selecione intervalos específicos usando o parâmetro --node.

Configuração de rede

Para garantir que as frações de TPU possam se comunicar entre si, siga estas etapas. Instale o JAX em cada uma das frações. Para mais informações, consulte Executar o código JAX em frações de TPU. Afirme que len(jax.devices()) é igual ao número de ícones no seu ambiente de várias frações. Para fazer isso, execute o seguinte em cada fração:

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

Se você executar esse código em quatro fatias de v4-16, haverá oito chips por fatia e quatro fatias. Um total de 32 chips (dispositivos) será retornado por jax.devices().

Listar recursos na fila

gcloud

Use o comando queued-resources list para conferir o estado dos recursos enfileirados:

$ gcloud compute tpus queued-resources list --zone=${ZONE}

A saída será assim:

NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central1-a  4           v5litepod-16             ACTIVE
...

Console

1. No console Google Cloud , acesse a página TPUs:

   Acessar TPUs

2. Clique na guia Recursos na fila.

Iniciar seu job em um ambiente provisionado

É possível executar manualmente as cargas de trabalho 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 ${TPU_NAME} \
    --zone=${ZONE} \
    --worker=all \
    --command="command-to-run"

Redefinir respostas rápidas

A API ResetQueuedResource pode ser usada para redefinir todas as VMs em um ACTIVE QR. A reinicialização forçada apaga 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 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 seguinte comando:

$ gcloud compute tpus queued-resources reset ${QR_ID} --zone=${ZONE}

Excluir recursos na fila

Para liberar recursos ao final da sessão de treinamento, exclua o recurso enfileirado. A exclusão leva de dois a cinco minutos para ser concluída. Se você estiver usando a CLI gcloud, execute esse comando em segundo plano com a flag opcional --async.

gcloud

$ gcloud compute tpus queued-resources \
    delete ${QR_ID} --zone=${ZONE} --force [--async]

Console

1. No console Google Cloud , acesse a página TPUs:

   Acessar TPUs

2. Clique na guia Recursos na fila.

3. Marque a caixa de seleção ao lado da solicitação de recurso na fila.

4. Clique em deleteExcluir.

Recuperação automática de falhas

Em caso de interrupção, o Multislice oferece reparo sem intervenção da fração afetada e redefinição de todas as frações depois. A fração afetada é substituída por uma nova, e as demais frações íntegras são redefinidas. Se não houver capacidade 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 pontos de verificação salvos. O script de inicialização é executado automaticamente sempre que uma fração é realocada ou uma VM é redefinida. Você especifica um script de inicialização no payload JSON enviado à API de solicitação de criação de QR.

O script de inicialização a seguir (usado em Criar QRs) permite a recuperação automática de falhas e a retomada do treinamento de 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 -m MaxText.train 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 isso.

Criação de perfil e depuração

O perfil é o mesmo em ambientes de fatia única e multifatia. Para mais informações, consulte Como criar perfis de programas JAX.

Como otimizar o treinamento

As seções a seguir descrevem como otimizar o treinamento Multislice.

Fragmentação com Multislice para performance máxima

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

Recomendamos começar com o paralelismo de dados puro. Usar o paralelismo de dados totalmente fragmentados é útil para liberar o uso da memória. A desvantagem é que a comunicação entre as frações usa a rede DCN e diminui a velocidade da 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 funciona bem em casos em que você tem uma carga de trabalho que está sendo executada corretamente, mas quer melhorar o desempenho dela escalonando em várias partes.

Para alcançar um escalonamento eficiente em várias fatias, o tempo necessário para realizar a redução total na DCN precisa ser menor do que o tempo necessário para realizar uma transmissão para trás. A DCN é usada para comunicação entre frações e é um fator limitante no throughput 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 vai usar de 32 a 64 bits de largura de banda de DCN para cada parâmetro por etapa. Se você usar duas fatias, o modelo vai usar 32 bits de largura de banda da DCN. Se você usar mais de duas fatias, o compilador vai realizar uma operação de redução completa de embaralhamento e usará até 64 bits de largura de banda de DCN para cada parâmetro por etapa. A quantidade de FLOPS necessária 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 uma transmissão direta e uma transmissão indireta é 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 a transmissão para trás é 4 * B.

Para garantir um escalonamento eficiente em várias frações, 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 a transmissão 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 transformador, se você estiver usando duas segmentações: Operational intensity = 4 * B / 32

Se você estiver usando mais de duas fatias: Operational intensity = 4 * B/64

Isso sugere um tamanho mínimo de lote 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, é necessário estimar o tempo de execução da transmissão para trás por fração (cronometrando com um criador de perfil ou contando FLOPS). Depois, compare isso com o tempo de execução esperado para a redução total na DCN e tenha uma boa estimativa de se o paralelismo de dados faz sentido para você.

Quando usar o paralelismo de dados totalmente fragmentados (FSDP)

O paralelismo de dados totalmente fragmentados (FSDP) combina o paralelismo de dados (fragmentação dos dados entre nós) com a fragmentação dos pesos entre nós. Para cada operação nas transmissões 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 all-reduce, eles são reduce-scattered à medida que são produzidos. Dessa forma, cada fatia recebe apenas os gradientes dos pesos de que é responsável.

Assim como o paralelismo de dados, o FSDP exige o escalonamento do tamanho do lote global de forma linear com o número de slices. O FSDP vai diminuir a pressão na memória à medida que você aumenta o número de slices. Isso acontece porque o número de pesos e o estado do otimizador por fração diminuem, mas ao custo de aumento do tráfego de rede e maior possibilidade de bloqueio devido a um coletivo atrasado.

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

As operações all-gather e all-reduce no FSDP funcionam de maneira semelhante às do DP. Assim, é possível determinar se a carga de trabalho do FSDP é limitada pela performance da DCN da mesma forma descrita na seção anterior.

Quando usar o paralelismo de pipeline

O paralelismo de pipeline se torna relevante quando se busca alta performance com outras estratégias de paralelismo que exigem um tamanho de lote global maior do que o tamanho máximo de lote preferido. O paralelismo de pipeline permite que as partes que compõem um pipeline "compartilhem" um lote. No entanto, o paralelismo de pipeline tem duas desvantagens significativas:

1. Isso causa o "pipeline bubble", em que os chips ficam ociosos porque estão aguardando dados.
2. Ele exige microlotes, o que diminui 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 exigirem um tamanho de lote global muito grande. Antes de tentar o paralelismo de pipeline, vale a pena fazer um experimento para verificar empiricamente se a convergência por amostra diminui no tamanho do lote necessário para alcançar um FSDP de alta performance. O FSDP tende a alcançar uma utilização maior de FLOP do modelo, mas se a convergência por amostra diminuir à medida que o tamanho do lote aumenta, o paralelismo de pipeline ainda pode ser a melhor opção. A maioria das cargas de trabalho pode tolerar tamanhos de lote suficientemente grandes 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 o FSDP. Isso permite minimizar a profundidade do pipeline e aumentar o tamanho do lote por pipeline até que a latência da DCN se torne menos importante na capacidade de transferência. Por exemplo, se você tiver N fatias, considere pipelines de profundidade 2 e N/2 réplicas de paralelismo de dados, depois pipelines de profundidade 4 e N/4 réplicas de paralelismo de dados, continuando da mesma maneira, até que o lote por pipeline fique grande o suficiente para que os coletivos DCN possam ser ocultados por trás da aritmética na transmissão para trás. Isso vai minimizar a lentidão introduzida pelo paralelismo de pipeline, permitindo que você dimensione além do limite global de tamanho do lote.

Práticas recomendadas para multislice

As seções a seguir descrevem as práticas recomendadas para o treinamento com várias fatias.

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 e eficiente que fragmenta o lote entre hosts é importante para evitar que as TPUs fiquem sem 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 exige um refragmento no modelo. Além disso, o MaxText ainda não oferece instantâneos determinísticos, o que permitiria que o iterador de dados carregasse os mesmos dados antes e depois da substituição.

Como estabelecer pontos de verificação

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

Configurações aceitas

As seções a seguir descrevem as formas de fração, a orquestração, os frameworks e o paralelismo compatíveis com o Multislice.

Formas

Todas as divisões precisam ter o mesmo formato (por exemplo, o mesmo AcceleratorType). Formatos heterogêneos não são aceitos.

Orquestração

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

Frameworks

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 o representante da sua contaGoogle Cloud .

Suporte e feedback

Queremos saber sua opinião. Para compartilhar feedback ou pedir suporte, entre em contato usando o formulário de suporte ou feedback da Cloud TPU.