Treinamento do Cloud TPU v5e
O Cloud TPU v5e é o acelerador de IA de última geração do Google Cloud. Com um menor consumo de 256 chips por pod, uma v5e é otimizada para ter o valor mais alto produto para transformador, texto para imagem e rede neural convolucional (CNN, na sigla em inglês) treinamento, ajuste e disponibilização. Para mais informações sobre como usar a Cloud TPU v5e para exibição, consulte Inferência usando a v5e.
Para mais informações sobre o hardware e as configurações da TPU 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
É preciso 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. Existem diferentes as cotas exigidas se você estiver usando sua TPU v5e para inferência. Para mais informações sobre cotas, consulte Cotas. Para solicitar a cota de TPU v5e, entre em contato com o Cloud Vendas.
Criar um projeto e uma conta do Google Cloud
Você precisa de uma conta e um projeto do Google Cloud para usar o Cloud TPU. Para mais mais informações, consulte Configurar um Cloud TPU de código aberto.
Criar uma Cloud TPU
A prática recomendada é provisionar Cloud TPU v5es na fila
recursos usando a queued-resource create
kubectl. Para mais informações, consulte Gerenciar filas
do Google Cloud.
Também é possível usar a API Create Node (gcloud alpha compute tpus tpu-vm create
)
para provisionar o Cloud TPU
v5es Para mais informações, consulte Gerenciar TPU
do Google Cloud.
Para mais informações sobre os papéis configurações da v5e para treinamento, consulte Tipos da Cloud TPU v5e para treinamento.
Configuração do framework
Nesta seção, descrevemos o processo geral de configuração do treinamento de modelo personalizado usando
JAX ou PyTorch com TPU v5e. O suporte ao TensorFlow está disponível no
TPU tpu-vm-tf-2.17.0-pjrt
e tpu-vm-tf-2.17.0-pod-pjrt
mais recentes do ambiente de execução.
Para instruções de configuração de inferência, consulte Inferência na v5e introdução.
Configuração para JAX
Se você tiver formas de fatia com mais de 8 ícones, terá várias VMs em uma
fatia Nesse caso, é necessário usar a flag --worker=all
para executar o
em todas as VMs de TPU em uma única etapa, sem usar SSH para fazer login em cada
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'
Execute o seguinte comando para verificar o número de dispositivos (as saídas mostradas aqui foram produzidos com uma fatia v5litepod-16). Esse código testa se tudo está instalado corretamente verificando se o JAX vê os TensorCores do Cloud TPU e pode executar operações básicas:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'
A saída será semelhante a esta:
SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4
jax.device_count()
mostra o número total de ícones em uma determinada fração.
jax.local_device_count()
indica o número de ícones acessíveis por um único
nessa fração.
# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'
A saída será semelhante a esta:
SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
Confira os Tutoriais do JAX neste documento para começar a treinamento da v5e usando JAX.
Configuração do PyTorch
A v5e só oferece suporte ao ambiente de execução PJRT. e o PyTorch 2.1+ 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
marcação "Gostei"
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
Executar um script com PJRT
unset LD_PRELOAD
O exemplo a seguir 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')
Veja os Tutoriais do PyTorch neste documento para começar a treinamento da v5e com o PyTorch.
Exclua a TPU e o recurso na fila no final da sessão. Para excluir um recurso na fila, exclua a fração e depois 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 enfileiradas que estão em
o estado FAILED
.
Exemplos de JAX/FLAX
As seções a seguir descrevem exemplos de como treinar modelos JAX e FLAX em TPU v5e
Treinar o ImageNet na v5e
Neste tutorial, descrevemos como treinar o ImageNet na v5e usando dados de entrada falsos. Se Se você quiser usar dados reais, consulte o arquivo README no GitHub.
Configurar
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
-
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}
Você poderá se conectar via SSH à VM da TPU quando o recurso na fila estiver o estado
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Quando o QueuedResource está no estado
ACTIVE
, a saída será semelhante a:state: ACTIVE
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'
Clone o modelo do 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'
Para gerar dados falsos, o modelo precisa de informações sobre as dimensões 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 que todas as etapas anteriores forem concluídas, você poderá 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 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 Hugging Face FLAX
Os modelos Hugging Face implementados no FLAX funcionam a partir do a caixa do Cloud TPU v5e. Nesta seção, fornecemos instruções para executar os modelos mais populares.
Treinar ViT no Imagenette
Neste tutorial, mostramos como treinar o Vision Transformer (ViT) do HuggingFace usando o Imagenette de IA rápida no Cloud TPU v5e.
O modelo ViT foi o primeiro a treinar com sucesso um codificador no ImageNet com resultados excelentes em comparação com redes convolucionais. Para mais , consulte os seguintes recursos:
Configurar
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
-
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 conexões SSH com sua VM da TPU quando o recurso estiver na fila está 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 para o seguinte:state: ACTIVE
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'
Faça o download do repositório Hugging Face e requisitos de instalação:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='git clone https://github.com/huggingface/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.17.0 && pip install -r examples/flax/vision/requirements.txt'
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'
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 de viT
O script de treinamento foi executado nas versões v5litepod-4, v5litepod-16 e v5litepod-64. O a tabela a seguir mostra as capacidades de processamento 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 de processamento (exemplos/segundo) | 263,40 | 429,34 | 470,71 |
Treine o Diffusion em Pokémon
Neste tutorial, mostramos como treinar o modelo de difusão estável HuggingFace usando o Pokémon no Cloud TPU v5e.
O modelo de difusão estável é um modelo latente de texto para imagem que gera imagens fotorrealistas a partir de qualquer entrada de texto. Para saber mais, acesse os recursos a seguir (links em inglês):
Configurar
Crie variáveis de ambiente:
export 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
-
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 uma conexão SSH com sua 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 estiver no estado
ACTIVE
, a saída será semelhante a esta:state: ACTIVE
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'
Faça o download do repositório HuggingFace e requisitos de instalação.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.17.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} \
--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 o recurso na fila
Exclua a TPU e o recurso na 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 comparativos de difusão
O script de treinamento foi executado nas versões 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 |
Degrau de trem | 1500 | 1500 | 1500 |
Tamanho global do lote | 32 | 64 | 128 |
Capacidade de processamento (exemplos/segundo) | 36,53 | 43,71 | 49,36 |
Treinar a 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 com 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
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
-
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 conexões SSH com sua VM da TPU quando o recurso estiver na fila está 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 para o seguinte:state: ACTIVE
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'
Faça o download do repositório HuggingFace e requisitos de instalação.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='git clone https://github.com/huggingface/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'
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 && gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt .'
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 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 nas versões 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 de processamento (exemplos/segundo) | 74,60 | 72,97 | 72,62 |
PyTorch/XLA
As seções a seguir descrevem exemplos de como treinar modelos PyTorch/XLA em TPU v5e
Treinar o ResNet usando o ambiente de execução PJRT
O PyTorch/XLA está migrando do PyTorch 2.0+ para o PjRt. Estes são os instruções atualizadas para configurar a v5e para cargas de trabalho de treinamento PyTorch/XLA.
Configurar
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
-
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 conexões SSH com sua VM da TPU assim que seu QueuedResource está 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
Instalar as 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 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 do comparativo.
Tipo de acelerador | Capacidade de processamento (exemplos/segundo) |
v5litepod-4 | 4240 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 em PyTorch/XLA usando o conjunto de dados wikitext.
Configurar
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
-
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}
Você poderá se conectar via SSH à VM da TPU quando o QueuedResource estiver
ACTIVE
estado: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:state: ACTIVE
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'
Faça o download do repositório HuggingFace e requisitos de instalação.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' git clone https://github.com/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 '
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/'
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 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 nas versões v5litepod-4, v5litepod-16 e v5litepod-64. O a tabela a seguir mostra as capacidades de processamento do comparativo para diferentes tipos de aceleradores.
v5litepod-4 | v5litepod-16 | v5litepod-64 | |
Período | 3 | 3 | 3 |
config | 600 milhões | 2 bilhões | 16 bi |
Tamanho global do lote | 64 | 128 | 256 |
Capacidade de processamento (exemplos/segundo) | 66 | 77 | 31 |
Treinar ViT na v5e
Este tutorial abordará como executar a VIT na v5e usando o repositório HuggingFace em PyTorch/XLA no conjunto de dados cifar10.
Configurar
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
-
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}
Você poderá se conectar via SSH à VM da TPU quando o QueuedResource estiver o 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:state: ACTIVE
Instalar 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 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'
Faça o download do repositório HuggingFace e requisitos de instalação.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=" git clone https://github.com/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 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 do comparativo de mercado para diferentes tipos de acelerador.
v5litepod-4 | v5litepod-16 | v5litepod-64 | |
Período | 3 | 3 | 3 |
Tamanho global do lote | 32 | 128 | 512 |
Capacidade de processamento (exemplos/segundo) | 201 | 657 | 2.844 |
TensorFlow 2.x
As seções a seguir descrevem exemplos de como treinar o TensorFlow 2.x modelos de ML na TPU v5e.
Treinar Resnet em um host único 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
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 serv5litepod-4
ouv5litepod-8
.-
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}
Você poderá se conectar via SSH à VM da TPU quando o recurso na fila estiver o estado
ACTIVE
. Para verificar o estado do recurso na fila, use este comando:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Conectar-se à TPU usando SSH
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
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
Altere para o diretório do repositório de modelos e os requisitos de instalação.
cd ${MODELS_REPO} && git checkout r2.15.0 pip install -r official/requirements.txt
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
Excluir a TPU
gcloud compute tpus tpu-vm delete ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet
Excluir a solicitação de recurso na fila
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet
Treinar Resnet em um v5e de vários hosts
Este tutorial descreve como treinar o ImageNet em v5litepod-16
ou maior usando
em um conjunto de dados fictício. Se você quiser usar um conjunto de dados diferente, consulte Como preparar o conjunto de dados.
Crie variáveis de ambiente:
export PROJECT_ID=your_project_ID export 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 serv5litepod-16
ou maior.-
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}
Você poderá se conectar via SSH à VM da TPU quando o recurso na fila estiver o estado
ACTIVE
. Para verificar o estado do seu recurso na fila, use o método seguinte comando:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Conecte-se à TPU (worker zero) usando SSH.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
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
Altere para o diretório do repositório de modelos e os requisitos de instalação.
cd $MODELS_REPO && git checkout r2.15.0 pip install -r official/requirements.txt
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
Excluir a TPU
gcloud compute tpus tpu-vm delete ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet
Excluir a solicitação de recurso na fila
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet