Visão geral da TPU multislice do Cloud

Importante: o Multislice só tem suporte para o Cloud TPU v4 e versões mais recentes.

O Cloud TPU multislice é uma tecnologia de escalonamento de desempenho de pilha completa que permite que um job de treinamento use várias fatias de TPU em um único pod ou em fatias em vários pods com paralelismo de dados simples. Com os chips TPU v4, isso significa que os jobs de treinamento podem usar mais de 4096 chips 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 fatias menores estão mais disponíveis, permitindo um tempo de inicialização mais rápido quando o recurso Multislice é usado com fatias menores.

Quando implantados em configurações Multislice, os chips TPU em cada fatia se comunicam por interconexões 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 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 a computação para o máximo desempenho.

Conceitos

Tipo de acelerador

A forma de cada fração de TPU que compreende uma fração múltipla. Cada pedaço em uma solicitação de vários pedaços é 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 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 fatia, a criação não será concluída até que o hardware esteja disponível. Depois que o novo slice é criado, todos os outros slices no ambiente Multislice são reiniciados 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 último ponto de verificação.

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 em grupo

Quando todos os slices de TPU são provisionados juntos, ao mesmo tempo, garantindo que todos ou nenhum dos slices sejam provisionados com sucesso.

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 sobre dados.

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 pela DCN.

Nó

No contexto de Multislice, 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 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 que é executado sempre que uma VM é inicializada ou reinicializada. Para o recurso Multislice, ele é especificado na solicitação de criação de QR. Para mais informações sobre scripts de inicialização do Cloud TPU, consulte Gerenciar recursos de 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 entre si usando a rede ICI.

VM de TPU

Uma máquina virtual que executa o Linux e tem acesso às TPUs. Para TPUs v4, cada VM de 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

As TPUs podem ser criadas com diferentes tipos de capacidade (consulte as opções de uso em Como funcionam os preços da TPU) :

• Reserva: para consumir uma reserva, você precisa ter um contrato de reserva com o Google. Use a flag --reserved ao criar seus recursos.
• Spot: define a cota preemptiva usando VMs spot. Seus recursos podem ser adiados para abrir espaço para solicitações de um job de maior prioridade. Use a flag --spot ao criar seus recursos.
• Sob demanda: define a cota sob demanda, que não precisa de uma reserva e não será substituída. 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.

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

Primeiros passos

Se você ainda não usou TPUs, comece instalando a Google Cloud CLI e configurando seu Cloud TPU Cloud. Para usar o recurso Multislice, os recursos de TPU precisam ser gerenciados como Recursos enfileirados.

Se você já é usuário da TPU v4 e tem uma reserva, talvez seja necessário migrar a reserva para um novo sistema. Para mais informações, entre em contato com o representante da conta do Google Cloud.

Exemplo introdutório

Este tutorial usa o código do repositório MaxText do GitHub. 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 maneira eficiente no Cloud TPU.

O código em shardings.py foi criado para ajudar você a começar a experimentar diferentes opções de paralelismo. Por exemplo, paralelismo de dados, paralelismo de dados totalmente fragmentado (FSDP, na sigla em inglês) e paralelismo de tensor. O código é dimensionado de uma única fatia para ambientes de várias fatias.

Paralelismo de ICI

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

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

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

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

Este exemplo pressupõe que você já sabe executar código em uma única fatia 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 fatia.

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 após executar o comando abaixo). Se uma mensagem informar 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 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 da Cloud TPU.

   node-id

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

   reserved

   Use uma reserva ao criar as fatias.

   spot

   Use 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.

   Console

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

      Acessar as TPUs

   2. Clique em Criar TPU.

   3. No campo Nome, insira um nome para o 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 aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.

   6. Na caixa Versão do software da TPU, selecione uma versão do software. 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 que será instalada. Para mais informações, consulte Imagens de VM de TPU.

   7. Clique no botão Ativar fila.

   8. No campo Nome do recurso em fila, digite um nome para a solicitação de recurso em 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 em fila começa, pode levar de um a cinco minutos para ser concluído, dependendo do tamanho do recurso em fila. É possível verificar o status de uma solicitação de recurso em fila usando a CLI gcloud ou o Console do Google Cloud:

   gcloud

   $ gcloud compute tpus queued-resources \
       list --filter=YOUR_QR_ID

   Console

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

      Acessar as TPUs

   2. Clique na guia Recursos em fila.

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

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 o JAX e outras dependências na fatia 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 no seu slice 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 com o ICI. Por exemplo, você pode tentar esta combinação:

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

