Cette page a été traduite par l'API Cloud Translation.

Entraînement Cloud TPU v5e

Cloud TPU v5e est l'accélérateur d'IA de dernière génération de Google Cloud. Avec une empreinte de 256 puces par pod, un v5e est optimisé pour être le produit le plus intéressant pour l'entraînement, l'ajustement et la diffusion des transformateurs, du texte en image et des réseaux de neurones convolutifs (RNN). Pour en savoir plus sur l'utilisation de Cloud TPU v5e pour la diffusion, consultez la section Inférence à l'aide de v5e.

Pour en savoir plus sur le matériel et les configurations des TPU Cloud TPU v5e, consultez la page TPU v5e.

Commencer

Les sections suivantes expliquent comment commencer à utiliser les TPU v5e.

Quota de requêtes

Vous avez besoin d'un quota pour utiliser des TPU v5e pour l'entraînement. Il existe différents types de quotas pour les TPU à la demande, les TPU réservés et les VM Spot TPU. Des quotas distincts sont requis si vous utilisez votre TPU v5e pour l'inférence. Pour en savoir plus sur les quotas, consultez la page Quotas. Pour demander un quota TPU v5e, contactez le service commercial Cloud.

Créer un compte et un projet Google Cloud

Vous avez besoin d'un compte et d'un projet Google Cloud pour utiliser Cloud TPU. Pour en savoir plus, consultez Configurer un environnement Cloud TPU.

Créer une instance Cloud TPU

Il est recommandé de provisionner des Cloud TPU v5 en tant que ressources mises en file d'attente à l'aide de la commande queued-resource create. Pour en savoir plus, consultez Gérer les ressources en file d'attente.

Vous pouvez également utiliser l'API Create Node (gcloud compute tpus tpu-vm create) pour provisionner des Cloud TPU v5. Pour en savoir plus, consultez Gérer les ressources TPU.

Pour en savoir plus sur les configurations v5e disponibles pour l'entraînement, consultez la section Types de Cloud TPU v5e pour l'entraînement.

Configuration du framework

Cette section décrit le processus de configuration général pour l'entraînement de modèles personnalisés à l'aide de JAX ou de PyTorch avec TPU v5e. La compatibilité avec TensorFlow est disponible dans les versions d'exécution tpu-vm-tf-2.18.0-pjrt et tpu-vm-tf-2.18.0-pod-pjrt du TPU.

Pour obtenir des instructions de configuration de l'inférence, consultez la section Présentation de l'inférence v5e.

Configuration pour JAX

Si vos formes de tranche contiennent plus de huit chips, vous aurez plusieurs VM dans une même tranche. Dans ce cas, vous devez utiliser l'indicateur --worker=all pour exécuter l'installation sur toutes les VM TPU en une seule étape, sans utiliser SSH pour vous connecter à chacune d'elles séparément:

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'

Description des options de commande

Variable	Description
TPU_NAME	ID de texte attribué par l'utilisateur du TPU créé lors de l'allocation de la requête de ressource mise en file d'attente.
PROJECT_ID	Nom du projet Google Cloud Utilisez un projet existant ou créez-en un à l'adresse Configurer votre projet Google Cloud.
ZONE	Pour connaître les zones compatibles, consultez le document Régions et zones TPU.
Worker [class name, see definition and ref site provided]	VM TPU ayant accès aux TPU sous-jacents.

Vous pouvez exécuter la commande suivante pour vérifier le nombre d'appareils (les sorties affichées ici ont été produites avec une tranche v5litepod-16). Ce code vérifie que tout est correctement installé en vérifiant que JAX voit les TensorCores Cloud TPU et qu'il peut exécuter des opérations de base:

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

Le résultat doit ressembler à ce qui suit :

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() indique le nombre total de chips dans la tranche donnée. jax.local_device_count() indique le nombre de puces accessibles par une seule VM dans cette tranche.

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

Le résultat doit ressembler à ce qui suit :

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

Suivez les tutoriels JAX de ce document pour commencer à utiliser l'entraînement v5e avec JAX.

Configuration pour PyTorch

Notez que la version v5e n'est compatible qu'avec l'environnement d'exécution PJRT et que PyTorch 2.1 et versions ultérieures utiliseront PJRT comme environnement d'exécution par défaut pour toutes les versions de TPU.

Cette section explique comment commencer à utiliser PJRT sur la version v5e avec PyTorch/XLA avec des commandes pour tous les nœuds de calcul.

Installer des dépendances

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'

Si vous recevez une erreur lors de l'installation des roues pour torch, torch_xla ou torchvision comme pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222, rétrogradez votre version avec cette commande:

pip3 install setuptools==62.1.0

Exécuter un script avec PJRT

unset LD_PRELOAD

Vous trouverez ci-dessous un exemple d'utilisation d'un script Python pour effectuer un calcul sur une 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)"'

Un résultat semblable à celui-ci doit s'afficher :

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

Suivez les tutoriels PyTorch de ce document pour commencer à entraîner v5e avec PyTorch.

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session. Pour supprimer une ressource en file d'attente, supprimez d'abord la tranche, puis la ressource en file d'attente en deux étapes:

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

Ces deux étapes peuvent également être utilisées pour supprimer les requêtes de ressources en file d'attente qui sont à l'état FAILED.

Exemples JAX/FLAX

Les sections suivantes décrivent des exemples d'entraînement de modèles JAX et FLAX sur un TPU v5e.

Entraîner ImageNet sur v5e

Ce tutoriel explique comment entraîner ImageNet sur v5e à l'aide de fausses données d'entrée. Si vous souhaitez utiliser des données réelles, consultez le fichier README sur GitHub.

Prérequis

Créez des variables d'environnement :

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

Créez une ressource TPU:

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

Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE:

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

Lorsque QueuedResource est dans l'état ACTIVE, le résultat ressemble à ce qui suit:

 state: ACTIVE

Installez la dernière version de JAX et de 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'

Clonez le modèle ImageNet et installez les exigences correspondantes:

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'

Pour générer de fausses données, le modèle a besoin d'informations sur les dimensions de l'ensemble de données. Vous pouvez le faire à partir des métadonnées de l'ensemble de données 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'

Entraîner le modèle

Une fois toutes les étapes précédentes terminées, vous pouvez entraîner le modèle.

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'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session.

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

Modèles Hugging Face FLAX

Les modèles Hugging Face implémentés dans FLAX fonctionnent immédiatement sur Cloud TPU v5e. Cette section fournit des instructions pour exécuter des modèles populaires.

Entraîner ViT sur Imagenette

Ce tutoriel vous explique comment entraîner le modèle Vision Transformer (ViT) de HuggingFace à l'aide de l'ensemble de données Imagenette de Fast AI sur Cloud TPU v5e.

Le modèle ViT a été le premier à entraîner avec succès un encodeur Transformer sur ImageNet, avec d'excellents résultats par rapport aux réseaux convolutifs. Pour en savoir plus, consultez les ressources suivantes:

Prérequis

Créez des variables d'environnement :

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

Créez une ressource TPU:

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

Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE:

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

Lorsque la ressource mise en file d'attente est dans l'état ACTIVE, le résultat ressemble à ce qui suit:

 state: ACTIVE

Installez JAX et sa bibliothèque:

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'

Téléchargez le dépôt Hugging Face et les exigences d'installation:

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

Téléchargez l'ensemble de données 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'

Entraîner le modèle

Entraînez le modèle avec un tampon prémappé de 4 Go.

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'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session.

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

Résultats du benchmark de la vidéo interactive

Le script d'entraînement a été exécuté sur v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant présente les débits avec différents types d'accélérateurs.

Type d'accélérateur	v5litepod-4	v5litepod-16	v5litepod-64
Époque	3	3	3
Taille du lot global	32	128	512
Débit (exemples/s)	263,40	429,34	470,71

Train Diffusion sur Pokémon

Ce tutoriel vous explique comment entraîner le modèle Stable Diffusion de HuggingFace à l'aide de l'ensemble de données Pokémon sur Cloud TPU v5e.

Le modèle Stable Diffusion est un modèle de texte vers image latent qui génère des images photoréalistes à partir de n'importe quelle entrée textuelle. Pour en savoir plus, consultez les ressources suivantes :

Prérequis

Configurez un bucket de stockage pour la sortie de votre modèle.

gcloud storage buckets create gs://your_bucket 

--project=your_project 

--location=us-west1 

export GCS_BUCKET_NAME=your_bucket

Créez des variables d'environnement

export GCS_BUCKET_NAME=your_bucket
export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west1-c
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

Description des options de commande

