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 un encombrement inférieur de 256 puces par pod, le modèle v5e est optimisé pour l'entraînement, l'affinage et l'inférence des données Transformer, de la génération de texte à l'image et des réseaux de neurones convolutifs (CNN). Pour en savoir plus sur l'utilisation de Cloud TPU v5e pour la diffusion, consultez la page Inférence avec v5e.

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

Démarrer

Les sections suivantes décrivent comment commencer à utiliser TPU v5e.

Quota de requêtes

Vous avez besoin d'un quota d'utilisation de 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 TPU Spot. Si vous utilisez votre TPU v5e pour l'inférence, des quotas distincts sont requis. Pour en savoir plus sur les quotas, consultez la page Quotas. Pour demander un quota de TPU v5e, contactez le service commercial Cloud.

Créer un projet et un compte Google Cloud

Vous devez disposer d'un compte et d'un projet Google Cloud pour utiliser Cloud TPU. Pour en savoir plus, consultez la section Configurer un environnement Cloud TPU.

Remarque :Pour utiliser TPU v5e, vous devez installer gcloud alpha components.

Créer une instance Cloud TPU

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

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

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

Configuration du framework

Cette section décrit le processus de configuration générale pour l'entraînement de modèle personnalisé à l'aide de JAX ou PyTorch avec TPU v5e. TensorFlow est compatible avec les versions d'exécution de TPU tpu-vm-tf-2.16.1-pjrt et tpu-vm-tf-2.16.1-pod-pjrt.

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

Configuration pour JAX

Si vous avez des formes de tranches supérieures à 8 chips, vous disposez de plusieurs VM dans une même tranche. Dans ce cas, vous devez utiliser l'option --worker=all pour exécuter l'installation sur toutes les VM TPU en une seule étape, sans avoir à vous connecter en SSH à chacune d'elles:

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'

Vous pouvez exécuter la commande suivante pour vérifier le nombre d'appareils (les résultats présentés ici ont été produits avec une tranche v5litepod-16). Ce code teste que tout est correctement installé en vérifiant que JAX voit les TensorCores de Cloud TPU et 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 chips 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 l'entraînement de la v5e à l'aide de 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+ utilisera PJRT comme environnement d'exécution par défaut pour toutes les versions de TPU.

Cette section explique comment commencer à utiliser PJRT sur 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, revenez à une version antérieure à l'aide de la commande suivante:

pip3 install setuptools==62.1.0

Exécuter un script avec PJRT

Remarque :Pour les modèles comportant des allocations importantes et fréquentes, l'utilisation de tcmalloc peut améliorer considérablement le temps d'entraînement par rapport à l'implémentation malloc par défaut. Par conséquent, la valeur malloc par défaut utilisée sur une VM TPU est tcmalloc. Cependant, en fonction de votre charge de travail (par exemple, avec DLRM qui dispose d'allocations très importantes pour ses tables de représentations vectorielles continues), tcmalloc peut entraîner un ralentissement. Dans ce cas, vous pouvez essayer d'annuler la définition de la variable suivante en utilisant la valeur par défaut malloc à la place:

unset LD_PRELOAD

Voici un exemple utilisant 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')

Essayez les tutoriels PyTorch de ce document pour commencer à entraîner le modèle v5e à l'aide de PyTorch.

Supprimez votre TPU et la ressource en file d'attente à la fin de votre session. Pour supprimer une ressource en file d'attente, supprimez 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 dont l'état est FAILED.

Exemples JAX/FLAX

Les sections suivantes décrivent des exemples d'entraînement des modèles JAX et FLAX sur 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éparation

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

2. Créez une ressource 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}
   

   Remarque :L'option 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. Consultez la section quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.

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

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

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

    state: ACTIVE
   

3. 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 "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

   Remarque :Si une erreur s'affiche vous suggérant d'utiliser jax>=0.4.8, vous pouvez ignorer le message en toute sécurité.

4. Clonez le modèle ImageNet et installez la configuration requise correspondante:

   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. Pour générer des données fictives, le modèle a besoin d'informations sur les dimensions de l'ensemble de données. Ces informations peuvent être collectées à 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 la 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 FLAX Hugging Face

Les modèles Hugging Face implémentés en FLAX sont prêts à l'emploi sur Cloud TPU v5e. Cette section fournit des instructions pour exécuter des modèles populaires.

Entraîner ViT sur Imagenette

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

Le modèle ViT a été le premier à avoir entraîné 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ésentation de ViT
• Ressource de code ViT

Préparation

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

2. Créez une ressource 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}
   

   Remarque :L'option 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. Consultez la section quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.

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

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

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

    state: ACTIVE
   

3. 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'
   

4. Téléchargez le dépôt Hugging Face et procédez comme suit:

   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. 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 la 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 de l'analyse comparative ViT

Le script d'entraînement a été exécuté sur v5litepod-4, v5litepod-16 et v5litepod-64. Le tableau suivant indique 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

