Treinamento do Cloud TPU v5e

O Cloud TPU v5e é o acelerador de IA de última geração do Google Cloud. Com uma menor pegada de 256 chips por pod, o v5e é otimizado para ser o produto de maior valor para transformador, treinamento, ajuste fino e veiculação de rede neural convolucional (CNN, na sigla em inglês) de texto para imagem. Para mais informações sobre como usar o Cloud TPU v5e para veiculação, consulte Inferência usando o v5e.

Para mais informações sobre o hardware e as configurações do Cloud TPU v5e, consulte TPU v5e.

Primeiros passos

As seções a seguir descrevem como começar a usar a TPU v5e.

Solicitação de cotas

É necessário ter cota para usar a TPU v5e no treinamento. Há diferentes tipos de cota para TPUs sob demanda, TPUs reservadas e VMs Spot da TPU. Há cotas separadas necessárias se você estiver usando a TPU v5e para inferência. Para mais informações sobre cotas, consulte Cotas. Para solicitar a cota da TPU v5e, entre em contato com a equipe de vendas do Google Cloud.

Criar uma conta e um projeto do Google Cloud

Você precisa de uma conta e um projeto do Google Cloud para usar o Cloud TPU. Para mais informações, consulte Configurar um ambiente do Cloud TPU.

Criar uma Cloud TPU

A prática recomendada é provisionar Cloud TPU v5 como recursos em fila usando o comando queued-resource create. Para mais informações, consulte Gerenciar recursos em fila.

Você também pode usar a API Create Node (gcloud compute tpus tpu-vm create) para provisionar v5s da Cloud TPU. Para mais informações, consulte Gerenciar recursos de TPU.

Para mais informações sobre as configurações v5e disponíveis para treinamento, consulte Tipos de Cloud TPU v5e para treinamento.

Configuração do framework

Esta seção descreve o processo de configuração geral para treinamento de modelo personalizados usando JAX ou PyTorch com TPU v5e. O suporte ao TensorFlow está disponível nas versões do ambiente de execução do TPU tpu-vm-tf-2.18.0-pjrt e tpu-vm-tf-2.18.0-pod-pjrt.

Para instruções de configuração de inferência, consulte Introdução à inferência v5e.

Configuração para JAX

Se você tiver formas de fatia maiores que 8 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 da TPU em uma única etapa sem usar o SSH para fazer login em cada uma delas:

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'

Descrições de sinalizações de comando

Variável Descrição
TPU_NAME O ID de texto atribuído pelo usuário do TPU, que é criado quando a solicitação de recurso em fila é alocada.
PROJECT_ID Nome do projeto do Google Cloud. Use um projeto existente ou crie um novo em Configurar seu projeto do Google Cloud.
ZONA Consulte o documento Regiões e zonas de TPU para conferir as zonas compatíveis.
worker A VM de TPU que tem acesso às TPUs subjacentes.

Você pode executar o comando abaixo para verificar o número de dispositivos. As saídas mostradas aqui foram produzidas com uma fatia v5litepod-16. 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 ícones na fração. 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.]

Teste os tutoriais do JAX neste documento para começar a treinamento da v5e usando o JAX.

Configuração para PyTorch

A v5e oferece suporte apenas ao ambiente de execução do PJRT, e o PyTorch 2.1 e versões mais recentes vão usar 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 na v5e com o PyTorch/XLA com comandos 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 -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

Se você receber um erro ao instalar as rodas para torch, torch_xla ou torchvision, como pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222, desinstale a versão com este comando:

pip3 install setuptools==62.1.0

Executar um script com o PJRT

unset LD_PRELOAD

Confira a seguir um exemplo de uso de um script Python para fazer um cálculo em uma VM v5e:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker all \
   --command='
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      unset LD_PRELOAD
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'

Isso gera um resultado semelhante ao seguinte:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

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

Exclua a TPU e o recurso na fila no fim da sessão. Para excluir um recurso em fila, exclua a fatia e, em seguida, o recurso em 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