Variable	Description
GCS_BUCKET_NAME	Affichée dans la console Google Cloud -> Cloud Storage -> Buckets
PROJECT_ID	Nom du projet Google Cloud Utilisez un projet existant ou créez-en un à l'adresse Configurer votre projet Google Cloud.
ACCELERATOR_TYPE	Consultez la page Versions de TPU pour connaître la version de votre TPU.
ZONE	Pour connaître les zones compatibles, consultez le document Régions et zones TPU.
RUNTIME_VERSION	Utilisez v2-alpha-tpuv5 pour RUNTIME_VERSION.
SERVICE_ACCOUNT	Il s'agit de l'adresse de votre compte de service, que vous pouvez trouver dans la console Google Cloud -> IAM -> Comptes de service. Par exemple: tpu-service-account@monIDdeprojet.iam.gserviceaccount.com
TPU_NAME	ID de texte attribué par l'utilisateur du TPU créé lors de l'allocation de la requête de ressource mise en file d'attente.
QUEUED_RESOURCE_ID	ID de texte attribué par l'utilisateur à la requête de ressource mise en file d'attente. Pour en savoir plus sur les ressources en file d'attente, consultez la documentation sur les ressources en file d'attente.
QUOTA_TYPE	Peut être `reserved` ou `spot`. Si aucune de ces valeurs n'est spécifiée, la valeur par défaut de QUOTA_TYPE est `on-demand`. Consultez la section Quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.
VALID_UNTIL_DURATION	Durée de validité de la requête. Pour en savoir plus sur les différentes durées valides, consultez la section Ressources mises en file d'attente.

Créez une ressource TPU:

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

Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE:

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

Lorsque la ressource mise en file d'attente est dans l'état ACTIVE, le résultat ressemble à ceci:

 state: ACTIVE

Installez JAX et sa bibliothèque.

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'

Téléchargez le dépôt HuggingFace et les exigences d'installation.

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

Entraîner le modèle

Entraînez le modèle avec un tampon prémappé de 4 Go.

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

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session.

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

Résultats des benchmarks pour la diffusion

Le script d'entraînement a été exécuté sur v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant présente les débits.

Type d'accélérateur	v5litepod-4	v5litepod-16	v5litepod-64
Étape d'entraînement	1500	1500	1500
Taille du lot global	32	64	128
Débit (exemples/s)	36,53	43,71	49,36

Entraîner GPT2 sur l'ensemble de données OSCAR

Ce tutoriel vous explique comment entraîner le modèle GPT2 de HuggingFace à l'aide de l'ensemble de données OSCAR sur Cloud TPU v5e.

GPT2 est un modèle de transformateur pré-entraîné sur des textes bruts sans étiquetage humain. Il a été entraîné pour prédire le mot suivant dans les phrases. Pour en savoir plus, consultez les ressources suivantes :

Prérequis

Créez des variables d'environnement :

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

Créez une ressource TPU:

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

Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource mise en file d'attente sera à l'état ACTIVE:

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

Lorsque la ressource mise en file d'attente est dans l'état ACTIVE, le résultat ressemble à ce qui suit:

 state: ACTIVE

Installez JAX et sa bibliothèque.

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'

Téléchargez le dépôt HuggingFace et les exigences d'installation.

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'

Téléchargez les configurations pour entraîner le modèle.

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

Entraîner le modèle

Entraînez le modèle avec un tampon prémappé de 4 Go.

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'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session.

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

Résultats des benchmarks pour GPT2

Le script d'entraînement a été exécuté sur v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant présente les débits.

	v5litepod-4	v5litepod-16	v5litepod-64
Époque	3	3	3
Taille du lot global	64	64	64
Débit (exemples/s)	74,60	72,97	72,62

PyTorch/XLA

Les sections suivantes décrivent des exemples d'entraînement de modèles PyTorch/XLA sur un TPU v5e.

Entraîner ResNet à l'aide de l'environnement d'exécution PJRT

PyTorch/XLA passe de XRT à PjRt à partir de PyTorch 2.0 et versions ultérieures. Voici les instructions mises à jour pour configurer la version 5e pour les charges de travail d'entraînement PyTorch/XLA.

Prérequis

Créez des variables d'environnement :

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

Créez une ressource TPU:

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

Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre QueuedResource sera dans l'état ACTIVE:

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

Lorsque la ressource mise en file d'attente est dans l'état ACTIVE, le résultat ressemble à ceci:

 state: ACTIVE

Installer des dépendances spécifiques à 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'

Entraîner le modèle 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'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session.

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

Résultat du benchmark

Le tableau suivant présente les débits de référence.

Type d'accélérateur	Débit (exemples/seconde)
v5litepod-4	4 240 ex/s
v5litepod-16	10 810 ex/s
v5litepod-64	46 154 ex/s

Entraîner GPT2 sur v5e

Ce tutoriel explique comment exécuter GPT2 sur v5e à l'aide du dépôt HuggingFace sur PyTorch/XLA à l'aide de l'ensemble de données wikitext.

