Visão geral de multislices do Cloud TPU

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

Quando implantados em configurações de multislices, os chips de TPU em cada fração se comunicarem pela Inter-chip-interconnect (ICI). Os chips de TPU em diferentes frações se comunicam transferindo dados para CPUs (hosts), que transmitem os dados pela rede de data center (DCN).

Os desenvolvedores não precisam escrever código para implementar a comunicação DCN entre fatias. O compilador XLA gera esse código para você e sobrepõe a comunicação com de computação para o melhor desempenho.

Conceitos

Tipo de acelerador

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

Reparo automático

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

Conjunto de dados

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

Rede de data center (DCN)

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

Programação de gangues

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

Host

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

Inferência

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

Interconexão de interconexão (ICI)

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

Multislice (em inglês)

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

Nó

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

Pod

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

Recurso em fila (QR)

Uma representação dos recursos da TPU, usada para enfileirar e gerenciar a solicitação de um em um ambiente de TPU de fração única ou Multislice.

Script de inicialização

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

Fração da TPU

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

VM da TPU

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

Tensor (link em inglês)

Uma estrutura de dados usada para representar dados multidimensionais em uma máquina modelo de machine learning.

Unidade de Processamento de Tensor (TPU)

Chip de aceleração de ML desenvolvido internamente pelo Google. Eles são projetados para oferecem computação rápida e eficiente para tarefas essenciais de machine learning, como multiplicação de matrizes.

Tipos de capacidade do Cloud TPU