Entraînement Diffusion sur Pokémon

Ce tutoriel explique comment entraîner le modèle Stable Diffusion à partir 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 texte-image latent qui génère des images photoréalistes à partir de n'importe quelle entrée de texte. Pour en savoir plus, consultez les ressources suivantes :

• Présentation de Stable Diffusion
• Source du code Stable Diffusion

Préparation

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

2. Créez une ressource 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}
   

   Remarque :L'option 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. Consultez la section quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.

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

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

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

    state: ACTIVE
   

3. 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'
   

4. Téléchargez le dépôt HuggingFace et configurez les conditions requises pour l'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.16.1 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} \
   --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'

Supprimer le TPU et la ressource en file d'attente

Supprimez votre TPU et la 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 de l'analyse comparative de la diffusion

Le script d'entraînement a été exécuté sur les versions 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 explique comment entraîner le modèle GPT2 depuis HuggingFace en utilisant l'ensemble de données OSCAR sur Cloud TPU v5e.

GPT2 est un modèle Transformer pré-entraîné sur du texte brut sans étiquetage manuel. Il a été entraîné pour prédire le mot suivant dans des phrases. Pour en savoir plus, consultez les ressources suivantes :

• Présentation de GPT2
• Source du code GPT2

Préparation

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

2. Créez une ressource 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}
   

   Remarque :L'option 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. Consultez la section quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.

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

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

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

    state: ACTIVE
   

3. 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'
   

4. Téléchargez le dépôt HuggingFace et procédez à l'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'
   

5. 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 && gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt .'
   

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 la 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 de l'analyse comparative pour GPT2

Le script d'entraînement a été exécuté sur les versions 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 des modèles PyTorch/XLA sur TPU v5e.

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

PyTorch/XLA migre de XRT vers PjRt depuis PyTorch 2.0 ou version ultérieure. Voici les instructions mises à jour pour configurer v5e pour les charges de travail d'entraînement PyTorch/XLA.

Préparation

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

2. Créez une ressource 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}
   

   Remarque :L'option 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. Consultez la section quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.

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

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

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

    state: ACTIVE
   

3. 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 la 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 de l'analyse comparative

Le tableau suivant présente les débits de benchmark.

Type d'accélérateur	Débit (exemples/seconde)
V5litepod-4	4240 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 la version v5e à l'aide du dépôt HuggingFace sur PyTorch/XLA à l'aide de l'ensemble de données wikitext.

Préparation

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

2. Créez une ressource 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}
   

   Remarque :L'option 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. Consultez la section quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.

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

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

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

   state: ACTIVE
   

3. 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'
   

4. Téléchargez le dépôt HuggingFace et procédez à l'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
         '
   

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

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

Entraîner le modèle

Entraînez le modèle 2B en utilisant 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 la 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 de l'analyse comparative

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

	V5litepod-4	V5litepod-16	V5litepod-64
Époque	3	3	3
config	600M	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 le VIT sur la version v5e à l'aide du dépôt HuggingFace sur PyTorch/XLA sur l'ensemble de données cifar10.

Préparation

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

2. Créez une ressource 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}
   

   Remarque :L'option 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. Consultez la section quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.

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

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

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

    state: ACTIVE
   

3. Installer des 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'
   

4. Téléchargez le dépôt HuggingFace et procédez à l'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 la 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 de l'analyse comparative

Le tableau suivant indique les débits de benchmark pour les 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 des modèles TensorFlow 2.x sur TPU v5e.

Entraîner Resnet sur une version v5e à hôte unique

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éparation

1. 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 être v5litepod-4 ou v5litepod-8.

2. Créez une ressource 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}
   

   Remarque :L'option 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. Consultez la section quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.

   Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente passe à l'état ACTIVE. Pour vérifier l'état de votre ressource en file d'attente, exécutez la commande suivante:

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

3. Se connecter à votre TPU à l'aide de SSH

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

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

5. Accédez au répertoire du dépôt de modèles et installez les conditions requises.

   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

1. Supprimer votre TPU

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

2. Supprimer votre requête 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 v5e à hôtes multiples

Ce tutoriel explique comment entraîner ImageNet sur v5litepod-16 ou version ultérieure à 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.

1. 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 de taille v5litepod-16 ou supérieure.

2. Créez une ressource 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}
   

   Remarque :L'option 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. Consultez la section quotas pour en savoir plus sur les différents types de quotas compatibles avec Cloud TPU.

   Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente passe à l'état ACTIVE. Pour vérifier l'état de votre ressource en file d'attente, exécutez la commande suivante:

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

3. Se connecter à votre TPU (nœud de calcul zéro) à l'aide de SSH

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

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

5. Accédez au répertoire du dépôt de modèles et installez les conditions requises.

    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

1. Supprimer votre TPU

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

2. Supprimer votre requête de ressource en file d'attente

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