10. Exclua o recurso na fila e a fatia de TPU quando terminar. Execute essas etapas de limpeza no ambiente em que você configurou a fatia. Primeiro, execute exit para sair da sessão SSH. A exclusão vai levar de dois a cinco minutos para ser concluída. Se você estiver usando a CLI gcloud, será possível executar esse comando em segundo plano com a flag --async opcional.

    Importante: os recursos em fila precisam ser excluídos após o uso. Caso contrário, eles vão continuar consumindo a cota de recursos em fila, que é monitorada separadamente da cota de TPU.

    gcloud

    $ gcloud compute tpus queued-resources
        delete YOUR_QR_ID --force (--async)

    Console

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

       Acessar as TPUs

    2. Clique na guia Recursos em fila.

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

    4. Clique em Excluir.

Fragmentação de várias fatias usando o paralelismo de 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	Nº de fatias
Paralelismo de dados bidirecional	2	1	1	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 de 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 fatias, mas isso depende do aplicativo.

À medida que você aumenta o número de fatias (assumindo que o tamanho e o lote por fatia sejam constantes), aumenta a quantidade de paralelismo de dados.

Como executar shardings.py em um ambiente Multislice

É possível executar shardings.py em um ambiente 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 do Começar a usar: experimentos rápidos em várias fatias do repositório MaxText, exceto 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 os mesmos TPUs. Como o script multihost_runner.py depende de conexões SSH de longa duração, não recomendamos o uso dele para jobs de longa duração. Se você quiser executar um job mais longo (por exemplo, horas ou dias), recomendamos usar multihost_job.py.

Neste tutorial, usamos o termo runner para indicar a máquina em que você executa o script multihost_runner.py. Usamos o termo workers para indicar as VMs de TPU que compõem suas fatias. É possível executar multihost_runner.py em uma máquina local ou em qualquer VM do Compute Engine no mesmo projeto que suas fatias. 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 fatias v4-16, um total de quatro VMs e 16 chips 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 de várias fatias usando o comando queued-resources create ou o console do Google Cloud.

   gcloud

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

   O comando a seguir mostra como criar uma TPU Multislice v4. Para usar a v5e, especifique uma accelerator-type v5e (por exemplo, v5litepod-16) e a runtime-version 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.

   accelerator-type

   O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. O suporte a várias frações está disponível apenas no Cloud TPU v4 e em versões mais recentes da TPU. 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 da 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 serão nomeadas como: mySlice-0, mySlice-1, continuando numericamente para cada fatia.

   reserved

   Use uma reserva ao criar as fatias.

   spot

   Use VMs Spot ao criar as fatias.

   Console

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

      Acessar as TPUs

   2. Clique em Criar TPU.

   3. No campo Nome, insira um nome para o 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 suporte para várias frações está disponível apenas no Cloud TPU v4 e em versões mais recentes. 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 da TPU, selecione uma versão do software. 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 que será instalada. Para mais informações, consulte Imagens de VM de TPU.

   7. Clique no botão Ativar fila.

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

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

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

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

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

   gcloud

   $ 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
   ...

   Console

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

      Acessar as TPUs

   2. Clique na guia Recursos em fila.

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

   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 as TPUs e o recurso na fila quando terminar. A exclusão vai levar 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 --async opcional.

Como dimensionar 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 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 fatias de v4-16, terá oito chips por fatia * 4 fatias, então len(jax.devices()) vai retornar 32.

Como escolher os tamanhos de fatia para ambientes com várias fatias

Para acelerar de forma linear, adicione novas fatias do mesmo tamanho que a atual. Por exemplo, se você usar uma fatia v4-512, o Multislice vai atingir aproximadamente o dobro da performance adicionando uma segunda fatia v4-512 e duplicando 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 várias fatias:

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 no README do MaxText (link em inglês).

Como o multihost_runner.py mantém conexões SSH persistentes, ele só é adequado para experimentos de tamanho modesto e relativamente curtos. É possível adaptar as etapas do tutorial multihost_runner.py à sua 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 outras preempções, é melhor fazer a integração diretamente com a API Create Queued Resource. Como exemplo, fornecemos multihost_job.py, que aciona a chamada da API Created Queued Resource com o script de inicialização adequado para executar o treinamento e retomar a preempção. O script multihost_job.py está documentado no README do MaxText (em inglês).

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

