Esta página foi traduzida pela API Cloud Translation.

Treinamento do Cloud TPU v5p

A Cloud TPU v5p é a TPU de quinta geração do Google Cloud e a sucessora da TPU v4. A v5p é otimizada para treinamento em grande escala e para ser uma plataforma líder no desenvolvimento de LLMs básicos, modelos de difusão e IA generativa. De modo geral, a v5p oferece até o dobro do desempenho da v4, além de agrupar o dobro de TPUs em um pod (6 mil fatias maiores em comparação com 3 mil na v4), resultando em até quatro vezes mais desempenho no nível do pod. Ele também é executado em uma frequência de relógio mais alta (1,75 GHz x 1,05 GHz), adiciona SparseCore para incorporações em grande escala e triplica a capacidade da memória de alta largura de banda (HBM).

Conceitos do Cloud TPU v5p

Se você não tem experiência com Cloud TPUs, confira a página inicial da documentação da TPU.

Os conceitos do Cloud TPU (por exemplo, fatias, 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 de acelerador específicos para treinamento ou inferência. Esses tipos de acelerador são descritos em configurações v5p.

Gerenciar recursos de TPU

Todos os comandos neste documento pressupõem que você está criando uma VM da TPU v5p. Para mais informações sobre os comandos para criar VMs de TPU, consulte Gerenciar TPUs ou Guia do usuário de recursos em fila para gerenciar recursos em fila. Para facilitar a execução dos comandos, os exemplos de código neste documento usam as seguintes variáveis de ambiente:

export PROJECT_ID=your-project
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export TPU_NAME=your-tpu-name

Descrições das variáveis de ambiente

PROJECT_ID: O projeto do Google Cloud em que você está criando a TPU.
ACCELERATOR_TYPE: O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. Para mais informações sobre os tipos de acelerador compatíveis com cada versão de TPU, consulte Versões de TPU.
ZONE: A zona em que você planeja criar a Cloud TPU.
RUNTIME_VERSION: A versão do software da TPU.
TPU_NAME: O nome definido pelo usuário da TPU que você está usando.

Configuração do framework

Esta seção descreve o processo de configuração geral para treinamento de modelos usando JAX ou PyTorch com TPU v5p.

Configuração para JAX

Se você tiver formas de fatia maiores que 4 chips, terá várias VMs em uma fatia. Nesse caso, você precisa usar a flag --worker=all para executar a instalação em todas as VMs de TPU com um único comando:

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

É possível executar o comando a seguir para verificar o número de dispositivos. As saídas mostradas aqui foram produzidas com uma fatia v5p-32. Esse código testa se tudo está instalado corretamente verificando se o JAX encontra 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 chips na faixa especificada. jax.local_device_count() indica a contagem de chips acessíveis por uma única VM nesta fatia.

# 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 de 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())"'

Teste os tutoriais do JAX neste documento para começar a treinar o v5p usando o JAX.

Configuração para PyTorch

O ambiente de execução do PJRT é o único ambiente de execução com suporte para a v5p, e o PyTorch 2.1+ usa o PJRT como o ambiente de execução padrão para todas as versões de TPU. Esta seção descreve 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
pip install numpy
pip install torch torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html
pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
'

Use um script Python com PJRT para validar a instalação. O script mostra os dispositivos TPU disponíveis (as saídas mostradas aqui foram produzidas com uma fatia 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 de 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\"))"
'

Teste os tutoriais do PyTorch neste documento para começar a treinar o v5p usando o PyTorch.

Monitorar e criar perfil

O Cloud TPU v5p oferece suporte a monitoramento e criação de perfil usando os mesmos métodos das gerações anteriores do Cloud TPU. Leia Crie o perfil do modelo com as ferramentas do Cloud TPU para saber mais sobre a criação de perfil e Monitorar VMs do Cloud TPU para saber mais sobre monitoramento.

Tutoriais de treinamento

Esta seção se concentra em tutoriais de treinamento para uma única fatia. Para adaptar esses tutoriais ao treinamento de várias fatias, adicione a flag --node=all aos comandos SSH. Para detalhes e práticas recomendadas, consulte a Introdução às fatias múltiplas.

Diffusion 2.1
MaxText
ResNet em um único host v5p
ResNet em um v5p com vários hosts

Tutoriais do JAX

Train Diffusion 2.1

Este tutorial mostra como treinar o modelo de difusão estável do HuggingFace usando o conjunto de dados Pokémon no Cloud TPU v5p.

O modelo Stable Diffusion é um modelo latente de texto para imagem que gera imagens fotorrealistas a partir de qualquer entrada de texto. Para saber mais, acesse os recursos a seguir (links em inglês):

Configurar

Crie variáveis de ambiente:

export GCS_BUCKET_NAME=your-bucket
export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5p-32
export ZONE=europe-west4-b
export LOCATION=europe-west4
export RUNTIME_VERSION=v2-alpha-tpuv5
export SERVICE_ACCOUNT=your-service-account
export TPU_NAME=your-tpu-name
export QUEUED_RESOURCE_ID=your-qr-name
export QUOTA_TYPE=spot
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 Versões de TPU para sua versão de TPU.
ZONA	Consulte o documento Regiões e zonas de TPU para saber quais zonas são compatíveis.
LOCAL	A região do Google Cloud em que o bucket do Cloud Storage será criado.
RUNTIME_VERSION	Para a v5p, use v2-alpha-tpuv5 para a RUNTIME_VERSION.
SERVICE_ACCOUNT	Esse é o endereço da sua conta de serviço, que pode ser encontrado em "Console do Google Cloud -> IAM -> Contas de serviço". Por exemplo: tpu-service-account@myprojectID.iam.gserviceaccount.com
TPU_NAME	O ID de texto atribuído pelo usuário do TPU, que é criado quando a solicitação de recurso em fila é alocada.
QUEUED_RESOURCE_ID	O ID de texto atribuído pelo usuário da solicitação de recurso em fila. Consulte o documento Recursos em fila para informações sobre recursos em fila.
QUOTA_TYPE	Pode ser `reserved` ou `spot`. Se nenhum deles for especificado, o padrão de QUOTA_TYPE será `on-demand`. Consulte cotas para informações sobre os diferentes tipos de cotas compatíveis com a Cloud TPU.
VALID_UNTIL_DURATION	O período em que a solicitação é válida. Consulte Recursos em fila para informações sobre as diferentes durações válidas.

Configure um bucket de armazenamento para a saída do modelo.

gcloud storage buckets create gs://$GCS_BUCKET_NAME \
 --project=$PROJECT_ID \
 --location=$LOCATION

Crie um recurso de TPU:

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

Você poderá usar o SSH na VM da TPU quando o recurso enfileirado estiver no estado ACTIVE. Para verificar o estado do recurso na fila, execute 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

Treine o modelo

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --project $PROJECT_ID --worker=all --command="
git clone https://github.com/google/maxdiffusion
cd maxdiffusion
git reset --hard 57629bcf4fa32fe5a57096b60b09f41f2fa5c35d # This identifies the GitHub commit to use.
pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html # Install the latest version of JAX
pip3 install -r requirements.txt
pip3 install .
export LIBTPU_INIT_ARGS=""
python -m src.maxdiffusion.models.train src/maxdiffusion/configs/base_2_base.yml run_name=my_run base_output_directory=gs://$GCS_BUCKET_NAME enable_profiler=False"

Limpar

Exclua a TPU e a solicitação de recurso na fila no final da sessão ou remova as solicitações de recurso na fila que estão no estado "FALHA". 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 em 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 taxa de transferência.

	v5p-8	v5p-32	v5p-128
Etapa de treinamento	150	150	150
Tamanho global do lote	32	64	64
Capacidade (exemplos/s)	12.10	18.08	19.10

MaxText

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

O MaxText é um LLM de alto desempenho, escalonável de forma arbitrária, de código aberto e bem testado, escrito em Python/JAX puro para Cloud TPUs. O MaxText oferece aos pesquisadores e desenvolvedores 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 seu ambiente do Cloud TPU.

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 a TPU.
ACCELERATOR_TYPE	Consulte a página Versões de TPU para sua versão de TPU.
ZONA	Consulte o documento Regiões e zonas de TPU para conferir as zonas compatíveis.
RUNTIME_VERSION	Para a v5p, use v2-alpha-tpuv5 para a versão do ambiente de execução.
RUN_NAME	Nome da 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 de várias fatias e quiser uma performance de rede ideal, considere criar uma rede dedicada com uma unidade máxima de transmissão (MTU) de 8.896 bytes e configurar as regras de firewall adequadas. Embora opcional, essa etapa pode melhorar significativamente o desempenho, especialmente ao aumentar o número de fatias na rede do data center (DCN, na sigla em inglês). 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}

Clonar o repositório MaxText

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

Treine 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 TPUs do Cloud e a instalação de dependências até a execução do modelo e a remoçã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"

Depois de iniciar o script, uma mensagem semelhante a esta vai aparecer no registro. O local do registro é referenciado na mensagem de saída. Clique no primeiro link para acessar os registros de todos os workers depois que o provisionamento do 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.

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 flags a seguir são necessárias apenas 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 também --node-id ao criar o recurso TPU na próxima etapa. Um erro será gerado se todas as três flags forem especificadas.
Criar um recurso de TPU.
```
gcloud 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 flag QUOTA_TYPE pode ser reserved ou spot. Se nenhum deles for especificado, o padrão será on-demand. Consulte cotas para informações sobre os diferentes tipos de cotas com suporte da Cloud TPU.

Você poderá se conectar às VMs da TPU usando SSH assim que o QueuedResource estiver no estado ACTIVE:

Use o comando describe para consultar o status do recurso em fila.
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}
--project ${PROJECT_ID} --zone ${ZONE}
```
Quando o recurso em fila estiver no estado ACTIVE, a saída será semelhante a esta:
```
 state: ACTIVE
```

Conectar-se ao TPU usando SSH

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

Instalar dependências

Observação: depois de usar ssh para se conectar à VM de TPU, é necessário definir as variáveis ${TPU_NAME} e ${MAXTEXT_OUTPUT_PATH} para usá-las na VM de 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'
```
Execute o modelo com vários scripts de configuração, como 32b.sh e 64b.sh. Se você estiver executando o script em uma VM TPU, adicione a flag --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 do MaxText foi executado de 32B a 1160B com precisão bf16. Os resultados dessas execuções são mostrados na tabela a seguir.

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

O modelo de parâmetro de 256B foi testado no v5p-512 e no v5p-1024 usando pesos bf16 e int8. A tabela a seguir mostra 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 do modelo (MFU)	66,98%	88,85%	67,09%	90,23%

Tutoriais do TensorFlow

Treinar o ResNet em um único host v5p

Este tutorial descreve como treinar o ImageNet em um TPU v5p-8 usando um conjunto de dados falso. Se você quiser usar um conjunto de dados diferente, consulte Como preparar o conjunto de dados.

Configurar

Crie variáveis de ambiente:

export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5p-32
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.18.0-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.

Crie um recurso de TPU:

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

Você poderá se conectar à VM da TPU usando SSH assim que o recurso em fila estiver no estado ACTIVE. Para verificar o estado do recurso em fila, use o seguinte comando:

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

Conectar-se ao TPU usando SSH

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

Definir 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

Mude para o diretório do repositório de modelos e instale os requisitos.

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

Treine o modelo

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 ResNet em um v5p com vários hosts

Este tutorial descreve como treinar o ImageNet em v5p-16 ou mais usando um conjunto de dados falso. Se você quiser usar um conjunto de dados diferente, consulte Como preparar o conjunto de dados.

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.18.0-pod-pjrt
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

ACCELERATOR_TYPE pode ser v5p-16 ou maior.

Crie um recurso de TPU:

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

Você poderá se conectar à VM da TPU usando SSH assim que o recurso em fila estiver no estado ACTIVE.

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

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

Conectar-se ao TPU (worker zero) usando SSH

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

Definir 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

Mude para o diretório do repositório de modelos e instale os requisitos.

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

Treine o modelo

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

Este tutorial vai abordar como treinar o modelo Llama 2 7B na v5p usando um fork do repositório HuggingFace no 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

Crie variáveis de ambiente.

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

Criar um recurso de TPU
```
gcloud 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 flag QUOTA_TYPE pode ser reserved ou spot. Se nenhum deles for especificado, o padrão será on-demand. Consulte cotas para informações sobre os diferentes tipos de cotas com suporte da Cloud TPU.

Você poderá se conectar à VM da TPU usando SSH assim que o QueuedResource estiver no estado ACTIVE:

Use o comando describe para consultar o status do recurso em fila.
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} \
--zone ${ZONE}
```
Quando o recurso em fila estiver no estado ACTIVE, a saída será semelhante a esta:
```
 state: ACTIVE
```
Observação: se você estiver executando em um ambiente de várias fatias, 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 TPU sejam usados corretamente.

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 torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

Faça o download do repositório do HuggingFace e instale os requisitos.

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'

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"

Treine 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 com várias fatias, defina a flag --spmd_dcn_parallelism como o número de fatias.

O SPMD_USER_GUIDE fornece um guia do usuário mais detalhado que explica todas as diferentes variáveis e alternâncias de ambiente do script de HF. 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 de todos os três tamanhos de modelo Llama 2 está incluída na tabela abaixo.

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

Suporte e feedback

Queremos saber sua opinião. Para compartilhar feedback ou solicitar suporte, preencha o formulário de suporte ou feedback da Cloud TPU.