Treinamento do Cloud TPU v5p

O Cloud TPU v5p é o Cloud TPU de quinta geração do Google Cloud e o sucessor da TPU v4. O v5p é otimizado para treinamento em grande escala e é uma plataforma líder no desenvolvimento de LLMs básicos, modelos de difusão e IA generativa. De modo geral, a v5p oferece até duas vezes mais desempenho do que a v4, além de empacotar duas vezes mais TPUs em um pod (6 mil fatias em comparação com 3 mil na v4), resultando em um desempenho de até quatro vezes no nível do pod. Ele também é executado a uma frequência de relógio maior (1,75 Ghz em vez de 1,05 Ghz), adiciona o SparseCore para embeddings de grande escala e triplica a capacidade de memória de alta largura de banda (HBM).

Conceitos do Cloud TPU v5p

Se você ainda não conhece Cloud TPUs, confira a página inicial da documentação da TPU.

Os conceitos do Cloud TPU (por exemplo, frações, hosts e TensorCores) e a arquitetura do sistema do Cloud TPU para todas as versões do Cloud TPU são descritos na página Arquitetura do sistema do Cloud TPU.

Cada versão do Cloud TPU requer tipos específicos de aceleradores para treinamento ou inferência. Esses tipos de aceleradores são descritos nas Configurações da v5p.

Gerenciar recursos da TPU

Todos os comandos que podem ser usados para gerenciar suas VMs de TPU estão descritos no Como gerenciar TPUs ou no Guia do usuário de recursos em fila para gerenciar recursos na fila.

Configuração do framework

Nesta seção, descrevemos o processo geral de configuração do treinamento de modelo usando JAX ou PyTorch com TPU v5p.

Configuração para JAX

Se você tiver formas de fração com mais de quatro ícones, terá várias VMs em uma fração. Nesse caso, é necessário usar a sinalização --worker=all para executar a instalação em todas as VMs da TPU usando um único comando:

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Execute o comando a seguir para verificar o número de dispositivos (as saídas mostradas aqui foram produzidas com uma fração v5p-32). Esse código testa se tudo está instalado corretamente, verificando se o JAX vê os TensorCores do Cloud TPU e pode executar operações básicas:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

A saída será semelhante a esta:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() mostra o número total de ícones na fatia especificada. jax.local_device_count() indica a contagem de ícones acessíveis por uma única VM nessa fração.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

A saída será semelhante a esta:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

Use --node=all para executar o comando em todos os workers multislice.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Veja os tutoriais do JAX neste documento para começar a usar o treinamento v5p.

Configuração do PyTorch

O ambiente de execução PJRT é o único compatível com a v5p, e o PyTorch 2.1 ou mais recente usa a PJRT como o ambiente de execução padrão para todas as versões de TPU. Nesta seção, descrevemos como começar a usar o PJRT em pods v5p com o PyTorch/XLA 2.2.0 para todos os workers.

Instalar dependências

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='
sudo apt-get update
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

Use um script Python com PJRT para validar a instalação e mostrar os dispositivos TPU disponíveis. As saídas mostradas aqui foram produzidas com uma fração v5p-32.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} --zone ${ZONE} --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']

Use --node=all para executar o comando em todos os workers multislice.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

Consulte os tutoriais do PyTorch neste documento para começar a usar o treinamento v5p.

Monitorar e criar perfil

O Cloud TPU v5p é compatível com monitoramento e criação de perfil usando os mesmos métodos das gerações anteriores do Cloud TPU. Para saber mais sobre monitoramento, leia Criar perfil do modelo com ferramentas do Cloud TPU para saber mais sobre criação de perfil e monitoramento de VMs do Cloud TPU.

Tutoriais de treinamento

Esta seção se concentra em tutoriais de treinamento para uma única fração. É possível adaptar esses tutoriais ao treinamento de multislices adicionando a sinalização --node=all aos comandos SSH. Para detalhes e práticas recomendadas, consulte Introdução ao multislice.

Observação: antes de executar qualquer um destes tutoriais, você precisa configurar seu ambiente do Cloud TPU. Depois de configurar o ambiente de TPU, é possível executar os tutoriais a seguir.

• Diffusion 2.1 (link em inglês)
• MaxText
• ResNet em um único host v5p
• ResNet em um v5p de vários hosts

Tutoriais do JAX

Train Diffusion 2.1 (link em inglês)

