Treinamento do Cloud TPU v5e

O Cloud TPU v5e é o acelerador de IA de última geração do Google Cloud. Com um tamanho menor de 256 chips por pod, um v5e é otimizado para ser o produto de maior valor para treinamento, ajuste e disponibilização de transformadores, texto para imagem e rede neural convolucional (CNN, na sigla em inglês).

Conceitos, arquitetura do sistema e configurações do Cloud TPU v5e

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

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

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

Inferência

Inferência é o processo de usar um modelo treinado para fazer previsões sobre novos dados. Ele é usado pelo processo de disponibilização.

Frações

Uma fração representa uma coleção de chips, todos localizados dentro do mesmo pod conectado por interconexões de alta velocidade entre chips (ICI). O v5e tem formatos de fatias 2D. Consulte a tabela na seção Configurações da v5e para conferir os formatos de fatia compatíveis.

A forma do chip e a topologia do chip também se referem a formas de fatia.

Exibição

Exibição é o processo de implantação de um modelo de machine learning treinado em um ambiente de produção onde ele pode ser usado para fazer previsões ou tomar decisões. A latência e a disponibilidade no nível de serviço são importantes para a exibição.

Host único versus host múltiplo

Um host é um computador físico (CPU) que executa VMs. Um host pode executar várias VMs ao mesmo tempo.

Slices que usam menos de 8 ícones usam no máximo um host. Slices com mais de 8 chips têm acesso a mais de um host e podem executar treinamento distribuído usando vários hosts. Consulte a página Arquitetura do sistema da TPU para ver detalhes sobre frações e chips.

A v5e oferece suporte para treinamento em vários hosts e inferência de vários hosts (usando SAX).

VM de TPU

Uma máquina virtual executando Linux que tem acesso às TPUs subjacentes. Para TPUs v5e, cada VM de TPU tem acesso direto a 1, 4 ou 8 chips, dependendo do tipo de acelerador especificado pelo usuário. Uma VM de TPU também é conhecida como worker.

Worker

Consulte VM da TPU.

Começar

Para mais informações sobre o hardware de TPU v5e, consulte Arquitetura do sistema.

Como garantir a capacidade

Entre em contato com a equipe de vendas do Cloud para começar a usar a Cloud TPU v5e nas suas cargas de trabalho de IA.

Preparar um projeto do Google Cloud

  1. Faça login na sua Conta do Google. Inscreva-se para uma nova conta caso ainda não tenha feito isso.
  2. No console do Google Cloud, selecione ou crie um projeto do Google Cloud na página do seletor de projetos.
  3. A configuração de faturamento é obrigatória para todos os usos do Google Cloud. Portanto, verifique se o faturamento está ativado para seu projeto.

  4. Instale o gcloud alpha components (link em inglês).

  5. Se você é um usuário da TPU que reutiliza componentes gcloud alpha atuais, atualize-os para garantir que os comandos e sinalizações relevantes sejam compatíveis:

    gcloud components update
    
  6. Ative a API TPU com o comando gcloud a seguir no Cloud Shell. Também é possível ativá-la no console do Google Cloud.

    gcloud services enable tpu.googleapis.com
    
  7. Ative a conta de serviço da TPU.

    As contas de serviço permitem que o serviço do Cloud TPU acesse outros serviços do Google Cloud. Uma conta de serviço gerenciado pelo usuário é uma prática recomendada do Google Cloud. Siga estes guias para criar e conceder papéis. Os seguintes papéis são necessários:

    • Administrador da TPU
    • Administrador do Storage: necessário para acessar o Cloud Storage.
    • Gravador de registros: necessário para gravar registros com a API Logging
    • Gravador de métricas do Monitoring: necessário para gravar métricas no Cloud Monitoring
  8. Configurar o projeto e a zona.

    O ID do projeto é o nome do projeto mostrado no console do Cloud.

    export PROJECT_ID=your-project-id
    export ZONE=us-west4-a
    
    gcloud compute tpus tpu-vm service-identity create --zone=${ZONE}
    
    gcloud auth login
    gcloud config set project ${PROJECT_ID}
    gcloud config set compute/zone ${ZONE}
    

Provisionar o ambiente da Cloud TPU

A prática recomendada é provisionar o Cloud TPU v5es como recursos na fila usando o comando queued-resource create. No entanto, também é possível usar a API Create Node (gcloud alpha compute tpus tpu-vm create) para provisionar o Cloud TPU v5es.

Criar variáveis de ambiente

Defina as variáveis de ambiente necessárias para a criação da TPU.

