Prepare um modelo com a TPU v5e

Com uma área de 256 chips por Pod, a TPU v5e está otimizada para ser um produto de alto valor para a preparação, o ajuste fino e a publicação de transformadores, texto para imagem e redes neurais convolucionais (CNN). Para mais informações sobre a utilização do Cloud TPU v5e para publicação, consulte o artigo Inferência com o v5e.

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

Começar

As secções seguintes descrevem como começar a usar a TPU v5e.

Quota de pedidos

Precisa de quota para usar a TPU v5e para preparação. Existem diferentes tipos de quotas para TPUs a pedido, TPUs reservadas e VMs Spot de TPU. Existem quotas separadas necessárias se estiver a usar a TPU v5e para inferência. Para mais informações sobre quotas, consulte o artigo Quotas. Para pedir quota de TPUs v5e, contacte as vendas do Google Cloud.

Crie uma Google Cloud conta e um projeto

Precisa de uma Google Cloud conta e um projeto para usar a Cloud TPU. Para mais informações, consulte o artigo Configure um ambiente de TPU na nuvem.

Crie uma Cloud TPU

A prática recomendada é aprovisionar Cloud TPUs v5es como recursos em fila usando o comando queued-resource create. Para mais informações, consulte o artigo Faça a gestão dos recursos em fila.

Também pode usar a API Create Node (gcloud compute tpus tpu-vm create) para aprovisionar TPUs em nuvem v5es. Para mais informações, consulte o artigo Faça a gestão dos recursos de TPU.

Para mais informações sobre as configurações v5e disponíveis para preparação, consulte o artigo Tipos de TPU na nuvem v5e para preparação.

Configuração da framework

Esta secção descreve o processo de configuração geral para a preparação de modelos personalizados com o JAX ou o PyTorch com a TPU v5e.

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

Defina algumas variáveis de ambiente:

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id

Configuração do JAX

Se tiver formas de divisão com mais de 8 chips, terá várias VMs numa divisão. Neste caso, tem de usar a flag --worker=all para executar a instalação em todas as VMs de TPU num único passo sem usar SSH para iniciar sessão em cada uma separadamente:

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 das flags de comando

Variável Descrição
TPU_NAME O ID de texto atribuído pelo utilizador da TPU que é criado quando o pedido de recurso em fila é atribuído.
PROJECT_ID Google Cloud Nome do projeto. Use um projeto existente ou crie um novo em Configure o seu Google Cloud projeto
ZONE Consulte o documento Regiões e zonas de TPU para ver as zonas suportadas.
trabalhadora A VM da TPU que tem acesso às TPUs subjacentes.

Pode executar o seguinte comando para verificar o número de dispositivos (os resultados apresentados aqui foram produzidos com uma fatia v5litepod-16). Este código testa se tudo está instalado corretamente, verificando se o JAX vê os TensorCores do Cloud TPU e consegue 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())"'

O resultado será semelhante ao seguinte:

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 fatia 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))"'

O resultado será semelhante ao seguinte:

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

Experimente os tutoriais do JAX neste documento para começar a usar o treino v5e com o JAX.

Configuração do PyTorch

Tenha em atenção que a v5e só suporta o tempo de execução PJRT e o PyTorch 2.1+ vai usar o PJRT como o tempo de execução predefinido para todas as versões de TPU.

Esta secção descreve como começar a usar o PJRT na v5e com o PyTorch/XLA com comandos para todos os trabalhadores.

Instale 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
      pip install mkl mkl-include
      pip install tf-nightly tb-nightly tbp-nightly
      pip install numpy
      sudo apt-get install libopenblas-dev -y
      pip install torch~=PYTORCH_VERSION torchvision torch_xla[tpu]~=PYTORCH_VERSION -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html'

Substitua PYTORCH_VERSION pela versão do PyTorch que quer usar. PYTORCH_VERSION é usado para especificar a mesma versão para o PyTorch/XLA. 2.6.0 é recomendado.

Para mais informações sobre as versões do PyTorch e do PyTorch/XLA, consulte os artigos PyTorch – Começar e Lançamentos do PyTorch/XLA.

Para mais informações sobre a instalação do PyTorch/XLA, consulte o artigo Instalação do PyTorch/XLA.

Se 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, reverter a versão com este comando:

pip3 install setuptools==62.1.0

Execute um script com o PJRT