Neste tutorial, mostramos como treinar o modelo de difusão estável da HuggingFace usando o conjunto de dados Pokémon no Cloud TPU v5p.

O modelo de difusão estável é um modelo latente de texto para imagem que gera imagens fotorrealistas de qualquer entrada de texto. Para saber mais, acesse os recursos a seguir (links em inglês):

• Visão geral da Difusão estável
• Fonte do código de difusão estável

Observação: antes de executar este tutorial, você precisa configurar seu ambiente do Cloud TPU.

Configurar

1. Crie variáveis de ambiente:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5p-32
   export ZONE=us-east5-a
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

   Descrições de sinalizações de comando

   Variável	Descrição
   PROJECT_ID	Nome do projeto do Google Cloud
   ACCELERATOR_TYPE	Consulte a página de versões de TPU para saber qual é sua versão de TPU.
   ZONA	Consulte o documento de regiões e zonas da TPU para conhecer as zonas compatíveis.
   RUNTIME_VERSION	Para v5p, use v2-alpha-tpuv5 para RUNTIME_VERSION.
   SERVICE_ACCOUNT	Esse é o endereço da sua conta de serviço que pode ser encontrado no console do Google Cloud -> IAM -> Contas de serviço. Por exemplo: [e-mail]
   TPU_NAME	O ID de texto atribuído pelo usuário da TPU, que é criado quando a solicitação de recurso na fila é alocada.
   QUEUED_RESOURCE_ID	O ID de texto atribuído pelo usuário da solicitação de recurso na fila. Consulte o documento Recursos na fila para ver informações sobre esses recursos.
   QUOTA_TYPE	Pode ser reserved ou best-effort. Se nenhum desses valores for especificado, on-demand será usado como padrão. Acesse quotas para saber mais sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.
   VALID_UNTIL_DURATION	Por quanto tempo a solicitação é válida. Consulte Recursos em fila para ver informações sobre as diferentes durações válidas.

2. Crie um recurso de TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --valid-until-duration ${VALID_UNTIL_DURATION} \
   --service-account ${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível estabelecer uma conexão SSH com sua VM da TPU quando o recurso na fila estiver no estado ACTIVE. Verifique o estado do seu recurso na fila executando o seguinte comando:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project ${PROJECT_ID} --zone ${ZONE}
   

   Quando o recurso na fila estiver no estado ACTIVE, a saída será semelhante a esta:

   state: ACTIVE
   

3. Instale o JAX e as dependências dele.

   # compatible with v5p: only jax version 0.4.19 and later \
   # jax 0.4.19 requires py 3.10 \
   
   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} --zone=${ZONE} --worker=all \
   --command='pip install "jax[tpu]==0.4.20" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Faça o download do repositório HuggingFace e os requisitos de instalação.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/huggingface/diffusers.git && cd diffusers && pip install . && pip install tensorflow clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
   

5. Treinar o modelo

   Treine o modelo com um buffer pré-mapeado de 4 GB.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='export PATH=$PATH:$HOME/.local/bin && cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=stabilityai/stable-diffusion-2-1 --dataset_name=lambdalabs/pokemon-blip-captions --resolution=256 --center_crop --random_flip --train_batch_size=1 --mixed_precision=bf16 --max_train_steps=150 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model --from_pt'
   

Limpar

Exclua a TPU e a solicitação de recurso na fila no final da sessão ou para remover as solicitações de recursos na fila que estão no estado "FAILED". Para excluir um recurso na fila, exclua as fatias e, em seguida, a solicitação de recurso na fila em duas etapas:

   gcloud compute tpus tpu-vm delete ${TPU_NAME} --project=${PROJECT_ID}
   --zone=${ZONE} --quiet

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
   --project ${PROJECT_ID} --zone ${ZONE} --quiet

Ou use --force para excluir as fatias e a solicitação de recurso na fila em uma única etapa:

# With --force
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
--project ${PROJECT_ID} --zone ${ZONE} --quiet --force

Resultados da comparação

O script de treinamento do Stable Diffusion foi executado nas versões v5p-8, v5p-32 e v5p-128. A tabela a seguir mostra a capacidade de processamento.

	v5p-8	v5p-32	v5p-128
Degrau de trem	150	150	150
Tamanho global do lote	32	64	64
Capacidade de processamento (exemplos/segundo)	12.10	18/08	19h10

MaxText