Essas duas etapas também podem ser usadas para remover solicitações de recursos em fila que estão no estado FAILED.

Exemplos de JAX/FLAX

As seções a seguir descrevem exemplos de como treinar modelos JAX e FLAX na TPU v5e.

Treinar o ImageNet na v5e

Este tutorial descreve como treinar o ImageNet na v5e usando dados de entrada falsos. Se você quiser usar dados reais, consulte o arquivo README no GitHub.

Configurar

  1. Crie variáveis de ambiente:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
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_UNTIL_DURATION=1d

  2. 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á fazer SSH na VM da TPU quando o recurso enfileirado estiver no estado ACTIVE:

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

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

     state: ACTIVE

  3. Instale a versão mais recente do JAX e do jaxlib:

    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'

  4. Clone o modelo ImageNet e instale os requisitos correspondentes:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'

  5. Para gerar dados falsos, o modelo precisa de informações sobre as dimensões do conjunto de dados. Isso pode ser coletado dos metadados do conjunto de dados do ImageNet:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'

Treine o modelo

Depois de concluir as etapas anteriores, você pode treinar o modelo.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no fim da sessão.

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

Modelos Hugging Face FLAX

Os modelos Hugging Face implementados no FLAX funcionam imediatamente na Cloud TPU v5e. Esta seção fornece instruções para executar modelos populares.

Treinar o ViT no Imagenette

Neste tutorial, mostramos como treinar o modelo Vision Transformer (ViT) do HuggingFace usando o conjunto de dados Imagenette do Fast AI no Cloud TPU v5e.

O modelo ViT foi o primeiro a treinar um codificador Transformer no ImageNet com resultados excelentes em comparação com as redes convolucionais. Para mais informações, consulte os seguintes recursos:

Configurar

  1. Crie variáveis de ambiente:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
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_UNTIL_DURATION=1d

  2. 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 do TPU quando o recurso enfileirado estiver no estado ACTIVE:

    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

  3. Instale o JAX e a biblioteca dele:

    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'

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

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.18.0 && pip install -r examples/flax/vision/requirements.txt'

  5. Faça o download do conjunto de dados Imagenette:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'

Treine 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='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no fim da sessão.

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

Resultados do comparativo de mercado da ViT

O script de treinamento foi executado em v5litepod-4, v5litepod-16 e v5litepod-64. A tabela a seguir mostra as taxas de transferência com diferentes tipos de acelerador.

Tipo de acelerador v5litepod-4 v5litepod-16 v5litepod-64
Período 3 3 3
Tamanho global do lote 32 128 512
Capacidade (exemplos/s) 263,40 429.34 470,71

Treinar a difusão em Pokémon

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

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

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

    gcloud storage buckets create gs://your_bucket 
    
--project=your_project 
    
--location=us-west1 
    
export GCS_BUCKET_NAME=your_bucket

  2. Criar variáveis de ambiente

    export GCS_BUCKET_NAME=your_bucket
export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west1-c
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
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
    GCS_BUCKET_NAME Exibida no console do Google Cloud -> Cloud Storage -> Buckets
    PROJECT_ID Nome do projeto do Google Cloud. Use um projeto existente ou crie um novo em Configurar seu 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 conferir as zonas compatíveis.
    RUNTIME_VERSION 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.

  3. 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á fazer SSH na VM da TPU quando o recurso enfileirado estiver no estado ACTIVE:

    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

  4. Instale o JAX e a biblioteca dele.

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

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

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

Treine o modelo

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

   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
   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=your_run base_output_directory=gs://${GCS_BUCKET_NAME}/ enable_profiler=False"

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no fim da sessão.

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

Resultados de comparação de mercado para difusão

O script de treinamento foi executado em v5litepod-4, v5litepod-16 e v5litepod-64. A tabela a seguir mostra as capacidades de processamento.

Tipo de acelerador v5litepod-4 v5litepod-16 v5litepod-64
Etapa de treinamento 1500 1500 1500
Tamanho global do lote 32 64 128
Capacidade (exemplos/s) 36,53 43,71 49,36