unset LD_PRELOAD

Segue-se um exemplo que usa um script do Python para fazer um cálculo numa 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
      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)"'

Isto 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')

Experimente os tutoriais do PyTorch neste documento para começar a usar o treino v5e com o PyTorch.

Eliminar a TPU e o recurso em fila no final da sessão. Para eliminar um recurso em fila, elimine a fatia e, em seguida, o recurso em fila em 2 passos:

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

Estes dois passos também podem ser usados para remover pedidos de recursos em fila que estejam no estado FAILED.

Exemplos de JAX/FLAX

As secções seguintes descrevem exemplos de como preparar modelos JAX e FLAX na TPU v5e.

Treine o ImageNet na v5e

Este tutorial descreve como preparar o ImageNet na v5e usando dados de entrada falsos. Se quiser usar dados reais, consulte o ficheiro README no GitHub.

Configurar

  1. Crie variáveis de ambiente:

    export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

    Descrições das variáveis de ambiente

    Variável Descrição
    PROJECT_ID O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
    TPU_NAME O nome da TPU.
    ZONE A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
    ACCELERATOR_TYPE O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
    RUNTIME_VERSION A versão do software do Cloud TPU.
    SERVICE_ACCOUNT O endereço de email da sua conta de serviço. Pode encontrá-la acedendo à página Contas de serviço na Google Cloud consola.

    Por exemplo: tpu-service-account@PROJECT_ID.iam.gserviceaccount.com

    QUEUED_RESOURCE_ID O ID de texto atribuído pelo utilizador do pedido de recurso em fila.

  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} \
   --service-account=${SERVICE_ACCOUNT}

    Pode estabelecer uma ligação SSH à sua VM da TPU assim que o recurso em fila estiver no estado ACTIVE:

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

    Quando o QueuedResource está no estado ACTIVE, o resultado é semelhante ao seguinte:

     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/coolkp/flax.git && cd flax && git checkout pmap-orbax-conversion && git pull"

  5. Para gerar dados falsos, o modelo precisa de informações sobre as dimensões do conjunto de dados. Estas informações podem ser recolhidas a partir dos metadados do conjunto de dados ImageNet:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command="cd flax/examples/imagenet && pip install -r requirements-cloud-tpu.txt"

Prepare o modelo

Depois de concluir todos os passos anteriores, pode preparar o modelo.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command="cd flax/examples/imagenet && bash ../../tests/download_dataset_metadata.sh && JAX_PLATFORMS=tpu python imagenet_fake_data_benchmark.py"

Elimine a TPU e o recurso em fila

Eliminar a TPU e o recurso em fila no final 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 FLAX da Hugging Face

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

Treine o ViT no Imagenette

Este tutorial mostra como preparar o modelo Vision Transformer (ViT) da HuggingFace usando o conjunto de dados Imagenette da Fast AI no Cloud TPU v5e.

O modelo ViT foi o primeiro a formar com êxito um codificador Transformer no ImageNet com excelentes resultados 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 TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-16
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

    Descrições das variáveis de ambiente

    Variável Descrição
    PROJECT_ID O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
    TPU_NAME O nome da TPU.
    ZONE A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
    ACCELERATOR_TYPE O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
    RUNTIME_VERSION A versão do software do Cloud TPU.
    SERVICE_ACCOUNT O endereço de email da sua conta de serviço. Pode encontrá-la acedendo à página Contas de serviço na Google Cloud consola.

    Por exemplo: tpu-service-account@PROJECT_ID.iam.gserviceaccount.com

    QUEUED_RESOURCE_ID O ID de texto atribuído pelo utilizador do pedido de recurso em fila.

  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} \
   --service-account=${SERVICE_ACCOUNT}

    Pode estabelecer uma ligação SSH à sua VM da TPU assim que 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, o resultado será semelhante ao seguinte:

     state: ACTIVE

  3. Instale o JAX e a respetiva biblioteca:

    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. Transfira o repositório do Hugging Face 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.19 && sed -i 's/torchvision==0.12.0+cpu/torchvision==0.22.1/' examples/flax/vision/requirements.txt && pip install -r examples/flax/vision/requirements.txt && pip install tf-keras'

  5. Transfira o 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'

Prepare o modelo

Prepare 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'

Elimine a TPU e o recurso em fila