Substitua as variáveis (em vermelho) na lista a seguir pelos valores que você usará no job de treinamento ou inferência.

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
Variável Descrição
PROJECT_ID Nome do projeto do Google Cloud
ACCELERATOR_TYPE Consulte a seção Tipos de acelerador para ver os tipos de aceleradores compatíveis.
ZONA Toda a capacidade está em us-west4-a.
RUNTIME_VERSION Use a v2-alpha-tpuv5-lite.
SERVICE_ACCOUNT Esse é o endereço da sua conta de serviço que pode ser encontrado no console do Google Cloud -> IAM -> Contas de serviço. Por exemplo: tpu-service-account@myprojectID.iam.gserviceaccount.com
TPU_NAME O ID do texto atribuído pelo usuário da TPU, que é criado quando a solicitação de recurso na fila é alocada.
QUEUED_RESOURCE_ID O ID do texto atribuído pelo usuário da solicitação de recurso na fila. Consulte o documento recursos na fila para informações sobre eles.
QUOTA_TYPE Pode ser reserved ou best-effort. Se nenhuma dessas opções for especificada, o valor padrão de QUOTA_TYPE será on-demand. Consulte as quotas para informações sobre os diferentes tipos de cotas compatíveis com o Cloud TPU.
VALID_UNTIL_DURATION A duração da solicitação. Consulte recursos na fila para informações sobre as diferentes durações válidas.

Criar um recurso de TPU

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

Se o recurso na fila for criado, o estado no campo response será WAITING_FOR_RESOURCES ou FAILED. Se o recurso na fila estiver no estado WAITING_FOR_RESOURCES, isso significa que ele passou por uma validação preliminar e está aguardando capacidade. Quando a capacidade estiver disponível, a solicitação vai passar para PROVISIONING. Estar no estado WAITING_FOR_RESOURCES não significa que você vai receber a cota alocada e pode levar algum tempo para mudar do status WAITING_FOR_RESOURCES para o status ACTIVE. Se o recurso na fila estiver no estado FAILED, o motivo da falha estará na saída. A solicitação expira se não for preenchida no --valid-until-duration e o estado se torna "FAILED".

Será possível acessar sua VM de TPU usando SSH quando o recurso na fila estiver no estado ACTIVE.

Use os comandos [list](/tpu/docs/managing-tpus-tpu-vm) ou [describe](/tpu/docs/managing-tpus-tpu-vm) para consultar o status do recurso na fila.

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

O state representa o status de um recurso na fila. Os estados são definidos como:

Estado Descrição
WAITING_FOR_RESOURCES O comando create do recurso na fila foi recebido e começará o provisionamento assim que a capacidade estiver disponível.
PROVISIONING As frações de TPU estão sendo provisionadas.
ACTIVE Todas as TPUs são provisionadas e estão prontas para uso. Se um script de inicialização for fornecido, ele começará a ser executado em todas as TPUs, quando o estado do recurso na fila mudar para ACTIVE.
FAILED Não foi possível provisionar as frações.
SUSPENDING Uma ou mais frações estão sendo excluídas.
SUSPENDED Todas as frações são excluídas, mas o recurso na fila permanece intacto até que seja explicitamente excluído. No momento, um recurso na fila suspenso não pode ser retomado e precisa ser excluído.
DELETING O recurso na fila está sendo excluído.

Conectar-se à VM da TPU usando SSH

Na seção a seguir, descrevemos como instalar binários em cada VM de TPU na fração de TPU e no código de execução. Nesse contexto, uma VM de TPU também é conhecida como worker.

Consulte a seção Tipos de VM para determinar quantas VMs sua fração terá.

Para instalar os binários ou executar o código, conecte-se à VM da TPU usando o comando tpu-vm ssh.

gcloud compute tpus tpu-vm ssh ${TPU_NAME}

Para acessar uma VM da TPU específica com SSH, use a sinalização --worker, que segue um índice baseado em 0:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} --worker=1

Se você tiver formas de fração maiores que 8 chips, terá várias VMs em uma fração. Nesse caso, use a sinalização --worker=all para executar a instalação em todas as VMs de TPU sem precisar se conectar a cada uma separadamente. Exemplo:

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'

Gerenciar

Todos os comandos que podem ser usados para gerenciar suas VMs de TPU são descritos em Como gerenciar TPUs.

Configuração do framework

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

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

Configuração para o JAX

Se você tiver formas de fração maiores que 8 chips, terá várias VMs em uma fração. Nesse caso, é preciso usar a sinalização --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 separadamente:

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'

