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 4096 ícones, uma única fatia 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.
- Reserva: segmenta a cota reservada. Para usar a cota reservada, você precisa ter um
de reserva com o Google. Use a sinalização
Primeiros passos
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.
Configure o ambiente:
$ gcloud auth login $ gcloud config set project your-project-id $ gcloud config set compute/zone your-zone
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 arquivogoogle_compute_engine
já existe, substitua a versão atual.$ ssh-keygen -f ~/.ssh/google_compute_engine
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.
Aguarde até que o QR esteja no estado
ACTIVE
, o que significa que os nós de trabalho estão no estadoREADY
. 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
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
Clone o MaxText (que inclui
shardings.py
) na VM da TPU.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
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).
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
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.$ 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
Clone MaxText no seu executor máquina virtual.
Acesse o diretório do repositório.
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 arquivogoogle_compute_engine
já existe, selecione para não manter seu versão atual.$ ssh-keygen -f ~/.ssh/google_compute_engine
Adicione uma variável de ambiente para definir a contagem de frações de TPU como
2
.$ export SLICE_COUNT=2
Crie um ambiente multislice usando
queued-resources create
.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 osruntime-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
comomySlice
, 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].
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
ouPROVISIONING
para mais de 15 minutos.Instale as dependências.
$ python3 multihost_runner.py \ --TPU_PREFIX=your-qr-id \ --COMMAND="bash setup.sh"
Execute
shardings.py
em cada worker usandomultihost_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. .
Limpe as TPUs e o QR code quando terminar. A exclusão vai levar de duas a cinco minutos para concluir e pode ser executado em segundo plano com o comando sinalização
--async
.
Como escalonar uma carga de trabalho para multislice
Antes de executar o modelo em um ambiente 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:
- Usando o script do executor de experimentos,
multihost_runner.py
- Usando o script do executor de produção,
multihost_job.py
- 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" } } ] } }
- 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.
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
- 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
- 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. A
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 TPU v4 tem um desempenho máximo de 275 * 1012 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 * 1012 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:
- Isso resulta no "bolha do pipeline" em que os chips estão inativos porque estão esperando para dados.
- 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.