Abordagem manual

Recomendamos usar ou adaptar 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 seu ambiente usando comandos de QR diretamente, consulte Gerenciar um ambiente com várias fatias.

Gerenciar um ambiente Multislice

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

Criar recursos na fila

gcloud

1. Defina as seguintes variáveis de ambiente 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

   Descrições das variáveis

   Entrada	Descrição
   YOUR_QR_ID	O ID atribuído pelo usuário do recurso em fila.
   PROJETO	Nome do projeto do Google 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 fatias.
   STARTUP_SCRIPT	Se você especificar um script de inicialização, ele será executado quando a fatia de TPU for provisionada ou reiniciada.

2. Crie uma solicitação de recurso em fila usando o seguinte comando:

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

   Descrições de sinalizações de comando

   YOUR_QR_ID

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

   project

   O projeto do Google Cloud em que você está criando a solicitação de recurso em fila.

   zone

   A zona do Google Cloud em que o recurso em fila será criado.

   node-count

   O número de fatias a serem criadas.

   accelerator-type

   O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. O suporte a várias frações está disponível apenas no Cloud TPU v4 e em versões mais recentes da TPU. 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 da Cloud TPU.

   network

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

   subnetwork

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

   reserved

   Use uma reserva ao criar as fatias.

   spot

   Use VMs Spot ao criar as fatias.

Confira se você tem a cota respectiva antes de selecionar --reserved, --spot ou a cota sob demanda padrão. 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 número do seu projeto do Google Cloud
   • your-zone: a zona em que você quer criar o recurso em fila
   • accelerator-type: a versão e o tamanho de uma única fatia. O recurso Multislice tem suporte apenas para o Cloud TPU v4 e versões mais recentes.
   • 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 em fila será anexado.
   • your-subnetwork-name: opcional, uma sub-rede à qual o recurso em fila 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 em fila 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 recurso em fila

   Para mais informações, consulte a documentação da API REST de recursos em fila para conferir todas as opções disponíveis.

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

   Para usar a capacidade do Spot, substitua:

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

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

2. Envie a solicitação de criação de recurso em 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_QR_ID

   Substitua os seguintes valores:

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

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 para receber informações sobre a solicitação de recurso em fila.

Console

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

   Acessar TPUs

2. Clique em Criar TPU.

3. No campo Nome, insira um nome para o 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 suporte para várias frações está disponível apenas no Cloud TPU v4 e versões mais recentes. 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 da TPU, selecione uma versão do software. 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 que será instalada. Para mais informações, consulte Imagens de VM de TPU.

7. Clique no botão Ativar fila.

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

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

10. No campo Número de fatias, insira o número de fatias 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 ${YOUR_QR_ID}

Para um recurso em fila que está no estado ACTIVE, a saída será 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 em fila que está no estado ACTIVE, a saída será esta:

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

Console

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

   Acessar as TPUs

2. Clique na guia Recursos em fila.

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

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

Em um cenário raro, seu recurso em fila pode estar no estado FAILED, enquanto algumas fatias estão ACTIVE. Se isso acontecer, exclua os recursos criados e tente novamente em alguns minutos ou entre em contato com o Suporte do Google Cloud.

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

Esse comando gcloud envia o comando especificado para todos os workers e nós no QR usando SSH. O comando é agrupado 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 em fila. Se o número de fatias que você está usando exceder o limite de linhas de execução do computador local (também chamado de limite de lote), você vai encontrar um deadlock. Como exemplo, suponha que o limite de lote na sua máquina local seja 64. Se você tentar executar um script de treinamento em mais de 64 fatias, digamos 100, o comando SSH vai dividir as fatias 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 fatias não pode ser concluído até que as 36 fatias restantes comecem a executar o script, causando um bloqueio.

Para evitar esse cenário, execute o script de treinamento em segundo plano em cada VM anexando um caractere & (&) ao comando do script especificado com a flag --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 o paralelismo no mesmo código QR, selecione fatias específicas usando o parâmetro --node.

Configuração de 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 recursos na fila

gcloud

É possível conferir o estado dos recursos em fila usando o comando queued-resources list:

$ gcloud compute tpus queued-resources list

A saída será assim:

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

Console

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

   Acessar as TPUs

2. Clique na guia Recursos em fila.