Neste tutorial, mostramos como treinar o modelo MaxText usando um conjunto de dados sintético no Cloud TPU.

O MaxText é um LLM bem testado, de alto desempenho, escalonável arbitrariamente e de código aberto, escrito em Python/JAX puro e voltado para Cloud TPUs. A MaxText capacita pesquisadores e desenvolvedores com uma ferramenta acessível e adaptável para avançar as fronteiras da pesquisa e do desenvolvimento de processamento de linguagem natural (PLN).

Antes de executar este tutorial, você precisa configurar o ambiente do Cloud TPU.

1. Configure as variáveis de ambiente

   export PROJECT_ID=your_project_ID
   export TPU_NAME=your_tpu_name # user defined TPU name
   export ACCELERATOR_TYPE=v5p-256
   export ZONE=us-east5-a
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export RUN_NAME=your_experiment_run_name # user defined name for this run
   export GCS_BUCKET_NAME=your_bucket_name # Output cloud folder. Should start with gs://
   export MAXTEXT_OUTPUT_PATH=${GCS_BUCKET_NAME}/your_experiment_output_path
   export NUM_SLICES=1 # Update the value to a number >1 for Multislice.
   

   Descrições de sinalizações de comando

   Variável	Descrição
   PROJECT_ID	Nome do projeto do Google Cloud
   TPU_NAME	Um nome definido pelo usuário para sua TPU.
   ACCELERATOR_TYPE	Consulte a página de versões de TPU para saber qual é sua versão de TPU.
   ZONA	Consulte o documento de regiões e zonas da TPU para conhecer as zonas compatíveis.
   RUNTIME_VERSION	Para v5p, use v2-alpha-tpuv5 para a versão do ambiente de execução.
   RUN_NAME	Nome de execução do experimento fornecido pelo usuário.

   Configuração opcional recomendada para o multislice:

   export NETWORK_NAME=your_network_name
   export FIREWALL_RULE_NAME=your_firewall_rule_name
   

   Se você estiver executando cargas de trabalho Multislice e quiser o desempenho de rede ideal, crie uma rede dedicada com uma unidade máxima de transmissão (MTU, na sigla em inglês) de 8.896 bytes e configure regras de firewall adequadas. Embora opcional, essa etapa pode melhorar significativamente o desempenho, especialmente ao aumentar o número de frações na rede de data center (DCN, na sigla em inglês). Observe que a criação de uma rede requer a permissão compute.networks.create no projeto. Os exemplos a seguir mostram como criar uma rede dedicada e uma regra de firewall.

   Crie uma rede dedicada:

   gcloud compute networks create ${NETWORK_NAME} \
   --mtu=8896 \
   --project=${PROJECT_ID} \
   --subnet-mode=auto \
   --bgp-routing-mode=regional
   

   Criar uma regra de firewall

   gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
   --network ${NETWORK_NAME} --allow tcp,icmp,udp --project=${PROJECT_ID}
   

2. Clonar o repositório MaxText

   git clone https://github.com/google/maxtext.git
   

3. Treinar o modelo

   As seções a seguir descrevem duas opções para treinar o MaxText.

   Opção 1

   Se você quiser que um script gerencie todo o fluxo de trabalho, desde o provisionamento de Cloud TPUs e a instalação de dependências até a execução do modelo e a eliminação de recursos, use multihost_job.py.

   cd maxtext && python3 multihost_job.py --PROJECT=${PROJECT_ID} --ZONE=${ZONE} \
   --NUM_SLICES=${NUM_SLICES} --TPU_TYPE=${ACCELERATOR_TYPE} \
   --VERSION=${RUNTIME_VERSION} --RUN_NAME=${RUN_NAME} #user defined run name \
   --BUCKET_NAME=${GCS_BUCKET_NAME} \ #used to store logs and configs
   --COMMAND="bash setup.sh && bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME} OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"
   

   Observação :adicione --CQR_EXTRA_ARGS="--network=${NETWORK_NAME}" se você estiver usando uma rede personalizada.

   Depois de iniciar o script, você verá uma mensagem semelhante à seguinte no registro: O local do registro é referenciado na mensagem de saída. Clique no primeiro link para acessar os registros de todos os workers quando o provisionamento de TPU for concluído.

   ------------------------------------
   
   multihost_job finished running, TPUs are starting up to run your job remotely.
   
   Logs for your job are displayed here:
   https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22_log%22%2529;?project=PROJECT_ID
   
   To see the output of a single host, you may edit the slice and worker
   number in the `log_file_path` property here:
   
   https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22RUN_NAME_log%22%2529%20AND%0Alabels.%22agent.googleapis.com%2Flog_file_path%22%3D%20%22%2FRUN_NAME%2Fmain_command_log_slice_0_worker_0%22;?project=PROJECT_ID
   
   When your job is finished, the main command log is in your Cloud Storage
   bucket:
   https://console.cloud.google.com/storage/browser/YOUR_BUCKET_NAME/RUN_NAME?project=PROJECT_ID
   
   View the status of the created TPUs using:
   gcloud compute tpus queued-resources list --filter=RUN_NAME
   --zone=ZONE --project=PROJECT_ID
   