É possível executar o seguinte comando para verificar o número de dispositivos. As saídas mostradas aqui foram produzidas com uma fração do v5litepod-16. Este código testa se tudo está instalado corretamente. Para isso, ele verifica se o JAX vê os TensorCores do Cloud TPU e pode executar operações básicas:

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

A saída será semelhante a esta:

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

jax.device\_count() mostra o número total de ícones na fatia especificada. jax.local\_device\_count() indica a contagem de ícones acessíveis por uma única VM nesta fração.

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

A saída será semelhante a esta:

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

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

Configuração do PyTorch

A v5e é compatível apenas com o ambiente de execução PJRT, e o PyTorch 2.1+ usará o PJRT como o ambiente de execução padrão em todas as versões da TPU.

Nesta seção, descrevemos 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ê não conseguir instalar rodas para torch/torch_xla/torchvision e vir um erro como pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222, faça downgrade da sua versão com este comando:

pip3 install setuptools==62.1.0

Execute um script com PJRT:

unset LD_PRELOAD

Veja a seguir um exemplo que usa 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')

Confira os Tutoriais do PyTorch neste documento para começar o treinamento da v5e usando o PyTorch.

Exclua a TPU e o recurso na fila ao final da sessão. Para excluir um recurso na fila, exclua a fração e, em seguida, o 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

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

Monitorar e criar perfil

O Cloud TPU v5e é compatível com monitoramento e criação de perfil usando os mesmos métodos das gerações anteriores do Cloud TPU. Leia Criar perfil do modelo com ferramentas do Cloud TPU para mais informações sobre criação de perfil e Como monitorar VMs do Cloud TPU.

Exemplos do JAX/FLAX