Treinar o GPT2 no conjunto de dados OSCAR

Este tutorial mostra como treinar o modelo GPT2 do HuggingFace usando o conjunto de dados OSCAR na Cloud TPU v5e.

O GPT2 é um modelo de transformador pré-treinado em textos brutos sem rotulagem humana. Ele foi treinado para prever a próxima palavra em frases. Para saber mais, acesse os recursos a seguir:

Configurar

  1. Crie variáveis de ambiente:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
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

  2. 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 em fila estiver no estado ACTIVE:

    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

  3. Instale o JAX e a biblioteca dele.

    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'

  4. 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 https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install TensorFlow && pip install -r examples/flax/language-modeling/requirements.txt'

  5. Faça o download das configurações para treinar o modelo.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt . --recursive'

Treine 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='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no fim da sessão.

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

Resultados de comparativos para a GPT2

O script de treinamento foi executado em v5litepod-4, v5litepod-16 e v5litepod-64. A tabela a seguir mostra as capacidades de processamento.

v5litepod-4 v5litepod-16 v5litepod-64
Período 3 3 3
Tamanho global do lote 64 64 64
Capacidade (exemplos/s) 74,60 72,97 72,62

PyTorch/XLA

As seções a seguir descrevem exemplos de como treinar modelos PyTorch/XLA na TPU v5e.

Treinar o ResNet usando o ambiente de execução do PJRT

O PyTorch/XLA está migrando do XRT para o PjRt do PyTorch 2.0+. Confira as instruções atualizadas para configurar a v5e para cargas de trabalho de treinamento do PyTorch/XLA.

Configurar

  1. Crie variáveis de ambiente:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. 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á fazer SSH na VM da TPU quando o QueuedResource estiver no estado ACTIVE:

    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. Instalar dependências específicas do Torch/XLA

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
Treinar o modelo ResNet
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      date
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export XLA_USE_BF16=1
      export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone https://github.com/pytorch/xla.git
      cd xla/
      git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no fim da sessão.

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
Resultado da comparação

A tabela a seguir mostra as capacidades de processamento do comparativo de mercado.

Tipo de acelerador Capacidade (exemplos/segundo)
v5litepod-4 4.240 ex/s
v5litepod-16 10.810 ex/s
v5litepod-64 46.154 ex/s

Treinar GPT2 na v5e

Este tutorial vai abordar como executar o GPT2 na v5e usando o repositório do HuggingFace no PyTorch/XLA usando o conjunto de dados wikitext.

Configurar

  1. Crie variáveis de ambiente:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
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

  2. 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á fazer SSH na VM da TPU quando o QueuedResource estiver no estado ACTIVE:

    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

  3. Instale as dependências do PyTorch/XLA.

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

  4. 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 https://github.com/pytorch/xla.git
      pip install --upgrade accelerate
      git clone https://github.com/huggingface/transformers.git
      cd transformers
      git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
      pip install .
      git log -1
      pip install datasets evaluate scikit-learn
      '

  5. Faça o download das configurações do modelo pré-treinado.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/ --recursive
      gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'

Treine o modelo

Treine o modelo 2B usando um tamanho de lote de 16.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      cd transformers/
      export LD_LIBRARY_PATH=/usr/local/lib/
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      python3 examples/pytorch/xla_spawn.py \
         --num_cores=4 \
         examples/pytorch/language-modeling/run_clm.py \
         --num_train_epochs=3 \
         --dataset_name=wikitext \
         --dataset_config_name=wikitext-2-raw-v1 \
         --per_device_train_batch_size=16 \
         --per_device_eval_batch_size=16 \
         --do_train \
         --do_eval \
         --logging_dir=./tensorboard-metrics \
         --cache_dir=./cache_dir \
         --output_dir=/tmp/test-clm \
         --overwrite_output_dir \
         --cache_dir=/tmp \
         --config_name=examples/pytorch/language-modeling/my_config_2.json \
         --tokenizer_name=gpt2 \
         --block_size=1024 \
         --optim=adafactor \
         --adafactor=true \
         --save_strategy=no \
         --logging_strategy=no \
         --fsdp=full_shard \
         --fsdp_config=examples/pytorch/language-modeling/fsdp_config.json'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no fim da sessão.

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