Opção 2

Para executar o script de treinamento várias vezes em um Cloud TPU provisionado, use o script multihost_runner.py para usar o recurso.

1. Configure variáveis para criar uma TPU.

   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export VALID_DURATION=1d
   export QUOTA_TYPE=quota_type
   

   Observação: as sinalizações a seguir só serão necessárias se você estiver usando o Multislice:

   --node-count ${NODE_COUNT} \
   --node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID
   

   Observação: se você especificar --node-count e --node-prefix, não especifique --node-id ao criar o recurso de TPU na etapa a seguir. Um erro será gerado se todas as três sinalizações forem especificadas.

2. Criar um recurso de TPU.

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --valid-until-duration ${VALID_DURATION} \
   --service-account ${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível se conectar às VMs da TPU usando SSH quando QueuedResource estiver no estado ACTIVE:

   Use o comando describe para consultar o status do recurso na fila.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}
   --project ${PROJECT_ID} --zone ${ZONE}
   

   Quando o recurso na fila está no estado ACTIVE, a saída é semelhante a esta:

    state: ACTIVE
   

3. Conectar-se à TPU usando SSH

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
     --project ${PROJECT_ID} \
     --zone ${ZONE}
   

4. Instalar dependências

   Observação: depois de usar ssh para se conectar à VM da TPU, é necessário definir as variáveis ${TPU_NAME} e ${MAXTEXT_OUTPUT_PATH} para usá-las na VM da TPU:

   export TPU_NAME=your_tpu_name
   export MAXTEXT_OUTPUT_PATH=output-path
   

   cd maxtext && python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
   --COMMAND='bash setup.sh'
   

5. Execute o modelo com vários scripts de configuração, como 32b.sh, 64b.sh. Se você estiver executando o script de uma VM da TPU, será necessário adicionar a sinalização --INTERNAL_IP=true.

   python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
   --COMMAND="bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME}
   OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"
   

   Observação: defina --TPU_PREFIX=${NODE_PREFIX} para Multislice em vez de --TPU_PREFIX=${TPU_NAME}.

Limpar

Exclua a TPU e os recursos na fila.

Resultados da comparação

O script de treinamento MaxText foi executado de 32B para 1160B com precisão de bf16. Os resultados dessas execuções são mostrados na tabela a seguir.

No de parâmetros	Tipo de acelerador	TFLOP/chip/s	Utilização de flops de modelo (MFU)
32 bi	v5p-128	3,28E+02	71,47%
64 bi	v5p-128	3,23E+02	70,31%
128 bi	v5p-256	3,15E+02	68,68%
128 bi	v5p-512	3,15E+02	68,53%
256 bi	v5p-1024	3,16E+02	68,82%
512 bi	v5p-1024	2,94E+02	63,99%
1.024 bi	v5p-2048	2,49E+02	64,05%
1.024 bi	v5p-4096	2,97E+02	64,80%
1.160 bi	v5p-7680	2,95E+02	64,27%
1.160 bi	v5p-12288	3,04E+02	66,23%

O modelo de parâmetros 256B foi testado em v5p-512 e v5p-1024 usando os pesos bf16 e int8. A tabela a seguir exibe os resultados desses testes.

	v5p-512	v5p-512	v5p-1024	v5p-1024
Tamanho global do lote (tokens)	5,24E+05	5,24E+05	1,05E+06	1,05E+06
Precisão	bf16	int8	bf16	int8
TFLOP/chip/s	307	408	308	414
Utilização de flops de modelo (MFU)	66,98%	88,85%	67,09%	90,23%

Tutoriais do TensorFlow

Treinar o ResNet em um único host v5p