Treinar o ImageNet no 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 alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
       --node-id=${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --valid-until-duration=${VALID_UNTIL_DURATION} \
       --service-account=${SERVICE_ACCOUNT} \
       --${QUOTA_TYPE}
    

    Será possível estabelecer o SSH para sua VM de TPU quando o recurso na fila 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 "jax[tpu]==0.4.16" -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 fictícios, 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'
    

Treinar o modelo

Depois de concluir todas as etapas anteriores, treine 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 ao 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 com rosto abraçado

Os modelos de hugging Face implementados no FLAX funcionam imediatamente no Cloud TPU v5e. Nesta seção, você verá instruções para executar os modelos mais usados.

Treinar o ViT no Imagenette

Neste tutorial, mostramos como treinar 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 treinar um codificador Transformer no ImageNet com resultados excelentes em comparação com redes convolucionais. Para mais informações, consulte os recursos abaixo:

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 alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
       --node-id=${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --valid-until-duration=${VALID_UNTIL_DURATION} \
       --service-account=${SERVICE_ACCOUNT} \
       --${QUOTA_TYPE}
    

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

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

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

     state: ACTIVE
    
  3. Instale o JAX e a biblioteca correspondente:

    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'
    
  4. Faça o download do repositório (link em inglês) 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.16.1 && pip install -r examples/flax/vision/requirements.txt'
    
  5. Faça o download do conjunto de dados do 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'
    

Treinar o modelo

Treine o modelo com um buffer pré-mapeado em 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 ao 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 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 capacidades com diferentes tipos de aceleradores.

Tipo de acelerador v5litepod-4 v5litepod-16 v5litepod-64
Época 3 3 3
Tamanho global do lote 32 128 512
Capacidade de processamento (exemplos/s) 263,40 429,34 470,71

Treine a Difusão no 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 de difusão estável é um modelo latente de texto para imagem que gera imagens realistas a partir de qualquer entrada de texto. Para saber mais, acesse os recursos a seguir (links em inglês):

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 alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
       --node-id=${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --valid-until-duration=${VALID_UNTIL_DURATION} \
       --service-account=${SERVICE_ACCOUNT} \
       --${QUOTA_TYPE}
    

    Será possível estabelecer o SSH para sua VM de TPU quando o recurso na fila estiver no estado ACTIVE:

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

    Quando o recurso na fila está no estado ACTIVE, a saída 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 "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  4. Faça o download do repositório 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 tensorflow==2.16.1 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
    

Treinar o modelo

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

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model'

Excluir a TPU e os recursos na fila

Exclua a TPU e o recurso na fila ao 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 do comparativo de mercado para difusão

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

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

Treinar o GPT2 no conjunto de dados do OSCAR

Neste tutorial, mostramos como treinar o modelo GPT2 da HuggingFace usando o conjunto de dados OSCAR no 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 alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
       --node-id=${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --valid-until-duration=${VALID_UNTIL_DURATION} \
       --service-account=${SERVICE_ACCOUNT} \
       --${QUOTA_TYPE}
    

    Será possível estabelecer o SSH para sua VM de TPU quando o recurso na fila estiver no estado ACTIVE:

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

    Quando o recurso na fila está no estado ACTIVE, a saída 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 "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  4. Faça o download do repositório 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. Fazer 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 && gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt .'
    

Treinar o modelo

Treine o modelo com um buffer pré-mapeado em 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 ao 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 do comparativo de mercado para GPT2

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

v5litepod-4 v5litepod-16 v5litepod-64
Época 3 3 3
Tamanho global do lote 64 64 64
Capacidade de processamento (exemplos/s) 74,60 72,97 72,62

PyTorch/XLA

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

O PyTorch/XLA está migrando do XRT para o PjRt do PyTorch 2.0 ou mais recente. Veja as instruções atualizadas para configurar a v5e em cargas de trabalho de treinamento do PyTorch/XLA.

Configuração
  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 alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
       --node-id=${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --valid-until-duration=${VALID_UNTIL_DURATION} \
       --service-account=${SERVICE_ACCOUNT} \
       --{QUOTA_TYPE}
    

    Será possível se conectar via SSH à 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 está 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 ao 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 do comparativo de mercado

A tabela a seguir mostra as capacidades de processamento de referência.

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

Treinar o GPT2 na v5e

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

Configuração

  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 alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
       --node-id=${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --valid-until-duration=${VALID_UNTIL_DURATION} \
       --service-account=${SERVICE_ACCOUNT} \
       --${QUOTA_TYPE}
    

    Será possível se conectar via SSH à 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 está no estado ACTIVE, a saída será semelhante a esta:

    state: ACTIVE
    
  3. Instale as dependências da lanterna/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. Fazer 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='
          gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/
          gsutil cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'
    

Treinar 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 ao 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 do comparativo de mercado

O script de treinamento foi executado no v5litepod-4, v5litepod-16 e v5litepod-64. A tabela a seguir mostra as capacidades de comparação para diferentes tipos de acelerador.

v5litepod-4 v5litepod-16 v5litepod-64
Época 3 3 3
config 600 milhões 2 bi 16 bi
Tamanho global do lote 64 128 256
Capacidade de processamento (exemplos/s) 66 77 31

Treinar o ViT na v5e

Neste tutorial, você verá como executar o VIT na v5e usando o repositório HuggingFace no PyTorch/XLA no conjunto de dados cifar10.

Configuração

  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 alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
       --node-id=${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --valid-until-duration=${VALID_UNTIL_DURATION} \
       --service-account=${SERVICE_ACCOUNT} \
       --${QUOTA_TYPE}
    

    Será possível se conectar via SSH à 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 está no estado ACTIVE, a saída será semelhante a esta:

     state: ACTIVE
    
  3. Instalar dependências da lanterna/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"
    

Treinar 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 ao 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 do comparativo de mercado

A tabela a seguir mostra as capacidades de comparação para diferentes tipos de aceleradores.

v5litepod-4 v5litepod-16 v5litepod-64
Época 3 3 3
Tamanho global do lote 32 128 512
Capacidade de processamento (exemplos/s) 201 657 2.844

TensorFlow 2.x

Treinar o Resnet em um único host v5e

Neste tutorial, descrevemos como treinar o ImageNet em v5litepod-4 ou v5litepod-8 usando um conjunto de dados fictício. 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 alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
       --node-id=${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --${QUOTA_TYPE}
    

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

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
       --project=${PROJECT_ID} \
       --zone=${ZONE}
    
  3. Conectar-se à TPU usando SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
       --project=${PROJECT_ID} \
       --zone=${ZONE}
    
  4. Defina algumas variáveis de ambiente

    export MODELS_REPO=/usr/share/tpu/models
    export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
    export MODEL_DIR=gcp-directory-to-store-model
    export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
    export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
    export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
    
  5. Altere para o diretório de repositório de modelos e os requisitos de instalação.

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

Treinar o modelo

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 sua 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 host multihost v5e

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

  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 alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
       --node-id=${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --${QUOTA_TYPE}
    

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

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
       --project=${PROJECT_ID} \
       --zone=${ZONE}
    
  3. Conectar-se à TPU (worker zero) usando SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
       --project=${PROJECT_ID} \
       --zone=${ZONE}
    
  4. Defina algumas variáveis de ambiente

    export 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. Altere para o diretório de repositório de modelos e os requisitos de instalação.

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

Treinar o modelo

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 sua solicitação de recurso na fila

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