Entraînement Cloud TPU v5e [version Preview publique]

Preview

Ce produit ou cette fonctionnalité est soumis aux "Conditions des offres de pré-DG" dans la section "Conditions d'utilisation générales" des Conditions spécifiques au service. Les produits et fonctionnalités pré-DG sont disponibles "en l'état" et peuvent bénéficier d'une assistance limitée. Pour en savoir plus, consultez les descriptions des étapes de lancement.

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, la version v5e est optimisée pour être le produit le plus performant en termes d'entraînement, d'ajustement et d'inférence des modèles de transformateur, texte-vers-image et de réseau de neurones convolutif (CNN).

Concepts de Cloud TPU v5e

Si vous ne connaissez pas les Cloud TPU, consultez la page d'accueil de la documentation sur les TPU.

Cette section définit des concepts spécifiquement liés à la version 5e. Les composants utilisés dans d'autres configurations de TPU (par exemple, les segments d'application, les hôtes, les puces et les TensorCores) sont décrits sur la page Architecture du système TPU.

Inférence

L'inférence consiste à utiliser un modèle entraîné pour réaliser des prédictions sur de nouvelles données. Il est utilisé par le processus de diffusion.

Diffusion

L'inférence est le processus de déploiement d'un modèle de machine learning entraîné dans un environnement de production où il peut être utilisé pour effectuer des prédictions ou prendre des décisions. La latence et la disponibilité au niveau du service sont importantes pour la diffusion.

Tranches

Une tranche représente un ensemble de chips qui se trouvent toutes dans le même pod et sont connectées par des interconnexions de puces à haut débit (ICI). La version v5e présente des formes de tranche 2D. Consultez le tableau de la section Types d'accélérateurs pour connaître les formes de tranche acceptées.

La forme de puce et la topologie de puce font également référence aux formes de tranche.

Hôte unique ou multi-hôte

Un hôte est un ordinateur physique (processeur) qui exécute des VM. Un hôte peut exécuter plusieurs VM à la fois.

Les segments utilisant moins de huit chips n'utilisent qu'un seul hôte. Les segments comportant plus de huit chips ont accès à plusieurs hôtes et peuvent exécuter un entraînement distribué à l'aide de plusieurs hôtes. Consultez la page Architecture du système TPU pour en savoir plus sur les tranches et les puces.

la v5e prend en charge l'entraînement multi-hôte et l'inférence sur un seul hôte. Cet aperçu n'est pas compatible avec l'inférence ou l'inférence multi-hôtes. Le nombre maximal de chips autorisés dans une tâche de diffusion est de huit.

VM TPU

Machine virtuelle exécutant Linux et ayant accès aux TPU sous-jacents Pour les TPU v5e, chaque VM TPU a un accès direct à 1, 4 ou 8 puces, selon le type d'accélérateur spécifié par l'utilisateur. Une VM TPU est également appelée nœud de calcul.

Nœud de calcul

Consultez la section VM TPU.

Architecture du système v5e

Chaque puce v5e contient un seul TensorCore. Chaque TensorCore comporte quatre unités de multiplication matricielle (MXU), une unité vectorielle et une unité scalaire. Le tableau suivant présente les spécifications des chips de clé et leurs valeurs pour une version v5e. Les spécifications des pods sont incluses dans le tableau qui suit les spécifications des puces.

Caractéristiques des puces de clés	Valeurs v5e
Pic de calcul par puce (bf16)	197 TFLOPs
Pic de calcul par puce (Int8)	393 TFLOPs
Capacité et bande passante HBM2	16 Go, 819 Go/s
Interchip Interconnect BW	1 600 Gbit/s

Spécifications clés des pods	Valeurs v5e
Taille du pod TPU	256 chips
Topologie d'interconnexion	2D Torus
Pic de calcul par pod	100 PetaOps(Int8)
Réduction totale de la bande passante par pod	51,2 To/s
Bande passante bissectionnelle par pod	1,6 To/s
Bande passante réseau des centres de données par pod	6,4 Tbit/s

Puce TPU v5e

Le schéma suivant illustre une puce TPU v5e.

Types d'accélérateurs

Cloud TPU v5e est un produit combiné d'entraînement et d'inférence (inférence). L'option AcceleratorType permet de différencier l'environnement TPU provisionné pour l'entraînement de l'environnement provisionné pour la diffusion.