TPUs podem ser criadas a partir de diferentes tipos de capacidade (consulte Opções de uso em Como funcionam os preços da TPU:

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

.

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

Começar

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

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

Exemplo introdutório

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

O código em shardings.py foi desenvolvido para ajudar você a começar a testar diferentes modelos de paralelismo . Por exemplo, paralelismo de dados, paralelismo de dados totalmente fragmentados (FSDP), e paralelismo de tensores. O código dimensiona de uma única fração para uma fração e ambientes de teste.

paralelismo da ICI

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

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

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

Essas entradas precisam ser restritas para que ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism é igual a com o número de ícones na fatia.

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

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

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

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

1. Configure o ambiente:

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

2. Crie chaves SSH para gcloud. Recomendamos deixar uma senha em branco (pressione digite duas vezes depois de executar o comando a seguir). Se aparecer a mensagem O arquivo google_compute_engine já existe, substitua a versão atual.

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

3. Provisione suas TPUs com o seguinte comando:

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

   Descrições de sinalizações de comando

   your-qr-id

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

   accelerator-type

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

   runtime-version

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

   node-id

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

   reserved

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

   best-effort

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

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

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

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

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

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

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

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

   $ bash setup.sh
   

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

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

   Confira os resultados nos registros. Suas TPUs devem alcançar cerca de 260 TFLOP por segundo ou uma utilização impressionante de mais de 90%do FLOP! Nesse caso, aproximadamente o lote máximo que se encaixa na Memória de largura de banda (HBM).

9. Fique à vontade para conhecer outras estratégias de fragmentação em vez de ICI. Por exemplo, é possível tentar a seguinte combinação:

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

10. Exclua a fração de QR e TPU quando terminar. Execute estas etapas de limpeza etapas do ambiente em que você configurou a fração (primeiro execute exit para saia da sessão SSH). A exclusão levará de dois a cinco minutos, e pode ser executado em segundo plano com a sinalização --async opcional.

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

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

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

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

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

Os valores desses parâmetros precisam ser restritos dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism é igual a o número de fatias.

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

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

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

Conforme você aumenta o número de fatias (supondo que você mantenha o tamanho da fatia e o lote por fatia), você aumenta a quantidade de paralelismo de dados.

Como executar shardings.py em um ambiente multislice

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

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

Neste tutorial, usamos o termo executor para indicar a máquina em que você execute o script multihost_runner.py. Usamos o termo trabalhadores para indicar VMs de TPU que compõem suas frações. É possível executar o multihost_runner.py em um máquina virtual ou qualquer VM do Compute Engine no mesmo projeto das frações. Em execução multihost_runner.py em um worker não é compatível.

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

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

Configurar o ambiente

1. Clone MaxText no seu executor máquina virtual.

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

3. Crie chaves SSH para gcloud. Recomendamos deixar uma senha em branco (pressione digite duas vezes depois de executar o comando a seguir). Se aparecer a mensagem o arquivo google_compute_engine já existe, selecione para não manter seu versão atual.

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

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

     $ export SLICE_COUNT=2
     

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

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

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

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

   Descrições de sinalizações de comando

   your-qr-id

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

   accelerator-type

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

   runtime-version

   A versão do software do Cloud TPU.

   node-count

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

   node-prefix

   O prefixo usado para gerar nomes para cada fatia. Um número é anexado ao prefixo de cada fatia. Por exemplo, se você definir node-prefix como mySlice, as fatias são nomeadas: mySlice-0, mySlice-1 e assim por diante.

   reserved

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

   best-effort

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

6. Pode levar até cinco minutos para que o provisionamento de QR code seja iniciado. dependendo do tamanho do QR code. Aguardar até que o recurso em fila (QR, na sigla em inglês) entre o estado ACTIVE. É possível verificar o status de uma solicitação de QR code usando o seguinte comando:

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

   Isso vai gerar uma saída semelhante a esta:

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

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

7. Instale as dependências.

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

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

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

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

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

Como escalonar uma carga de trabalho para multislices

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

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

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

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

Como escolher tamanhos de fatia para ambientes multislice

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

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

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

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

Script do executor do experimento

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

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

Script do executor Production

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

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

Abordagem manual

Recomendamos que você use ou adapte multihost_runner.py ou multihost_job.py para executar a carga de trabalho personalizada a 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 QRs de forma manual sem usar as ferramentas fornecido no repositório MaxText, leia o nas seções a seguir.

Criar QR codes

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

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

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

Criar uma solicitação de QR code usando gcloud

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

Descrições de sinalizações de comando

your-qr-id

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

project

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

zone

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

node-count

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

accelerator-type

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

runtime-version

A versão do software do Cloud TPU.

network

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

subnetwork

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

reserved

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

spot

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

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

Criar uma solicitação de QR code usando curl

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

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

Observação :substitua os seguintes valores:

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

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

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

Para usar a capacidade do Spot, substitua:

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

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

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

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

Observação :substitua os seguintes valores:

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

A resposta será semelhante a esta:

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

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

Recuperar o status de um QR code

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

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

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

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

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

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

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

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

Em casos raros, seu QR code pode ser encontrado no estado FAILED, enquanto alguns frações são ACTIVE. Se isso acontecer, exclua os recursos criados e tente de novo em alguns minutos ou entre em contato para a equipe do Cloud TPU resolver o problema.

SSH e instalar dependências

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

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

Este comando gcloud envia o comando especificado para todos os workers e nós em QR usando SSH. O comando é reunido em grupos de quatro e enviado ao mesmo tempo. O próximo lote de comandos é enviado quando o lote atual conclui a execução. Se houver falha em um dos comandos, o processamento paradas e nenhum outro lote é enviado. Para mais informações, consulte a Referência da API de recursos em fila. Se o número de frações que você está usando exceder a linha de execução do seu computador local (também chamado de limite de lotes), você encontrará um impasse. Por exemplo, presume que o limite de lotes na sua máquina local é de 64. Se você tentar executar uma script de treinamento em mais de 64 frações, digamos 100, o comando SSH quebrará as fatias em lotes. Ele vai executar o script de treinamento no primeiro lote de 64 frações e aguardar a conclusão dos scripts antes de executar o script na lote restante de 36 frações. No entanto, o primeiro lote de 64 frações não pode concluir até que as 36 frações restantes comecem a executar o script, o que resulta um impasse.

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

Configuração da rede

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

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

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

Listar QRs

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

$ gcloud compute tpus queued-resources list

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

Descrever QRs

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

Utilizando gcloud:

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

Utilizando curl:

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

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

Iniciar o job em um ambiente provisionado

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

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

Redefinindo QRs

A API ResetQueuedResource pode ser usada para redefinir todas as VMs em um QR code ACTIVE. Redefinir as VMs à força apaga a memória de a máquina e a redefine para o estado inicial dela. Todos os dados armazenados localmente permaneçam intactos e o script de inicialização será invocado após uma redefinição. O A API ResetQueuedResource pode ser útil quando você quer reiniciar todas as TPUs. Para exemplo, quando o treinamento trava e redefinir todas as VMs é mais fácil do que depurar.

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

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

Excluindo QRs

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

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

Recuperação automática de falhas

Em caso de interrupção, o Multislice oferece serviços sem intervenção reparação da fração afetada e redefinição de todas as frações posteriormente. Os afetados fração é substituída por uma nova e as fatias restantes saudáveis são redefinidas. Se nenhuma capacidade for para alocar uma fração de substituição, o treinamento é interrompido.

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

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

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

Clone o repositório MaxText antes de tentar para fora.

Criação de perfil e depuração

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

Como otimizar o treinamento

Como fragmentar com o multislice para o desempenho máximo

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

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

Quando usar o paralelismo de dados

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

Para alcançar um escalonamento forte em vários frações, o tempo necessário para realizar a redução total em DCN precisa ser menor que o tempo necessário para dar um passe para trás. O DCN é usado para comunicação entre frações e é um fator limitante na capacidade de processamento 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 máxima de rede de 50 Gbps.

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

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

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

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

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

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

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

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

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

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

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

Assim como o paralelismo de dados, o FSDP exige o escalonamento do tamanho do lote global linearmente com o número de fatias. O FSDP diminui a pressão sobre a memória você aumenta o número de fatias. Isso ocorre porque o número de pesos e o estado do otimizador por fração diminui, mas o faz ao preço de um aumento o tráfego de rede e a maior possibilidade de bloqueio devido a um atraso coletivo.

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

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

Quando usar o paralelismo de pipeline

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

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

O paralelismo de pipeline só deve ser usado se as outras estratégias de paralelismo que exigem um tamanho de lote global muito grande. Antes de tentarmos o paralelismo de pipeline, testar empiricamente se a convergência por amostra diminuir o tamanho do lote necessário para alcançar um FSDP de alto desempenho. o FSDP tende a alcançar maior utilização de FLOP do modelo, mas se a convergência por amostra diminuir à medida que aumentar o tamanho do lote, o paralelismo de pipeline ainda pode ser a melhor escolha. Mais frequentes as cargas de trabalho tolerem tamanhos de lote grandes o suficiente para não se beneficiarem paralelismo de pipeline, mas sua carga de trabalho pode ser diferente.

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

Práticas recomendadas de multislices

Carregamento de dados

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

Como estabelecer pontos de verificação

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

Configurações aceitas

Formas

Todas as fatias precisam ter o mesmo formato (por exemplo, o mesmo AcceleratorType). Não há suporte para formatos heterogêneos de fração.

Orquestração

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

Frameworks

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

Paralelismo

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

Suporte e feedback

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