Resultado da comparação

O script de treinamento foi executado em v5litepod-4, v5litepod-16 e v5litepod-64. A tabela a seguir mostra as taxas de transferência de referência para diferentes tipos de acelerador.

v5litepod-4 v5litepod-16 v5litepod-64
Período 3 3 3
config 600 milhões 2B 16B
Tamanho global do lote 64 128 256
Capacidade (exemplos/s) 66 77 31

Treinar ViT na v5e

Este tutorial vai abordar como executar o VIT na v5e usando o repositório do HuggingFace no PyTorch/XLA no conjunto de dados cifar10.

Configurar

  1. Crie variáveis de ambiente:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  2. 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á fazer SSH na VM da TPU quando o QueuedResource estiver no estado ACTIVE:

     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

  3. Instalar dependências do PyTorch/XLA

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

  4. 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 https://github.com/suexu1025/transformers.git vittransformers; \
      cd vittransformers; \
      pip3 install .; \
      pip3 install datasets; \
      wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"

Treine o modelo

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export TF_CPP_MIN_LOG_LEVEL=0
      export XLA_USE_BF16=1
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers
      python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 \
      --remove_unused_columns=False \
      --label_names=pixel_values \
      --mask_ratio=0.75 \
      --norm_pix_loss=True \
      --do_train=true \
      --do_eval=true \
      --base_learning_rate=1.5e-4 \
      --lr_scheduler_type=cosine \
      --weight_decay=0.05 \
      --num_train_epochs=3 \
      --warmup_ratio=0.05 \
      --per_device_train_batch_size=8 \
      --per_device_eval_batch_size=8 \
      --logging_strategy=steps \
      --logging_steps=30 \
      --evaluation_strategy=epoch \
      --save_strategy=epoch \
      --load_best_model_at_end=True \
      --save_total_limit=3 \
      --seed=1337 \
      --output_dir=MAE \
      --overwrite_output_dir=true \
      --logging_dir=./tensorboard-metrics \
      --tpu_metrics_debug=true'

Excluir a TPU e o recurso na fila

Exclua a TPU e o recurso na fila no fim da sessão.

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

Resultado da comparação

A tabela a seguir mostra a taxa de transferência de referência para diferentes tipos de acelerador.

v5litepod-4 v5litepod-16 v5litepod-64
Período 3 3 3
Tamanho global do lote 32 128 512
Capacidade (exemplos/s) 201 657 2.844

TensorFlow 2.x

As seções a seguir descrevem exemplos de como treinar modelos do TensorFlow 2.x na TPU v5e.

Treinar o Resnet em uma v5e de host único

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

Configurar

  1. Crie variáveis de ambiente:

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

    ACCELERATOR_TYPE pode ser v5litepod-4 ou v5litepod-8.

  2. 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á fazer SSH na VM da TPU quando o recurso enfileirado 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}

  3. Conectar-se ao TPU usando SSH

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

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

  5. 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"

Excluir a TPU e o recurso na fila

  1. Excluir a TPU

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

  2. Excluir a solicitação de recurso na fila

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

Treinar o Resnet em um v5e com vários hosts

Este tutorial descreve como treinar o ImageNet em v5litepod-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.

  1. Crie variáveis de ambiente:

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

    ACCELERATOR_TYPE pode ser v5litepod-16 ou maior.

  2. 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á fazer SSH na VM da TPU quando o recurso enfileirado estiver no estado ACTIVE. Para verificar o estado do recurso em fila, use o comando a seguir:

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

  3. Conectar-se ao TPU (worker zero) usando SSH

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

  4. 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 TPU_LOAD_LIBRARY=0
export TPU_NAME=your_tpu_name

  5. 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 \
   --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*"

Excluir a TPU e o recurso na fila

  1. Excluir a TPU

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

  2. Excluir a solicitação de recurso na fila

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