Les tâches d'entraînement sont optimisées pour le débit et la disponibilité, tandis que les tâches de diffusion sont optimisées pour la latence. Ainsi, un job d'entraînement sur des TPU provisionnés pour la diffusion peut avoir une mauvaise expérience (disponibilité inférieure) et, de la même manière, un job de diffusion exécuté sur des TPU provisionnés pour l'entraînement peut s'avérer médiocre (latence plus élevée).

AcceleratorType utilise le nombre de cœurs de Tensor dans le nœud pour décrire la taille de la tranche. AcceleratorType est une chaîne au format v$VERSION_NUMBER-$CORES_COUNT.

Pour en savoir plus sur les types d'accélérateurs, consultez la section Configurations TPU compatibles.

Un v5e a un cœur par puce. Les formes de tranche 2D suivantes sont compatibles avec la version 5es:

Topology	Nombre de puces TPU	Nombre d'hôtes
1x1	1	1/8
2x2	4	1/2
2x4	8	1
4x4	16	2
4x8	32	4
8x8	64	8
8x16	128	16
16x16	256	32

Pour connaître le nombre de VM par topologie, consultez la section Types de VM.

Types de Cloud TPU v5e

L'entraînement est pris en charge jusqu'à 256 chips.

Pour provisionner des TPU pour une tâche d'entraînement v5e, utilisez l'un des types d'accélérateurs suivants dans votre requête de création de TPU ou de CLI:

• V5E-1
• V5E-4
• V5E-8
• V5E-16
• V5E-32
• V5E-64
• V5E-128
• V5E-256

Le provisionnement pour l'inférence est décrit dans l'introduction de l'inférence v5e.

Types de VM TPU

Chaque VM TPU d'une tranche de TPU contient une, quatre ou huit puces. Les segments à quatre puces ou plus petites auront le même nœud NUMA (Non Uniform Memory Access). Pour en savoir plus sur les nœuds NUMA, consultez le document sur l'architecture du système. Pour les VM TPU à 8 puces, la communication CPU-TPU est plus efficace au sein des partitions NUMA. Par exemple, dans la figure suivante, la communication CPU0-Chip0 est plus rapide que la communication CPU0-Chip4.

Premiers pas

Sécuriser la capacité

Cloud TPU v5e est désormais en version Preview publique. Pour commencer à utiliser Cloud TPU v5e pour vos charges de travail d'IA, veuillez contacter le service commercial Cloud.

Préparer un projet Google Cloud

1. Connectez-vous à votre compte Google. Si ce n'est pas déjà fait, créez un compte.
2. Dans la console Google Cloud, sélectionnez ou créez un projet Cloud sur la page du sélecteur de projet.

3. La configuration de la facturation est requise pour toute utilisation de Google Cloud. Assurez-vous donc que la facturation est activée pour votre projet.

   La facturation de l'utilisation de la version Preview publique respecte les tarifs régionaux standards de la page Tarifs de Cloud TPU.

4. Installez les composants gcloud alpha.

5. Si vous utilisez TPU et réutilisez des composants gcloud alpha existants, veuillez les mettre à jour pour vous assurer que les commandes et les indicateurs pertinents sont pris en charge:

   gcloud components update
   

6. Activez l'API TPU via la commande gcloud suivante dans Cloud Shell. (Vous pouvez également l'activer à partir de la console Google Cloud.)

   gcloud services enable tpu.googleapis.com
   

7. Activez le compte de service TPU.

   Les comptes de service permettent au service Cloud TPU d'accéder à d'autres services Google Cloud. L'utilisation d'un compte de service géré par l'utilisateur est recommandée pour Google Cloud. Suivez ces guides pour créer et attribuer des rôles. Les rôles suivants sont nécessaires:

   • Administrateur TPU
   • Administrateur de l'espace de stockage: autorisation requise pour accéder à Cloud Storage
   • Rédacteur de journaux: nécessaire pour écrire des journaux avec l'API Logging
   • Rédacteur de métriques Monitoring: nécessaire pour écrire des métriques dans Cloud Monitoring.

8. Configurez le projet et la zone.

   Votre ID de projet est le nom de votre projet affiché dans la console Cloud.

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

Provisionner l'environnement Cloud TPU