Eliminar a TPU e o recurso em fila no final 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 dos testes de referência da ViT

O script de preparação foi executado em v5litepod-4, v5litepod-16 e v5litepod-64. A tabela seguinte mostra as taxas de transferência com diferentes tipos de aceleradores.

Tipo de acelerador v5litepod-4 v5litepod-16 v5litepod-64
Época 3 3 3
Tamanho do lote global 32 128 512
Tráfego transmitido (exemplos/seg) 263,40 429,34 470,71

Treine a difusão no Pokémon

Este tutorial mostra como preparar o modelo Stable Diffusion a partir do HuggingFace usando o conjunto de dados Pokémon no Cloud TPU v5e.

O modelo Stable Diffusion é um modelo de texto para imagem latente que gera imagens fotorrealistas a partir de qualquer entrada de texto. Para mais informações, consulte os seguintes recursos:

Configurar

  1. Defina uma variável de ambiente para o nome do seu contentor de armazenamento:

    export GCS_BUCKET_NAME=your_bucket_name

  2. Configure um contentor de armazenamento para a saída do modelo:

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

  3. Crie variáveis de ambiente:

    export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west1-c
export ACCELERATOR_TYPE=v5litepod-16
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

    Descrições das variáveis de ambiente

    Variável Descrição
    PROJECT_ID O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
    TPU_NAME O nome da TPU.
    ZONE A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
    ACCELERATOR_TYPE O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
    RUNTIME_VERSION A versão do software do Cloud TPU.
    SERVICE_ACCOUNT O endereço de email da sua conta de serviço. Pode encontrá-la acedendo à página Contas de serviço na Google Cloud consola.

    Por exemplo: tpu-service-account@PROJECT_ID.iam.gserviceaccount.com

    QUEUED_RESOURCE_ID O ID de texto atribuído pelo utilizador do pedido de recurso em fila.

  4. 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} \
   --service-account=${SERVICE_ACCOUNT}

    Pode estabelecer uma ligação SSH à sua VM da TPU assim que 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, o resultado será semelhante ao seguinte:

     state: ACTIVE

  5. Instale o JAX e a respetiva biblioteca.

    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'

  6. Transfira o 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/RissyRan/diffusers.git && cd diffusers && pip install . && pip install -U -r examples/text_to_image/requirements_flax.txt && pip install tensorflow==2.17.1 clu && pip install tensorboard==2.17.1'

Prepare o modelo

Prepare 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
    pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
    pip3 install -r requirements.txt
    pip3 install .
    pip3 install gcsfs
    export LIBTPU_INIT_ARGS=''
    python -m src.maxdiffusion.train src/maxdiffusion/configs/base_2_base.yml run_name=my_run \
    jax_cache_dir=gs://${GCS_BUCKET_NAME} activations_dtype=bfloat16 weights_dtype=bfloat16 \
    per_device_batch_size=1 precision=DEFAULT dataset_save_location=gs://${GCS_BUCKET_NAME} \
    output_dir=gs://${GCS_BUCKET_NAME}/ attention=flash"

Limpar

Elimine a TPU, o recurso em fila e o contentor do Cloud Storage no final da sessão.

  1. Elimine a TPU:

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

  2. Elimine o recurso em fila:

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

  3. Elimine o contentor do Cloud Storage:

    gcloud storage rm -r gs://${GCS_BUCKET_NAME}

Resultados de testes de referência para difusão

O script de preparação foi executado em v5litepod-4, v5litepod-16 e v5litepod-64. A tabela seguinte mostra as taxas de transferência.

Tipo de acelerador v5litepod-4 v5litepod-16 v5litepod-64
Etapa de comboio 1500 1500 1500
Tamanho do lote global 32 64 128
Tráfego transmitido (exemplos/seg) 36,53 43,71 49,36

PyTorch/XLA

As secções seguintes descrevem exemplos de como preparar modelos PyTorch/XLA na TPU v5e.

Faça o treino da ResNet com o tempo de execução da PJRT

O PyTorch/XLA está a migrar do XRT para o PjRt a partir do PyTorch 2.0 e superior. Seguem-se as instruções atualizadas para configurar a v5e para cargas de trabalho de preparação do PyTorch/XLA.