Neste tutorial, descrevemos como treinar o ImageNet em uma TPU v5p-8 usando um conjunto de dados fictício. Se você quiser usar um conjunto de dados diferente, consulte Como preparar o conjunto de dados.

Observação: antes de executar este tutorial, é preciso configurar o ambiente do Cloud TPU.

Configurar

1. Crie variáveis de ambiente:

   export PROJECT_ID=your-project-ID
   export ACCELERATOR_TYPE=v5p-8
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.16.1-pjrt
   export TPU_NAME=your-tpu-name
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type
   

   Neste tutorial, use v5p-8 como ACCELERATOR_TYPE.

2. Crie um recurso de TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
     --node-id ${TPU_NAME} \
     --project ${PROJECT_ID} \
     --zone ${ZONE} \
     --accelerator-type ${ACCELERATOR_TYPE} \
     --runtime-version ${RUNTIME_VERSION} \
     --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível se conectar à VM da TPU usando SSH quando o recurso na fila estiver no estado ACTIVE. Para verificar o estado do recurso na fila, use o seguinte comando:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
     --project ${PROJECT_ID} \
     --zone ${ZONE}
   

3. Conectar-se à TPU usando SSH

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
     --project ${PROJECT_ID} \
     --zone ${ZONE}
   

4. Defina algumas variáveis de ambiente

   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
   export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
   

5. Altere para o diretório do repositório de modelos e os requisitos de instalação.

   cd ${MODELS_REPO} && git checkout r2.15.0
   pip install -r official/requirements.txt
   

Treinar o modelo

1. Execute o script de treinamento.

   python3 official/vision/train.py \
     --tpu=local \
     --experiment=resnet_imagenet \
     --mode=train_and_eval \
     --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
     --model_dir=${MODEL_DIR} \
     --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"
   

Limpar

Exclua a TPU e os recursos na fila.

Treinar o ResNet em um v5p de vários hosts

Neste tutorial, descrevemos como treinar o ImageNet em v5p-16 ou maior usando um conjunto de dados fictício. Se você quiser usar um conjunto de dados diferente, consulte Como preparar o conjunto de dados.

Observação: antes de executar este tutorial, é preciso configurar o ambiente do Cloud TPU.

1. Crie variáveis de ambiente:

   export PROJECT_ID=your_project_ID
   export TPU_NAME=your_tpu_name
   export ZONE=us-east1-c
   export ACCELERATOR_TYPE=v5p-16
   export RUNTIME_VERSION=tpu-vm-tf-2.16.1-pod-pjrt
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type
   

   ACCELERATOR_TYPE pode ser v5p-16 ou maior.

2. Crie um recurso de TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
     --node-id ${TPU_NAME} \
     --project ${PROJECT_ID} \
     --zone ${ZONE} \
     --accelerator-type ${ACCELERATOR_TYPE} \
     --runtime-version ${RUNTIME_VERSION} \
     --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível se conectar à VM da TPU usando SSH quando o recurso na fila estiver no estado ACTIVE.

   Use o comando describe para consultar o status do recurso na fila:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
     --project ${PROJECT_ID} \
     --zone ${ZONE}
   

3. Conecte-se à TPU (worker zero) usando SSH.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
     --project ${PROJECT_ID} \
     --zone ${ZONE}
   

4. Defina algumas variáveis de ambiente

   export TPU_NAME=your_tpu_name
   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export TPU_LOAD_LIBRARY=0
   

5. Altere para o diretório do repositório de modelos e os requisitos de instalação.

   cd $MODELS_REPO && git checkout r2.15.0
   pip install -r official/requirements.txt
   

Treinar o modelo

1. Execute o script de treinamento.

   python3 official/vision/train.py \
     --tpu=${TPU_NAME} \
     --experiment=resnet_imagenet \
     --mode=train_and_eval \
     --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
     --model_dir=${MODEL_DIR} \
     --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"
   

Limpar

Exclua a TPU e os recursos na fila.

PyTorch/XLA

Llama 2

Neste tutorial, abordaremos como treinar o modelo Llama 2 7B na v5p usando uma bifurcação do repositório HuggingFace em PyTorch/XLA com paralelização geral e escalonável para gráficos de computação de ML (GSPMD, na sigla em inglês).

Configuração

1. Crie variáveis para ID do projeto, tipo de acelerador, zona, versão do ambiente de execução e nome da TPU.

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5p-8
   export ZONE=us-east5-a
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_DURATION=1d
   