Les Cloud TPU v5es sont gérés en tant que ressources en file d'attente. Vous pouvez provisionner de la capacité à l'aide de la commande queued-resource create.

Créez des variables d'environnement

Définissez les variables d'environnement nécessaires à la création de TPU:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-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
   

Variable	Description
PROJECT_ID	Nom du projet Google Cloud
ACCELERATOR_TYPE	Consultez la section Types d'accélérateurs pour connaître les formes de tranche acceptées.
ZONE	Toutes les capacités de Version preview sont situées dans la région us-west4-a.
RUNTIME_VERSION	Utilisez la version v2-alpha-tpuv5-lite pour la version Preview publique.
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. Exemple: [e-mail]
TPU_NAME	ID de texte attribué par l'utilisateur au TPU, créé lorsque la requête de ressources en file d'attente est allouée.
ID_RESSOURCE_QUEUED	ID de texte attribué par l'utilisateur pour la requête de ressources en file d'attente. Consultez le document Ressources en file d'attente pour en savoir plus sur les ressources en file d'attente.

Créer 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} \
   --service-account ${SERVICE_ACCOUNT} \
   --reserved

Remarque :L'option --reserved n'est nécessaire que pour les ressources réservées, et non pour les ressources à la demande.

Si la ressource en file d'attente est correctement créée, l'état dans le champ response sera ACCEPTED ou FAILED. Si la ressource en file d'attente est à l'état ACCEPTED, cela signifie qu'elle a passé une validation préliminaire et qu'elle est en attente de capacité. Une fois la capacité disponible, la demande passera à PROVISIONING. L'état ACCEPTED ne signifie pas que vous recevrez le quota alloué. Le passage de l'état ACCEPTED à l'état ACTIVE peut prendre un certain temps. Si la ressource en file d'attente est à l'état FAILED, le motif de l'échec est indiqué dans la sortie.

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

Utilisez les commandes list ou describe pour interroger l'état de votre ressource en file d'attente.

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

La valeur state représente l'état d'une ressource en file d'attente. Les états sont définis comme suit:

État	Description
ACCEPTÉ	La commande create de la ressource en file d'attente a été reçue et commencera le provisionnement dès que de la capacité sera disponible.
PROVISIONING	Les tranches de TPU sont en cours de provisionnement.
ACTIVE	Tous les TPU sont provisionnés et prêts à être utilisés. Si un script de démarrage est fourni, il commence à s'exécuter sur tous les TPU lorsque l'état de la ressource en file d'attente passe à ACTIVE.
FAILED	Impossible de provisionner les segments.
SUSPENDING	Un ou plusieurs secteurs sont en cours de suppression.
SUSPENDED	Toutes les tranches sous-jacentes sont supprimées, mais la ressource en file d'attente reste intacte, jusqu'à ce qu'elle soit explicitement supprimée. Pour le moment, une ressource suspendue en file d'attente ne peut pas être réactivée et doit être supprimée.
DELETING	La ressource en file d'attente est en cours de suppression.

Important :L'état ACCEPTED va être abandonné. Elle sera remplacée par WAITING_FOR_RESOURCES.

Se connecter en SSH aux VM TPU

La section suivante explique comment installer des binaires sur chaque VM TPU de votre tranche de TPU et exécuter du code. Dans ce contexte, une VM TPU est également appelée nœud de calcul.

Consultez la section Types de VM pour calculer le nombre de VM de votre tranche.

Pour installer les binaires ou exécuter du code, vous devez vous connecter en SSH à la VM TPU à l'aide de tpu-vm ssh command.

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}

Pour vous connecter en SSH à une VM TPU spécifique, utilisez l'option --worker qui suit un index basé sur 0:

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

Si vos formes de tranche sont supérieures à huit chips, vous aurez plusieurs VM dans une tranche. Dans ce cas, utilisez 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 séparément. Exemple :

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} /
--zone  ${ZONE} \
--worker=all \
--command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Gérer

Toutes les commandes que vous pouvez utiliser pour gérer vos VM TPU sont décrites dans la section Gérer des TPU.

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 PyTorch avec TPU v5e. La compatibilité avec TensorFlow sera bientôt activée.

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

Configuration pour JAX