Prérequis

Créez des variables d'environnement :

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

Créez une ressource TPU:

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

Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE:

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

Lorsque la ressource mise en file d'attente est dans l'état ACTIVE, le résultat ressemble à ceci:

state: ACTIVE

Installez les dépendances 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'

Téléchargez le dépôt HuggingFace et les exigences d'installation.

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
      '

Téléchargez les configurations du modèle pré-entraîné.

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

Entraîner le modèle

Entraînez le modèle 2B avec une taille de lot 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'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session.

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

Résultat du benchmark

Le script d'entraînement a été exécuté sur v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant présente les débits de référence pour différents types d'accélérateurs.

	v5litepod-4	v5litepod-16	v5litepod-64
Époque	3	3	3
config	600 M	2 milliards	16 Mrds
Taille du lot global	64	128	256
Débit (exemples/s)	66	77	31

Entraîner ViT sur v5e

Ce tutoriel explique comment exécuter VIT sur v5e à l'aide du dépôt HuggingFace sur PyTorch/XLA sur l'ensemble de données cifar10.

Prérequis

Créez des variables d'environnement :

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

Créez une ressource TPU:

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

Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre QueuedResource sera dans l'état ACTIVE:

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

Lorsque la ressource mise en file d'attente est dans l'état ACTIVE, le résultat ressemble à ceci:

 state: ACTIVE

Installer les dépendances 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'

Téléchargez le dépôt HuggingFace et les exigences d'installation.

   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"

Entraîner le modèle

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'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et votre ressource en file d'attente à la fin de votre session.

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

Résultat du benchmark

Le tableau suivant présente les débits de référence pour différents types d'accélérateurs.

	v5litepod-4	v5litepod-16	v5litepod-64
Époque	3	3	3
Taille du lot global	32	128	512
Débit (exemples/s)	201	657	2 844

TensorFlow 2.x

Les sections suivantes décrivent des exemples d'entraînement de modèles TensorFlow 2.x sur TPU v5e.

Entraîner Resnet sur un hôte unique v5e

Ce tutoriel explique comment entraîner ImageNet sur v5litepod-4 ou v5litepod-8 à l'aide d'un ensemble de données factice. Si vous souhaitez utiliser un autre ensemble de données, consultez la section Préparer l'ensemble de données.

Prérequis

Créez des variables d'environnement :

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 peut correspondre à v5litepod-4 ou v5litepod-8.

Créez une ressource TPU:
```
gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --${QUOTA_TYPE}
```
Remarque :L'indicateur QUOTA_TYPE peut être reserved ou best-effort. Si aucune de ces options n'est spécifiée, la valeur par défaut est on-demand. Pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU, consultez la section Quotas.

Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource mise en file d'attente sera à l'état ACTIVE. Pour vérifier l'état de votre ressource mise en file d'attente, utilisez la commande suivante:
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}
```

Se connecter à votre TPU à l'aide de SSH

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

Définir des variables d'environnement

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

Accédez au répertoire du dépôt de modèles et aux exigences d'installation.

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

Entraîner le modèle

Exécutez le script d'entraînement.

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"

Supprimer le TPU et la ressource en file d'attente

Supprimer votre TPU

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

Supprimer votre demande de ressource en file d'attente

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

Entraîner Resnet sur une version v5e multi-hôte

Ce tutoriel explique comment entraîner ImageNet sur v5litepod-16 ou plus à l'aide d'un ensemble de données factice. Si vous souhaitez utiliser un autre ensemble de données, consultez Préparer l'ensemble de données.

Créez des variables d'environnement :

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 peut être égal à v5litepod-16 ou supérieur.

Créez une ressource TPU:
```
gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --${QUOTA_TYPE}
```
Remarque :L'indicateur QUOTA_TYPE peut être reserved ou best-effort. Si aucune de ces options n'est spécifiée, la valeur par défaut est on-demand. Pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU, consultez la section Quotas.

Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource mise en file d'attente sera à l'état ACTIVE. Pour vérifier l'état de votre ressource mise en file d'attente, utilisez la commande suivante:
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}
```

Se connecter à votre TPU (nœud de calcul 0) à l'aide de SSH

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

Définir des variables d'environnement

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

Accédez au répertoire du dépôt de modèles et aux exigences d'installation.

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

Entraîner le modèle

Exécutez le script d'entraînement.

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*"

Supprimer le TPU et la ressource en file d'attente

Supprimer votre TPU

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

Supprimer votre demande de ressource en file d'attente

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