Iniciar o trabalho em um ambiente provisionado

É possível executar manualmente os workloads conectando-se a todos os hosts em cada fatia por SSH e executando o comando abaixo 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 permanecem intactos, e o script de inicialização é 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 YOUR_QR_ID

Como excluir recursos na fila

Para liberar os recursos no final da sessão de treinamento, exclua o recurso na fila. A exclusão vai levar de dois a cinco minutos. Se você estiver usando a CLI gcloud, será possível executar esse comando em segundo plano com a flag --async opcional.

gcloud

$ gcloud compute tpus queued-resources \
    delete YOUR_QR_ID --force (--async)

Console

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

   Acessar as TPUs

2. Clique na guia Recursos em fila.

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

4. Clique em Excluir.

Recuperação automática de falhas

Em caso de interrupção, o Multislice oferece reparo sem intervenção do segmento afetado e redefinição de todos os segmentos depois. O segmento afetado é substituído por um novo, e os segmentos restantes são redefinidos. 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 a partir de pontos de controle armazenados em um bucket do Cloud Storage durante o treinamento do MaxText:

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

Clone o repositório MaxText antes de tentar 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

Fragmentação com o recurso Multislice para desempenho máximo

Para alcançar o desempenho máximo em ambientes Multislice, é necessário considerar como fragmentar as várias fatias. Normalmente, há três opções (paralelismo de dados, paralelismo de dados totalmente fragmentado e paralelismo de pipeline). Não recomendamos o particionamento 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. 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 de dados puro funciona bem nos casos em que você tem uma carga de trabalho que está funcionando bem, mas gostaria de melhorar a performance escalonando em várias fatias.

Para alcançar uma escalação forte em várias fatias, o tempo necessário para realizar a redução total em DCN precisa ser menor do que o tempo necessário para realizar um passe reverso. A DCN é usada para a comunicação entre as frações e é um fator limitante no throughput da carga de trabalho.

Cada chip 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 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 de DCN. Se você usar mais de duas fatias, o compilador vai realizar uma operação de redução total de shuffle e você vai 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 transformador, o número de FLOPS necessário para uma transmissão direta e reversa é 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 passagem reversa é 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 a transmissão 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 fatias: 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 significativa para erros, implantando o paralelismo de dados somente se o tamanho do lote por pod for de pelo menos 350 mil (dois pods) a 700 mil (muitos pods).

Para outras arquiteturas de modelo, você precisará estimar o tempo de execução da transmissão reversa por fatia (cronometrando-a usando um perfilador ou contando FLOPS). Em seguida, compare esse tempo com o tempo de execução esperado para reduzir o tempo de execução da DCN e ter uma boa estimativa de se o paralelismo de dados faz sentido para você.

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

O paralelismo de dados totalmente fragmentado (FSDP, na sigla em inglês) combina o paralelismo de dados (fragmentação 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 reunidos para que cada fatia tenha os pesos necessários. Em vez de sincronizar os gradientes usando a redução total, os gradientes são espalhados de forma reduzida conforme são 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 vai diminuir a pressão sobre a memória à medida que 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 de um aumento no tráfego de rede e uma 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. Assim, 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 do pipeline se torna relevante ao alcançar um alto desempenho com outras estratégias de paralelismo que exigem um tamanho de lote global maior que o tamanho máximo preferido. 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 gera o "balão do pipeline", em que os chips ficam inativos porque estão aguardando dados.
2. Ele exige 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. A FSDP tende a alcançar uma utilização de FLOP de modelo mais alta, 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 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 o paralelismo de dados ou o FSDP. Isso vai permitir que você minimize a profundidade do pipeline e aumente o tamanho do lote por pipeline até que a latência do DCN se torne menos um fator na capacidade. Especificamente, 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 e assim por diante, até que o lote por pipeline fique grande o suficiente para que os coletivos de DCN possam ser ocultados pela aritmética no passo reverso. 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 o modelo. Ter um carregador de dados eficiente e assíncrono que fragmenta o lote entre hosts é importante para evitar a falta de trabalho nas TPUs. 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 que o iterador de dados carregasse os mesmos dados antes e depois da preempção.

Como estabelecer pontos de verificação

A biblioteca de pontos de verificação Orbax fornece primitivas para pontos de verificação de PyTrees JAX no armazenamento local ou no Google Cloud Storage. 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 conosco usando o formulário de suporte ou feedback do Cloud TPU.