Si vos formes de tranche sont supérieures à huit chips, vous aurez plusieurs VM dans une 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 passer par SSH pour vous connecter à chacune d'elles séparément:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} \
--zone  ${ZONE} \
--worker=all \
--command='pip install "jax[tpu]==0.4.6" -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 sorties présenté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 peut exécuter des opérations de base:

gcloud alpha 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 ce segment.

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

Essayez les tutoriels JAX de ce document pour vous familiariser avec l'entraînement v5e à l'aide de JAX.

Configuration pour PyTorch

Notez que la version 5e 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 dans la version 5e avec PyTorch/XLA avec des commandes pour tous les nœuds de calcul.

Installer des dépendances

gcloud alpha 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 numactl libopenblas-dev -y
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
   pip3 install torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl
   pip3 install torch_xla[tpuvm]
'

Si vous n'avez pas réussi à installer wheels for torch/torch_xla/torchvision et qu'un message d'erreur tel que pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222 s'affiche, passez à une version antérieure à l'aide de la commande suivante:

pip3 install setuptools==62.1.0

Exécutez un script simple avec PJRT:

Pour les modèles disposant d'allocations importantes et fréquentes, l'utilisation de tcmalloc peut considérablement améliorer la durée 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. Toutefois, en fonction de votre charge de travail (par exemple, avec DLRM, qui dispose d'allocations très importantes pour ses tables d'intégration), tcmalloc peut entraîner un ralentissement. Dans ce cas, vous pouvez essayer d'annuler la définition de la variable suivante à l'aide de la valeur par défaut malloc:

unset LD_PRELOAD

Voici un exemple utilisant un script Python pour effectuer un calcul simple sur une VM v5e:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker all --command='
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/$USER/.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/$USER/.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 faire vos premiers pas avec l'entraînement v5e à l'aide de 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 la tranche, puis la ressource en file d'attente en deux étapes:

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

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

Ces deux étapes peuvent également servir à supprimer les demandes de ressources en file d'attente qui sont à l'état FAILED.

Surveiller et profiler

Cloud TPU v5e permet la surveillance et le profilage à l'aide des mêmes méthodes que celles des générations précédentes de Cloud TPU. Pour en savoir plus sur le profilage, consultez la page Profiler votre modèle avec les outils Cloud TPU. Pour en savoir plus sur la surveillance, consultez la page sur la surveillance des VM Cloud TPU.

Exemples JAX/FLAX

Entraîner ImageNet sur la version 5e

Ce tutoriel explique comment entraîner ImageNet dans la version 5e à 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

1. Créez des variables d'environnement :

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-test-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   

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

   Remarque:L'option --reserved n'est nécessaire que pour les ressources réservées, et non pour les ressources à la demande.
   Remarque:Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE:

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

   When the QueuedResource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. Installez la dernière version de JAX et jaxlib:

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

   Remarque :Si vous voyez une erreur vous suggérant d'utiliser jax>=0.4.8, vous pouvez ignorer ce message.

4. Clonez le modèle ImageNet et installez les éléments requis:

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='git clone -b throughput_testing https://github.com/gkroiz/flax.git && cd flax && pip install -e . && cd examples/imagenet && pip install -r requirements.txt && pip install --upgrade git+https://github.com/google/flax.git'
   

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 données peuvent être recueillies à partir des métadonnées de l'ensemble de données ImageNet:

   gcloud alpha 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. Le dépôt GitHub contient plusieurs fichiers de configuration. Vous pouvez également personnaliser les configurations en les remplaçant au niveau de la ligne de commande.

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='cd flax/examples/imagenet && JAX_USE_PJRT_C_API_ON_TPU=1 JAX_PLATFORMS=tpu TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0 TPU_MIN_LOG_LEVEL=0 python3 imagenet_fake_data_benchmark.py --config=configs/fake_data_16_chips.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 alpha compute tpus tpu-vm delete $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --quiet
   

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

Modèles FLAX pour le visage

Modèles de faces Hugging implémentés en fichier FLAX et 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 imagenette de Fast AI sur Cloud TPU v5e.

Le modèle ViT est 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érequis

1. Créez des variables d'environnement :

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-test-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   

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

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

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

   When the queued resource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. Installez JAX et sa bibliothèque:

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Téléchargez le dépôt Hugging Face et installez la configuration requise:

   gcloud alpha 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/vision/requirements.txt'
   