2. Criar um recurso de TPU

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --valid-until-duration ${VALID_DURATION} \
   --service-account ${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}
   

   Observação: a sinalização QUOTA_TYPE pode ser reserved ou best-effort. Se nenhum deles for especificado, o padrão será on-demand. Acesse quotas para mais informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.

   Será possível se conectar à VM da TPU usando SSH quando QueuedResource estiver no estado ACTIVE:

   Use o comando describe para consultar o status do recurso na fila.

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project ${PROJECT_ID} \
   --zone ${ZONE}
   

   Quando o recurso na fila estiver no estado ATIVO, a saída será semelhante a esta:

    state: ACTIVE
   

   Observação: se você estiver executando em um ambiente de várias parcelas, use gcloud compute tpus queued-resources ssh --node=all ... em vez de gcloud compute tpus tpu-vm ssh ... para seus comandos. Isso garante que todos os chips de TPU sejam utilizados corretamente.

3. Instale o Pytorch/XLA e as dependências necessárias.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project ${PROJECT_ID} \
   --zone  ${ZONE} \
   --worker=all \
   --command='
   sudo apt-get update
   sudo apt-get install libopenblas-dev -y
   pip3 install numpy
   pip3 install typing-extensions
   pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
   '
   

4. Faça o download do repositório HuggingFace e os requisitos de instalação.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
   git clone -b llama2-google-next-training https://github.com/pytorch-tpu/transformers.git
   cd transformers
   pip3 install git+file://$PWD
   pip3 install datasets accelerate evaluate scikit-learn'
   

5. Faça o download da configuração do modelo 7B.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command="curl https://huggingface.co/TheBloke/Llama-2-7B-fp16/raw/main/config.json --output ~/config.json"
   

6. Treinar o modelo

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
   export PJRT_DEVICE=TPU
   export XLA_USE_BF16=1
   export XLA_IR_DEBUG=1
   export XLA_HLO_DEBUG=1
   
   export LIBTPU_INIT_ARGS="--xla_enable_async_collective_permute=true
   --xla_tpu_enable_async_collective_fusion_multiple_steps=true
   --xla_tpu_enable_async_collective_fusion=true
   --xla_tpu_overlap_compute_collective_tc=true
   --xla_enable_async_all_gather=true
   --xla_jf_spmd_threshold_for_windowed_einsum_mib=0"
   
   export PROFILE_EPOCH=0
   export PROFILE_STEP=3
   export PROFILE_DURATION_MS=20000
   export PROFILE_LOGDIR=/tmp/home/
   
   cd transformers
   python examples/pytorch/language-modeling/run_clm.py \
    --tokenizer_name hf-internal-testing/llama-tokenizer \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --per_device_train_batch_size 96 \
    --per_device_eval_batch_size 8 \
    --num_train_epochs 1 \
    --do_train \
    --output_dir /tmp/output \
    --overwrite_output_dir \
    --config_name ~/config.json \
    --save_strategy no \
    --logging_strategy no \
    --remove_unused_columns no \
    --optim adafactor \
    --torch_dtype bfloat16 \
    --dataloader_drop_last yes \
    --block_size 2048 \
    --spmd_2d_sharding 1 \
    --spmd_grad_chkpt
   '
   

Se você estiver executando em um ambiente de várias fatias, defina a flag --spmd_dcn_parallelism como o número de frações.

O SPMD_USER_GUIDE fornece um guia do usuário mais detalhado que explica todas as diferentes variáveis de ambiente e alternâncias do script HF. Observação: o LIBTPU_INIT_ARGS será incorporado ao PyTorch/XLA e ativado por padrão em versões futuras.

Limpar

Exclua a TPU e os recursos na fila.

Resultados da comparação

A capacidade de processamento dos três tamanhos de modelo do Llama 2 está incluída na tabela abaixo.

	v5p-8	v5p-128	v5p-128
Tamanho do modelo	7 bi	13 bi	70 bi
Tamanho global do lote	96	1024	128
Como fragmentar o formato da malha	(4 e 1)	(64 e 1)	(16 e 4)
Utilização de flops de modelo (MFU)	56,67%	55,80%	51,85%

Suporte e feedback

Seu feedback é muito bem-vindo! Para compartilhar feedback ou solicitar suporte, preencha o formulário de suporte ou feedback do Cloud TPU.