Configurar

  1. Crie variáveis de ambiente:

    export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-16
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

    Descrições das variáveis de ambiente

    Variável Descrição
    PROJECT_ID O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
    TPU_NAME O nome da TPU.
    ZONE A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
    ACCELERATOR_TYPE O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
    RUNTIME_VERSION A versão do software do Cloud TPU.
    SERVICE_ACCOUNT O endereço de email da sua conta de serviço. Pode encontrá-la acedendo à página Contas de serviço na Google Cloud consola.

    Por exemplo: tpu-service-account@PROJECT_ID.iam.gserviceaccount.com

    QUEUED_RESOURCE_ID O ID de texto atribuído pelo utilizador do pedido de recurso em fila.

  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} \
   --service-account=${SERVICE_ACCOUNT}

    Pode estabelecer uma ligação SSH à sua VM da TPU assim que 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, o resultado será semelhante ao seguinte:

     state: ACTIVE

  3. Instale 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
     pip install torch==PYTORCH_VERSION torchvision torch_xla[tpu]==PYTORCH_VERSION -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html'

    Substitua PYTORCH_VERSION pela versão do PyTorch que quer usar. PYTORCH_VERSION é usado para especificar a mesma versão para o PyTorch/XLA. 2.6.0 é recomendado.

    Para mais informações sobre as versões do PyTorch e do PyTorch/XLA, consulte os artigos PyTorch – Começar e Lançamentos do PyTorch/XLA.

    Para mais informações sobre a instalação do PyTorch/XLA, consulte o artigo Instalação do PyTorch/XLA.

Prepare o modelo ResNet
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      date
      export PJRT_DEVICE=TPU
      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 checkout release-r2.6
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

Elimine a TPU e o recurso em fila

Eliminar a TPU e o recurso em fila no final 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 de referência

A tabela seguinte mostra os rendimentos de referência.

Tipo de acelerador Tráfego transmitido (exemplos/segundo)
v5litepod-4 4240 ex/s
v5litepod-16 10 810 ex/s
v5litepod-64 46 154 ex/s

Forme o ViT no v5e

Este tutorial aborda a execução do VIT na v5e através do 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 TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-16
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

    Descrições das variáveis de ambiente

    Variável Descrição
    PROJECT_ID O seu Google Cloud ID do projeto. Use um projeto existente ou crie um novo.
    TPU_NAME O nome da TPU.
    ZONE A zona na qual criar a VM da TPU. Para mais informações sobre as zonas suportadas, consulte o artigo Regiões e zonas de TPUs.
    ACCELERATOR_TYPE O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que quer criar. Para mais informações sobre os tipos de aceleradores suportados para cada versão da TPU, consulte o artigo Versões da TPU.
    RUNTIME_VERSION A versão do software do Cloud TPU.
    SERVICE_ACCOUNT O endereço de email da sua conta de serviço. Pode encontrá-la acedendo à página Contas de serviço na Google Cloud consola.

    Por exemplo: tpu-service-account@PROJECT_ID.iam.gserviceaccount.com

    QUEUED_RESOURCE_ID O ID de texto atribuído pelo utilizador do pedido de recurso em fila.

  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} \
   --service-account=${SERVICE_ACCOUNT}

    Pode estabelecer uma ligação SSH à sua VM da TPU assim que 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, o resultado será semelhante ao seguinte:

     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 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
  pip install torch==PYTORCH_VERSION torchvision torch_xla[tpu]==PYTORCH_VERSION -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html
  pip install jax==0.4.38 jaxlib==0.4.38 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/

    Substitua PYTORCH_VERSION pela versão do PyTorch que quer usar. PYTORCH_VERSION é usado para especificar a mesma versão para o PyTorch/XLA. 2.6.0 é recomendado.

    Para mais informações sobre as versões do PyTorch e do PyTorch/XLA, consulte os artigos PyTorch – Começar e Lançamentos do PyTorch/XLA.

    Para mais informações sobre a instalação do PyTorch/XLA, consulte o artigo Instalação do PyTorch/XLA.

  4. Transfira o 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"

Prepare o modelo

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU
      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'

Elimine a TPU e o recurso em fila

Eliminar a TPU e o recurso em fila no final 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 de referência

A tabela seguinte mostra os rendimentos de referência para diferentes tipos de aceleradores.

v5litepod-4 v5litepod-16 v5litepod-64
Época 3 3 3
Tamanho do lote global 32 128 512
Tráfego transmitido (exemplos/seg) 201 657 2844