5. Téléchargez l'ensemble de données Imagenette:

   gcloud alpha 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 alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='cd transformers && TPU_PREMAPPED_BUFFER_SIZE=4294967296 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 alpha compute tpus tpu-vm delete $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --quiet
   

   gcloud alpha 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 les versions v5e-4, v5e-16 et v5e-64. Le tableau suivant présente les débits avec différents types d'accélérateurs.

Type d'accélérateur	V5E-4	V5E-16	V5E-64
Epoch	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 explique comment entraîner le modèle de diffusion stable à partir de HuggingFace à l'aide de l'ensemble de données Pokémon sur Cloud TPU v5e.

Le modèle de diffusion stable est un modèle latent texte-image 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 la diffusion stable
• Source du code de diffusion stable

Prérequis

1. Créez des variables d'environnement :

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-test-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   

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

   Remarque:L'option --reserved n'est nécessaire que pour les ressources réservées, et non pour les ressources à la demande.
   Remarque:Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE:

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

   When the queued resource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. Installez JAX et sa bibliothèque.

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Téléchargez le dépôt HuggingFace et la configuration requise.

   gcloud alpha 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 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 alpha compute tpus tpu-vm ssh $TPU_NAME --=$PROJECT=${PROJECT_ID} --zone=$ZONE --worker=all --command='cd diffusers/examples/text_to_image && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --fake_data=1 --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 votre ressource en file d'attente à la fin de votre session.

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --quiet
   

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

Résultats de l'analyse comparative pour la diffusion

Le script d'entraînement était exécuté sur les versions v5e-4, v5e-16 et v5e-64. Le tableau suivant indique les débits.

Type d'accélérateur	V5E-4	V5E-16	V5E-64
Étape d'entraînement	1 500	1 500	1 500
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 à partir de HuggingFace en utilisant l'ensemble de données OSCAR sur Cloud TPU v5e.

GPT2 est un modèle Transformer pré-entraîné sur des textes bruts 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érequis

1. Créez des variables d'environnement :

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-test-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   

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

   Remarque:L'option --reserved n'est nécessaire que pour les ressources réservées, et non pour les ressources à la demande.
   Remarque:Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre ressource en file d'attente sera à l'état ACTIVE:

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

   When the queued resource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. Installez JAX et sa bibliothèque.

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Téléchargez le dépôt HuggingFace et la configuration requise.

   gcloud alpha 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 alpha 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-v2-images-artifacts/config/v5litepod-preview/jax/gpt .'
   

Entraîner le modèle

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

   gcloud alpha 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 alpha compute tpus tpu-vm delete $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --quiet
   

   gcloud alpha 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 était exécuté sur les versions v5e-4, v5e-16 et v5e-64. Le tableau suivant indique les débits.

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

PyTorch/XLA

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

PyTorch/XLA migre de XRT vers PjRt depuis PyTorch 2.0+. Voici les instructions mises à jour pour configurer la version 5e pour les charges de travail d'entraînement PyTorch/XLA.

Prérequis

1. Créez des variables d'environnement :

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-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=your_queued_resource_id
   

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

   Remarque:L'option --reserved n'est nécessaire que pour les ressources réservées, et non pour les ressources à la demande.
   Remarque:Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre QueuedResource sera à l'état ACTIVE:

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

   When the queued resource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. Cloner PyTorch/XLA

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command='cd /usr/share/
   sudo git clone --recursive https://github.com/pytorch/pytorch
   cd pytorch/
   sudo git clone --recursive https://github.com/pytorch/xla.git
   cd xla/
   sudo git checkout 3b10663
   sudo pip3 install numpy'
   

4. Installer Torch, Torch/xla, Torchvision

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} --project ${PROJECT_ID} --worker=all --command='
   pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch-2.0%2B20230222-cp38-cp38-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-2.0%2B20230222-cp38-cp38-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torchvision-2.0%2B20230222-cp38-cp38-linux_x86_64.whl
   sudo rm -rf /usr/local/lib/python3.8/dist-packages/libtpu*
   sudo pip3 install torch_xla[tpuvm]
   '
   

5. Installer des dépendances spécifiques à Torch/XLA

   gcloud alpha 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 numactl libopenblas-dev -y
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
   pip3 install torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl
   pip3 install torch_xla[tpuvm]
   '
   

6. Télécharger libtpu

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command="
   sudo gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/20230216/libtpu.so /home/$USER/.local/lib/python3.8/site-packages/libtpu/libtpu.so
   "
   

7. Mettre à jour tpu.py avec la configuration v5e

   gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command="
   sudo gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/20230216/tpu.py /home/$USER/.local/lib/python3.8/site-packages/torch_xla/experimental/tpu.py
   "
   

Entraîner le modèle ResNet

      gcloud alpha 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/$USER/.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 alpha compute tpus tpu-vm delete $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --quiet

   gcloud alpha 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 l'analyse comparative.

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

Entraîner GPT2 sur la version 5e

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

Prérequis

1. Créez des variables d'environnement :

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-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
   

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

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

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

   When the queued resource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. Installez les dépendances Torch/xla.

   gcloud alpha 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 numactl libopenblas-dev -y
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
   pip3 install torchvision-0.16.0a0+fc838ad-cp310-cp310-linux_x86_64.whl
   pip3 install torch_xla[tpuvm]
   '
   

4. Téléchargez le dépôt HuggingFace et la configuration requise.

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command="
   rm -rf transformers/
   git clone https://github.com/huggingface/transformers.git; \
   cd transformers; \
   sudo pip3 install -e .; \
   git checkout -q ebdb185befaa821304d461ed6aa20a17e4dc3aa2; \
   pip3 install datasets; \
   pip3 install evaluate; \
   pip3 install scikit-learn;"
   

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

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command="
   cd transformers &&
   gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/GPT2/my_config_2.json examples/pytorch/language-modeling/
   gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/GPT2/fsdp_config.json examples/pytorch/language-modeling/
   "
   

Entraîner le modèle

Entraînez le modèle 2B avec une taille de lot de 16.

   gcloud alpha 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 XLA_USE_BF16=1
   export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
   export TPU_LIBRARY_PATH=/home/$USER/.local/lib/python3.10/site-packages/libtpu/libtpu.so
   cd transformers
   numactl --cpunodebind=0 python3 -u 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
    --output_dir /tmp/test-clm
    --overwrite_output_dir
    --config_name examples/pytorch/language-modeling/my_config_2.json --cache_dir /tmp
    --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 alpha compute tpus tpu-vm delete $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --quiet
   

   gcloud alpha 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 était exécuté sur les versions v5e-4, v5e-16 et v5e-64. Le tableau suivant présente les débits de benchmark pour différents types d'accélérateurs.

	V5E-4	V5E-16	V5E-64
Epoch	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 VIT dans la version 5e à l'aide du dépôt HuggingFace sur PyTorch/XLA sur l'ensemble de données cifar10.

Prérequis

1. Créez des variables d'environnement :

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-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
   

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

   Remarque:L'option --reserved n'est nécessaire que pour les ressources réservées, et non pour les ressources à la demande.
   Remarque:Vous pourrez vous connecter en SSH à votre VM TPU une fois que votre QueuedResource sera à l'état ACTIVE:

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

   When the queued resource is in the ACTIVE state, the output will be similar to the following:
   …
   state: ACTIVE
   …
   

3. Installer les dépendances Torch/xla

   gcloud alpha 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 numactl libopenblas-dev -y
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly+20230726-cp310-cp310-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly+20230726-cp310-cp310-linux_x86_64.whl
   gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
   pip3 install torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl
   pip3 install torch_xla[tpuvm]
   '
   

4. Téléchargez le dépôt HuggingFace et la configuration requise.

   gcloud alpha 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 alpha 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/$USER/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers;
      numactl --cpunodebind=0 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 alpha compute tpus tpu-vm delete $TPU_NAME 

   --project=${PROJECT_ID} 

   --zone=$ZONE 

   --quiet
   

   gcloud alpha 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 pour différents types d'accélérateurs.

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

Assistance et commentaires

Tous vos commentaires sont les bienvenus ! Pour nous envoyer vos commentaires ou demander de l'aide, cliquez ici ou envoyez un e-mail à [e-mail].

Conditions d'utilisation

Toutes les informations que Google vous a fournies à propos de cette Version preview sont des informations confidentielles de Google et sont soumises aux dispositions de confidentialité énoncées dans les Conditions d'utilisation de Google Cloud Platform (ou tout autre accord régissant votre utilisation de Google Cloud Platform).